tao-run-inference-service

TAO 推理微服务

说明

启动推理服务：

1. 收集所需输入（第 1 节）并解析容器镜像（第 2 节）。
2. 构建作业负载和内部命令（第 3–4.1 节）；使用 references/code-templates.yaml → job_payload_builder。
3. 阅读 skills/platform/<platform>/SKILL.md 并启动容器（第 4.2 节）。
4. 写入服务注册表并轮询就绪状态（第 4.3 节）；使用 references/code-templates.yaml → registry_write.<platform> 和 readiness_check。

发送推理请求：

1. 按照第 6.0 节解析接收请求的服务（按 job_id、按 network_arch，或在运行多个服务时由用户明确选择——服务不止一个时，绝不能默默默认为 "latest"），然后使用解析出的 job_id，从 references/code-templates.yaml → request.registry_read 读取端点。
2. 构建请求正文前，先询问用户采用 vLLM 风格的采样参数（第 6.1 节）。 提供 max_tokens、top_p、temperature（以及任何架构专属的额外参数）及其默认值；允许用户逐项覆盖或跳过以接受默认值。绝不能默默使用默认值。
3. 按第 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 模板。

---

密钥规则（适用于此技能中的每个生成代码块）

绝不要求用户在提示中输入密钥值。 对于每个密钥值：

1. 告诉用户需要设置哪个环境变量（例如 export HF_TOKEN=...）。
2. 生成代码，通过 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、lepton、slurm 或 kubernetes。
num_gpus	默认值为 1；推理至少需要 1。

---

2. 镜像解析

每个 network_arch 都有一个名为 {network_arch}.config.json 的旁置配置文件。按以下步骤解析容器镜像：

1. 读取 {network_arch}.config.json，获取 api_params.image（例如 COSMOS_RL）。这是 references/service.yaml 中 docker_image_defaults.mapping 的一个键。
2. 在映射中查找该键。如果设置了主机环境变量 IMAGE_<KEY>（例如 IMAGE_COSMOS_RL），则该值会覆盖映射的默认值。
3. 映射值通常是仓库根目录 versions.yaml 清单中的一个点分隔键（例如 tao_toolkit.cosmos_rl）。通过查找 versions.yaml → images.<group>.<name>，将其解析为具体的 nvcr.io/... 镜像 URI。绝对 URI 会原样传递，因此，包含完整 URI 的 IMAGE_<KEY> 环境变量覆盖值仍然有效。此处的 Python 辅助程序位于 references/code-templates.yaml。
4. 如果配置文件缺失或 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
lepton	lepton
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、lepton、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。推理服务不依赖云存储；唯一可能适用的凭据环境变量是 HF_TOKEN（用于受限访问的 HuggingFace 模型）。
• 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
lepton	resource_shape（GPU 形状 ID，例如 gpu.8xh100-sxm）；dedicated_node_group（可选）
slurm	partition 和 account——检查 SLURM_PARTITION / SLURM_ACCOUNT 环境变量；如果未设置，请询问用户
kubernetes	namespace（默认值：default）；image_pull_secret（使用 nvcr.io 镜像时必需）

端口绑定（local-docker 和 brev）： 使用直接运行 docker（而非 DockerSDK），以便传入 -p <host_port>:8080，并确保容器名称与 job_id 完全相同。

端口分配规则（local-docker 和 brev；并发服务必需）： 启动服务前，读取注册表（/tmp/tao-inf-ms-state.json），收集同一平台上所有现有条目的 host_port 值（对于 brev，还需限定为同一 instance_id）。选择从 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 而失败。Lepton、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 无效）
lepton	Lepton 端点 URL	job.id	轮询 sdk.get_job_status，直到状态为 Running；从控制台或运行 lep job get <job.id> 获取端点
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 — 容器名称	无
lepton	entry["platform_job_id"] — Lepton 作业 ID	无
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 时，可以按架构解析目标。按以下顺序应用这些规则：

1. 用户提供了明确的 job_id → 使用它。验证它存在于 state 中。
2. 用户指定了 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；不要自动选择。
   • 未匹配到 → 停止并告知用户没有正在运行的该架构服务。
3. 既没有 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 风格采样参数。不得悄悄应用默认值。使用结构化提问（例如在 Claude Code 中使用 AskUserQuestion，每个字段单独提问），并且：

1. 列出每个适用字段及其类型和默认值。
2. 允许用户跳过或接受任意字段的默认值——无需输入任何值。
3. 一次性收集所有字段。

提问后，逐字应用用户输入的值；任何跳过的字段均使用默认值。不得编造值，也不得悄悄限制取值范围。

字段列表、默认值及各架构的适用性： 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 聊天补全）；完整字段架构和内容项形状（文本 / 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。