TAO Inference Microservice

Instructions

To start an inference service:

1. Collect required inputs (Section 1) and resolve the container image (Section 2).
2. Build the job payload and inner command (Sections 3–4.1); use references/code-templates.yaml → jobpayloadbuilder.
3. Read skills/platform/<platform>/SKILL.md and start the container (Section 4.2).
4. Write the service registry and poll readiness (Section 4.3); use references/code-templates.yaml → registrywrite.<platform> and readinesscheck.

To send an inference request:

1. Resolve which service receives the request per Section 6.0 (by jobid, by networkarch, or by explicit user choice when multiple services run — never silently default to "latest" when more than one service exists), then read the endpoint from references/code-templates.yaml → request.registryread with the resolved jobid.
2. Before building the request body, prompt the user for the vLLM-style sampling parameters (Section 6.1). Present maxtokens, topp, temperature (and any per-arch extras) with their defaults; let the user override or skip each one to accept the default. Never silently use defaults.
3. Build and send the body per Section 6.2; handle the response per Section 6.3.

To stop a service: Read references/code-templates.yaml → stop.registryread to resolve the jobid, read skills/platform/<platform>/SKILL.md, then follow Section 5.

Reference data (schemas, mappings, valid values — no instructions):

• references/service.yaml — image mappings, valid network_arch names, job payload schema, env var names, secrets classification.
• references/request.yaml — endpoint definition, request field schema, response shapes, code examples.
• references/code-templates.yaml — Python templates for payload building, registry writes, readiness checks, and stop/request flows.

Secrets rule (applies to every generated code block in this skill)

Never ask the user to type a secret value into a prompt. For every secret value:

1. Tell the user which environment variable to set (e.g. export HF_TOKEN=...).
2. Generate code that reads it with os.environ["VAR_NAME"] — never hard-code, interpolate, or prompt for the value.

Secret env vars (full list in references/service.yaml → secretshandling): HFTOKEN, WANDBAPIKEY, CLEARMLAPIACCESSKEY, CLEARMLAPISECRETKEY, TAOAPIKEY, TAOUSERKEY.

Safe to collect in the prompt: networkarch, modelpath, numgpus, prompt text, WANDB config URLs, CLEARMLHOST URLs.

1. What to collect from the user

2. Image resolution

Each networkarch has a sidecar config file named {networkarch}.config.json. Resolve the container image as follows:

1. Read {networkarch}.config.json and take apiparams.image (e.g. COSMOSRL). This is a key into dockerimage_defaults.mapping in references/service.yaml.
2. Look up that key in the mapping. If the host env var IMAGE<KEY> is set (e.g. IMAGECOSMOS_RL), it overrides the mapped default.
3. The mapped value is normally a dotted key into the repo-root versions.yaml manifest (e.g. taotoolkit.cosmosrl). Resolve it to a concrete nvcr.io/... image URI by looking up versions.yaml → images.<group>.<name>. Absolute URIs pass through unchanged, so an IMAGE_<KEY> env-var override that contains a full URI still works. The Python helper for this lives in references/code-templates.yaml.
4. If the config file is missing or apiparams.image is empty, fall back to the COSMOSRL key.

The config file also has specparams.inference.modelpath which drives folder vs file path semantics: if the value contains the substring folder, the container treats the path as a directory.

3. Environment variables (no callbacks)

Set these in envpayload before encoding envjson. Do not set TAOLOGGINGSERVERURL or TAOADMIN_KEY.

TAOEXECUTIONBACKEND — must match the platform:

CLOUDBASED — always "False" for this skill (disables callback posting to TAOLOGGINGSERVERURL).

GPU env vars — only needed when the platform skill does not handle GPU injection automatically:

• Tegra / Jetson: --runtime=nvidia with NVIDIADRIVERCAPABILITIES=all and NVIDIAVISIBLEDEVICES=<ids>.
• Standard x86 + nvidia-container-toolkit: use Docker device_requests. The platform skill handles this.

4. Executing across platforms

The job payload and inner command (Sections 1–3) are platform-agnostic. For each platform, read skills/platform/<name>/SKILL.md for preflight checks and credentials before generating any execution code.

4.1 Build the inner command (per arch)

The inner-command shape is per networkarch — there is no uniform template. Look up the per-arch entry in references/service.yaml → containercommands.<networkarch>; if not present, the arch is unsupported — stop and ask. Pick the matching sub-block in references/code-templates.yaml → jobpayloadbuilder.<networkarch>. Prefix the command with umask 0 && and keep it identical across platforms (local-docker, brev, slurm, kubernetes).

Common across arches:

• job_id: fresh uuid.uuid4() — becomes the container name and registry key.
• image: resolve per Section 2.
• Secrets (accesskey, secretkey, HF_TOKEN, etc.) are read from env vars at runtime — never hard-code, never log or print.

Arch-specific notes (full details in references/service.yaml → container_commands):

