tao-run-on-lepton

Lepton

基于 DGX Cloud 的托管式 GPU 计算平台。作业以容器工作负载的形式提交，在专用或共享 GPU 节点组上运行。Lepton 负责调度、拉取镜像、收集日志以及管理作业生命周期。

如果你需要基于云的 GPU 计算，又不想直接管理 Kubernetes 或 SLURM 基础设施，请使用 Lepton。

预检

Lepton 以 API 为先——没有 docker-run 替代方案。此技能需要安装带有 Lepton 扩展的 TAO SDK。nvidia-tao-sdk 位于公共 PyPI 上；固定版本记录在 versions.yaml（wheels.tao_sdk_lepton）中，并通过 scripts/resolve_versions_key.py 解析：

PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_lepton)
python -c "import tao_sdk" 2>/dev/null || {
  echo "MISSING: nvidia-tao-sdk not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}
python -c "import leptonai" 2>/dev/null || {
  echo "MISSING: lepton extra not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

如果缺少该 SDK，Agent 会提示用户授权通过 Bash 安装，然后重新运行预检再继续。

凭据

• LEPTON_WORKSPACE_ID（必需）：决定作业运行在哪个集群和计费账户下。
• LEPTON_AUTH_TOKEN（必需）：用于向 Lepton 控制平面进行身份验证的 API 令牌。
• NGC_KEY（可选）：用于创建镜像拉取密钥，以从 nvcr.io 拉取 TAO 容器镜像。
• ACCESS_KEY / SECRET_KEY（可选）：用于数据集和检查点 URI 的兼容 S3 存储密钥。
• S3_ENDPOINT_URL（可选）：自定义 S3 端点（例如用于 MinIO 或非 AWS S3）。
• S3_BUCKET_NAME（可选）：用于存储作业输出构件的存储桶。
• CLOUD_REGION（可选）：存储区域（例如 us-east-1）。

启动预检

生成脚本或提交作业前：

1. 确认已设置 LEPTON_WORKSPACE_ID 和 LEPTON_AUTH_TOKEN。
2. 使用随附的辅助程序确认工作区 API 可访问： scripts/check_tao_launch_preflight.py --platform lepton ...。
3. 对于 s3:// 数据集/结果，确认已设置 ACCESS_KEY 和 SECRET_KEY， 并且可通过 aws s3 ls 读取确切路径。
4. 对于已挂载的 NFS/Lustre 路径，必须提供 Lepton 卷/存储权限的证明， 确认该路径会挂载到作业中。对于 Lepton 作业，不要将 Agent 主机上的本地 文件系统 test -e 视为证明。
5. 在启动前确认模型专用凭据（例如 HF_TOKEN）。

后端详情

除平台无关的参数（image、command、gpu_count、env_vars、inputs、outputs、hooks）外，LeptonSDK.create_job 还接受以下 Lepton 专用关键字参数：

• resource_shape：显式 GPU 资源形状 ID（例如 "gpu.8xh100-sxm"）。设置后，将跳过根据 gpu_count 自动解析的步骤。其格式不透明（取决于 Lepton API 返回的 instance metadata.id）——可通过 sdk.list_resource_shapes() 查找有效 ID。
• dedicated_node_group：用于保证 GPU 分配的节点组 ID（不会被抢占）。省略则使用共享资源。
• num_nodes：分布式训练的节点数。默认值为 1。大于 1 时，会启用作业内通信和 PyTorch 分布式初始化（参见多节点训练）。
• mounts：用于 NFS / Lustre 的预构建 Mount 对象。未设置时，会根据节点组自动检测。

查找工作区的形状 / 卷

shapes = sdk.list_resource_shapes()
# {<platform_id>: {"cluster": ..., "gpu_type": "gpu.8xh100-sxm",
#                   "gpu_count": 8, "instance_type": ..., ...}, ...}

volumes = sdk.get_volumes(node_group_id="my-h100-pool")
# [{"name": "lustre", "from_path": "/lustre", "type": "Lustre"}, ...]

prefixes = sdk.get_storage_permissions("lustre", "my-h100-pool")
# ["/lustre/fsw/portfolios/edgeai/...", ...]

多节点训练（分布式）

向 create_job 传入 num_nodes > 1，即可进行多节点分布式训练。Lepton 处理器（tao_sdk/platforms/lepton/handler.py）通过设置 intra_job_communication=True（启用 Pod 间网络）、parallelism=num_nodes 和 completions=num_nodes（由 Lepton 调度 N 个副本），并将 WORLD_SIZE=num_nodes 导出为容器环境变量，来配置底层的 LeptonJob。

Lepton 原生的每副本环境变量使用 Lepton 专属名称（LEPTON_JOB_WORKER_INDEX、LEPTON_JOB_TOTAL_WORKERS、LEPTON_JOB_WORKER_PREFIX、LEPTON_SUBDOMAIN），因此处理器会添加一段引导代码，用于加载 Lepton 官方转换脚本：

wget -O init.sh https://raw.githubusercontent.com/leptonai/scripts/main/lepton_env_to_pytorch.sh
chmod +x init.sh
source init.sh
# user command runs here

加载后会设置以下环境变量：

环境变量	来源	值
MASTER_ADDR	脚本	${LEPTON_JOB_WORKER_PREFIX}-0.${LEPTON_SUBDOMAIN}
MASTER_PORT	脚本	29400
NNODES	脚本	${LEPTON_JOB_TOTAL_WORKERS}
NODE_RANK	脚本	${LEPTON_JOB_WORKER_INDEX}
WORKER_ADDRS	脚本	以逗号分隔的非主节点工作器主机名列表
WORLD_SIZE	TAO SDK 处理器	num_nodes（TAO 容器的约定值，与 NNODES 相同）
NUM_GPU_PER_NODE	TAO SDK 处理器	gpu_count（由 TAO 容器的入口点读取）

job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:6.26.3-pyt',
    command='dino train -e /tmp/spec.yaml',  # TAO entrypoint reads WORLD_SIZE + NUM_GPU_PER_NODE
    gpu_count=8,                          # GPUs per node
    num_nodes=4,                          # 4 × 8 = 32 GPUs total
    dedicated_node_group='my-h100-pool',
    inputs={'/data/train.json': 's3://bucket/coco/train.json'},
    outputs=['/results/'],
)

对于基于 torchrun 的原始命令（非 TAO 容器）：

command='torchrun --nnodes=$NNODES --nproc-per-node=8 --node-rank=$NODE_RANK '
        '--master-addr=$MASTER_ADDR --master-port=$MASTER_PORT train.py'

在 Lepton 上运行分布式任务的两种方式

路径	适用场景
TAO SDK create_job(num_nodes=N)（此技能）	从 Agent 代码以编程方式提交；需要 SDK 提供的 S3 封装、监控、故障分析和 JobStore。
Lepton“Torchrun”任务类型（Lepton UI / lep CLI）	通过 Lepton 控制台手动构建并提交。Lepton UI 提供一流的“Torchrun”模式，可为你配置会合过程，无需引导脚本。请参阅官方示例。

参考资料

• NVIDIA 的 Lepton 多节点 PyTorch 示例（UI / Torchrun 模式）：https://docs.nvidia.com/dgx-cloud/lepton/examples/batch-job/distributed-training-with-pytorch/
• SDK 加载的转换脚本：https://github.com/leptonai/scripts/blob/main/lepton_env_to_pytorch.sh
• PyTorch 分布式（环境变量会合）：https://pytorch.org/docs/stable/elastic/run.html
• NCCL 网络调优：https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html

注意事项

• 多节点场景优先使用 dedicated_node_group，以便将副本安排在同一低延迟互连（NVLink / InfiniBand）上。
• 如果共享节点组中的某个副本被抢占，整个任务都会失败——Lepton 在 v1 中不会弹性重启。长时间运行的任务请使用专用节点组。
• 对于由 Lustre 支持的数据集，每个副本都能访问同一个挂载点，无需为每个副本单独封装 I/O。

云存储

尽管平台是 Lepton，存储层兼容 S3。数据集和 results_dir 的 cloud_metadata 键始终使用 aws，URI 协议始终使用 s3://。

• 正确：s3://bucket-name/path
• 错误：lepton://bucket-name/path

容器的 get_cloud_storage_class_object() 会解析 URI 协议，以便从 CLOUD_METADATA[protocol][bucket] 查找凭据。

共享存储（NFS/Lustre）

节点组可以挂载 NFS 或 Lustre 卷。SDK 会自动检测这些卷，并将其挂载到容器中，以便在任务之间持久共享数据。

SDK 函数

• sdk.get_volumes(node_group_id=None) — 返回节点组规范中的可用卷（名称、from_path、类型）
• sdk.get_storage_permissions(volume_name, node_group_id) — 返回卷允许的路径前缀

LeptonSDK.create_job() 会自动调用这些函数来检测挂载，并为任务规范构建相应的 Mount 对象。

脚本运行器如何使用挂载

有可用的 Lustre 挂载时：

• 输入：S3 路径会映射到 Lustre（s3://bucket/path → /mnt/lustre/bucket/path）。如果文件已存在于 Lustre，则直接使用（无需下载）。如果不存在，则从 S3 下载到 Lustre，并保留供后续作业使用。
• 输出：结果先写入 Lustre（速度快且持久），然后上传到 S3（持久可靠）。下游作业（例如，差异分析）可以直接从 Lustre 读取结果，无需经过 S3 中转。

卷优先级顺序

lustre > filestore > first available

Lustre 缓存失效

Lustre 会在多个作业之间持久缓存文件。它没有内置的失效机制。如果上游数据发生变化，但 S3 路径保持不变，Lustre 就会提供过期的缓存版本。要强制缓存未命中：

• 重命名 S3 上的文件（例如，使用 prompt_v2.txt，而不是覆盖 prompt.txt）
• 在不同迭代之间使用新的 storage_root，以避免跨迭代读取过期数据
• 为任何重新生成的制品使用新路径

监控

作业状态

使用 sdk.get_job_status(job_id) 获取高级状态（Pending、Running、Complete、Error）。

副本状态

在启动期间使用 sdk.get_job_replicas(job_id) 获取详细的副本级信息。每个副本都是一个字典：

replicas = sdk.get_job_replicas(job_id)
for r in replicas:
    node = r["status"]["node"]["name"]           # e.g., "node-ip-10-50-111-24"
    node_group = r["status"]["node"]["node_group_id"]
    cpu = r["status"]["cpu"]                      # e.g., 2
    memory_mb = r["status"]["memory_in_mb"]       # e.g., 8192
    readiness = r["status"].get("readiness_issue")
    if readiness:
        reason = readiness["reason"]   # "InProgress", "Failed", "ConfigError"
        message = readiness["message"] # "Pulling image", "Mount point not found", etc.

常见的 readiness_issue 模式：

• reason="InProgress"，message="Pulling image" — 正在拉取镜像（大型镜像出现这种情况很正常）
• reason="Failed" — 镜像拉取失败（检查 NGC_KEY）
• reason="ConfigError" — 节点问题（挂载失败、GPU 错误）
• 没有 readiness_issue — 副本正在运行

当作业卡在 Pending 状态时，副本状态尤其有用——它能显示问题是出在镜像拉取、资源调度还是节点健康状况上。

作业日志

使用 sdk.get_job_logs(job_id, tail=N) 获取最近的 N 行日志。日志从 Lepton 的日志收集服务中获取。

并行作业

对于并行运行的工作流阶段（例如，视频生成 x8）：

1. 启动： 对每个拆分项调用 execute_step(plan, step_id, extra_args={"split_id": i})。每次调用都会立即返回一个 job_id。
2. 监控： 轮询所有作业：对每个作业调用 sdk.get_job_status(job_id)。使用 get_job_replicas(job_id) 进行启动诊断。
3. 完成： 所有作业的状态均为 Complete 或 Error 时，即表示全部完成。
4. 部分失败： 只重试失败的拆分项——成功的拆分项无需重新运行。向 execute_step 传入相同的 split_id。

失败分析

作业失败时，使用 sdk.get_failure_analysis(job_id) 自动检测根本原因：

analysis = sdk.get_failure_analysis(job_id)
if analysis:
    print(analysis["err_class"])    # e.g., "ERR_PROGRAM"
    print(analysis["suggestion"])   # Human-readable fix
    for event in analysis.get("job_failure_by_node_event", []):
        print(event["node_event_name"], event["message"])
        # e.g., "OOM", "OOM encountered, victim process: cosmos-rl-evalu, pid: 3368483"

返回：

• err_class：错误分类（ERR_PROGRAM、ERR_INFRA 等）
• suggestion：可能的问题及修复方法
• job_failure_by_node_event：节点级事件（OOM 终止、GPU 错误、挂载失败）
• log_streams：包含错误上下文的相关日志片段

重试前务必对失败的作业调用此方法——它可以区分用户错误（配置错误、OOM）和基础设施问题（节点故障、驱逐）。

故障模式

OOM 被终止：容器超出 GPU 或系统内存限制。检测方法：get_failure_analysis() 返回 node_event_name: "OOM"。常见原因：evaluation.batch_size 过高，或 max_length 超出可用 KV 缓存的容量。恢复方法：减小 batch_size、通过张量并行添加 GPU，或减小 max_length。

镜像拉取失败：无法从 nvcr.io 拉取 TAO 容器镜像。通常是因为缺少或已过期的镜像拉取密钥。SDK 会根据 NGC_KEY 自动配置密钥，但如果 NGC_KEY 无效，作业就会失败。检测方法：检查 get_job_replicas()——如果 readiness_issue.reason 长时间显示 InProgress，且 message = "Pulling image"，则表示正在拉取镜像；如果拉取失败，则会显示 Failed。恢复方法：确认 NGC_KEY 有效。

资源不可用：请求的 GPU 规格不可用。作业会无限期处于 Queueing 状态。检测方法：Pending 超过 15 分钟，且副本没有节点分配信息。恢复方法：尝试其他 resource_shape 或 dedicated_node_group，或等待资源可用。

身份验证失败：LEPTON_AUTH_TOKEN 无效或已过期。所有 API 调用都会失败并返回 401/403。检测方法：创建作业时立即引发异常。恢复方法：刷新令牌并重新初始化 SDK。

节点不健康：分配的节点存在基础设施问题（挂载失败、GPU 错误、网络问题）。检测方法：检查 get_job_replicas()——readiness_issue.reason = "ConfigError"，并显示类似 "Mount point not found" 的消息。作业会在故障节点上无限期处于 Pending 状态。恢复方法：取消作业并重新提交，Lepton 会将其调度到其他节点。如果问题再次出现，请尝试其他 dedicated_node_group 或 resource_shape。

作业被驱逐：在共享节点组上，Lepton 可能会在资源紧张时驱逐作业。检测方法：作业意外地从 Running 状态转为 Error。恢复方法：重试，或使用 dedicated_node_group。