feat(orchestration): return and persist a record of every submission

nemo_gym/cli/main.py

@@ -567,9 +567,11 @@ def _reject_scratch_namespace_additions(overrides: list[str]) -> None:
 def _eval_submit(args: argparse.Namespace, overrides: list[str]) -> None:
     from pathlib import Path
 
+    import rich
     from hydra import compose, initialize_config_dir
     from hydra.core.global_hydra import GlobalHydra
     from omegaconf import OmegaConf
+    from rich.markup import escape
 
     from nemo_gym.orchestration.api import SubmitConfig
     from nemo_gym.orchestration.submit import submit
@@ -585,7 +587,31 @@ def _eval_submit(args: argparse.Namespace, overrides: list[str]) -> None:
     resolved = OmegaConf.to_container(composed, resolve=True)
     scratch_keys = {key for key in resolved if key.startswith("_")}
     config = SubmitConfig.model_validate({key: value for key, value in resolved.items() if key not in scratch_keys})
-    submit(config, dry_run=args.dry_run)
+
+    record = submit(config, dry_run=args.dry_run)
+    if record is None:
+        return
+
+    if args.json:
+        # The record alone, so the output parses.
+        print(record.model_dump_json(indent=2))
+    else:
+        rich.print(f"Run directory: [bold]{record.run_dir}[/bold]")
+        for benchmark in record.benchmarks:
+            if benchmark.job_id is None:
+                # sbatch's message is not ours to format: `escape` disables markup
+                # for it, so an error carrying square brackets is printed as
+                # written instead of being swallowed or raising MarkupError.
+                rich.print(f"[red]failed[/red] {benchmark.benchmark}: {escape(benchmark.error or '')}")
+            else:
+                rich.print(
+                    f"[green]submitted[/green] {benchmark.benchmark} → Slurm job [bold]{benchmark.job_id}[/bold]"
+                )
+
+    if record.failed:
+        names = ", ".join(b.benchmark for b in record.failed)
+        print(f"Error: {len(record.failed)} benchmark(s) failed to submit: {names}", file=sys.stderr)
+        sys.exit(1)
 
 
 def _eval_run(args: argparse.Namespace, overrides: list[str]) -> None:
@@ -1113,6 +1139,11 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None:
                     "--dry-run", action="store_true", help="Print generated job scripts without submitting."
                 ),
             ),
