关于 tao-run-inference-service
用于在不同计算平台启动、检查、请求和停止模型推理服务的 TAO 技能。它规定了按架构生成命令、登记服务、检查就绪状态及处理密钥的流程。
主要功能
- 01跨平台启动推理容器
- 02解析镜像和架构专属命令
- 03登记服务并检查就绪状态
- 04将请求路由到指定服务
- 05通过环境变量处理凭据
适用场景
- 01为指定模型架构启动推理服务
- 02向多个运行中服务之一发送请求
- 03根据任务 ID 停止服务
安装
复制下面的命令,在终端中运行。
npx skills add https:/ / github.com/ nvidia/ skills --skill tao-run-inference-service安装后在 Claude Code 中输入/tao-run-inference-service确认已加载。
包含的文件
除 SKILL.md 外,安装时一起下载的参考资料和脚本。
BENCHMARK.mdskill-card.mdskill.oms.sigconfig/skillspector-baseline.yamlevals/evals.jsonreferences/code-templates.yamlcosmos-predict2.5.config.jsonrequest.yamlservice.yamlskill_info.yamltao-dataservices.config.json详细介绍
TAO 推理微服务
独立安装? 如果此会话不是由 TAO skill bank 插件初始化的,请先运行
tao-setup技能(主机预检、凭据、跨技能发现)。
说明
启动推理服务:
- 收集必需的输入(第 1 节)并解析容器映像(第 2 节)。
- 构建作业负载和内部命令(第 3–4.1 节);使用
references/code-templates.yaml→job_payload_builder。 - 阅读
skills/platform/<platform>/SKILL.md并启动容器(第 4.2 节)。 - 写入服务注册表并轮询就绪状态(第 4.3 节);使用
references/code-templates.yaml→registry_write.<platform>和readiness_check。
发送推理请求:
- 根据第 6.0 节确定由哪个服务接收请求(通过
job_id、通过network_arch,或在运行多个服务时由用户明确选择——服务不止一个时,绝不能默认使用"latest"),然后使用解析出的job_id,通过references/code-templates.yaml→request.registry_read读取端点。 - 构建请求正文前,先向用户询问 vLLM 风格的采样参数(第 6.1 节)。 列出
max_tokens、top_p、temperature(以及任何架构专属参数)及其默认值;允许用户逐项覆盖或跳过以接受默认值。绝不能默默使用默认值。 - 按照第 6.2 节构建并发送正文;按照第 6.3 节处理响应。
停止服务: 通过 references/code-templates.yaml → stop.registry_read 解析 job_id,阅读 skills/platform/<platform>/SKILL.md,然后按照第 5 节操作。
参考数据(架构、映射、有效值——不含操作说明):
references/service.yaml— 映像映射、有效的network_arch名称、作业负载架构、环境变量名称、密钥分类。references/request.yaml— 端点定义、请求字段架构、响应格式、代码示例。references/code-templates.yaml— 用于构建负载、写入注册表、检查就绪状态以及停止/请求流程的 Python 模板。
密钥规则(适用于此技能中生成的每个代码块)
绝不要要求用户在提示词中输入密钥值。 对于每个密钥值:
- 告知用户要设置哪个环境变量——在其 Shell 中运行
export HF_TOKEN=...,或在经用户批准、通过set -a; source /path/to/.env; set +a加载的环境文件中添加一行KEY=value。 - 绝不要打印、执行
cat或Read此类文件——只验证其是否存在:[ -n "$HF_TOKEN" ] && echo SET || echo UNSET。 - 生成的代码应使用
os.environ["VAR_NAME"]读取该值——绝不要硬编码、插入或提示用户输入该值。
密钥环境变量(完整列表见 references/service.yaml → secrets_handling):
HF_TOKEN、WANDB_API_KEY、CLEARML_API_ACCESS_KEY、CLEARML_API_SECRET_KEY、TAO_API_KEY、TAO_USER_KEY。
可在提示词中安全收集: network_arch、model_path、num_gpus、提示词文本、WANDB_* 配置 URL、CLEARML_*_HOST URL。
1. 向用户收集的信息
| 输入 | 作用 |
|---|---|
network_arch | 选择容器镜像、每种架构对应的内部命令格式(references/service.yaml → container_commands.<network_arch>),以及适用时作业 JSON 中的 neural_network_name。必须与 references/service.yaml 中 valid_network_arch_config_basenames 的某个基名匹配(例如 cosmos-rl、cosmos-predict2.5)。 |
model_path | 训练好的模型检查点。有效格式:hf_model://<org>/<model>(HuggingFace Hub — 对受限模型设置 HF_TOKEN),或本地容器文件系统路径。不支持云 URI(s3://、gs://、az://)——推理服务没有云存储依赖。始终询问用户;切勿用占位符代替。参见 references/service.yaml → model_path_protocols。 |
platform | 计算平台:local-docker、brev、slurm 或 kubernetes。 |
num_gpus | 默认值为 1;推理至少需要 1。 |
2. 镜像解析
每个 network_arch 都有一个名为 {network_arch}.config.json 的侧车配置文件。按以下步骤解析容器镜像:
- 读取
{network_arch}.config.json并获取api_params.image(例如COSMOS_RL)。这会从references/service.yaml中的docker_image_defaults选择一个条目。 - 如果设置了主机环境变量
IMAGE_<KEY>(例如IMAGE_COSMOS_RL),则会覆盖打包的默认值。 - 对于
model_skill_defaults下的键,使用其模型、操作和后端调用scripts/resolve_tao_image.py。这样可以确保使用 Cosmos 模型技能的references/skill_info.yaml中指定的确切 Cosmos 镜像。 - 对于
mapping下的键,使用scripts/resolve_versions_key.py,根据仓库根目录的versions.yaml解析点分隔的值。绝对环境变量覆盖值会原样传递。Python 示例位于references/code-templates.yaml。 - 如果配置文件缺失或
api_params.image为空,则回退到COSMOS_RL键。
配置文件还包含 spec_params.inference.model_path,它决定目录与文件路径语义:如果该值包含子字符串 folder,容器会将此路径视为目录。
3. 环境变量(无回调)
在编码 env_json 之前,将这些变量设置到 env_payload 中。不要设置 TAO_LOGGING_SERVER_URL 或 TAO_ADMIN_KEY。
TAO_EXECUTION_BACKEND — 必须与平台匹配:
| 平台 | TAO_EXECUTION_BACKEND 值 |
|---|---|
| local-docker | local-docker |
| brev | local-docker |
| slurm | slurm |
| kubernetes | local-k8s |
CLOUD_BASED — 此技能中始终为 "False"(禁用向 TAO_LOGGING_SERVER_URL 发布回调)。
GPU 环境变量 — 仅当平台技能不会自动处理 GPU 注入时才需要:
- Tegra / Jetson:使用
--runtime=nvidia,并设置NVIDIA_DRIVER_CAPABILITIES=all和NVIDIA_VISIBLE_DEVICES=<ids>。 - 标准 x86 + nvidia-container-toolkit:使用 Docker
device_requests。平台技能会处理此项。
4. 跨平台执行
作业负载和内部命令(第 1–3 节)与平台无关。对于每个平台,在生成任何执行代码之前,请阅读 skills/platform/<name>/SKILL.md,了解预检和凭据要求。
4.1 构建内部命令(按架构)
内部命令的形式因 network_arch 而异——没有统一模板。查阅 references/service.yaml → container_commands.<network_arch> 中对应架构的条目;如果不存在,则表示该架构不受支持——停止并询问。选择 references/code-templates.yaml → job_payload_builder.<network_arch> 中匹配的子区块。在命令前加上 umask 0 &&,并确保命令在所有平台(local-docker、brev、slurm、kubernetes)上完全一致。
所有架构的共同要求:
job_id:使用全新的uuid.uuid4()——它将成为容器名称和注册表键。image:按第 2 节解析。- 密钥(
access_key、secret_key、HF_TOKEN等)在运行时从环境变量中读取——绝不硬编码,也绝不记录或打印。
架构特定说明(详情见 references/service.yaml → container_commands):
cosmos-rl— 使用单个--job '<JOB_JSON>' --docker_env_vars '<ENV_JSON>'数据块;使用json.dumps(...)+shlex.quote(...)。env_payload包含TAO_EXECUTION_BACKEND(参见第 3 节表格)、TAO_API_JOB_ID、CLOUD_BASED=False。推理服务不依赖云存储;唯一可能适用的凭据环境变量是用于受限访问 HuggingFace 模型的HF_TOKEN。cosmos-predict2.5— 使用标志形式的cosmos_predict inference_microservice start ... --port 8080(没有setup.前缀;使用tyro.conf.OmitArgPrefixes)。不接受--job/--docker_env_vars。将model_path转换为--checkpoint-path(本地路径)或--model <registered_key>(hf_model://);拒绝云 URI。唯一可能适用的凭据环境变量是用于受限访问 HuggingFace 模型的HF_TOKEN。每个请求的参数(prompt、inference_type、num_output_frames、guidance、seed、num_steps、negative_prompt)应放在请求体中,而不是在启动时传入。TAO_EXECUTION_BACKEND/TAO_API_JOB_ID/CLOUD_BASED不会使用,可以省略。
4.2 将执行委托给平台技能
阅读 skills/platform/<platform>/SKILL.md 并按照其中的说明启动容器。
基础参数(适用于所有平台):
| 参数 | 值 |
|---|---|
image | 已解析的容器镜像(第 2 节) |
command | inner — 在第 4.1 节中构建的 shell 字符串 |
gpu_count | num_gpus |
env_vars | env_payload |
| 作业 / 容器名称 | job_id — 必须等于第 4.1 节中的 UUID,以便注册表可以引用它 |
host_port (local-docker、brev) | 绑定到容器端口 8080 的主机端口。默认值为 8080,但 每个并发服务的端口必须唯一 — 请参阅下方的端口分配规则。 |
平台特定的附加输入:
| 平台 | 附加输入 |
|---|---|
| local-docker | 除基础参数外,无 |
| brev | instance_id(可选 — 重用现有实例);对于多凭证 / 多工作区帐户,首次创建时还需要 cloud_cred_id 和 workspace_group_id — 请参阅 skills/platform/tao-run-on-brev/SKILL.md |
| slurm | partition 和 account — 检查 SLURM_PARTITION/SLURM_ACCOUNT 环境变量;如果未设置,请询问用户 |
| kubernetes | namespace(默认值:default);image_pull_secret(nvcr.io 镜像必需) |
端口绑定(local-docker 和 brev): 使用 直接执行 docker run,以便传入 -p <host_port>:8080,并确保容器名称与 job_id 完全一致。
端口分配规则(local-docker 和 brev,并发服务必需): 启动服务前,读取注册表(/tmp/tao-inf-ms-state.json),并收集同一平台(对于 brev,还需是同一个 instance_id)下所有现有条目的 host_port 值。选择从 8080 开始的 最小空闲端口,即该端口不在已占用端口集合中 — 例如 host_port = next(p for p in range(8080, 8200) if p not in used_ports)。只有在没有其他服务运行时,默认端口 8080 才适用。这项规则可确保“启动 3 个服务,每个服务都能通过不同的 host_url 访问”正常工作;否则,第 2 和第 3 个服务会因 bind: address already in use 而失败。SLURM 和 kubernetes 通过各自的平台机制分配不同的端点,无需执行此步骤。
4.3 启动后:服务注册表和端点
平台确认容器正在运行后,立即写入服务注册表。注册表(/tmp/tao-inf-ms-state.json)以 job_id 为键;"latest" 始终指向最近启动的服务。
请参阅 references/code-templates.yaml → registry_write.<platform> 中的 Python 模板。
| 平台 | host_url | platform_job_id | 写入前的额外步骤 |
|---|---|---|---|
| local-docker | http://localhost:{host_port} | — | 无 |
| brev | http://{brev_ip}:{host_port} | — | 运行 brev ls → 获取实例 IP(远程 VM 上的 localhost 无效) |
| slurm | http://localhost:{host_port} | SLURM 调度器作业 ID | 等待状态变为 Running;通过 SSH 将端口转发 localhost:{host_port}→{node}:8080 |
| kubernetes | http://{external_ip}:8080 | k8s 作业名称 | 执行 kubectl expose job … --type=LoadBalancer;等待外部 IP |
写入注册表后,打印 job_id 和 URL:
print(f"Inference service started.")
print(f" Job ID : {job_id}")
print(f" Arch : {network_arch}")
print(f" URL : {state[job_id]['host_url']}/v1/chat/completions")
print(f"Use this Job ID to send requests or stop the service.")
然后轮询检查就绪状态 — 请参阅 references/code-templates.yaml → readiness_check。容器会在后台加载模型;在收到 200 响应前不要发送请求。
5. 停止推理服务
询问用户要停止的 job_id。如果用户未提供,则默认使用 state["latest"],并确认要停止的是哪个 job_id。使用 references/code-templates.yaml → stop.registry_read 读取注册表,然后读取 skills/platform/<platform>/SKILL.md,并使用其中的取消 / 停止机制。
| 平台 | 要传入的标识符 | 额外清理 |
|---|---|---|
| local-docker | job_id_to_stop — 容器名称 | 无 |
| brev | job_id_to_stop — 容器名称 | 无 |
| slurm | entry["platform_job_id"] — SLURM 作业 ID | pkill -f "ssh.*-L.*{entry['host_port']}" |
| kubernetes | entry["platform_job_id"] — k8s 作业名称 | kubectl delete svc {entry["platform_job_id"]} -n <namespace> |
其中 entry = state[job_id_to_stop]。停止后,清理注册表:references/code-templates.yaml → stop.registry_cleanup。
6. 发送推理请求
6.0 确定接收此请求的服务(必需)
每个请求都必须路由到运行匹配模型的特定服务。路由依据是 job_id — 注册表会为每个条目存储 network_arch,因此用户提供模型名称而非 job_id 时,可按架构查找目标。按以下顺序应用这些规则:
- 用户提供了明确的
job_id→ 使用该 ID。确认它存在于state中。 - 用户指定了
network_arch(例如“将此请求发送到 cosmos-rl 服务”)→ 查找匹配的条目:candidates = [j for j, e in state.items() if j != "latest" and isinstance(e, dict) and e["network_arch"] == arch]。- 恰好匹配一项 → 使用该项。
- 匹配多项 → 询问用户,列出候选
job_id及其started_at;不要自动选择。 - 没有匹配项 → 停止,并告知用户当前没有运行该架构的服务。
- 既未提供
job_id,也未提供network_arch→ 统计state中非"latest"条目的数量:- 恰好有一个运行中的服务 → 使用该服务。
- 有两个或更多 → 不要默默默认使用
state["latest"]。向用户提供完整列表(job_id、network_arch、host_url),并要求用户明确选择。"latest"指针适用于单服务工作流,但不能作为多个服务共存时的路由回退选项。 - 没有 → 停止,并告知用户先启动一个服务。
确定目标后,从注册表读取端点(references/code-templates.yaml → request.registry_read),并将解析出的 job_id 作为 user_provided_job_id 传入。向用户确认:“正在发送到 job_id=… arch=… url=…”。如果服务可能仍在加载,先轮询检查就绪状态(references/code-templates.yaml → readiness_check)。
**发送前交叉检查:**如果用户提供的请求体包含特定架构的字段(例如 guidance / num_steps / seed / negative_prompt → cosmos-predict2.5;必需的 image_url/video_url 内容项 → cosmos-rl),请确认这些字段与 state[job_id]["network_arch"] 一致。如果不匹配,请停止并询问用户 — 将 cosmos-predict2.5 请求体发送到 cosmos-rl 服务会导致容器返回 4xx/5xx 错误,比起在此处提前发现问题更难排查。
6.1 采样参数 — 每次请求前都必须提示用户
构造请求正文之前,必须明确提示用户提供 vLLM 风格的采样参数。不要悄悄应用默认值。使用结构化提示,每个字段单独提问,并且:
- 列出所有适用字段及其类型和默认值。
- 允许用户跳过或接受任意字段,以采用该字段的默认值——无需必须输入值。
- 在一轮提示中收集所有字段。
提示后,逐字应用用户输入的值,并为所有跳过的字段代入默认值。不要编造值,也不要悄悄限制取值范围。
字段列表、默认值和各架构的适用性:references/request.yaml → chat_completions_request_body(基础采样字段:max_tokens、top_p、temperature)以及 network_arch_constraints.<network_arch>(各架构的覆盖项和额外字段,例如 cosmos-predict2.5 的 guidance/num_steps/seed/negative_prompt)。如果某字段标记为当前架构不支持,不要提示用户填写,也不要将其包含在请求正文中。
6.2 请求格式
向 {BASE_URL}/v1/chat/completions 发送 POST 请求,设置 Content-Type: application/json,超时时间至少为 300 s。请求正文与 OpenAI 兼容(vLLM 聊天补全);完整字段架构和内容项格式(text / image_url / video_url)请参阅 references/request.yaml → chat_completions_request_body,可直接运行的 Python 和 curl 示例请参阅 code_examples。
**约束:**只处理第一条用户消息。请求正文中不得包含机密值。各网络的约束(例如,cosmos-rl 要求每个请求都包含图像或视频;cosmos-rl 会拒绝 data: URI)请参阅 references/request.yaml → network_arch_constraints。
6.3 响应处理
| HTTP 状态 | 含义 | 操作 |
|---|---|---|
| 200 | 成功 — choices[0].message.content 包含生成的文本 | 读取结果 |
| 202 | 服务仍在初始化,或模型仍在加载 | 延迟后重试 |
| 503 | 初始化失败、模型加载失败,或模型尚未就绪 | 检查 error.type:model_not_ready → 重试;initialization_error / model_load_error → 放弃并检查日志 |
| 400 | JSON 正文缺失或为空 | 修正请求 |
| 500 | 推理期间发生未处理的异常 | 检查容器日志 |
对于 202 和 503,正文包含 {"error": {"type": "<error_type>", "message": "<reason>"}}。错误类型字符串请参阅 references/request.yaml 中的 container_response_shapes。