multigpu_executor.py

# Copyright (C) 2025 Arcee AI
# SPDX-License-Identifier: BUSL-1.1
"""
Implementation of multi-GPU parallel task execution.

Handles distribution of parallelizable tasks across multiple GPUs while respecting:
- Main-thread-only task requirements
- Task dependency graphs
- GPU assignment of connected task components
- Intermediate result storage locations
"""

import concurrent.futures
import logging
import queue
import threading
from collections import defaultdict
from typing import Any, Dict, Iterator, List, Optional, Set, Tuple

import networkx as nx
import torch
import tqdm

from .graph import Executor, Task

logger = logging.getLogger(__name__)


class MultiGPUExecutor:
    """
    Execute tasks across multiple GPUs.

    Attributes:
        num_gpus: Number of GPUs to utilize (None = all available)
        storage_device: Device for storing tensors between stages
        targets: Final output tasks to retain results for
    """

    def __init__(
        self,
        tasks: List[Task],
        num_gpus: Optional[int] = None,
        storage_device: Optional[torch.device] = None,
    ):
        """
        Initialize the executor with a list of tasks.

        Args:
            tasks: List of tasks to execute
            num_gpus: Number of GPUs to utilize (None = all available)
            storage_device: Device for storing tensors between stages
        """
        self.results = {}
        self.targets = set(tasks)
        self.storage_device = storage_device

        if num_gpus is None:
            num_gpus = torch.cuda.device_count()

        # Create temp executor to get full schedule
        temp_exec = Executor(tasks)
        ordered_tasks = temp_exec._make_schedule(tasks)
        self.dependencies = temp_exec.dependencies
        self.total_tasks = len(ordered_tasks)

        leading_tasks = self._find_leading_tasks(ordered_tasks)
        trailing_tasks = self._find_trailing_tasks(ordered_tasks)
        self.trailing_main_tasks = [t for t in ordered_tasks if t in trailing_tasks]
        self.leading_main_tasks = [t for t in ordered_tasks if t in leading_tasks]

        self.trailing_dependencies = set()
        for task in self.trailing_main_tasks:
            self.trailing_dependencies.update(self.dependencies[task])

        parallel_tasks = [
            t
            for t in ordered_tasks
            if (t not in trailing_tasks and t not in leading_tasks)
        ]
        logger.info(
            f"Task breakdown: {len(self.leading_main_tasks)} leading, "
            f"{len(parallel_tasks)} parallel, "
            f"{len(self.trailing_main_tasks)} trailing"
        )
        if any(t.main_thread_only() for t in parallel_tasks):
            raise RuntimeError(
                "Main-thread-only tasks must be either leading or trailing"
            )
        self.gpu_assignments = self._assign_islands_to_gpus(parallel_tasks, num_gpus)

        self.task_completion_queue = queue.Queue()
        self.done_event = threading.Event()

    def run(self, quiet: bool = False) -> Iterator[Tuple[Task, Any]]:
        """
        Execute all tasks and yield target results.

        Yields:
            Iterator[Tuple[Task, Any]]: Task and result pairs
        """
        with tqdm.tqdm(
            total=self.total_tasks, disable=quiet, desc="Executing graph"
        ) as pbar:
            if self.leading_main_tasks:
                exec = Executor(
                    self.leading_main_tasks,
                    math_device=self.storage_device or torch.device("cpu"),
                    storage_device=self.storage_device or torch.device("cpu"),
                )
                for task, result in exec.run(quiet=True):
                    pbar.update()
                    self.results[task] = result

                logger.debug("Leading tasks complete, beginning parallel execution")

            def update_progress():
                while not self.done_event.is_set():
                    try:
                        task, result = self.task_completion_queue.get(timeout=0.1)
                        self.results[task] = result
                        pbar.update()
                    except queue.Empty:
                        continue

            progress_thread = threading.Thread(target=update_progress)
            progress_thread.start()

            # Run parallel tasks
            with concurrent.futures.ThreadPoolExecutor() as executor:
                futures = []
                for device, island_tasks in self.gpu_assignments.items():
                    futures.append(
                        executor.submit(
                            self._device_worker,
                            task_list=island_tasks,
                            cached_values=dict(self.results),
                            device=device,
                            quiet=True,
                        )
                    )

                for future in concurrent.futures.as_completed(futures):
                    if future.exception():
                        self.done_event.set()
                        executor.shutdown(wait=False)
                        raise future.exception()

            self.done_event.set()
            progress_thread.join()

            logger.debug("Parallel tasks complete")

            # Run main thread tasks
            if self.trailing_main_tasks:
                exec = Executor(
                    self.trailing_main_tasks,
                    math_device=self.storage_device or torch.device("cpu"),
                    storage_device=self.storage_device or torch.device("cpu"),
                    cached_values=dict(self.results),
                )
                for task, result in exec.run(quiet=True):
                    pbar.update()
                    if task in self.targets:
                        self.results[task] = result

        # Yield final results
        for task, result in self.results.items():
            if task in self.targets:
                yield task, result

    def execute(self) -> None:
        """Execute all tasks and discard results"""
        for _ in self.run(quiet=False):
            pass

    def _find_trailing_tasks(self, tasks: List[Task]) -> Set[Task]:
        """
        Identify tasks that must execute AFTER parallel GPU tasks complete.

        Trailing tasks must:
        - Require main thread execution
        - Not have non-trailing dependants
        """
        dependants = defaultdict(set)
        for task, deps in self.dependencies.items():
            for dep in deps:
                dependants[dep].add(task)

        trailing_tasks = set()
        to_explore = set([t for t in tasks if not dependants[t]])
        while to_explore:
            task = to_explore.pop()
            if not task.main_thread_only():
                continue
            if all(d in trailing_tasks for d in dependants[task]):
                trailing_tasks.add(task)
                to_explore.update(self.dependencies[task])
        return trailing_tasks

    def _find_leading_tasks(self, tasks: List[Task]) -> Set[Task]:
        """Identify tasks that must execute BEFORE parallel GPU tasks.

        Leading tasks must:
        - Require main thread execution
        - Not have non-leading dependencies
        """
        leading_tasks = set()
        for task in tasks:
            if not task.main_thread_only():
                continue
            if self.dependencies[task] and any(
                dep not in leading_tasks for dep in self.dependencies[task]
            ):
                continue
            leading_tasks.add(task)
        return leading_tasks

    def _assign_islands_to_gpus(
        self, tasks: List[Task], num_gpus: int
    ) -> Dict[torch.device, List[Task]]:
        """
        Assign task islands to GPUs.

        Task islands (weakly connected components) are groups of tasks that
        can execute independently. This method identifies islands in the
        non-trailing, non-leading task graph and assigns them to devices.
        """

        island_graph = nx.DiGraph()
        island_graph.add_nodes_from(tasks)

        # Add edges only between parallel tasks
        for task in tasks:
            for dep in self.dependencies[task]:
                if dep in tasks:
                    island_graph.add_edge(dep, task)

        islands = list(nx.weakly_connected_components(island_graph))
        logger.info(f"Found {len(islands)} islands in parallel task graph")
        assignments = {}
        for island in islands:
            # Borrow orderings from original task list
            island_tasks = [t for t in tasks if t in island]
            # assign to GPU with fewest tasks
            device_idx = min(
                range(num_gpus),
                key=lambda i: len(assignments.get(torch.device(f"cuda:{i}"), [])),
            )
            device = torch.device(f"cuda:{device_idx}")
            assignments[device] = assignments.get(device, []) + island_tasks
        return assignments

    def _device_worker(
        self,
        task_list: List[Task],
        cached_values: Dict[Task, Any],
        device: torch.device,
        quiet: bool,
    ):
        """
        Execute a set of tasks on a single GPU.

        Args:
            island_tasks: List of tasks to execute
            cached_values: Values of previously-executed dependent tasks
            device: Device to execute tasks on
            quiet: Suppress progress bar output
        """
        stream = torch.cuda.Stream(device=device)
        with torch.cuda.stream(stream):
            exec = Executor(
                tasks=task_list,
                math_device=device,
                storage_device=self.storage_device or device,
                cached_values=cached_values,
            )
            count = 0
            for task, result in exec.run(quiet=quiet):
                count += 1
                if not (task in self.targets or task in self.trailing_dependencies):
                    result = None
                self.task_completion_queue.put((task, result))
        torch.cuda.synchronize(device=device)