+            Flag(
+                register=lambda p: p.add_argument(
+                    "--json", action="store_true", help="Emit the submission record as JSON."
+                ),
+            ),
         ),
     ),
     "eval compare": Command(

nemo_gym/decorators.py

@@ -15,15 +15,18 @@
 
 import functools
 
-import rich
+from rich.console import Console
+
+
+_stderr = Console(stderr=True)
 
 
 def experimental(fn):
-    """Decorator that prints an experimental warning before the function runs."""
+    """Decorator that warns on stderr that a function is experimental, then runs it."""
 
     @functools.wraps(fn)
     def wrapper(*args, **kwargs):
-        rich.print(
+        _stderr.print(
             f"[yellow]Warning:[/yellow] [bold]{fn.__name__}[/bold] is experimental and may change or be removed without notice."
         )
         return fn(*args, **kwargs)

nemo_gym/global_config.py

@@ -1297,12 +1297,11 @@ def parse(self, parse_config: Optional[GlobalConfigDictParserConfig] = None) ->
         )
 
         with open_dict(global_config_dict):
-            # Populate head server defaults
-            if not global_config_dict.get(HEAD_SERVER_KEY_NAME):
-                global_config_dict[HEAD_SERVER_KEY_NAME] = {
-                    "host": default_host,
-                    "port": DEFAULT_HEAD_SERVER_PORT,
-                }
+            # Head server defaults, filled per key so a config may pin just one.
+            head_server = global_config_dict.get(HEAD_SERVER_KEY_NAME) or {}
+            head_server.setdefault("host", default_host)
+            head_server.setdefault("port", DEFAULT_HEAD_SERVER_PORT)
+            global_config_dict[HEAD_SERVER_KEY_NAME] = head_server
 
             # Store final list of disallowed ports.
             global_config_dict[DISALLOWED_PORTS_KEY_NAME] = disallowed_ports

nemo_gym/orchestration/api.py

@@ -156,6 +156,15 @@ class DriverConfig(_StrictModel):
     # Name of a service in `services:` to use as the policy model. When set, injects
     # policy_base_url/policy_model_name/policy_api_key into each benchmark's run config.
     policy_model: str | None = None
+    # Which responses_api_models asset serves as the policy, passed as
+    # `--model-type`. Not every benchmark wants the same one: Gym permits exactly
+    # one entry under `policy_model.responses_api_models`, so composing
+    # openai_model against a benchmark that ships its own vllm_model policy (e.g.
+    # lmarena_v3) fails validation with "Dictionary should have at most 1 item
+    # after validation, not 2", and overrides keyed on `vllm_model.*` land on a
+    # server that was never composed. Set to "" to compose no policy model config
+    # at all, for a benchmark whose own config already declares a complete one.
+    policy_model_type: str = "openai_model"
     benchmarks: dict[str, BenchmarkRunConfig]
     env: dict[str, str] = {}
     # Pyxis-style bind mounts passed as --container-mounts.

nemo_gym/orchestration/executors/base.py

@@ -14,10 +14,48 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable
+from pathlib import Path
 
 from nemo_gym.orchestration.api import SubmitConfig
+from nemo_gym.orchestration.jobs import MANIFEST_NAME, SubmissionRecord
 
 
 class BaseExecutor(ABC):
     @abstractmethod
-    def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: ...
+    def run(self, config: SubmitConfig, *, dry_run: bool = False) -> SubmissionRecord | None:
+        """Submit `config` and return the record describing it.
+
+        Returns None on a dry run, which renders the scripts and stops before
+        anything is submitted; every other path either returns a record or
+        raises.
+        """
+
+    def persist(self, record: SubmissionRecord, write_manifest: Callable[[Path, str], None]) -> None:
+        """Store the record, in the order that survives a partial failure.
+
+        Shared by every executor because the ordering and the failure handling
+        are policy rather than transport. The machine-local index goes first
+        precisely because it cannot fail the submit, so if the manifest write
+        does fail there is still a parseable record for the by-hand recovery the
+        error asks for.
+
+        Only the manifest's transport differs per executor, so it arrives as
+        `write_manifest` -- `Connection.write_text` already has this signature --
+        rather than this class owning a connection it cannot know how to open.
+        Call it while that transport is still open: reopening one here would pay
+        a second connection on every submit.
+
+        The jobs are queued by the time this runs, so a failure has to name them
+        or they are stranded with no record anywhere.
+        """
+        record.write_local_index()
+        manifest = Path(record.run_dir) / MANIFEST_NAME
+        try:
+            write_manifest(manifest, record.dumps())
+        except Exception as error:
+            queued = ", ".join(f"{b.benchmark}={b.job_id}" for b in record.benchmarks if b.job_id)
+            raise RuntimeError(
+                f"Submitted jobs but could not write the manifest to {manifest}: {error}. "
+                f"Already queued: {queued or 'nothing'}. Record these by hand before collecting."
+            ) from error

nemo_gym/orchestration/executors/connection.py

@@ -36,6 +36,9 @@ def copy(self, local: Path, remote: Path) -> None: ...
     @abstractmethod
     def run(self, commands: list[str]) -> str: ...
 
+    @abstractmethod
+    def write_text(self, remote: Path, content: str) -> None: ...
+
     def close(self) -> None:
         pass
 
@@ -47,7 +50,14 @@ def copy(self, local: Path, remote: Path) -> None:
         shutil.copytree(local, remote)
 
     def run(self, commands: list[str]) -> str:
-        return "\n".join(_checked(shlex.split(cmd)) for cmd in commands)
+        # Piped to bash, not shlex.split into argv: callers send compound bash
+        # (`out=$(...); rc=$?; ...`), and this has to mean the same thing here
+        # as it does over SSH.
+        return _checked(["bash", "-s"], input="\n".join(commands), context="local commands")
+
+    def write_text(self, remote: Path, content: str) -> None:
+        remote.parent.mkdir(parents=True, exist_ok=True)
+        remote.write_text(content, encoding="utf-8")
 
 
 class SSHConnection(Connection):
@@ -128,6 +138,21 @@ def run(self, commands: list[str]) -> str:
             context=f"ssh commands on {self._hostname}",
         )
 
