[tab] nemo_run/run/experiment.py (python)
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextvars
import copy
import importlib.util
import inspect
import json
import os
import pprint
import shutil
import sys
import time
import traceback
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Optional, Type, Union

import fiddle as fdl
import networkx as nx
import rich
from fiddle._src import daglish, diffing
from rich.console import Group
from rich.live import Live
from rich.panel import Panel
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskID, TimeElapsedColumn
from rich.progress import Task as RichTask
from rich.syntax import Syntax
from torchx.specs.api import AppState

import nemo_run as run
from nemo_run.config import (
    Config,
    ConfigurableMixin,
    Partial,
    Script,
    get_nemorun_home,
    get_type_namespace,
)
from nemo_run.core.execution.base import Executor
from nemo_run.core.execution.dgxcloud import DGXCloudExecutor
from nemo_run.core.execution.docker import DockerExecutor
from nemo_run.core.execution.lepton import LeptonExecutor
from nemo_run.core.execution.local import LocalExecutor
from nemo_run.core.execution.skypilot import SkypilotExecutor
from nemo_run.core.execution.skypilot_jobs import SkypilotJobsExecutor
from nemo_run.core.execution.slurm import SlurmExecutor
from nemo_run.core.frontend.console.api import CONSOLE, configure_logging, deconfigure_logging
from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
from nemo_run.core.tunnel.client import SSHTunnel, Tunnel
from nemo_run.core.tunnel.rsync import rsync
from nemo_run.run.job import Job, JobGroup
from nemo_run.run.plugin import ExperimentPlugin
from nemo_run.run.torchx_backend.runner import get_runner
from nemo_run.run.utils import TeeStdoutStderr

_current_experiment: contextvars.ContextVar["Experiment"] = contextvars.ContextVar(
    "nemo_current_experiment"
)


class DummyConsole:
    """A dummy console that mimics rich.console.Console but does nothing."""

    def __getattr__(self, name):
        """Return a no-op function for any attribute access."""

        def no_op(*args, **kwargs):
            pass

        return no_op


