执行 NeMo Run#

配置 NeMo-Run 后,下一步是执行它。NeMo-Run 将配置与执行分离,允许您配置一次函数或任务,然后在多个环境中执行它。使用 NeMo-Run,您可以选择在不同的远程集群上同时执行单个任务或多个任务,并在实验下管理它们。 这就引出了执行的核心构建块: run.Executor 和 run.Experiment

单个配置任务的每次执行都需要一个执行器。NeMo-Run 提供了 run.Executor,这些 API 用于配置您的远程执行器并设置代码的打包。目前我们支持

  • run.LocalExecutor

  • run.DockerExecutor

  • run.SlurmExecutor 以及可选的 SSHTunnel,用于从本地计算机在 Slurm 集群上执行

  • run.SkypilotExecutor(在 python 包的可选功能 skypilot 下可用)。

任务和执行器的元组构成一个执行单元。NeMo-Run 的一个关键目标是允许您混合和匹配任务和执行器,以任意定义执行单元。

创建执行单元后,下一步是运行它。run.run 函数执行单个任务,而 run.Experiment 提供更精细的控制来定义复杂的实验。run.run 使用单个任务包装 run.Experimentrun.Experiment 是一个 API,用于启动和管理全部使用纯 Python 的多个任务。run.Experiment 负责存储运行元数据、在指定的集群上启动它以及同步日志等。此外,run.Experiment 还提供管理工具,以便轻松检查和重现过去的实验。run.Experiment 的灵感来自 xmanager,并在底层使用 TorchX 来处理执行。

注意: NeMo-Run 假定您熟悉 Docker,并使用 Docker 镜像作为远程执行的环境。这意味着当使用远程执行器时,您必须提供包含所有必要依赖项和配置的 Docker 镜像。

注意: 所有实验元数据都存储在启动实验的机器上的 NEMORUN_HOME 环境变量下。默认情况下,NEMORUN_HOME 值为 ~/.run。请务必根据您的需要更改此值。

执行器#

执行器是数据类,用于配置您的远程执行器并设置代码的打包。所有受支持的执行器都继承自基类 run.Executor,但具有特定于其执行环境的配置参数。理解执行器的具体细节并进行设置最初需要一定的成本,但随着时间的推移,这种努力很容易被摊销。

每个 run.Executor 都有两个属性: packager 和 launcher。 packager 指定如何打包代码以供执行,而 launcher 确定使用哪个工具来启动任务。

启动器#

我们支持以下 launchers

  • defaultNone:这将直接启动您的任务,而无需使用任何特殊的启动器。如果您不想使用特定的启动器,请设置 executor.launcher = None(这是默认值)。

  • torchrunrun.Torchrun:这将使用 torchrun 启动任务。有关配置选项,请参阅 Torchrun 类。您可以使用 executor.launcher = "torchrun"executor.launcher = Torchrun(...) 来使用它。

  • ftrun.core.execution.FaultTolerance:这将使用 NVIDIA 的容错启动器启动任务。有关配置选项,请参阅 FaultTolerance 类。您可以使用 executor.launcher = "ft"executor.launcher = FaultTolerance(...) 来使用它。

注意: 启动器可能无法与 run.Script 很好地配合使用。请在 https://github.com/NVIDIA/NeMo-Run/issues 报告任何问题。

打包器#

打包器支持矩阵如下所述

执行器

打包器

LocalExecutor

run.Packager

DockerExecutor

run.Packager, run.GitArchivePackager, run.PatternPackager

SlurmExecutor

run.Packager, run.GitArchivePackager, run.PatternPackager

SkypilotExecutor

run.Packager, run.GitArchivePackager, run.PatternPackager

run.Packager 是一个直通基础打包器。

run.GitArchivePackager 使用 git archive 来打包您的代码。有关使用 git archive 打包的确切机制,请参阅 run.GitArchivePackager 的 API 参考。从高层次上讲,它的工作方式如下

  1. base_path = git rev-parse --show-toplevel

  2. 可以选择通过在 GitArchivePackager 上设置 subpath 属性,将子路径定义为 base_path/GitArchivePackager.subpath

  3. cd base_path && git archive --format=tar.gz --output={output_file} {GitArchivePackager.subpath}:{subpath}

这个提取的 tar 文件将成为您作业的工作目录。例如,给定以下目录结构,其中 subpath="src"

- docs
- src
  - your_library
- tests

执行时,您的工作目录将如下所示

- your_library

