Basic SLURM orchestration

[prokotg]

gym eval submit — Slurm/Pyxis orchestration (initial implementation)

Introduces gym eval submit, a new CLI command that takes a declarative YAML config and submits one sbatch job per benchmark to a Pyxis-enabled Slurm cluster, handling staging, SSH transfer, service startup, health checks, and driver execution end-to-end.

What's in this PR

Config model (nemo_gym/orchestration/api.py)

• SubmitConfig — top-level config with services, compute, driver, job sections; all models use extra="forbid" so typos in YAML surface immediately
• Discriminated unions for ServiceConfig (vllm, ray) and ComputeConfig (slurm) — unknown types are rejected at parse time
• BaseModelServiceConfig — base for services that can be wired as the policy model; currently VllmServiceConfig
• NodePool with structured allocation fields (nodes, ntasks_per_node, gpus_per_node) the executor uses for deployment decisions, plus extra_args for arbitrary #SBATCH directives
• driver.policy_model — syntactic sugar: naming a service auto-injects policy_base_url, policy_model_name, policy_api_key into every benchmark's run args, with conflict detection if those keys are already set
• driver.gym_install — optional {repo, ref} to install gym from source at runtime via uv
• benchmarks as a dict (name is the key); each benchmark has prepare and run dicts of Hydra overrides

Executor (nemo_gym/orchestration/executors/)

• BaseExecutor — thin ABC; SlurmExecutor is the only implementation
• SlurmExecutor — stages job dirs locally, rsync's to remote, submits all benchmarks in one SSH session, prints Slurm job IDs on success; --dry-run prints generated scripts without touching the cluster
• Connection — LocalConnection (shutil + subprocess) vs SSHConnection (ControlMaster socket, single session for copy + all sbatch calls); dispatched based on whether hostname matches socket.gethostname()
• build_sbatch_script — generates a Pyxis sbatch script per benchmark: #SBATCH directives, backgrounded srun per service with --overlap --no-container-mount-home, curl-based health checks with dead-process detection, then a single driver srun that optionally runs gym eval prepare before gym eval run in the same container step
• Hydra overrides use + (not ++) so config_paths participates in gym's _merge_config_paths coalescing and doesn't clobber model-type configs
• Output artifacts land at <remote_bench_dir>/artifacts/rollouts.jsonl; logs/ and artifacts/ dirs pre-created during staging

CLI

• gym eval submit --config <yaml> [--dry-run]
• @experimental decorator warns users on invocation

Known gaps / next steps

1. Local submit (login node) — LocalConnection is wired when hostname matches, but untested end-to-end
2. Env var injection — per-service and driver env vars (literal values and host env var references) not yet supported
3. Rollout output path — artifacts/rollouts.jsonl is set but downstream profiling commands haven't been validated against it
4. Multi-instance vLLM — single vLLM instance per job; multi-instance requires allocation shape reasoning (gpus_per_node is modeled but not acted on yet)
5. Mounts — per-service and driver container mounts (--container-mounts) not yet supported
6. RayService — placeholder only; needs a real Ray cluster spanning the full allocation wired into Gym's Ray backend
7. Multi-node vLLM — --distributed-executor-backend (mp or ray) not yet generated
8. Broader benchmark coverage — only gsm8k tested; more complex benchmarks (e.g. MCP-Atlas) may surface config or timing issues
9. Startup parallelism — currently: start services → wait for health → prepare → run; prepare could overlap with service startup to reduce wall time

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[iguehring-nvidia]

I guess this is not in the scope of the first version. But would be great to see this tested against the vllm configs that we have in EFB to see that we achieve parity. E.g., sth like