+    def write_text(self, remote: Path, content: str) -> None:
+        # A quoted heredoc delimiter: the payload reaches the file byte for byte,
+        # with no parameter or command substitution applied to it on the way.
+        # The heredoc supplies the newline before the delimiter, so a payload
+        # that already ends in one (jobs.dumps does) must shed it -- otherwise
+        # the remote manifest gains a blank line the local index does not have,
+        # and the two stores stop being byte-identical.
+        payload = content.removesuffix("\n")
+        script = f"cat > {shlex.quote(str(remote))} <<'GYM_EOF'\n{payload}\nGYM_EOF\n"
+        _checked(
+            ["ssh", *self._ssh_opts(), self._hostname, "bash", "-s"],
+            input=script,
+            context=f"writing {remote} on {self._hostname}",
+        )
+
     def close(self) -> None:
         subprocess.run(
             ["ssh", *self._ssh_opts(), "-O", "exit", self._hostname],

nemo_gym/orchestration/executors/script_templates.py

@@ -97,6 +97,25 @@ def render_ray_prelude() -> str:
     return _RAY_PRELUDE
 
 
+def escape_for_single_quoted_block(body: str) -> str:
+    """Make `body` safe to embed inside a single-quoted `bash -c '...'` block.
+
+    POSIX shells do not nest single quotes: an inner quote ENDS the outer string
+    rather than nesting in it. Every value that had to be quoted -- a Hydra
+    override containing a space, or any of the JSON blobs vLLM flags take
+    (`--hf-overrides '{"architectures":[...]}'`) -- would otherwise break out of
+    the block and word-split. Both failure modes have been observed on real
+    submissions: Hydra rejecting `+multistage.stages=[{num_tasks:` on its own,
+    and `/usr/bin/env: Argument list too long` from a multi-node vLLM command
+    whose JSON flags reopened the quoting.
+
+    `'"'"'` is the standard end-quote / literal-quote / reopen-quote sequence.
+    It leaves `$VAR` and `$(( ))` untouched, which matters: the inner shell is
+    the one meant to expand them.
+    """
+    return body.replace("'", "'\"'\"'")
+
+
 def render_vllm_ray_symmetric_run(inner_cmd: str, total_nodes: int, resource_flags: str) -> str:
     """Render the Ray head/worker bootstrap that wraps a single vLLM instance's TP/PP command so
     it spans multiple Slurm nodes.
@@ -106,7 +125,13 @@ def render_vllm_ray_symmetric_run(inner_cmd: str, total_nodes: int, resource_fla
     pin fall back to manually starting head/worker Ray processes, keyed on Slurm's per-node task
     rank ($SLURM_NODEID).
     """
-    return _VLLM_RAY_SYMMETRIC_RUN.format(total_nodes=total_nodes, resource_flags=resource_flags, inner_cmd=inner_cmd)
+    # Only the interpolated values are escaped; the template's own structure is
+    # what the quoting is meant to preserve.
+    return _VLLM_RAY_SYMMETRIC_RUN.format(
+        total_nodes=total_nodes,
+        resource_flags=escape_for_single_quoted_block(resource_flags),
+        inner_cmd=escape_for_single_quoted_block(inner_cmd),
+    )
 
 
 def render_health_check(name: str, port: int, path: str, timeout: int) -> str:
@@ -125,13 +150,18 @@ def render_gym_cmd(subcommand: str, var_name: str, args: list[str]) -> str:
     return f"{var_name}=(\n    " + "\n    ".join(entries) + "\n)"
 
 
-def render_repo_checkout(repo: str, ref: str) -> str:
-    """Render an &&-chained command that installs git if missing, then clones and checks out `ref`."""
+def render_repo_checkout(repo: str, ref: str, dest: str | None = None) -> str:
+    """Render an &&-chained command that installs git if missing, then clones and checks out `ref`.
+
+    `dest` is emitted verbatim so it may be a shell expression (the driver passes
+    `"$GYM_SRC/gym"` to clone outside the job directory); omit it to clone into a
+    directory named after the repo in the current one.
+    """
     repo_name = repo.rstrip("/").split("/")[-1].removesuffix(".git")
+    target = dest if dest is not None else shlex.quote(repo_name)
     ensure_git = "command -v git >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq git)"
     return (
-        f"({ensure_git})"
-        f" && git clone {shlex.quote(repo)} && cd {shlex.quote(repo_name)} && git checkout {shlex.quote(ref)}"
+        f"({ensure_git}) && git clone {shlex.quote(repo)} {target} && cd {target} && git checkout {shlex.quote(ref)}"
     )
 
 
@@ -154,10 +184,16 @@ def render_driver_entrypoint(
     preamble: list[str] = []
 
     if repo and ref:
+        # Clone to /tmp rather than the job directory: a checkout plus its .venv
+        # inside every benchmark's rundir is slow to write on lustre and noise in
+        # the artifacts. `cd` into it because a benchmark's prepare_script and
+        # jsonl_fpath resolve against cwd; the driver's output path is absolute,
+        # so nothing depends on the clone being reachable afterwards.
         preamble += [
             "curl -LsSf https://astral.sh/uv/install.sh | sh",
             'source "$HOME/.local/bin/env"',
-            render_repo_checkout(repo, ref),
+            'GYM_SRC="$(mktemp -d /tmp/gym-install-XXXXXX)"',
+            render_repo_checkout(repo, ref, dest='"$GYM_SRC/gym"'),
             # A real venv, not --system: --system targets whatever interpreter happens to be on
             # the container's PATH, sidestepping uv's own project-aware Python selection - `uv
             # venv` instead reads requires-python from pyproject.toml and auto-downloads a
@@ -176,4 +212,5 @@ def render_driver_entrypoint(
 
     preamble.append('exec "$@"')
     body = "\n    ".join(["set -euo pipefail", *preamble])
+    body = escape_for_single_quoted_block(body)
     return f"bash -c '\n    {body}\n' -- \"${{GYM_CMD[@]}}\""