如果您正在执行 Python 函数,则此工作目录将自动包含在您的 Python 路径中。

注意: git archive 不打包未提交的更改。将来,我们可能会添加对包含未提交更改的支持,同时遵守 .gitignore

run.PatternPackager 是一个使用模式来打包代码的打包器。它对于打包未受版本控制的代码非常有用。例如,如果您有如下目录结构

- docs
- src
  - your_library

您可以使用 run.PatternPackager 通过将 include_pattern 指定为 src/** 并将 relative_path 指定为 os.getcwd() 来打包您的代码。这将打包整个 src 目录。用于获取要打包的文件列表的命令是

# relative_include_pattern = os.path.relpath(self.include_pattern, self.relative_path)
cd {relative_path} && find {relative_include_pattern} -type f

定义执行器#

接下来,我们将详细介绍如何设置下面的每个执行器。

LocalExecutor#

LocalExecutor 是最简单的执行器。它在本地执行您的任务,与您当前的工作目录位于不同的进程或组中。

定义它的最简单方法是调用 run.LocalExecutor()

DockerExecutor#

DockerExecutor 允许在您的本地计算机上使用 docker 启动任务。它需要预先安装并运行 docker

DockerExecutor 使用 docker python 客户端,并且大多数选项都直接传递给客户端。

以下是配置 Docker 执行器的示例

run.DockerExecutor(
    container_image="python:3.12",
    num_gpus=-1,
    runtime="nvidia",
    ipc_mode="host",
    shm_size="30g",
    volumes=["/local/path:/path/in/container"],
    env_vars={"PYTHONUNBUFFERED": "1"},
    packager=run.Packager(),
)

SlurmExecutor#

SlurmExecutor 允许在带有 Pyxis 的 Slurm 集群上启动配置的任务。  此外,您可以配置 run.SSHTunnel,这使您可以从本地计算机在 Slurm 集群上执行任务,同时 NeMo-Run 为您管理 SSH 连接。此设置支持诸如在多个 Slurm 集群上启动同一任务的用例。

以下是配置 Slurm 执行器的示例

def your_slurm_executor(nodes: int = 1, container_image: str = DEFAULT_IMAGE):
    # SSH Tunnel
    ssh_tunnel = run.SSHTunnel(
        host="your-slurm-host",
        user="your-user",
        job_dir="directory-to-store-runs-on-the-slurm-cluster",
        identity="optional-path-to-your-key-for-auth",
    )
    # Local Tunnel to use if you're already on the cluster
    local_tunnel = run.LocalTunnel()

    packager = GitArchivePackager(
        # This will also be the working directory in your task.
        # If empty, the working directory will be toplevel of your git repo
        subpath="optional-subpath-from-toplevel-of-your-git-repo"
    )

    executor = run.SlurmExecutor(
        # Most of these parameters are specific to slurm
        account="your-account",
        partition="your-partition",
        ntasks_per_node=8,
        gpus_per_node=8,
        nodes=nodes,
        tunnel=ssh_tunnel,
        container_image=container_image,
        time="00:30:00",
        env_vars=common_envs(),
        container_mounts=mounts_for_your_hubs(),
        packager=packager,
    )

# You can then call the executor in your script like
executor = your_slurm_cluster(nodes=8, container_image="your-nemo-image")

从本地计算机启动时使用 SSH 隧道,如果您已在 Slurm 集群上,则使用本地隧道。

SkypilotExecutor#

此执行器用于配置 Skypilot。确保已安装 Skypilot,并且至少使用 sky check 配置了一个云。

这是 Kubernetes 的 SkypilotExecutor 的示例

def your_skypilot_executor(nodes: int, devices: int, container_image: str):
    return SkypilotExecutor(
        gpus="RTX5880-ADA-GENERATION",
        gpus_per_node=devices,
        nodes = nodes
        env_vars=common_envs()
        container_image=container_image,
        cloud="kubernetes",
        # Optional to reuse Skypilot cluster
        cluster_name="tester",
        setup="""
    conda deactivate
    nvidia-smi
    ls -al ./
    """,
    )

# You can then call the executor in your script like
executor = your_skypilot_cluster(nodes=8, devices=8, container_image="your-nemo-image")

正如示例中所示,在 Python 中定义执行器提供了极大的灵活性。您可以轻松地混合和匹配诸如通用环境变量之类的内容,并且将任务与执行器分离使您可以在任何受支持的执行器上运行相同的配置任务。