• cosmos-rl — single --job '<JOBJSON>' --dockerenvvars '<ENVJSON>' blob; json.dumps(...) + shlex.quote(...). envpayload carries TAOEXECUTIONBACKEND (per Section 3 table), TAOAPIJOBID, CLOUDBASED=False. The inference service has no cloud-storage dependency; HFTOKEN is the only cred env var that ever applies (for gated HuggingFace models).
• cosmos-predict2.5 — flag-style cosmospredict inferencemicroservice start ... --port 8080 (no setup. prefix; uses tyro.conf.OmitArgPrefixes). --job/--dockerenvvars are not accepted. Translate modelpath to --checkpoint-path (local path) or --model <registeredkey> (hfmodel://); cloud URIs are rejected. The only cred env var that ever applies is HFTOKEN for gated HuggingFace models. Per-request params (prompt, inferencetype, numoutputframes, guidance, seed, numsteps, negativeprompt) go in the request body, not at startup. TAOEXECUTIONBACKEND/TAOAPIJOBID/CLOUD_BASED are unused and may be omitted.

4.2 Delegate execution to the platform skill

Read skills/platform/<platform>/SKILL.md and follow it to start the container.

Base parameters (all platforms):

Platform-specific additional inputs:

Port binding (local-docker and brev): use direct docker run (not DockerSDK) so that -p <hostport>:8080 can be passed and the container name equals jobid exactly.

Port allocation rule (local-docker and brev, REQUIRED for concurrent services): Before starting a service, read the registry (/tmp/tao-inf-ms-state.json) and collect the set of hostport values from every existing entry on the same platform (and, for brev, the same instanceid). Pick the lowest free port starting from 8080 that is not in that set — e.g. hostport = next(p for p in range(8080, 8200) if p not in usedports). The default 8080 only applies when no other service is running. This is what makes "start 3 services, each reachable at a distinct host_url" work; without it, services 2 and 3 fail with bind: address already in use. SLURM and kubernetes get distinct endpoints from their own platform mechanisms and do not need this step.

4.3 After start: service registry and endpoint

Write the service registry immediately after the platform confirms the container is running. The registry (/tmp/tao-inf-ms-state.json) is keyed by job_id; "latest" always points to the most recently started service.

See references/code-templates.yaml → registry_write.<platform> for the Python template.

After writing the registry, print the job_id and 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.")

Then poll for readiness — see references/code-templates.yaml → readiness_check. The container loads the model in the background; do not send requests before it returns 200.

5. Stopping the inference service

Ask the user for the jobid to stop. If they don't provide one, default to state["latest"] and confirm which jobid is being stopped. Read the registry using references/code-templates.yaml → stop.registry_read, then read skills/platform/<platform>/SKILL.md and use its cancellation / stop mechanism.

where entry = state[jobidtostop]. After stopping, clean up the registry: references/code-templates.yaml → stop.registrycleanup.

6. Sending inference requests

6.0 Resolve which service receives this request (REQUIRED)

Each request must be routed to the specific service that runs the matching model. Routing happens by jobid — the registry stores networkarch per entry, so you can resolve a target by arch when the user names a model instead of a job_id. Apply these rules in order:

1. User provided an explicit job_id → use it. Verify it exists in state.
2. User named a networkarch (e.g. "send this to the cosmos-rl service") → look up matching entries: candidates = [j for j, e in state.items() if j != "latest" and isinstance(e, dict) and e["networkarch"] == arch].

- Exactly one match → use it. - Multiple matches → prompt the user with the candidate jobids and their startedat; do not auto-pick. - No match → stop and tell the user no service for that arch is running.

1. No jobid and no networkarch → count non-"latest" entries in state:

- Exactly one running service → use it. - Two or more → do not silently default to state["latest"]. Prompt the user with the full list (jobid, networkarch, host_url) and require an explicit choice. The "latest" pointer is a convenience for single-service workflows, not a routing fallback when multiple services coexist. - Zero → stop and tell the user to start a service first.

After resolving, read the endpoint from the registry (references/code-templates.yaml → request.registryread), passing the resolved jobid as userprovidedjobid. Confirm to the user: "Sending to jobid=… arch=… url=…". If the service may still be loading, poll readiness first (references/code-templates.yaml → readiness_check).

Cross-check before sending: if the user-supplied request body contains arch-specific fields (e.g. guidance / numsteps / seed / negativeprompt → cosmos-predict2.5; required imageurl/videourl content items → cosmos-rl), verify they are consistent with state[jobid]["networkarch"]. On mismatch, stop and ask — sending a cosmos-predict2.5 body to a cosmos-rl service will fail at the container with a 4xx/5xx that is harder to diagnose than catching it here.

6.1 Sampling parameters — REQUIRED user prompt before each request

Before constructing the request body, you MUST explicitly prompt the user for the vLLM-style sampling parameters. Do not silently apply defaults. Use a structured prompt, one question per field, that:

1. Lists every applicable field with its type and default value.
2. Lets the user skip / accept any field to take that field's default — entering a value is never required.
3. Collects all fields in one round.

After the prompt, apply each user-entered value verbatim and substitute the default for any skipped field. Do not invent values or silently clamp.

Field list, defaults, and per-arch applicability: references/request.yaml → chatcompletionsrequestbody (base sampling fields: maxtokens, topp, temperature) and networkarchconstraints.<networkarch> (per-arch overrides and extras such as guidance/numsteps/seed/negativeprompt for cosmos-predict2.5). If a field is marked unsupported for the active arch, do not prompt for it and do not include it in the body.

6.2 Request format

Send a POST to {BASEURL}/v1/chat/completions with Content-Type: application/json and a timeout of at least 300 s. The body is OpenAI-compatible (vLLM chat completions); see references/request.yaml → chatcompletionsrequestbody for the full field schema and content-item shapes (text / imageurl / videourl), and code_examples for ready-to-run Python and curl samples.

Constraints: only the first user message is processed. No secret values in request bodies. Per-network constraints (e.g. cosmos-rl requires every request to include an image or video; cosmos-rl rejects data: URIs) are in references/request.yaml → networkarchconstraints.

6.3 Response handling

For 202 and 503, the body contains {"error": {"type": "<errortype>", "message": "<reason>"}}. See containerresponse_shapes in references/request.yaml for error type strings.