model:
    type: vllm
    model: path/to/weights
    protocol: chat_completions
    port: 5000
    tensor_parallel_size: 4
    data_parallel_size: 16
    num_nodes: 8
    image: /path/to/image.sqsh
    startup_timeout: 3600.0
    extra_args:
      - "--trust-remote-code"
      - "--enable-auto-tool-choice"
      - "--tool-call-parser=x"
      - "--reasoning-parser=x"
      - "--compilation-config={\"pass_config\":{\"fuse_allreduce_rms\":false}}"
      - "--load-format=dummy"
      - "--data-parallel-size-local=2"
      - "--data-parallel-backend=ray"
      - "--api-server-count=1"
      - "--enable-log-requests"
      - "--enable-prefix-caching"
      - "--enable-chunked-prefill"
      - "--kv-cache-dtype=fp8"
      - "--enable-expert-parallel"
    extra_env:
      HF_HOME: /cache/huggingface
      HF_TOKEN: ${HF_TOKEN}
      VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
      VLLM_ENGINE_READY_TIMEOUT_S: "3600"
      RAY_raylet_start_wait_time_s: "120"
      VLLM_FLASHINFER_ALLREDUCE_BACKEND: trtllm
      VLLM_ALLREDUCE_USE_SYMM_MEM: "0"
      FLASHINFER_WORKSPACE_BASE: /tmp
    container_mounts:
      - /cache/huggingface:/cache/huggingface
      - /cache/vllm:/cache/vllm
    generation:
      temperature: 1.0
      top_p: 1.0
    node_pool: gpu

[iguehring-nvidia]

nit: I didn't check the complete repo, but I saw logger being used in other parts of the code. That seems preferably to me. I still hope that some day, we will configure gym logging to output structured (json) logs and then collect all of them and be able to analyse runs (e.g., in Grafana) by querying the logs.

Same for all other rich.print statements.

[iguehring-nvidia]

nit: For my subjective, personal taste, Gym is much too flat. Wondering if we should start here by having a utils package?

[prokotg]

I'll leave this to a wider audience

[iguehring-nvidia]

nit: This function looks a bit old-style / manual to me with argv, string parsing, and error message about missing args. I didn't really check below LLM output, but maybe sth like this?

def _load_submit_config() -> tuple[SubmitConfig, bool]:
    parser = argparse.ArgumentParser(prog="ng-submit")
    parser.add_argument("--config", "-c", required=True, help="Path to submit config YAML")
    parser.add_argument("--dry-run", action="store_true")
    args, overrides = parser.parse_known_args()

    if malformed := [t for t in overrides if "=" not in t]:
        parser.error(f"overrides must be key=value, got: {' '.join(malformed)}")

    merged = OmegaConf.merge(
        OmegaConf.load(args.config),
        OmegaConf.from_dotlist([t.lstrip("+") for t in overrides]),
    )
    return SubmitConfig.model_validate(OmegaConf.to_container(merged, resolve=True)), args.dry_run

[prokotg]

You're right, addressed in 9308727

[iguehring-nvidia]

nit: maybe log that? I see its getting logged in _dry_run later, but not sure if also for non-dry-run.

[iguehring-nvidia]

nit: logging.info?

[iguehring-nvidia]

Mutates benchmark.run during validation, injecting policy_base_url / policy_model_name. api.py:135-139 raises if those keys are already present. So round-tripping a config (SubmitConfig.model_validate(cfg.model_dump())) raises "already sets" on a config that just validated cleanly.

[prokotg]

I couldn't reproduce. In my case this mechanism works well, can you give an example where this fails?

[iguehring-nvidia]

Not sure if this is in the initial scope, but these are quite helpful: --container-mounts / --container-env, so host paths can be remapped into the container (HF_HOME=/cache/huggingface, vLLM's /cache/vllm compile cache) and environment reaches the process (HF_TOKEN).

[AI feedback] srun with no --nodes/--ntasks inherits the full allocation, so on a multi-node pool this launches one independent vLLM per node, all binding the same port — and the driver srun at line 105 fans out the same way into N duplicate benchmark runs. It fails silently: the health check curls localhost, gets a 200 from the node-0 copy, and the job proceeds using a fraction of the allocation.

[oyilmaz-nvidia]

Ran a few tests and LGTM. Will create 1-2 more PRs for the remaining features.