class Experiment(ConfigurableMixin):
    """
    A context manager to launch and manage multiple runs, all using pure Python.

    run.Experiment provides researchers with
    a simple and flexible way to create and manage their ML experiments.

    Building on the core blocks of nemo_run,
    the Experiment can be used as an umbrella under which a user can
    launch different configured functions on multiple remote clusters.

    The Experiment takes care of storing the run metadata,
    launching it on the specified cluster, and syncing the logs and artifacts.

    Additionally, the Experiment also provides management tools to easily inspect and reproduce past experiments.
    Some of the use-cases that it enables are listed below:

    1. Check the status and logs of a past experiment
    2. Reconstruct a past experiment and relaunch it after some changes
    3. Compare different runs of the same experiment.

    This API allows users to programmatically define their experiments.
    To get a glance of the flexibility provided, here are some use cases
    which can be supported by the Experiment in just a few lines of code.

    1. Launch a benchmarking run on different GPUs at the same time in parallel
    2. Launch a sequential data processing pipeline on a CPU heavy cluster
    3. Launch hyperparameter grid search runs on a single cluster in parallel
    4. Launch hyperparameter search runs distributed across all available clusters

    The design is heavily inspired from `XManager <https://github.com/google-deepmind/xmanager/blob/main/docs/xm_launch_api_principles.md>`_.

    Under the hood, the Experiment metadata is stored in the local filesystem
    inside a user specified directory controlled by get_nemorun_home() env var.
    We will explore making the metadata more persistent in the future.

    .. note::
        `Experiment.add` and `Experiment.run` methods inside Experiment can currently only be used within its context manager.

    Examples
    --------
    .. code-block:: python

        # An experiment that runs a pre-configured training example
        # on multiple GPU specific clusters (A100 and H100 shown here) in parallel using torchrun
        # Assumes that example_to_run is pre-configured using run.Partial
        with run.Experiment("example-multiple-gpus", executor="h100_cluster") as exp:
            # Set up the run on H100
            # Setting up a single task is identical to setting up a single run outside the experiment
            h100_cluster: run.SlurmExecutor = exp.executor.clone()
            h100_cluster.nodes = 2

            # torchrun manages the processes on a single node
            h100_cluster.ntasks_per_node = 1
            h100_cluster.gpus_per_task = 8

            h100_cluster.packager.subpath = "subpath/to/your/code/repo"
            h100_cluster.launcher = "torchrun"

            exp.add(
                "example_h100",
                fn=example_to_run,
                tail_logs=True,
                executor=h100_cluster,
            )

            # Set up the run on A100
            a100_cluster: run.Config[SlurmExecutor] = h100_cluster.clone()
            a100_cluster.tunnel = run.Config(
                SSHTunnel,
                host=os.environ["A100_HOST"],
                user="your_user_in_cluster",
                identity="path_to_your_ssh_key"
            )

            exp.add(
                "example_a100",
                fn=example_to_run,
                tail_logs=True,
                executor=a100_cluster,
            )

            # Runs all the task in the experiment.
            # By default, all tasks will be run in parallel if all different executors support parallel execution.
            # You can set sequential=True to run the tasks sequentially.
            exp.run()

        # Upon exiting the context manager, the Experiment will automatically wait for all tasks to complete,
        # and optionally tail logs for tasks that have tail_logs=True.
        # A detach mode (if the executors support it) will be available soon.
        # Once all tasks have completed,
        # the Experiment will display a status table and clean up resources like ssh tunnels.

        # You can also manage the experiment at a later point in time
        exp = run.Experiment.from_title("example-multiple-gpus")
        exp.status()
        exp.logs(task_id="example_a100")

    """

    GOODBYE_MESSAGE_PYTHON = """
# The experiment was run with the following tasks: {tasks}
# You can inspect and reconstruct this experiment at a later point in time using:
experiment = run.Experiment.from_id("{exp_id}")
experiment.status() # Gets the overall status
experiment.logs("{tasks[0]}") # Gets the log for the provided task
experiment.cancel("{tasks[0]}") # Cancels the provided task if still running
"""

    GOODBYE_MESSAGE_BASH = """
# You can inspect this experiment at a later point in time using the CLI as well:
nemo experiment status {exp_id}
nemo experiment logs {exp_id} 0
nemo experiment cancel {exp_id} 0
"""
    _PARALLEL_SUPPORTED_EXECUTORS = (
        SlurmExecutor,
        LocalExecutor,
        SkypilotExecutor,
        SkypilotJobsExecutor,
        DockerExecutor,
        DGXCloudExecutor,
        LeptonExecutor,
    )
    _DETACH_SUPPORTED_EXECUTORS = (
        SlurmExecutor,
        SkypilotExecutor,
        SkypilotJobsExecutor,
        DGXCloudExecutor,
        LeptonExecutor,
    )
    _DEPENDENCY_SUPPORTED_EXECUTORS = (SlurmExecutor,)
    _RUNNER_DEPENDENT_EXECUTORS = (LocalExecutor,)
    _CONFIG_FILE = "_CONFIG"
    _VERSION_FILE = "_VERSION"
    _TASK_FILE = "_TASKS"
    _DONE_FILE = "_DONE"
    _TUNNELS_FILE = "_TUNNELS"
    _current_experiment_token: Optional[contextvars.Token]

    @classmethod
    def catalog(
        cls: Type["Experiment"],
        title: str = "",
    ) -> list[str]:
        """
        List all experiments inside get_nemorun_home(), optionally with the provided title.
        """
        parent_dir = os.path.join(get_nemorun_home(), "experiments", title)
        return _get_sorted_dirs(parent_dir)

    @classmethod
    def _from_config(cls: Type["Experiment"], exp_dir: str) -> "Experiment":
        id = os.path.basename(exp_dir)
        with open(os.path.join(exp_dir, cls._CONFIG_FILE), "r") as f:
            config = f.read()

        if not config:
            raise ValueError(f"Experiment {id} not found.")

        serializer = ZlibJSONSerializer()
        cfg: Config["Experiment"] = fdl.cast(Config, serializer.deserialize(config))
        if "id" not in cfg.__arguments__:
            cfg.id = id

        cfg._reconstruct = True

        exp: "Experiment" = fdl.build(cfg)
        exp._jobs = exp._load_jobs()
        try:
            exp.tunnels = exp._load_tunnels()
        except Exception as e:
            exp.console.log(
                f"Exception {e} loading tunnels for experiment {id}, will continue without loading tunnels."
            )

        return exp

    @classmethod
    def from_id(
        cls: Type["Experiment"],
        id: str,
    ) -> "Experiment":
        """
        Reconstruct an experiment with the specified id.
        """
        title, _, _ = id.rpartition("_")
        parent_dir = os.path.join(get_nemorun_home(), "experiments", title)
        exp_dir = os.path.join(parent_dir, id)

        assert os.path.isdir(exp_dir), f"Experiment {id} not found."

        exp = cls._from_config(exp_dir)
        return exp

    @classmethod
    def from_title(
        cls: Type["Experiment"],
        title: str,
    ) -> "Experiment":
        """
        Reconstruct an experiment with the specified title.
        """
        parent_dir = os.path.join(get_nemorun_home(), "experiments", title)
        exp_dir = _get_latest_dir(parent_dir)

        assert os.path.isdir(exp_dir), f"Experiment {id} not found."

        exp = cls._from_config(exp_dir)
        return exp

    def __init__(
        self,
        title: str,
        executor: Executor | None = None,  # type: ignore
        id: str | None = None,
        log_level: str = "INFO",
        _reconstruct: bool = False,
        jobs: list[Job | JobGroup] | None = None,
        base_dir: str | None = None,
        clean_mode: bool = False,
        enable_goodbye_message: bool = True,
        threadpool_workers: int = 16,
        skip_status_at_exit: bool = False,
        serialize_metadata_for_scripts: bool = True,
    ) -> None:
        """
        Initializes an experiment run by creating its metadata directory and saving the experiment config.

        Args:
            title: Title or name for the experiment
            executor: Any executor that subclasses run.Executor and is supported by NeMo-Run.
                This will be used as the default executor for tasks if an explicit one is not specified.
                Users can also clone this and make task specific executor changes.
            id (Optional): Unique id for the experiment run.
                If not specified, will be set automatically based on the current timestamp.
            log_level: Set log level for the experiment. Defaults to WARN.
            _reconstruct: Generally, the user does not need to specify this flag.
                This is only set to True when using run.Experiment.from_dir.
            clean_mode: If True, disables all console output (logs, progress bars, etc.). Defaults to False.
            enable_goodbye_message: if True, prints goodbye message after submitting job. Defaults to True.
        """
        configure_logging(level=log_level)
        self._reconstruct = _reconstruct
        if _reconstruct:
            assert id, "Cannot reconstruct an experiment without id."

        self._title = title
        self._id = id or f"{title}_{int(time.time())}"
        self._enable_goodbye_message = enable_goodbye_message
        self._threadpool_workers = threadpool_workers
        self._skip_status_at_exit = skip_status_at_exit
        self._serialize_metadata_for_scripts = serialize_metadata_for_scripts

        base_dir = str(base_dir or get_nemorun_home())
        self._exp_dir = os.path.join(base_dir, "experiments", title, self._id)

        self.log_level = log_level
        self._runner = get_runner(component_defaults=None, experiment=self)

        if not _reconstruct:
            self.executor = executor if executor else LocalExecutor()
        else:
            assert isinstance(executor, Executor)
            self.executor = executor

        self._jobs: list[Job | JobGroup] = jobs or []
        self.tunnels: dict[str, Tunnel] = {}
        self.console = CONSOLE
        self.clean_mode = clean_mode
        if self.clean_mode:
            self.console = DummyConsole()
        self._launched = False
        self._live_progress = None
        self._current_experiment_token = None

    def to_config(self) -> Config:
        return Config(
            self.__class__,
            title=self._title,
            id=self._id,
            executor=self.executor.to_config(),
            log_level=self.log_level,
            clean_mode=self.clean_mode,
            threadpool_workers=self._threadpool_workers,
            enable_goodbye_message=self._enable_goodbye_message,
            skip_status_at_exit=self._skip_status_at_exit,
            serialize_metadata_for_scripts=self._serialize_metadata_for_scripts,
        )

    def _save_experiment(self, exist_ok: bool = False):
        os.makedirs(self._exp_dir, exist_ok=exist_ok)
        self._save_config()

    def _save_config(self):
        with open(os.path.join(self._exp_dir, self.__class__._CONFIG_FILE), "w+") as f:
            f.write(ZlibJSONSerializer().serialize(self.to_config()))

        with open(os.path.join(self._exp_dir, self.__class__._VERSION_FILE), "w+") as f:
            f.write(f"{run.__version__}
")

    def _save_tunnels(self):
        serializer = ZlibJSONSerializer()
        serialized_tunnels = {
            k: serializer.serialize(v.to_config()) for k, v in self.tunnels.items()
        }
        with open(os.path.join(self._exp_dir, self.__class__._TUNNELS_FILE), "w+") as f:
            json.dump(serialized_tunnels, f)

    def _load_tunnels(self) -> dict[str, Tunnel]:
        with open(os.path.join(self._exp_dir, self.__class__._TUNNELS_FILE)) as f:
            serialized_tunnels = json.load(f)
        serializer = ZlibJSONSerializer()
        return {k: fdl.build(serializer.deserialize(v)) for k, v in serialized_tunnels.items()}

    def _save_jobs(self):
        serialized_jobs = list(map(lambda job: job.serialize(), self.jobs))
        with open(os.path.join(self._exp_dir, self.__class__._TASK_FILE), "w+") as f:
            json.dump(serialized_jobs, f)

        if "__main__" in sys.modules:
            main_module = sys.modules["__main__"]
            try:
                with open(os.path.join(self._exp_dir, "__main__.py"), "w+") as f:
                    f.write(inspect.getsource(main_module))
            except TypeError:
                ...

    def _load_jobs(self) -> list[Job | JobGroup]:
        with open(os.path.join(self._exp_dir, self._TASK_FILE)) as f:
            serialized_jobs = json.load(f)

        serializer = ZlibJSONSerializer()
        jobs = []
        for job_cfg, task_cfg in serialized_jobs:
            job_cfg = serializer.deserialize(job_cfg)

            job: Job | JobGroup = fdl.build(job_cfg)
            if isinstance(job, Job):
                job.task = task_cfg  # type: ignore
            elif isinstance(job, JobGroup):
                job.tasks = task_cfg  # type: ignore
            else:
                raise ValueError(f"Unknown task type: {task_cfg.__fn_or_cls__}")

            jobs.append(job)

        return jobs

    def _prepare(self, exist_ok: bool = False):
        self._save_experiment(exist_ok=exist_ok)

        for job in self.jobs:
            job.prepare(serialize_metadata_for_scripts=self._serialize_metadata_for_scripts)

        self._save_jobs()

    def _add_single_job(
        self,
        task: Union[Partial, Script],
        executor: Executor,
        name: str = "",
        plugins: Optional[list[ExperimentPlugin]] = None,
        tail_logs: bool = False,
        dependencies: Optional[list[str]] = None,
    ) -> str:
        if isinstance(task, Script):
            default_name = task.get_name()
        else:
            default_name = get_type_namespace(task.__fn_or_cls__)

        reuse_job_dir = True if name else False
        name = name or default_name
        if any(map(lambda job: job.id == name, self.jobs)):
            task_id = f"{name}_{len(self.jobs)}"
        else:
            task_id = name

        self._validate_task(task_info=task_id, task=task)

        executor = executor.clone()
        executor.assign(
            self._id,
            self._exp_dir,
            task_id=task_id,
            task_dir=name if reuse_job_dir else task_id,
        )

        cloned = copy.deepcopy(task) if isinstance(task, Script) else task.clone()
        job = Job(
            id=task_id,
            task=cloned,
            executor=executor,
            plugins=plugins,
            tail_logs=tail_logs,
            dependencies=dependencies or [],
        )
        plugins = plugins or []
        for plugin in plugins:
            plugin.assign(self._id)
            plugin.setup(cloned, executor)

        self._jobs.append(job)
        return job.id

    def _add_job_group(
        self,
        tasks: list[Partial | Script],
        executor: list[Executor] | Executor,
        name: str,
        plugins: Optional[list[ExperimentPlugin]] = None,
        tail_logs: bool = False,
        dependencies: Optional[list[str]] = None,
    ) -> str:
        if any(map(lambda task: task.id == name, self.jobs)):
            task_id = f"{name}_{len(self.jobs)}"
        else:
            task_id = name

        for i, _task in enumerate(tasks):
            self._validate_task(task_info=f"Job Group: {task_id}, job index: {i}", task=_task)

        executors = executor if isinstance(executor, list) else [executor]
        cloned_executors = []
        for executor in executors:
            new_executor = executor.clone()
            cloned_executors.append(new_executor)
            new_executor.assign(self._id, self._exp_dir, task_id, task_dir=name)

        cloned_tasks = []
        for task in tasks:
            cloned_task = copy.deepcopy(task) if isinstance(task, Script) else task.clone()
            cloned_tasks.append(cloned_task)

        job_group = JobGroup(
            id=task_id,
            tasks=cloned_tasks,
            executors=cloned_executors,
            plugins=plugins,
            tail_logs=tail_logs,
            dependencies=dependencies or [],
        )
        plugins = plugins or []
        for plugin in plugins:
            for i, task in enumerate(cloned_tasks):
                _executor = job_group.executors if job_group._merge else job_group.executors[i]  # type: ignore
                assert isinstance(_executor, Executor)
                plugin.setup(task, _executor)

        self._jobs.append(job_group)
        return job_group.id

    def _validate_task(self, task_info: str, task: Union[Partial, Script]) -> None:
        valid = True
        message = ""
        if isinstance(task, Partial):
            serializer = ZlibJSONSerializer()
            serialized = serializer.serialize(task)
            deserialized = serializer.deserialize(serialized)
            diff = diffing.build_diff(deserialized, task)
            diff = {
                daglish.path_str(d.target): (d.new_value if hasattr(d, "new_value") else None)  # type: ignore
                for d in diff.changes
            }
            if deserialized != task:
                valid = False
                message += f"""
Deserialized task does not match original task. The following paths in your task need to be wrapped in `run.Config` or `run.Partial`:

{pprint.PrettyPrinter(indent=4).pformat(diff)}

For more information about `run.Config` and `run.Partial`, please refer to https://github.com/NVIDIA-NeMo/Run/blob/main/docs/source/guides/configuration.md.
"""
        if not valid:
            raise RuntimeError(f"Failed to validate task {task_info}.
{message}")

    def add(
        self,
        task: Union[Partial, Script] | list[Union[Partial, Script]],
        executor: Executor | list[Executor] | None = None,
        name: str = "",
        plugins: Optional[list[ExperimentPlugin]] = None,
        tail_logs: bool = False,
        dependencies: Optional[list[str]] = None,
    ) -> str:
        """
        Add a configured function along with its executor config to the experiment.
        """
        assert _current_experiment.get(None) == self, (
            "Using Experiment without it's context manager is not permitted."
        )

        job_ids = set([job.id for job in self.jobs])
        for dep in dependencies or []:
            assert dep in job_ids, f"Dependency {dep} not found."

        executor = executor or self.executor
        if not isinstance(task, list):
            assert executor and isinstance(executor, Executor)
            job_id = self._add_single_job(
                task,
                executor,
                name,
                plugins=plugins,
                tail_logs=tail_logs,
                dependencies=dependencies.copy() if dependencies else None,
            )
        else:
            assert name, "name is required for task group."
            job_id = self._add_job_group(
                task,
                executor,
                name,
                plugins=plugins,
                tail_logs=tail_logs,
                dependencies=dependencies.copy() if dependencies else None,
            )

        return job_id

    def dryrun(self, log: bool = True, exist_ok: bool = False, delete_exp_dir: bool = True):
        """
        Logs the raw scripts that will be executed for each task.
        """
        if log:
            self.console.log(f"[bold magenta]Experiment {self._id} dryrun...")

        self._prepare(exist_ok=exist_ok)

        for job in self.jobs:
            if isinstance(job, Job):
                if log:
                    self.console.log(f"[bold magenta]Task {job.id}
")
            elif isinstance(job, JobGroup):
                if log:
                    self.console.log(f"[bold magenta]Task Group {job.id}
")
            job.launch(wait=False, runner=self._runner, dryrun=True, direct=False, log_dryrun=log)

        if delete_exp_dir:
            shutil.rmtree(self._exp_dir)

    def run(
        self,
        sequential: bool = False,
        detach: bool = False,
        tail_logs: bool = False,
        direct: bool = False,
    ):
        """
        Runs all the tasks in the experiment.

        By default, all tasks are run in parallel.

        If sequential=True, all tasks will be run one after the other.
        The order is based on the order in which they were added.

        Parallel mode only works if all executors in the experiment support it.
        Currently, all executors support parallel mode.

        In sequential mode, if all executor supports dependencies, then all tasks will be scheduled at once
        by specifying the correct dependencies to each task.
        Otherwise, the experiment.run call will block and each task that is scheduled will be executed sequentially.
        In this particular case, we cannot guarantee the state of the experiment if the process exits in the middle.

        Currently, only the slurm executor supports dependencies.

        Args:
            sequential: If True, runs all tasks sequentially in the order they were added. Defaults to False.
            detach: If True, detaches from the process after launching the tasks. Only supported for Slurm and Skypilot. Defaults to False.
            tail_logs: If True, tails logs from all tasks in the experiment. If False, relies on task specific setting. Defaults to False.
            direct: If True, runs all tasks in the experiment sequentially in the same process. Note that if direct=True, then sequential also will be True. Defaults to False.
        """
        assert _current_experiment.get(None) == self, (
            "Using Experiment without it's context manager is not permitted."
        )

        if self._launched:
            self.console.log("[bold magenta]Experiment already running...")
            return

        if self._reconstruct:
            self.console.log("[bold magenta]Experiment in inspection mode...")
            return

        # Prepare experiment before running

        # in case of multi-node execution with LocalExecutor+torchrun+slurm, run only on first rank
        if int(os.getenv("SLURM_PROCID", 0)) == 0:
            self._prepare()

        if direct:
            self.console.log(
                "[bold magenta]Running the experiment with direct=True. "
                "This will launch all jobs sequentially in the same process."
            )
            if not self.jobs:
                self.console.log("[bold red]No jobs to run in this experiment.")
                return

            assert all(map(lambda job: isinstance(job, Job), self.jobs)), (
                "Jobs in this experiment contain JobGroup which cannot be run directly for now."
            )

            assert all(map(lambda job: not job.dependencies, self.jobs)), (
                "Jobs in this experiment contain dependencies which cannot be run directly for now."
            )

            for job in self.jobs:
                assert isinstance(job, Job)
                with TeeStdoutStderr(
                    os.path.join(job.executor.job_dir, f"log_{job.id}_direct_run.out")
                ):
                    job.launch(wait=True, direct=True, runner=self._runner)

            self._save_jobs()
            self._launched = any(map(lambda job: job.launched, self.jobs))
            self._direct = True
            return

        executors = set()
        for job in self.jobs:
            if isinstance(job, Job):
                executors.add(job.executor.__class__)
            elif isinstance(job, JobGroup):
                if isinstance(job.executors, list):
                    for executor in job.executors:
                        executors.add(executor.__class__)
                else:
                    executors.add(job.executors.__class__)

        if detach and any(map(lambda x: x not in self._DETACH_SUPPORTED_EXECUTORS, executors)):
            self.console.log(
                "[bold red] Cannot detach from this experiment. Please keep it running until completion."
            )
            detach = False

        is_dag = any(map(lambda job: len(job.dependencies) > 0, self.jobs))
        assert not (is_dag and sequential), (
            "Jobs in this experiment have dependencies, they cannot be run sequentially. Set sequential=False."
        )

        if sequential:
            for i in range(1, len(self.jobs)):
                self.jobs[i].dependencies.append(self.jobs[i - 1].id)

        self.dryrun(log=False, exist_ok=True, delete_exp_dir=False)
        for tunnel in self.tunnels.values():
            if isinstance(tunnel, SSHTunnel):
                tunnel.connect()
                assert tunnel.session, f"SSH tunnel {tunnel.key} failed to connect."
                rsync(tunnel.session, self._exp_dir, os.path.dirname(tunnel.job_dir))

            symlink_cmds = []
            for packaging_job in tunnel.packaging_jobs.values():
                if packaging_job.symlink:
                    symlink_cmds.append(packaging_job.symlink_cmd())

            if symlink_cmds:
                tunnel.run(" && ".join(symlink_cmds))

        self._save_tunnels()

        return self._run_dag(detach=detach, tail_logs=tail_logs, executors=executors)

    def _run_dag(self, detach: bool, tail_logs: bool, executors: set[Executor]):
        job_map = {job.id: job for job in self._jobs}
        graph = nx.DiGraph()
        job_ids = set([job.id for job in self.jobs])
        for job in self.jobs:
            graph.add_node(job.id, job=job)
            for dep in job.dependencies:
                assert dep in job_ids, f"Dependency {dep} not found in job list {job_ids}."
                graph.add_edge(dep, job.id)

        assert nx.is_directed_acyclic_graph(graph), "Jobs have cyclic dependencies."
        order = [sorted(generation) for generation in nx.topological_generations(graph)]
        add_deps = False
        if len(order) > 1:
            if all(map(lambda x: x in self._DEPENDENCY_SUPPORTED_EXECUTORS, executors)):
                wait = False
                add_deps = True
                self.detach = detach
            else:
                wait = True
                if len(self.jobs) > 1:
                    self.console.log(
                        f"[bold cyan]Dependencies not supported for atleast one of {executors}."
                        "All jobs will be run one after the other based on their dependencies, please keep the process alive."
                    )
                if detach:
                    self.console.log(
                        "[bold red] Cannot detach from this experiment. Please keep it running until completion."
                    )
        else:
            # All jobs will be executed in parallel
            assert all(map(lambda x: x in self._PARALLEL_SUPPORTED_EXECUTORS, executors)), (
                f"Parallel mode not supported for atleast one of {executors}. Set sequential=True."
            )
            wait = False
            self.detach = detach

        for level in order:
            # Launch jobs in this level concurrently since they are independent

            def _set_context(ctx: contextvars.Context):
                for var, value in ctx.items():
                    var.set(value)

            ctx = contextvars.copy_context()

            def _launch(node: str):
                job: Job | JobGroup = job_map[node]
                self.console.log(f"[bold cyan]Launching job {job.id} for experiment {self._title}")
                if tail_logs:
                    job.tail_logs = True

                try:
                    if add_deps:
                        deps = []
                        for dep_id in job.dependencies:
                            dep = job_map[dep_id]
                            handle = dep.handle
                            assert dep.launched and handle, (
                                f"Dependency {dep.id} for {job.id} not yet launched."
                            )
                            deps.append(handle)

                        job.executor.dependencies = deps  # type: ignore

                    job.launch(wait=False, runner=self._runner)
                    return job

                except Exception as e:
                    self.console.log(f"Error running job {job.id}: {e}")
                    raise e

            launched_jobs: list[Job | JobGroup] = []
            with ThreadPoolExecutor(
                initializer=_set_context, initargs=(ctx,), max_workers=self._threadpool_workers
            ) as pool:
                futures = [pool.submit(_launch, node) for node in level]
                for future in as_completed(futures):
                    launched_jobs.append(future.result())

            if wait:
                self._wait_for_jobs(jobs=launched_jobs)

        self._save_jobs()
        self._launched = any(map(lambda job: job.launched, self.jobs))
        self._waited = wait

    def _wait_for_jobs(self, jobs: list[Job | JobGroup]):
        def set_context(context: contextvars.Context):
            for var, value in context.items():
                var.set(value)

        context = contextvars.copy_context()
        with ThreadPoolExecutor(initializer=set_context, initargs=(context,)) as executor:
            futures: dict[Future, Job | JobGroup] = {}
            for job in jobs:
                if isinstance(job, Job):
                    handle_exists = job.handle
                else:
                    handle_exists = len(job.handles) > 0 and all(job.handles)

                if job.launched and handle_exists:
                    self._initialize_live_progress()
                    self._add_progress(job=job)
                    future = executor.submit(
                        job.wait,
                        runner=self._runner
                        if isinstance(
                            job.executor,
                            self._RUNNER_DEPENDENT_EXECUTORS,
                        )
                        else get_runner(),
                    )
                    futures[future] = job

            for future in as_completed(futures.keys()):
                job = futures[future]
                try:
                    future.result()
                    self._update_progress(job, job.state)
                except Exception as e:
                    self.console.log(f"Exception while waiting for Job {job.id}: {e}")
                    self.console.log(*traceback.format_exception(e))
                    self._update_progress(job, AppState.UNKNOWN)
                finally:
                    job.cleanup()

    def _initialize_tunnels(self, extract_from_executors: bool = False):
        if extract_from_executors:
            for job in self.jobs:
                if (
                    isinstance(job.executor, SlurmExecutor)
                    and job.executor.tunnel.key not in self.tunnels
                ):
                    self.tunnels[job.executor.tunnel.key] = job.executor.tunnel

        for tunnel in self.tunnels.values():
            if isinstance(tunnel, SSHTunnel):
                tunnel.connect()
                assert tunnel.session, f"SSH tunnel {tunnel.key} failed to connect."

    def status(self, return_dict: bool = False) -> Optional[dict[str, dict[str, str]]]:
        """
        Prints a table specifying the status of all tasks.

        .. note::
            status is not supported for local executor
            and the status for a task using the local executor
            will be listed as UNKNOWN in most cases
        """
        _set_current_experiment = False
        if not self._current_experiment_token:
            _current_experiment.set(self)
            _set_current_experiment = True

        def _get_job_info_and_dict(
            idx: int, job: Job | JobGroup
        ) -> tuple[list[str], dict[str, str]]:
            job_info = []
            job_info.append(f"[bold green]Task {idx}[/bold green]: [bold orange1]{job.id}")
            job_info.append(
                f"- [bold green]Status[/bold green]: {str(job.status(runner=self._runner))}"
            )
            job_info.append(f"- [bold green]Executor[/bold green]: {job.executor.info()}")

            try:
                _, _, path_str = job.handle.partition("://")
                path = path_str.split("/")
                app_id = path[1]
            except Exception:
                app_id = ""

            job_info.append(f"- [bold green]Job id[/bold green]: {app_id}")
            directory_info = [
                "- [bold green]Local Directory[/bold green]: " + job.executor.job_dir,
            ]
            job_dict = {
                "name": job.id,
                "status": job.status(runner=self._runner),
                "executor": job.executor.info(),
                "job_id": app_id,
                "handle": job.handle,
                "local_dir": job.executor.job_dir,
            }

            if isinstance(job.executor, SlurmExecutor) and isinstance(
                job.executor.tunnel, SSHTunnel
            ):
                directory_info.extend(
                    [
                        "- [bold green]Remote Directory[/bold green]: "
                        + os.path.join(
                            job.executor.tunnel.job_dir,
                            Path(job.executor.job_dir).name,
                        ),
                    ]
                )
                job_dict["remote_dir"] = os.path.join(
                    job.executor.tunnel.job_dir,
                    Path(job.executor.job_dir).name,
                )
            job_info.extend(directory_info)
            return job_info, job_dict

        self._initialize_tunnels(extract_from_executors=True)
        try:
            result_dict = {}
            job_infos: list[Group | None] = [None] * len(self.jobs)

            # Parallelize IO-bound status retrieval across jobs
            def _collect(arg):
                idx, job = arg
                job_info, job_dict = _get_job_info_and_dict(idx, job)
                return idx, job.id, job_info, job_dict

            # Propagate context variables to worker threads so helpers that rely on them keep working
            def _set_context(ctx: contextvars.Context):
                for var, value in ctx.items():
                    var.set(value)

            ctx = contextvars.copy_context()
            with ThreadPoolExecutor(
                initializer=_set_context, initargs=(ctx,), max_workers=self._threadpool_workers
            ) as pool:
                futures = [pool.submit(_collect, (idx, job)) for idx, job in enumerate(self.jobs)]
                for future in as_completed(futures):
                    idx, job_id, job_info, job_dict = future.result()
                    job_infos[idx] = Group(*job_info)
                    result_dict[job_id] = job_dict

            # Remove potential None slots (should not occur)
            job_infos = [ji for ji in job_infos if ji is not None]

            if return_dict:
                return result_dict

            self.console.print()
            self.console.print(
                f"[bold green]Experiment Status for[/bold green] [bold orange1]{self._id}",
                new_line_start=True,
            )
            for job_info in job_infos:
                self.console.print(job_info, soft_wrap=True, new_line_start=True, highlight=False)
            self.console.print()
        finally:
            if _set_current_experiment and self._current_experiment_token:
                _current_experiment.reset(self._current_experiment_token)
                self._current_experiment_token = None

    def cancel(self, job_id: str):
        """
        Cancels an existing job if still running.
        """
        _set_current_experiment = False
        if not self._current_experiment_token:
            _current_experiment.set(self)
            _set_current_experiment = True

        self.console.log(f"[bold cyan]Cancelling {job_id} if still running")
        try:
            job = next(filter(lambda x: x.id == job_id, self.jobs))
            job.cancel(runner=self._runner)
        except StopIteration:
            self.console.log(f"[bold red]Job {job_id} not found")
        except Exception as e:
            self.console.log(f"[bold red]Failed to cancel {job_id}
Error: {e}
")
            self.console.log(*traceback.format_exception(e))
        finally:
            if _set_current_experiment and self._current_experiment_token:
                _current_experiment.reset(self._current_experiment_token)
                self._current_experiment_token = None

    def logs(self, job_id: str, regex: str | None = None):
        """
        Prints the logs of the specified job_id, optionally filtered by regex.
        """
        _set_current_experiment = False
        if not self._current_experiment_token:
            _current_experiment.set(self)
            _set_current_experiment = True

        self.console.log(f"[bold cyan]Fetching logs for {job_id}")
        try:
            job = next(filter(lambda x: x.id == job_id, self.jobs))
            if isinstance(job, Job) and job.handle.endswith("direct_run"):
                self.console.log("This job was run with direct=True.")
                self.console.log(
                    f"Logs may be present in task directory at:
[bold]{job.executor.job_dir}."
                )
                return

            try:
                job.logs(runner=self._runner, regex=regex)
            except Exception as e:
                self.console.log(f"[bold red]Failed to get logs for {job_id}
Error: {e}
")
                self.console.log(
                    f"Logs may be present in job directory at:
[bold]{job.executor.job_dir}."
                )
        except StopIteration:
            self.console.log(f"[bold red]Job {job_id} not found")
        finally:
            if _set_current_experiment and self._current_experiment_token:
                _current_experiment.reset(self._current_experiment_token)
                self._current_experiment_token = None

    def reset(self) -> "Experiment":
        """
        Resets an experiment to make it ready for a relaunch.
        Only works if the current experiment run has already been launched.
        """
        if not self._reconstruct and not os.path.isfile(
            os.path.join(self._exp_dir, self._DONE_FILE)
        ):
            self.console.log(
                f"[bold magenta]Experiment {self._id} has not run yet, skipping reset..."
            )
            return self

        old_id, old_exp_dir, old_launched = self._id, self._exp_dir, self._launched
        self._id = f"{self._title}_{int(time.time())}"
        self._exp_dir = os.path.join(get_nemorun_home(), "experiments", self._title, self._id)
        self._launched = False
        self._live_progress = None

        jobs = self._jobs
        self._jobs = []
        serializer = ZlibJSONSerializer()
        _set_current_experiment = False
        if not self._current_experiment_token:
            _current_experiment.set(self)
            _set_current_experiment = True

        try:
            if "__external_main__" not in sys.modules:
                maybe_load_external_main(old_exp_dir)

            for job in jobs:
                if isinstance(job, Job):
                    if isinstance(job.task, str):
                        _task = serializer.deserialize(job.task)
                        if _task.__fn_or_cls__ == Script:
                            job.task = fdl.build(_task)
                        else:
                            job.task = _task  # type: ignore

                    self.add(
                        job.task,
                        job.executor,
                        name=job.id,
                        tail_logs=job.tail_logs,
                    )
                else:
                    if isinstance(job.tasks, str):
                        tasks = serializer.deserialize(job.tasks)
                        job.tasks = [
                            fdl.build(task) if task.__fn_or_cls__ == Script else task
                            for task in tasks
                        ]

                    self.add(
                        job.tasks,
                        job.executors,
                        name=job.id,
                        tail_logs=job.tail_logs,
                    )
        except Exception as e:
            self.console.log(
                f"[bold magenta]Failed resetting Experiment {self._id} due to error: {e}"
            )
            # Double check exp dir is unchanged
            new_path = os.path.join(get_nemorun_home(), "experiments", self._title, self._id)
            if self._exp_dir == new_path and new_path != old_exp_dir:
                shutil.rmtree(self._exp_dir)

            self._id = old_id
            self._exp_dir = old_exp_dir
            self._launched = old_launched
            self._jobs = self._load_jobs()
        finally:
            if _set_current_experiment and self._current_experiment_token:
                _current_experiment.reset(self._current_experiment_token)
                self._current_experiment_token = None

        self._reconstruct = False
        return self

    def _initialize_live_progress(self):
        if not self._live_progress:
            # Disable live progress if we are tailing logs for any task
            # as tty output consistency can not be guaranteed as of now
            if self.clean_mode or any(map(lambda job: job.tail_logs, self.jobs)):
                return

            assert isinstance(self.console, rich.console.Console)
            self._progress = Progress(
                "{task.description}",
                SpinnerColumn(),
                BarColumn(bar_width=None),
                TimeElapsedColumn(),
            )
            self._exp_panel = Panel(
                self._progress,
                title=f"[b]{self._id}",
                padding=(1, 3),
            )
            self._task_progress: dict[str, TaskID] = {}
            self._live_progress = Live(self._exp_panel, console=self.console, refresh_per_second=10)
            self._live_progress.start(refresh=True)

    def _add_progress(self, job: Job | JobGroup):
        if self._live_progress:
            self._task_progress[job.id] = self._progress.add_task(
                f"[bold green]{job.id}", total=None
            )

    def _update_progress(self, job: Job | JobGroup, state: AppState):
        if self._live_progress:
            color = "[bold green]" if state == AppState.SUCCEEDED else "[bold red]"
            task_progress_id = self._task_progress[job.id]
            self._progress.stop_task(task_progress_id)
            self._progress.update(
                task_progress_id,
                description=f"{color}{job.id} {state}",
            )
            progress_task: RichTask = self._progress._tasks[task_progress_id]
            progress_task.finished_time = progress_task.elapsed
            progress_task.completed = progress_task.elapsed or 0.0
            progress_task.total = progress_task.elapsed

            self._progress.refresh()

    def _cleanup(self, tunnels: bool = True):
        if tunnels and hasattr(self, "tunnels"):
            for tunnel in self.tunnels.values():
                try:
                    tunnel.cleanup()
                except Exception:
                    ...

        self._runner.close()

        if (
            _current_experiment is not None
            and _current_experiment.get(None)
            and self._current_experiment_token
        ):
            _current_experiment.reset(self._current_experiment_token)
            self._current_experiment_token = None

    def __enter__(self) -> "Experiment":
        self._current_experiment_token = _current_experiment.set(self)
        self.console.rule(
            f"[bold magenta]Entering Experiment {self._title} with id: {self._id}",
        )
        return self

    def __exit__(self, exc_type, exc_value, tb):
        try:
            if hasattr(self, "detach") and self.detach:
                self.console.rule(f"[bold magenta]Detaching from Experiment {self._id}.")
                self.console.log(
                    "Task specific cleanup won't be run.
"
                    "Ephemeral logs and artifacts may be lost.",
                )

                if self._launched and not self._skip_status_at_exit:
                    self.status()
                return

            if self._launched:
                if hasattr(self, "_direct") and self._direct:
                    self.console.rule(
                        f"[bold magenta]Direct run Experiment {self._id}",
                    )
                    if not self._skip_status_at_exit:
                        self.status()
                    return

                if hasattr(self, "_waited") and self._waited:
                    self.console.rule(
                        f"[bold magenta]Done waiting for Experiment {self._id}",
                    )
                    if not self._skip_status_at_exit:
                        self.status()
                    return

                self.console.rule(
                    f"[bold magenta]Waiting for Experiment {self._id} to finish",
                )
                if not self._skip_status_at_exit:
                    self.status()

                self._wait_for_jobs(jobs=self.jobs)
        finally:
            if self._live_progress:
                self._live_progress.stop()

            self._cleanup(tunnels=False)
            if self._launched:
                Path(os.path.join(self._exp_dir, self._DONE_FILE)).touch()
                if self._enable_goodbye_message:
                    self.console.print(
                        Syntax(
                            self.GOODBYE_MESSAGE_PYTHON.format(
                                exp_id=self._id,
                                tasks=list(map(lambda job: job.id, self.jobs)),
                            ),
                            "python",
                            theme=os.environ.get("NEMO_RUN_CODE_THEME", "monokai"),
                        )
                    )
                    self.console.print(
                        Syntax(
                            self.GOODBYE_MESSAGE_BASH.format(
                                exp_id=self._id,
                                tasks=list(map(lambda job: job.id, self.jobs)),
                            ),
                            "shell",
                            theme=os.environ.get("NEMO_RUN_CODE_THEME", "monokai"),
                        )
                    )

    def _repr_svg_(self):
        return self.to_config()._repr_svg_()

    def __del__(self):
        try:
            deconfigure_logging()
            self._cleanup()
        except Exception:
            pass

    @property
    def jobs(self) -> list[Job | JobGroup]:
        return Jobs(self._jobs)

    @jobs.setter
    def jobs(self, jobs: list[Job | JobGroup]):
        self._jobs = jobs

    @property
    def tasks(self) -> list[Config]:
        serializer = ZlibJSONSerializer()

        for job in self._jobs:
            if isinstance(job, Job):
                if isinstance(job.task, str):
                    _task = serializer.deserialize(job.task)
                    if _task.__fn_or_cls__ == Script:
                        job.task = fdl.build(_task)
                    else:
                        job.task = _task  # type: ignore
            else:
                if isinstance(job.tasks, str):
                    tasks = serializer.deserialize(job.tasks)
                    job.tasks = [
                        fdl.build(task) if task.__fn_or_cls__ == Script else task for task in tasks
                    ]

        return Tasks((job.task if isinstance(job, Job) else job.tasks) for job in self._jobs)


class Tasks(list, ConfigurableMixin): ...


class Jobs(list, ConfigurableMixin): ...


def _get_latest_dir(path) -> str:
    dirs = [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]
    latest_dir = max(dirs, key=lambda d: os.path.getctime(os.path.join(path, d)))
    return os.path.join(path, latest_dir)


def _get_sorted_dirs(path: str) -> list[str]:
    if not os.path.exists(path):
        return []
    dirs = [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]
    dirs = sorted(dirs, key=lambda d: os.path.getctime(os.path.join(path, d)))
    return list(dirs)


_LOADED_MAINS = set()


def maybe_load_external_main(exp_dir: str):
    main_file = Path(exp_dir) / "__main__.py"
    if main_file.exists() and main_file not in _LOADED_MAINS:
        _LOADED_MAINS.add(main_file)

        spec = importlib.util.spec_from_file_location("__external_main__", main_file)
        if spec is not None and spec.loader is not None:
            new_main_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(new_main_module)

            if "__external_main__" not in sys.modules:
                sys.modules["__external_main__"] = new_main_module
            else:
                external = sys.modules["__external_main__"]
                for attr in dir(new_main_module):
                    if not attr.startswith("__"):
                        setattr(external, attr, getattr(new_main_module, attr))

            existing_main = sys.modules["__main__"]
            for attr in dir(new_main_module):
                if not attr.startswith("__"):
                    setattr(existing_main, attr, getattr(new_main_module, attr))