tao-run-on-slurm

SLURM

适用于由 SLURM 管理的集群的远程 GPU 计算平台。作业从 TAO 服务或 SDK 主机通过 SSH 提交到登录节点，在共享文件系统上暂存，使用 sbatch 提交，并通过支持容器的 srun 执行。

当用户可以访问受管理的 GPU 集群、共享 Lustre 存储和由调度器分配的 GPU 时，请使用 SLURM。对于仅存在于 Agent 机器上的本地文件，请勿使用 SLURM；数据和输出必须能从集群访问。

预检

# 1. SSH to the login node works without a password prompt
SLURM_HOST="${SLURM_HOSTNAME%%,*}"
[ -n "$SLURM_USER" ] && [ -n "$SLURM_HOST" ] || {
  echo "MISSING: set SLURM_USER and SLURM_HOSTNAME (comma-separated for failover) in your env (~/.config/tao/.env)."
  exit 1
}
ssh -o BatchMode=yes -o ConnectTimeout=10 "${SLURM_USER}@${SLURM_HOST}" "true" 2>/dev/null || {
  echo "MISSING: passwordless SSH to ${SLURM_USER}@${SLURM_HOST} not working. See references/ssh-setup.md."
  exit 1
}

# 2. Optional: TAO SDK wrapper for Job handles + S3 wrapping.
# nvidia-tao-sdk is on public PyPI; pin lives in versions.yaml (wheels.tao_sdk_slurm).
PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_slurm)
python -c "import tao_sdk" 2>/dev/null || {
  echo "MISSING: nvidia-tao-sdk not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

如果检查失败，Agent 会提示用户通过 Bash 授权安装或修复。

第三个预检步骤仅适用于私有 nvcr.io 镜像：计算节点上的 Pyxis 需要在集群的 ~/.config/enroot/.credentials 中保存持久化的 enroot 凭据（它不会从作业环境读取 NGC_KEY）。如果没有这些凭据，受身份验证保护的拉取操作会在作业启动时失败，并显示“Could not process JSON input”。每个（集群、用户）只需执行一次。完整检查步骤以及使用 printf | ssh 安装、避免 NGC_KEY 出现在历史记录、文件和聊天输出中的模式，请参阅 references/ssh-setup.md。公有镜像无需执行此步骤。

先决条件

提交任何作业之前，运行 TAO 服务或 SDK 的主机必须能够通过 SSH 登录 SLURM_HOSTNAME 中至少一台主机，且不会出现交互式密码提示。处理程序会以非交互方式运行 sbatch、squeue、sacct、scancel 和日志尾随命令，因此密码或双重身份验证提示会导致作业在提交或查看状态时失败。

每个（主机、登录节点、用户）组合只需配置一次：创建 SSH 密钥对，在每台登录主机上安装公钥，信任主机密钥，将私钥权限设为 chmod 600，并使用 ssh -o BatchMode=yes ... 验证。完整分步说明（包括 ~/.ssh/config 别名、容器密钥挂载说明，以及双重身份验证 / SSH_AUTH_SOCK 备用方案）请参阅 references/ssh-setup.md。该文件还包含密码免提示 SSH 失败时向用户显示的 SSH 故障修复提示词。

凭据

• SLURM_USER（必需）：登录节点的 SSH 用户名。在微服务工作区元数据中，此字段为 cloud_specific_details.slurm_user。
• SLURM_HOSTNAME（必需）：以逗号分隔的登录主机名，用于故障转移。微服务架构将其存储为列表字段 cloud_specific_details.slurm_hostname。
• SLURM_PARTITION（必需）：用于提交 GPU 作业的分区列表。请在必需的 SLURM 信息收集清单中询问此项。打包的默认值为 polar,polar3,polar4,grizzly，这些分区按 4 小时队列处理。
• SSH_KEY_PATH（首选，且启动前应已设置）：用于向登录节点进行非交互式公钥身份验证的私钥路径。如果密码免提示 SSH 失败，请向用户询问 SSH_KEY_PATH=/path/to/private_key，并展示 references/ssh-setup.md 中的设置步骤；不要把这个选项藏在多个备用选项之后。
• SSH_AUTH_SOCK（高级备用方案）：已加载可接受密钥的 SSH Agent 套接字。在面向用户的故障修复提示中，优先使用 SSH_KEY_PATH。
• SLURM_BASE_RESULTS_DIR（可选）：共享文件系统的基础路径。tao-core 的默认约定为 /lustre/fsw/portfolios/edgeai/<your-dir>，其中 <your-dir> 是你在集群上的个人目录。
• SLURM_ACCOUNT（通常由站点策略要求）：通过 #SBATCH --account 计费的账户。

除非用户说明其站点要求提供账户、想要自定义结果根目录，或工作流必须覆盖默认值才能继续，否则不要在初始信息收集中询问 SLURM_ACCOUNT 或 SLURM_BASE_RESULTS_DIR。

后端详情

将作业路由到此平台时，使用 backend_details.backend_type = "slurm"。微服务架构支持的后端详情如下：

{
  "backend_type": "slurm",
  "partition": "polar,polar3,polar4,grizzly",
  "cluster_name": "optional-name"
}

运行时元数据存储在 backend_details.slurm_metadata 下，尤其是 slurm_job_id 和 job_dir。不要自行编造这些值。它们会在 sbatch 返回调度器作业 ID 后写入。

存储

SLURM 作业在集群上运行，因此 API 主机上的本地路径不能用作有效的数据集路径。请优先使用共享文件系统 URI：

• 对于 Lustre 上由用户提供的数据集，请使用 lustre:///absolute/path。
• 微服务元数据中可能会出现 slurm:// 路径，这些路径会在容器启动前转换为实际的 Lustre 路径。
• 对于 SLURM，请避免使用裸路径 /local/path 和 file:// 数据集 URI。tao-core 中的验证会拒绝远程后端使用本地路径和 file 路径。

可以使用数据集根目录，也可以使用直接的 spec 键路径：

• 根目录模式：/lustre/.../<model>/train，模型技能会将其映射到所需文件，例如 <root>/annotations.json，并将 <root> 用作媒体路径。
• 直接 spec 模式：精确字段，例如 custom.train_dataset.annotation_path=/lustre/.../train.json 和 custom.train_dataset.media_path=/lustre/.../videos.tar.gz。

无密码 SSH 成功后、生成脚本前，请从登录主机验证每个所需的数据集文件/路径：

ssh -o BatchMode=yes <SLURM_USER>@<working-login-host> \
  'test -e /lustre/.../annotations.json && test -e /lustre/.../media_or_archive'

如果远程 test -e 失败，请停止并要求提供修正后的路径，或将数据暂存到集群共享存储中。不要创建会在第一个训练作业中运行失败的运行脚本。

结果默认保存到：

/lustre/fsw/portfolios/edgeai/<your-dir>/results/<job_id>

<your-dir> 是你在集群上的个人目录。

运行器会将 TAO_API_RESULTS_DIR 设置为结果父目录，因为容器代码在写入状态和产物时会追加作业 ID。

SLURM 作业输入请使用 Lustre，而非 S3。 SLURM 的调度器会强制执行 GPU 空闲超时——如果脚本开头执行耗时很长的 s3:// 下载，可能会在训练开始前耗尽分配时间，调度器也可能终止作业。 请先将训练数据暂存到 Lustre；只有较小的辅助输入（检查点、配置）可以从 S3 / HF / NGC 预取。完整原因请参阅 references/sdk-usage.md。

容器执行

tao-core 使用 SLURM 处理器，通过 Pyxis/Enroot 运行 TAO 容器：

1. 在 <job_dir>/specs、<job_dir>/env 和 <job_dir>/meta 下暂存 spec、环境和云元数据的紧凑 JSON 文件。
2. 可选：使用 srun -n1 -p <conversion_partition> enroot import 将 Docker 镜像转换为缓存的 SQSH 镜像。
3. 在 <job_dir>/sbatch/job_<job_id>.sbatch 下写入 sbatch 脚本。
4. 提交 sbatch --export=ALL <script>。
5. 使用 srun --container-image=<image> --container-mounts=/lustre 运行容器。

处理器接受的镜像格式：

• /path/to/image.sqsh
• registry#image:tag
• docker://registry#image:tag
• 普通的 registry/image:tag，需要时会将其转换为 Pyxis 格式

SQSH 转换会按镜像名称缓存。对于 :latest 镜像，除非启用了 force_reconvert_latest，否则会使用缓存的 SQSH 镜像。

资源映射

tao-core 的默认值：

• num_nodes：1
• num_gpus：4
• max_num_gpus_per_node：8
• cpus_per_task：16
• time_hours：4
• timeout_hours：3.8
• max_time_hours：4
• container_mounts：/lustre
• use_requeue：true
• use_sqsh：true

为 SLURM 生成启动器或包装脚本时，请显式使用打包的平台资源默认值设置墙钟时间默认值：

export SLURM_TIME_HOURS="${SLURM_TIME_HOURS:-4}"
export SLURM_TIMEOUT_HOURS="${SLURM_TIMEOUT_HOURS:-3.8}"

在 SLURM 上，不要将默认值设为 12 小时。如果用户提供的 SLURM_TIME_HOURS 更长，请在提交前确认所选分区支持该时长。对于打包的默认分区列表 polar,polar3,polar4,grizzly，拒绝超过 4 小时的请求；仅当用户确实需要更长的墙钟时间时，才要求更换分区。

当 num_gpus 大于或等于 max_num_gpus_per_node 时，处理程序会将请求视为每个节点独占，并在必要时根据 GPU 总数计算额外节点数。

对于多节点作业（num_nodes > 1），sbatch 脚本会导出 WORLD_SIZE、MASTER_ADDR、MASTER_PORT、NODE_RANK 和 NUM_GPU_PER_NODE；Cosmos-RL 对控制器、策略和推理工作进程有特殊的多节点角色处理逻辑。完整的 sbatch 指令、会合环境变量表及其约定和集群要求，请参阅 references/multi-node.md。

监控

• 调度器状态通过存储的 SLURM 作业 ID 使用 squeue 或 sacct 获取。
• TAO 终端状态来自共享结果文件夹中的 status.json。
• 如果用户启用了聊天监控，当作业处于 PENDING、RUNNING 或其他非终止状态时，请按请求的间隔继续轮询。不要在固定的经过时间（例如 30 分钟）后停止；在共享 GPU 分区上，长时间排队很常见。
• 启用聊天监控时，对于仍处于非终止状态的 SLURM 作业，不要发送最终回复。最终回复会使监控分离；只有在用户要求分离／停止或作业进入终止状态时才使用。
• 通过 SSH 从以下位置读取日志：

<job_dir>/slurm-logs/<slurm_job_name>-<slurm_job_id>/main.out
<job_dir>/slurm-logs/<slurm_job_name>-<slurm_job_id>/main.err

状态映射：

• PENDING -> Pending
• RUNNING 或 COMPLETING -> Running
• COMPLETED -> 检查 status.json
• FAILED、BOOT_FAIL、DEADLINE、OUT_OF_MEMORY、NODE_FAIL -> 如果日志符合可重试的基础设施问题模式，则重试；否则为 Error
• CANCELLED、PREEMPTED、REVOKED -> Canceled
• TIMEOUT -> Error
• SUSPENDED、STOPPED -> Paused

取消

查找 backend_details.slurm_metadata.slurm_job_id，并通过 SSH 运行 scancel <slurm_job_id> 进行取消。对于缺失或已终止的 SLURM 作业，将取消视为成功。

多节点训练（分布式）

SLURM 是大型多节点运行的首选平台——传入 num_nodes > 1，SDK 会自动处理 sbatch 指令和 PyTorch 分布式环境变量。一个完整的 create_job 示例、生成的 sbatch 指令、会合环境变量表（WORLD_SIZE、NUM_GPU_PER_NODE、NODE_RANK、MASTER_ADDR、MASTER_PORT）、Cosmos-RL 角色说明、集群要求（Pyxis/Enroot、InfiniBand/NVLink、Lustre）以及上游参考链接，请参阅 references/multi-node.md。

通过 TAO SDK 运行

SDK 的安装步骤已在预检中介绍 — pip install 'nvidia-tao-sdk[slurm]'。 如果你希望由 SDK 处理 Job 句柄、sbatch/squeue/sacct 流程，通过 ActionWorkflow 确保运行文件夹持久化，或便捷地进行云存储 I/O（s3://、hf_model://、ngc://），请使用 SDK。不使用 SDK 时，请自行运行 sbatch 和 srun。

自动重试完全自动：后台监控程序会轮询 squeue/sacct，遇到看似由基础设施引起的故障时，最多会重新对暂存脚本执行 sbatch，重试次数上限为 MAX_JOB_RETRIES = 10；普通训练故障则会立即报告。此外，默认会设置 #SBATCH --requeue（SLURM_USE_REQUEUE 默认为 true）。有关 SlurmSDK / build_entrypoint 代码示例、不得使用 S3 而应使用 Lustre 的规则、可重试故障的分类，以及自动重试和重新排队的完整行为，请参阅 references/sdk-usage.md。

故障模式

常见故障包括：SSH 身份验证失败、本地数据集路径被拒绝、SQSH 转换超时、Pyxis/Enroot 不可用，以及节点异常或临时 GPU 故障（处理程序会按配置的上限重试）。有关每种故障的诊断和解决方法，请参阅 references/troubleshooting.md。