Handle side-effects from Command class

[activatedgeek]

A typical side-effect for the Command class is the value of ports at which the additional container service is exposed.

Another side-effect can be waiting for another container to start the service at a given port, e.g. the nemo-skills container waiting for the model server.

We need a standardized way to handle such scenarios. Two most common scenarios:

1. Passing the exposed service port values to the nemo-skills client container.
2. The client container waiting on other services before launching evaluation.

[activatedgeek]

See #793 for an example of the side-effect of passing around environment variables.

[activatedgeek]

cc @gwarmstrong @Kipok , just tagging for visibility.

[gwarmstrong]

This is pretty much the crux of supporting auxiliary containers. @activatedgeek I'm going to dump some design ideas here. Let me know what you think and whether this handles your primary concerns here.

Auxiliary Container Support Design

Overview

Let's extend the multi-model generation system we are introducing in #1018 to support pipeline-level auxiliary containers (like retrieval services) with automatic port coordination, health check waiting, and deployment flexibility (co-located or dedicated het groups).

Key Principles:

1. Container names are keys to cluster config (like server types)
2. Auto-inferred environment variable naming: NEMO_SKILLS_<UPPERNAME>_<RESOURCE>
3. All containers get all aux resources in their environment
4. Standard nemo-skills exports (env vars, mounts) apply to aux containers
5. Single HTTP port per aux container (if requested)
6. Client waits for aux services to be healthy before starting
7. Aux containers can run on rank 0 (co-located) or dedicated het groups

Exposing Service Variables

Each aux container automatically gets:

• Primary Port: NEMO_SKILLS_<NAME>_PORT (e.g., NEMO_SKILLS_FAISS_PORT)
• Hostname: NEMO_SKILLS_<NAME>_HOST (resolved from het_group)

CLI Examples

Single Port (Simple Case - Co-located)

ns generate \
  --model model1 --model model2 \
  --aux-container faiss \
  --aux-cmd "python server.py --port \$NEMO_SKILLS_FAISS_PORT --data-dir /data"

What happens:

• System allocates faiss_port = get_free_port(strategy="random")
• Container image from cluster config: cluster_config["containers"]["faiss"]
• FAISS runs in Group 0 (same node as model 0 server and client)
• All containers get: NEMO_SKILLS_FAISS_PORT=34521, NEMO_SKILLS_FAISS_HOST=batch-node-0075

Dedicated Het Group (GPU-heavy aux container)

ns generate \
  --model model1 \
  --aux-container embedder \
  --aux-gpus 8 \
  --aux-nodes 1 \
  --aux-dedicated-group \
  --aux-cmd "python embedder_server.py --port \$NEMO_SKILLS_EMBEDDER_PORT"

What happens:

• Embedder runs in its own het group (Group 1)
• Gets dedicated 8 GPUs on separate node(s)
• Client can still access via NEMO_SKILLS_EMBEDDER_HOST and NEMO_SKILLS_EMBEDDER_PORT

Multiple Aux Containers (Mixed Deployment)

ns generate \
  --model model1 \
  --server-gpus 4 \
  --aux-container faiss embedder feedback \
  --aux-gpus - 8 - \
  --aux-nodes - 1 - \
  --aux-dedicated-group false true false \
  --aux-cmd "python faiss.py --port \$NEMO_SKILLS_FAISS_PORT" \
           "python embedder.py --port \$NEMO_SKILLS_EMBEDDER_PORT" \
           "python feedback.py --port \$NEMO_SKILLS_FEEDBACK_PORT"

What happens:

• Group 0: Model server, FAISS, Feedback, Client
• Group 1: Embedder
• Client can wait for all services to be healthy before starting

Health Check and Service Waiting

Problem

Client needs to wait for aux services to be ready before starting generation. Each service might have different health check endpoints.

Solution

Add --aux-health-endpoint parameter (list syntax):

python -m nemo_skills.pipeline.generate \
  --model model1 \
  --aux-container faiss \
  --aux-cmd "python server.py --port \$NEMO_SKILLS_FAISS_PORT" \
  --aux-health-endpoint /health

Default Health Endpoints:

• If --aux-health-endpoint not specified, default to /health
• If service doesn't have health endpoint, set to none to skip waiting

Example with Custom Health Checks

python -m nemo_skills.pipeline.generate \
  --model model1 \
  --aux-container faiss embedder \
  --aux-health-endpoint - /api/status \
  --aux-cmd "python faiss.py --port \$NEMO_SKILLS_FAISS_PORT" \
           "python embed.py --port \$NEMO_SKILLS_EMBEDDER_PORT"

Generated generation plus wait command:

<usual front matter> &&
wait_for_server http://faiss-host:34521/health && \
wait_for_server http://embedder-host:34522/api/status && \
python -m nemo_skills.inference.generate ...