diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 7b68aff3b2..464b86c225 100644 --- a/nemo_gym/cli/main.py +++ b/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( diff --git a/nemo_gym/decorators.py b/nemo_gym/decorators.py index 39a2bfcfac..b7a75e7598 100644 --- a/nemo_gym/decorators.py +++ b/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) diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 3a77533d45..0237d4783d 100644 --- a/nemo_gym/global_config.py +++ b/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 diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 0975fbb7a1..fa08af9ad9 100644 --- a/nemo_gym/orchestration/api.py +++ b/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. diff --git a/nemo_gym/orchestration/executors/base.py b/nemo_gym/orchestration/executors/base.py index 8eb95dce6d..f6108261e7 100644 --- a/nemo_gym/orchestration/executors/base.py +++ b/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 diff --git a/nemo_gym/orchestration/executors/connection.py b/nemo_gym/orchestration/executors/connection.py index 3282cd2871..ac1b45e6b9 100644 --- a/nemo_gym/orchestration/executors/connection.py +++ b/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], diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index 918b506682..ce9fb5d52c 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/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[@]}}\"" diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index ae96686cf3..1ee7411a36 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -13,21 +13,103 @@ # See the License for the specific language governing permissions and # limitations under the License. +import getpass import re import shlex import tempfile from datetime import datetime from pathlib import Path -import rich - +from nemo_gym import __version__ from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.base import BaseExecutor -from nemo_gym.orchestration.executors.connection import Connection, LocalConnection, get_connection +from nemo_gym.orchestration.executors.connection import Connection, get_connection from nemo_gym.orchestration.executors.slurm_script import build_sbatch_script +from nemo_gym.orchestration.jobs import ( + BenchmarkJob, + SubmissionRecord, + new_gym_job_id, + utc_now, + utc_timestamp, +) + + +# Each sbatch reports its own result on a line that names its benchmark, so a +# failure cannot shift the benchmarks after it onto the wrong job ids. +_MARKER = "__GYM_JOB:" + +# Benchmark names are interpolated into the marker line, so they must not carry +# its delimiter or anything the shell would act on. +_VALID_BENCHMARK_NAME = re.compile(r"^[A-Za-z0-9._-]+$") + +# Real `sbatch --parsable` output is a bare "" or a federated ";". +# stderr is merged onto the same line (see `_sbatch_command`), and a *successful* +# sbatch can still print a warning (job_submit plugin notices, QOS/ntasks +# adjustments are routine on production Slurm) — so only the payload's last +# whitespace-delimited token is checked against this, not the whole line. +_JOB_ID_RE = re.compile(r"^\d+(;\S+)?$") + + +def _validate_benchmark_names(benchmarks: list[str]) -> None: + """Fail before anything is staged or copied. + + `_sbatch_command` raises this same check, but only once submission is + already underway — after staging, connecting, mount validation and the + rsync. Checking here first means a bad name fails before any of that, and + fails on `--dry-run` too, instead of the dry run printing a clean script + listing for a benchmark that could never actually be submitted. + """ + bad = [name for name in benchmarks if not _VALID_BENCHMARK_NAME.match(name)] + if bad: + raise ValueError( + f"Invalid benchmark name(s) {', '.join(map(repr, bad))}: names must match " + f"{_VALID_BENCHMARK_NAME.pattern} so they can be reported back from the submit script." + ) -_SBATCH_JOB_ID_RE = re.compile(r"Submitted batch job (\d+)") +def _sbatch_command(benchmark: str, script: Path) -> str: + """One `sbatch` that reports its own benchmark, exit status and output. + + `rc` is captured immediately: `$?` after the `tr` pipeline would be *tr's* + status, which is zero however badly sbatch failed. `tr` flattens sbatch's + multi-line error messages so the whole result stays on one marker line. + """ + if not _VALID_BENCHMARK_NAME.match(benchmark): + raise ValueError( + f"Invalid benchmark name {benchmark!r}: names must match {_VALID_BENCHMARK_NAME.pattern} " + "so they can be reported back from the submit script." + ) + return ( + f"out=$(sbatch --parsable {shlex.quote(str(script))} 2>&1); rc=$?; " + "out=$(echo \"$out\" | tr '\\n' ' '); " + f'echo "{_MARKER}{benchmark}:$rc:$out"' + ) + + +def _parse_sbatch_results(output: str) -> dict[str, tuple[str | None, str | None]]: + """Benchmark name to `(job_id, error)`, exactly one of which is set.""" + results: dict[str, tuple[str | None, str | None]] = {} + for line in output.splitlines(): + if not line.startswith(_MARKER): + continue + benchmark, _, rest = line[len(_MARKER) :].partition(":") + status, _, payload = rest.partition(":") + payload = payload.strip() + if status == "0": + tokens = payload.split() + candidate = tokens[-1] if tokens else "" + if _JOB_ID_RE.match(candidate): + # A federated sbatch answers "jobid;cluster"; the ledger wants the id. + results[benchmark] = (candidate.split(";")[0], None) + else: + # Exit 0 but the last token isn't an id: no output at all, or a + # warning with nothing that looks like a job id after it. Either + # way there is no id to trust, so this is a failure, not a + # success with garbage (or an empty string) in job_id. + results[benchmark] = (None, payload or "sbatch exited 0 with no output") + else: + results[benchmark] = (None, payload or f"sbatch exited {status}") + return results def _validate_mounts(config: SubmitConfig, conn: Connection) -> None: @@ -38,13 +120,12 @@ def _validate_mounts(config: SubmitConfig, conn: Connection) -> None: if not srcs_by_label: return - if isinstance(conn, LocalConnection): - missing = {src for _, src in srcs_by_label if not Path(src).exists()} - else: - # SSHConnection.run() pipes commands as a bash script, so || works fine. - checks = [f'test -e {shlex.quote(src)} || echo "__GYM_MISSING:{src}"' for _, src in srcs_by_label] - output = conn.run(checks) - missing = {line[len("__GYM_MISSING:") :] for line in output.splitlines() if line.startswith("__GYM_MISSING:")} + # Both connections pipe commands to bash, so one shell program checks the + # mounts wherever the submit is going -- and the local path exercises the + # same code the SSH path runs. + checks = [f'test -e {shlex.quote(src)} || echo "__GYM_MISSING:{src}"' for _, src in srcs_by_label] + output = conn.run(checks) + missing = {line[len("__GYM_MISSING:") :] for line in output.splitlines() if line.startswith("__GYM_MISSING:")} if missing: bad = [(label, src) for label, src in srcs_by_label if src in missing] @@ -60,14 +141,18 @@ class SlurmExecutor(BaseExecutor): bash inside the sbatch script (no container needed — they just poll HTTP). """ - def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: + def run(self, config: SubmitConfig, *, dry_run: bool = False) -> SubmissionRecord | None: compute = next(iter(config.compute.values())) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - remote_run_dir = Path(config.job.output_path) / f"gym-job-{timestamp}" + cluster = next(iter(config.compute)) + benchmark_names = list(config.driver.benchmarks) + _validate_benchmark_names(benchmark_names) + now = utc_now() + gym_job_id = new_gym_job_id(now) + remote_run_dir = Path(config.job.output_path) / gym_job_id if dry_run: self._dry_run(config, compute, remote_run_dir) - return + return None with tempfile.TemporaryDirectory(prefix="gym-submit-") as staging_str: staging = self._stage(config, compute, remote_run_dir, Path(staging_str)) @@ -75,18 +160,52 @@ def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: _validate_mounts(config, conn) conn.copy(staging, remote_run_dir) output = conn.run( - [ - f"sbatch {shlex.quote(str(remote_run_dir / name / 'job.sh'))}" - for name in config.driver.benchmarks - ] + [_sbatch_command(name, remote_run_dir / name / "job.sh") for name in benchmark_names] ) - - benchmark_names = list(config.driver.benchmarks) - job_ids = _SBATCH_JOB_ID_RE.findall(output) - for name, job_id in zip(benchmark_names, job_ids): - rich.print(f"[green]submitted[/green] {name} → Slurm job [bold]{job_id}[/bold]") - for name in benchmark_names[len(job_ids) :]: - rich.print(f"[green]submitted[/green] {name} (job ID unavailable)") + record = self._build_record(cluster, compute, gym_job_id, now, remote_run_dir, benchmark_names, output) + # Inside the connection, because that is the transport persist() + # needs and reopening one would cost a second connection per + # submit. Ordering and failure handling live in the base class. + self.persist(record, conn.write_text) + + return record + + def _build_record( + self, + cluster: str, + compute: SlurmComputeConfig, + gym_job_id: str, + now: datetime, + remote_run_dir: Path, + benchmark_names: list[str], + output: str, + ) -> SubmissionRecord: + results = _parse_sbatch_results(output) + benchmarks = [] + for name in benchmark_names: + # A benchmark absent from the output produced no marker line at all — + # the shell died before reaching it, or the transport truncated. That + # is a failure, not a success with a missing id. + job_id, error = results.get(name, (None, "sbatch produced no result for this benchmark")) + benchmarks.append( + BenchmarkJob( + benchmark=name, + job_dir=str(remote_run_dir / name), + job_id=job_id, + error=error, + ) + ) + return SubmissionRecord( + gym_job_id=gym_job_id, + gym_version=__version__, + submitted_at=utc_timestamp(now), + run_dir=str(remote_run_dir), + cluster=cluster, + executor="slurm", + hostname=compute.hostname, + submitted_by=getpass.getuser(), + benchmarks=benchmarks, + ) def _dry_run(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> None: print(f"[dry-run] remote run dir: {remote_run_dir}") diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index d22be589ee..2dcaed08e4 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -31,6 +31,7 @@ from nemo_gym.orchestration.executors.script_templates import ( ENSURE_RAY_INSTALLED, bash_var, + escape_for_single_quoted_block, render_driver_entrypoint, render_gym_cmd, render_health_check, @@ -63,7 +64,8 @@ def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benc lines.append(f"#SBATCH --account={compute.account}") if compute.walltime: lines.append(f"#SBATCH --time={compute.walltime}") - # --chdir makes relative paths (logs/, artifacts/) resolve correctly inside the job. + # --chdir sets the batch script's cwd on the HOST, so srun --output=logs/... resolves there. + # Container-side cwd is set separately per step (see driver_workdir_flag). lines.append(f"#SBATCH --chdir={remote_bench_dir}") for key, val in compute.extra_args.items(): lines.append(f"#SBATCH --{key}={val}") @@ -173,6 +175,17 @@ def _build_vllm_single_instance_multi_node_command(service: VllmServiceConfig, t return render_vllm_ray_symmetric_run(inner_cmd, total_nodes, resource_flags) +# vLLM refuses `--api-server-count` in headless mode ("no API servers are started in headless +# mode") and exits before loading anything. The flag is legitimate on the head node and reaches us +# through a service's own extra_args, so it is stripped from the worker command rather than +# rejected: Gym decides which nodes run headless, so Gym keeps their command valid. +_HEADLESS_INCOMPATIBLE_FLAG = re.compile(r"\s--api-server-count(?:[= ]\S+)?") + + +def _strip_headless_incompatible_flags(cmd: str) -> str: + return _HEADLESS_INCOMPATIBLE_FLAG.sub("", cmd) + + def _build_vllm_multi_instance_multi_node_command(service: VllmServiceConfig, total_nodes: int) -> str: # Data-parallel replicas span nodes. vLLM's Ray-based DP auto-placement doesn't spread ranks # across physical nodes - launching a single `vllm serve --data-parallel-size N` from one node @@ -194,18 +207,23 @@ def _build_vllm_multi_instance_multi_node_command(service: VllmServiceConfig, to trust_flag = " --trust-remote-code" if service.trust_remote_code else "" head_cmd = common + dp_flags + trust_flag worker_cmd = ( - common + _strip_headless_incompatible_flags(common) + dp_flags + trust_flag + " --headless" + f" --data-parallel-start-rank $(( SLURM_NODEID * {dp_size_local} ))" ) + # Both branches go inside a single-quoted `bash -lc '...'`, and this service's + # command carries JSON flags that are themselves single-quoted + # (--hf-overrides, --limit-mm-per-prompt, --media-io-kwargs). Unescaped they + # end the block early and the whole invocation word-splits; mmlu-prox died + # that way with "/usr/bin/env: Argument list too long". return ( "bash -lc '\n" ' if [ "$SLURM_NODEID" = "0" ]; then\n' - f" {head_cmd}\n" + f" {escape_for_single_quoted_block(head_cmd)}\n" " else\n" - f" {worker_cmd}\n" + f" {escape_for_single_quoted_block(worker_cmd)}\n" " fi\n" "'" ) @@ -369,8 +387,15 @@ def build_sbatch_script( if benchmark.prepare: prepare_cmd = "gym eval prepare " + " ".join(flatten_run_args(benchmark.prepare)) - output_path = "+output_jsonl_fpath=artifacts/rollouts.jsonl" - extra_flags = ["--model-type openai_model"] if config.driver.policy_model else [] + # ABSOLUTE, not relative. The driver `cd`s into the Gym checkout so that a + # benchmark's own relative `prepare_script` / `jsonl_fpath` resolve, which + # means a relative output path would write every artifact inside that + # checkout instead of the job directory -- the run completes, exits 0, and + # leaves nothing behind. Making the OUTPUT absolute is what keeps artifacts + # in the job directory without constraining cwd. + output_path = f"+output_jsonl_fpath={remote_bench_dir}/artifacts/rollouts.jsonl" + policy_type = config.driver.policy_model_type + extra_flags = [f"--model-type {shlex.quote(policy_type)}"] if config.driver.policy_model and policy_type else [] run_args = _with_default_capture_dir(benchmark.run, remote_bench_dir) gym_cmd = render_gym_cmd("eval run", "GYM_CMD", [output_path] + extra_flags + flatten_run_args(run_args)) entrypoint = render_driver_entrypoint( @@ -381,12 +406,21 @@ def build_sbatch_script( prepare_command = "" driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" driver_node_flags = " --nodes=1 --ntasks=1" if is_multi_node else "" - driver_mounts_flag = ( - f" --container-mounts={','.join(shlex.quote(m) for m in config.driver.mounts)}" if config.driver.mounts else "" - ) + # The driver writes everything relative to the job directory -- `output_path` + # above is `artifacts/rollouts.jsonl`. `#SBATCH --chdir` sets the cwd of the + # BATCH script on the host, but inside a Pyxis container the cwd is whatever + # the image declares and the job directory is not visible at all unless it is + # mounted. Without both of these the run completes cleanly, exits 0, and + # writes every artifact into the container's ephemeral overlay, which is + # discarded on exit: no rollouts, no metrics, no preprocessed data, and + # nothing to say so. Logs survive only because srun resolves `--output` on + # the host, which is what makes the loss so easy to miss. + driver_mounts = [*config.driver.mounts, f"{remote_bench_dir}:{remote_bench_dir}"] + driver_mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in driver_mounts)}" driver_command = ( f"{gym_cmd}\n" - f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_node_flags}{driver_mounts_flag} --container-image={shlex.quote(config.driver.container)} " + f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_node_flags}{driver_mounts_flag}" + f" --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) diff --git a/nemo_gym/orchestration/jobs.py b/nemo_gym/orchestration/jobs.py new file mode 100644 index 0000000000..34328944e7 --- /dev/null +++ b/nemo_gym/orchestration/jobs.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The record a submission leaves behind. + +`api.py` is the input schema (what to submit); this is the output schema (what +was submitted). + +This module is executor-agnostic on purpose, and that is a rule rather than a +coincidence: it must describe a submission made by ANY executor -- Slurm today, +k8s or a local runner later. Two consequences, both enforced by +`test_jobs_module_is_executor_agnostic`: + +* Nothing here imports from `executors/`, so a reader -- today EFB's collect, + tomorrow `gym eval status` -- can load a record without pulling Slurm, SSH and + the sbatch templates into its import path. +* No field is typed or named after one executor's vocabulary. Where a docstring + explains a value by example, it says which executor the example comes from. + +Compatibility. A record outlives the Gym that wrote it, so `schema_version` is +read in one direction only: a reader accepts anything at or below its own +version and refuses only what is newer than it understands. Upgrading nemo-gym +therefore never strands records already on disk. + +That is only sound if the schema keeps its side of the bargain, so within one +schema version a field may only be ADDED WITH A DEFAULT, or removed. A field +must never be repurposed or have its meaning changed -- an older record would +then be read as if it meant the new thing. Anything of that kind needs a +`SCHEMA_VERSION` bump and an explicit migration, not a silent reinterpretation. +""" + +import os +import secrets +import sys +from datetime import datetime, timezone +from pathlib import Path + +from pydantic import BaseModel + + +SCHEMA_VERSION = 1 + +# The manifest's name inside the run directory, wherever that directory lives. +# Readers look for exactly this. +MANIFEST_NAME = "gym-job.json" + + +class BenchmarkJob(BaseModel): + """One benchmark's submission. + + `job_id` is None exactly when the executor failed to enqueue this benchmark + (for the Slurm executor, a failed `sbatch`), in which case `error` says why. + The other benchmarks in the same submission are unaffected -- one failure + does not discard the record for the ones that did start. + """ + + benchmark: str + job_dir: str + job_id: str | None = None + error: str | None = None + + +class SubmissionRecord(BaseModel): + """Everything needed to find a submitted run again. + + `executor` names the executor that produced this record, and is the field a + reader branches on before interpreting the executor-shaped parts of the rest + -- `job_id`, for instance, is a Slurm job ID under the Slurm executor and + need not be numeric under another. It is deliberately a plain `str` rather + than an enum of the executors that happen to exist today, so that adding one + does not require a schema version bump on every reader. + + `hostname` None means the submission was made from the machine that runs the + workload manager's client directly, rather than reaching it over SSH -- not + that the host is unknown. Executors with no remote-submission concept at all + leave it None. + """ + + gym_job_id: str + gym_version: str + submitted_at: str + run_dir: str + cluster: str + executor: str + submitted_by: str + benchmarks: list[BenchmarkJob] + hostname: str | None = None + # Whatever the executor needs to find this submission again that the fields + # above cannot express -- a k8s namespace, say. Free-form because the shape + # is the executor's business; a reader keys off `executor` before reading it. + executor_metadata: dict[str, str] = {} + schema_version: int = SCHEMA_VERSION + + @property + def failed(self) -> list[BenchmarkJob]: + return [b for b in self.benchmarks if b.job_id is None] + + def dumps(self) -> str: + """The manifest's on-disk bytes; one spelling, so every store matches.""" + return self.model_dump_json(indent=2) + "\n" + + def write_local_index(self) -> Path | None: + """Record the submission on this machine, or report why not and carry on. + + Best-effort by design: this runs after the jobs are queued, so raising + here would report a failure for work that is really running. The durable + copy is the manifest in the run directory. + """ + path = local_index_dir() / f"{self.gym_job_id}.json" + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(self.dumps(), encoding="utf-8") + except OSError as error: + print(f"Could not write the local job index at {path}: {error}", file=sys.stderr) + return None + return path + + @classmethod + def load(cls, payload: dict) -> "SubmissionRecord": + """Parse a record, refusing only one written by a NEWER Gym. + + Older records stay readable: fields added since take their defaults, which + the module docstring's compatibility rule is there to guarantee. Refusing + them instead would mean an upgrade silently orphaned every job already + submitted. + + A newer record is the one case that cannot be read safely -- it may carry + fields whose meaning this version does not know -- so it fails, and says + which way round the mismatch is. + """ + version = payload.get("schema_version") + if not isinstance(version, int): + raise ValueError( + f"Job record has no usable schema_version (got {version!r}); it was not written by `gym eval submit`." + ) + if version > SCHEMA_VERSION: + raise ValueError( + f"Job record schema_version {version} was written by a newer nemo-gym; this one understands " + f"up to {SCHEMA_VERSION}. Upgrade nemo-gym to read it." + ) + return cls.model_validate(payload) + + +def utc_now() -> datetime: + """The clock the run directory is named from. + + Here rather than in an executor because run-directory identity is this + module's concern, and every executor needs the same answer. A seam: tests + freeze it to prove two submits in the same second still get distinct + directories. + """ + return datetime.now(timezone.utc) + + +def utc_timestamp(now: datetime) -> str: + """ISO 8601, seconds, explicit Z. A run directory is read on a cluster whose + timezone need not match the submitter's.""" + return now.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def new_gym_job_id(now: datetime) -> str: + """The submission's primary key, and the run directory's name. + + The random suffix is what keeps two submits in the same second against the + same `job.output_path` from sharing a directory. Sharing one is not merely + untidy: an executor that stages by mirroring a directory will delete what it + does not recognise (the Slurm executor copies with `rsync --delete`), so the + second submit silently erases the first's staged scripts. + """ + return f"gym-job-{now.strftime('%Y%m%dT%H%M%SZ')}-{secrets.token_hex(3)}" + + +def local_index_dir() -> Path: + """Where this machine remembers its own submissions.""" + base = os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache") + return Path(base) / "nemo-gym" / "jobs" diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index b15f778039..582da9fe59 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -16,6 +16,7 @@ from nemo_gym.decorators import experimental from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.slurm import SlurmExecutor +from nemo_gym.orchestration.jobs import SubmissionRecord _EXECUTORS = { @@ -24,6 +25,6 @@ @experimental -def submit(config: SubmitConfig, *, dry_run: bool = False) -> None: # pragma: no cover +def submit(config: SubmitConfig, *, dry_run: bool = False) -> SubmissionRecord | None: # pragma: no cover compute = next(iter(config.compute.values())) - _EXECUTORS[type(compute)]().run(config, dry_run=dry_run) + return _EXECUTORS[type(compute)]().run(config, dry_run=dry_run) diff --git a/tests/unit_tests/test_cli_eval_submit.py b/tests/unit_tests/test_cli_eval_submit.py index c546ad1173..58a03ec8a3 100644 --- a/tests/unit_tests/test_cli_eval_submit.py +++ b/tests/unit_tests/test_cli_eval_submit.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import argparse +import json +import sys import pytest import yaml @@ -21,7 +23,9 @@ from pytest import MonkeyPatch import nemo_gym.orchestration.submit as submit_module -from nemo_gym.cli.main import _eval_submit +from nemo_gym.cli.main import _eval_submit, main +from nemo_gym.orchestration.api import SlurmComputeConfig +from nemo_gym.orchestration.jobs import BenchmarkJob, SubmissionRecord COMPUTE = {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}} @@ -30,8 +34,8 @@ JOB = {"output_path": "/tmp/gym-jobs"} -def _args(config_path, *, dry_run: bool = False) -> argparse.Namespace: - return argparse.Namespace(config=str(config_path), dry_run=dry_run) +def _args(config_path, *, dry_run: bool = False, json_output: bool = False) -> argparse.Namespace: + return argparse.Namespace(config=str(config_path), dry_run=dry_run, json=json_output) def _capture_submit(monkeypatch: MonkeyPatch) -> dict: @@ -246,3 +250,144 @@ def test_repeated_calls_do_not_leak_global_hydra_state(self, tmp_path, monkeypat _eval_submit(_args(config_path), overrides=[]) assert captured["config"].job.output_path == "/tmp/gym-jobs" + + +def _record(*, failed: bool = False) -> SubmissionRecord: + return SubmissionRecord( + gym_job_id="gym-job-20260909T100203Z-abc123", + gym_version="0.6.0", + submitted_at="2026-09-09T10:02:03Z", + run_dir="/jobs/gym-job-20260909T100203Z-abc123", + cluster="hsg", + executor="slurm", + submitted_by="wprazuch", + hostname="login-01", + benchmarks=[ + BenchmarkJob( + benchmark="gsm8k", + job_dir="/jobs/gym-job-20260909T100203Z-abc123/gsm8k", + job_id=None if failed else "12345", + error="sbatch: error: bad account" if failed else None, + ) + ], + ) + + +def _returning(monkeypatch: MonkeyPatch, record): + monkeypatch.setattr(submit_module, "submit", lambda config, *, dry_run=False: record) + + +def _config_file(tmp_path): + path = tmp_path / "submit.yaml" + path.write_text(yaml.dump({"services": {"svc": SERVICE}, "compute": COMPUTE, "driver": DRIVER, "job": JOB})) + return path + + +class TestEvalSubmitOutput: + def test_human_output_names_each_benchmark_and_job(self, tmp_path, monkeypatch, capsys): + _returning(monkeypatch, _record()) + + _eval_submit(_args(_config_file(tmp_path)), overrides=[]) + + out = capsys.readouterr().out + assert "gsm8k" in out and "12345" in out + assert "/jobs/gym-job-20260909T100203Z-abc123" in out + + def test_json_output_is_the_record_and_nothing_else(self, tmp_path, monkeypatch, capsys): + record = _record() + _returning(monkeypatch, record) + + _eval_submit(_args(_config_file(tmp_path), json_output=True), overrides=[]) + + assert json.loads(capsys.readouterr().out) == json.loads(record.model_dump_json()) + + def test_a_failed_benchmark_exits_non_zero(self, tmp_path, monkeypatch): + _returning(monkeypatch, _record(failed=True)) + + with pytest.raises(SystemExit) as exit_info: + _eval_submit(_args(_config_file(tmp_path)), overrides=[]) + + assert exit_info.value.code == 1 + + def test_a_failed_benchmark_still_emits_json(self, tmp_path, monkeypatch, capsys): + _returning(monkeypatch, _record(failed=True)) + + with pytest.raises(SystemExit): + _eval_submit(_args(_config_file(tmp_path), json_output=True), overrides=[]) + + assert json.loads(capsys.readouterr().out)["benchmarks"][0]["job_id"] is None + + def test_dry_run_prints_nothing_extra_and_does_not_exit(self, tmp_path, monkeypatch, capsys): + _returning(monkeypatch, None) + + _eval_submit(_args(_config_file(tmp_path), dry_run=True, json_output=True), overrides=[]) + + assert capsys.readouterr().out == "" + + +class TestEvalSubmitThroughTheRealCli: + """Drive `gym eval submit` the way a caller does: `main()` with argv, and + nothing between it and the code under test but a fake executor. + + Every other test in this file monkeypatches `submit_module.submit`, which + replaces the function the `@experimental` decorator wraps -- so the + decorator never runs, and anything it writes to stdout is invisible. That + is how a warning printed in front of the JSON shipped: EFB does + `json.loads(result.stdout)` on the whole stream and gets a JSONDecodeError, + so no run is ever recorded. These tests parse the *entire* stdout. + """ + + def _fake_executor(self, monkeypatch, record): + class _FakeExecutor: + def run(self, config, *, dry_run: bool = False): + return record + + monkeypatch.setattr(submit_module, "_EXECUTORS", {SlurmComputeConfig: _FakeExecutor}) + + def _argv(self, monkeypatch, config_path, *extra): + monkeypatch.setattr(sys, "argv", ["gym", "eval", "submit", "--config", str(config_path), *extra]) + + def test_json_stdout_parses_whole(self, tmp_path, monkeypatch, capsys): + record = _record() + self._fake_executor(monkeypatch, record) + self._argv(monkeypatch, _config_file(tmp_path), "--json") + + main() + + captured = capsys.readouterr() + assert json.loads(captured.out) == json.loads(record.model_dump_json()) + + def test_the_experimental_warning_goes_to_stderr(self, tmp_path, monkeypatch, capsys): + self._fake_executor(monkeypatch, _record()) + self._argv(monkeypatch, _config_file(tmp_path), "--json") + + main() + + captured = capsys.readouterr() + assert "experimental" in captured.err + assert "experimental" not in captured.out + + def test_json_stdout_parses_whole_when_a_benchmark_failed(self, tmp_path, monkeypatch, capsys): + # The partial-failure path still has to hand EFB a parseable record: + # that is what keeps the siblings that did queue from being stranded. + self._fake_executor(monkeypatch, _record(failed=True)) + self._argv(monkeypatch, _config_file(tmp_path), "--json") + + with pytest.raises(SystemExit) as exit_info: + main() + + assert exit_info.value.code == 1 + assert json.loads(capsys.readouterr().out)["benchmarks"][0]["job_id"] is None + + def test_human_output_survives_an_error_containing_markup(self, tmp_path, monkeypatch, capsys): + # An sbatch message with square brackets is markup to rich: without + # escaping it is either eaten or raises MarkupError mid-report. + record = _record(failed=True) + record.benchmarks[0].error = "sbatch: error: Invalid account [dev] for user" + self._fake_executor(monkeypatch, record) + self._argv(monkeypatch, _config_file(tmp_path)) + + with pytest.raises(SystemExit): + main() + + assert "[dev]" in capsys.readouterr().out diff --git a/tests/unit_tests/test_connection.py b/tests/unit_tests/test_connection.py new file mode 100644 index 0000000000..c6f0d9eba0 --- /dev/null +++ b/tests/unit_tests/test_connection.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import subprocess +from pathlib import Path + +from pytest import MonkeyPatch + +from nemo_gym.orchestration.executors import connection as connection_module +from nemo_gym.orchestration.executors.connection import LocalConnection, SSHConnection +from nemo_gym.orchestration.jobs import BenchmarkJob, SubmissionRecord + + +def test_local_connection_runs_a_compound_bash_command(tmp_path): + # Task 2's sbatch command is bash, not a simple argv. shlex.split would + # mangle it, so local submits depend on this going through a shell. + marker = tmp_path / "ran" + LocalConnection().run([f'out=hello; rc=$?; echo "$out:$rc" > {marker}']) + + assert marker.read_text().strip() == "hello:0" + + +def test_local_connection_runs_every_command_in_one_shell(): + # Both connections must agree: a failing command does not abort the rest. + output = LocalConnection().run(["echo first", "false", "echo third"]) + + assert "first" in output and "third" in output + + +def test_local_connection_writes_the_file(tmp_path): + target = tmp_path / "nested" / "gym-job.json" + + LocalConnection().write_text(target, '{"a": 1}\n') + + assert target.read_text() == '{"a": 1}\n' + + +def _script_for(monkeypatch: MonkeyPatch, remote: Path, content: str) -> str: + """The bash `SSHConnection.write_text` would send, without opening a socket.""" + captured = {} + + def fake_checked(cmd, *, input=None, context=""): + captured["cmd"] = cmd + captured["input"] = input + return "" + + monkeypatch.setattr(connection_module, "_checked", fake_checked) + SSHConnection("login-01").write_text(remote, content) + + assert captured["cmd"][-2:] == ["bash", "-s"] + return captured["input"] + + +def test_ssh_connection_writes_the_same_bytes_the_local_index_holds(monkeypatch: MonkeyPatch, tmp_path): + # Run the generated script through a real bash and compare the file it + # produces, byte for byte, against `dumps`. Asserting on substrings of the + # command instead is what let a stray trailing newline ship: the remote + # manifest and the local index have to be the same bytes, and only the + # file the shell actually writes can show that. + # + # The destination has a space in it, so a build that dropped `shlex.quote` + # fails here rather than passing on a path that never needed quoting. In + # production the run directory is `/gym-job-...`, and + # `job.output_path` comes from a config file. + record = SubmissionRecord( + gym_job_id="gym-job-20260909T100203Z-abc123", + gym_version="0.6.0", + submitted_at="2026-09-09T10:02:03Z", + run_dir="/jobs/gym-job-20260909T100203Z-abc123", + cluster="hsg", + executor="slurm", + submitted_by="wprazuch", + hostname="login-01", + benchmarks=[ + BenchmarkJob( + benchmark="gsm8k", + job_dir="/jobs/gym-job-20260909T100203Z-abc123/gsm8k", + job_id="12345", + ) + ], + ) + target = tmp_path / "run dir" / "gym-job.json" + target.parent.mkdir() + + script = _script_for(monkeypatch, target, record.dumps()) + subprocess.run(["bash", "-s"], input=script, text=True, check=True) + + assert target.read_bytes() == record.dumps().encode() + # And the store this is supposed to match, written the other way. + local = tmp_path / "local.json" + LocalConnection().write_text(local, record.dumps()) + assert target.read_bytes() == local.read_bytes() + + +def test_ssh_connection_quotes_the_heredoc_delimiter(monkeypatch: MonkeyPatch, tmp_path): + # A quoted delimiter stops the shell expanding anything inside the payload: + # a manifest carrying `$HOME` or a backtick must land as written. + target = tmp_path / "gym-job.json" + content = '{"note": "$HOME and `id` and ${PATH}"}\n' + + script = _script_for(monkeypatch, target, content) + assert "<<'GYM_EOF'" in script + subprocess.run(["bash", "-s"], input=script, text=True, check=True) + + assert target.read_text() == content diff --git a/tests/unit_tests/test_global_config.py b/tests/unit_tests/test_global_config.py index 6ebd390ded..ad14b571d8 100644 --- a/tests/unit_tests/test_global_config.py +++ b/tests/unit_tests/test_global_config.py @@ -2773,3 +2773,35 @@ def test_rejects_a_field_the_incoming_agent_does_not_declare(self, monkeypatch: with raises(ConfigKeyError): self._parse_with_cli(self._config(), self._cli_override(renamed, no_such_field=1), monkeypatch) + + +def test_partial_head_server_inherits_the_resolved_host(monkeypatch): + """Pinning only the port must not suppress the host default. + + A caller that constrains the head server to an allocated port range cannot + also supply the host: it is the address of whichever node the job lands on, + which only `use_absolute_ip` resolves, and only here. + """ + from omegaconf import OmegaConf + + from nemo_gym.global_config import ( + HEAD_SERVER_KEY_NAME, + USE_ABSOLUTE_IP, + GlobalConfigDictParser, + GlobalConfigDictParserConfig, + ) + + monkeypatch.setattr("nemo_gym.global_config.gethostname", lambda: "node-17") + monkeypatch.setattr("nemo_gym.global_config.gethostbyname", lambda _h: "10.1.2.3") + + initial = OmegaConf.create( + { + **GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT, + USE_ABSOLUTE_IP: True, + HEAD_SERVER_KEY_NAME: {"port": 63000}, + } + ) + parsed = GlobalConfigDictParser().parse_no_environment(initial_global_config_dict=initial) + + assert parsed[HEAD_SERVER_KEY_NAME]["port"] == 63000, "explicit port must survive" + assert parsed[HEAD_SERVER_KEY_NAME]["host"] == "10.1.2.3", "host must be filled in" diff --git a/tests/unit_tests/test_orchestration_jobs.py b/tests/unit_tests/test_orchestration_jobs.py new file mode 100644 index 0000000000..a0dcba0491 --- /dev/null +++ b/tests/unit_tests/test_orchestration_jobs.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from pytest import MonkeyPatch + +from nemo_gym.orchestration.executors.base import BaseExecutor +from nemo_gym.orchestration.jobs import ( + SCHEMA_VERSION, + BenchmarkJob, + SubmissionRecord, + local_index_dir, + new_gym_job_id, +) + + +NOW = datetime(2026, 9, 9, 10, 2, 3, tzinfo=timezone.utc) + + +def _record(**overrides) -> SubmissionRecord: + defaults = dict( + gym_job_id="gym-job-20260909T100203Z-abc123", + gym_version="0.6.0", + submitted_at="2026-09-09T10:02:03Z", + run_dir="/jobs/gym-job-20260909T100203Z-abc123", + cluster="cluster", + executor="slurm", + hostname="login-01", + submitted_by="wprazuch", + benchmarks=[ + BenchmarkJob( + benchmark="gsm8k", + job_id="12345", + job_dir="/jobs/gym-job-20260909T100203Z-abc123/gsm8k", + ) + ], + ) + return SubmissionRecord(**{**defaults, **overrides}) + + +def test_record_round_trips_through_json(): + record = _record() + restored = SubmissionRecord.load(json.loads(record.model_dump_json())) + assert restored == record + + +def test_record_defaults_to_current_schema_version(): + assert _record().schema_version == SCHEMA_VERSION + + +def test_load_record_accepts_an_older_record_after_an_upgrade(): + """The upgrade path: a record written by an older Gym must stay readable. + + Simulated the way it really happens -- the payload predates a field, so the + key is simply absent and the model's default fills it. Refusing this is what + would strand every job already on disk the moment nemo-gym is upgraded. + """ + payload = json.loads(_record().dumps()) + payload["schema_version"] = SCHEMA_VERSION - 1 + payload.pop("hostname") + + record = SubmissionRecord.load(payload) + + assert record.schema_version == SCHEMA_VERSION - 1 + assert record.hostname is None + + +def test_load_record_rejects_a_payload_with_no_version(): + payload = json.loads(_record().dumps()) + del payload["schema_version"] + with pytest.raises(ValueError, match="no usable schema_version"): + SubmissionRecord.load(payload) + + +def test_load_record_refuses_a_newer_schema_version(): + payload = json.loads(_record().model_dump_json()) + payload["schema_version"] = SCHEMA_VERSION + 1 + with pytest.raises(ValueError, match="written by a newer nemo-gym"): + SubmissionRecord.load(payload) + + +def test_gym_job_id_embeds_the_utc_timestamp(): + assert new_gym_job_id(NOW).startswith("gym-job-20260909T100203Z-") + + +def test_gym_job_ids_minted_in_the_same_second_differ(): + assert new_gym_job_id(NOW) != new_gym_job_id(NOW) + + +def test_local_index_dir_honours_xdg_cache_home(tmp_path, monkeypatch: MonkeyPatch): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + assert local_index_dir() == tmp_path / "nemo-gym" / "jobs" + + +def test_local_index_dir_falls_back_to_home_cache(tmp_path, monkeypatch: MonkeyPatch): + monkeypatch.delenv("XDG_CACHE_HOME", raising=False) + monkeypatch.setattr(Path, "home", lambda: tmp_path) + assert local_index_dir() == tmp_path / ".cache" / "nemo-gym" / "jobs" + + +def test_write_local_index_writes_the_record(tmp_path, monkeypatch: MonkeyPatch): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + record = _record() + + path = record.write_local_index() + + assert path == tmp_path / "nemo-gym" / "jobs" / f"{record.gym_job_id}.json" + assert SubmissionRecord.load(json.loads(path.read_text())) == record + + +def test_write_local_index_returns_none_when_it_cannot_write(tmp_path, monkeypatch: MonkeyPatch, capsys): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + + def boom(*_args, **_kwargs): + raise OSError("read-only file system") + + monkeypatch.setattr(Path, "mkdir", boom) + + assert _record().write_local_index() is None + assert "read-only file system" in capsys.readouterr().err + + +def test_write_local_index_writes_exactly_what_dumps_produces(tmp_path, monkeypatch: MonkeyPatch): + # The local index and the remote manifest must be byte-identical; both go + # through dumps(), and this is what holds that true. + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + record = _record() + + path = record.write_local_index() + + assert path.read_text() == record.dumps() + + +def test_executor_metadata_carries_whatever_an_executor_needs(): + """The record cannot grow a field per executor, so anything one needs that + the shared fields cannot express goes here -- a k8s namespace, say. It round + trips like any other field and defaults to empty for executors with nothing + to add.""" + assert _record().executor_metadata == {} + + record = _record() + record.executor_metadata = {"namespace": "frontier-eval", "context": "prod"} + + assert SubmissionRecord.load(json.loads(record.dumps())).executor_metadata == { + "namespace": "frontier-eval", + "context": "prod", + } + + +def test_persist_writes_the_local_index_even_when_the_manifest_fails(tmp_path, monkeypatch: MonkeyPatch): + """The ordering is the whole point of persist() living in the base class. + + The index is written before the manifest precisely so that a failed manifest + write still leaves a parseable record behind -- the by-hand recovery the + error message asks for has to have something to recover from. + """ + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + + class _Executor(BaseExecutor): + def run(self, config, *, dry_run: bool = False): # pragma: no cover - unused + raise NotImplementedError + + def _explode(path, text): + raise OSError("remote is read-only") + + record = _record() + with pytest.raises(RuntimeError, match="Record these by hand"): + _Executor().persist(record, _explode) + + index = tmp_path / "nemo-gym" / "jobs" / f"{record.gym_job_id}.json" + assert SubmissionRecord.load(json.loads(index.read_text())) == record + + +def test_persist_names_the_queued_jobs_when_the_manifest_fails(tmp_path, monkeypatch: MonkeyPatch): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + + class _Executor(BaseExecutor): + def run(self, config, *, dry_run: bool = False): # pragma: no cover - unused + raise NotImplementedError + + record = _record() + with pytest.raises(RuntimeError, match=r"Already queued: .*gsm8k=12345"): + _Executor().persist(record, lambda path, text: (_ for _ in ()).throw(OSError("nope"))) + + +def test_jobs_module_is_executor_agnostic(): + """`jobs.py` must not drag any executor into a reader's import path. + + A reader that only wants to parse a record -- EFB's collect, `gym eval + status` -- should not end up importing Slurm, SSH and the sbatch templates. + Checked in a SUBPROCESS on purpose: this test module imports executors + itself, so asserting against the already-loaded sys.modules here would pass + no matter what jobs.py does. + """ + probe = ( + "import sys;" + "import nemo_gym.orchestration.jobs;" + "leaked = sorted(m for m in sys.modules if 'orchestration.executors' in m);" + "print(','.join(leaked))" + ) + result = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, check=True) + + assert result.stdout.strip() == "", f"jobs.py pulled in executor modules: {result.stdout.strip()}" + + +def test_submission_record_executor_is_not_closed_over_todays_executors(): + """A future k8s or local executor must be able to write a record without + editing this schema, so `executor` is an open str rather than an enum of the + executors that exist today.""" + record = SubmissionRecord( + gym_job_id="gym-job-20260101T000000Z-abc123", + gym_version="0.0.0", + submitted_at="2026-01-01T00:00:00Z", + run_dir="/runs/gym-job-20260101T000000Z-abc123", + cluster="some-cluster", + executor="kubernetes", + submitted_by="someone", + benchmarks=[], + ) + assert record.executor == "kubernetes" + assert SubmissionRecord.load(json.loads(record.dumps())) == record diff --git a/tests/unit_tests/test_slurm_executor.py b/tests/unit_tests/test_slurm_executor.py new file mode 100644 index 0000000000..649c69073e --- /dev/null +++ b/tests/unit_tests/test_slurm_executor.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +from nemo_gym.orchestration.api import SubmitConfig +from nemo_gym.orchestration.executors import slurm as slurm_module +from nemo_gym.orchestration.executors.connection import LocalConnection +from nemo_gym.orchestration.executors.slurm import ( + SlurmExecutor, + _parse_sbatch_results, + _sbatch_command, + _validate_mounts, +) +from nemo_gym.orchestration.jobs import MANIFEST_NAME, SubmissionRecord + + +def test_sbatch_command_captures_the_status_of_sbatch_not_the_pipeline(): + command = _sbatch_command("gsm8k", Path("/jobs/run/gsm8k/job.sh")) + # rc must be captured immediately after sbatch; taking $? after a pipe + # would report tr's status and call every failure a success. + assert "rc=$?" in command + assert command.index("rc=$?") < command.index("tr ") + assert "sbatch --parsable /jobs/run/gsm8k/job.sh" in command + assert "__GYM_JOB:gsm8k:$rc:$out" in command + + +@pytest.mark.parametrize("name", ["has space", "has:colon", "has$dollar", "", "a;b"]) +def test_sbatch_command_rejects_a_benchmark_name_that_would_break_the_marker(name): + with pytest.raises(ValueError, match="benchmark name"): + _sbatch_command(name, Path("/jobs/run/job.sh")) + + +@pytest.mark.parametrize("name", ["gsm8k", "gpqa-no-tools", "tau2.airline", "aime_24"]) +def test_sbatch_command_accepts_ordinary_benchmark_names(name): + assert f"__GYM_JOB:{name}:$rc:$out" in _sbatch_command(name, Path("/jobs/run/job.sh")) + + +def test_parse_sbatch_results_reads_a_successful_submission(): + assert _parse_sbatch_results("__GYM_JOB:gsm8k:0:12345 ") == {"gsm8k": ("12345", None)} + + +def test_parse_sbatch_results_drops_the_federation_cluster_suffix(): + assert _parse_sbatch_results("__GYM_JOB:gsm8k:0:12345;hsg ") == {"gsm8k": ("12345", None)} + + +def test_parse_sbatch_results_ignores_unrelated_output(): + output = "Loading modules\n__GYM_JOB:gsm8k:0:12345 \nsome trailing chatter" + assert _parse_sbatch_results(output) == {"gsm8k": ("12345", None)} + + +def test_parse_sbatch_results_records_a_failure_with_its_message(): + output = "__GYM_JOB:gsm8k:1:sbatch: error: Invalid account 'nope' " + assert _parse_sbatch_results(output) == {"gsm8k": (None, "sbatch: error: Invalid account 'nope'")} + + +def test_parse_sbatch_results_falls_back_to_the_exit_code_when_sbatch_said_nothing(): + assert _parse_sbatch_results("__GYM_JOB:gsm8k:1: ") == {"gsm8k": (None, "sbatch exited 1")} + + +def test_parse_sbatch_results_treats_a_silent_success_as_a_failure(): + # Exit 0 with no output at all is not a job id; recording it as a "success" + # with job_id="" would slip past `SubmissionRecord.failed` (which only + # checks `job_id is None`). + job_id, error = _parse_sbatch_results("__GYM_JOB:gsm8k:0: ")["gsm8k"] + assert job_id is None + assert error is not None + + +def test_parse_sbatch_results_treats_a_non_numeric_success_payload_as_a_failure(): + # A warning on a successful sbatch (job_submit plugin notices, QOS/ntasks + # adjustments) merges onto the same line via 2>&1. If it doesn't end in + # something that looks like a job id, there is no id to trust. + output = "__GYM_JOB:gsm8k:0:sbatch: Warning: blah " + assert _parse_sbatch_results(output) == {"gsm8k": (None, "sbatch: Warning: blah")} + + +def test_parse_sbatch_results_takes_the_trailing_id_off_a_warning_line(): + # The actual C1 regression: a successful sbatch that also warns must not + # let the warning text end up in job_id. + output = "__GYM_JOB:gsm8k:0:sbatch: Warning: can't honor --ntasks-per-node 12345 " + assert _parse_sbatch_results(output) == {"gsm8k": ("12345", None)} + + +def test_a_failure_mid_list_does_not_shift_the_benchmarks_after_it(): + # The regression test for the positional-zip bug: bench_b fails, and bench_c + # must still get its own job id rather than inheriting the next one along. + output = "\n".join( + [ + "__GYM_JOB:bench_a:0:111 ", + "__GYM_JOB:bench_b:1:sbatch: error: Invalid account ", + "__GYM_JOB:bench_c:0:333 ", + ] + ) + + results = _parse_sbatch_results(output) + + assert results["bench_a"] == ("111", None) + assert results["bench_b"][0] is None + assert results["bench_c"] == ("333", None) + + +def _submit_config(tmp_path, benchmarks): + return SubmitConfig.model_validate( + { + "services": {}, + "compute": {"hsg": {"type": "slurm", "account": "my-account", "hostname": None}}, + "driver": {"container": "gym:latest", "benchmarks": {name: {} for name in benchmarks}}, + "job": {"output_path": str(tmp_path / "jobs")}, + } + ) + + +class _FakeConnection(LocalConnection): + """A local connection whose `run` answers as a scheduler would.""" + + def __init__(self, replies): + self._replies = replies + self.commands = [] + + def run(self, commands): + self.commands.append(commands) + return self._replies.pop(0) + + +def _install(monkeypatch, conn): + monkeypatch.setattr(slurm_module, "get_connection", lambda hostname: conn) + monkeypatch.setattr(slurm_module, "_validate_mounts", lambda config, connection: None) + + +def test_run_returns_a_record_naming_every_benchmark(tmp_path, monkeypatch): + conn = _FakeConnection(["__GYM_JOB:bench_a:0:111 \n__GYM_JOB:bench_b:0:222 "]) + _install(monkeypatch, conn) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a", "bench_b"])) + + assert record is not None + assert [(b.benchmark, b.job_id) for b in record.benchmarks] == [("bench_a", "111"), ("bench_b", "222")] + assert record.cluster == "hsg" + assert record.executor == "slurm" + assert record.hostname is None + assert record.run_dir.endswith(record.gym_job_id) + assert record.benchmarks[0].job_dir == f"{record.run_dir}/bench_a" + + +def test_run_writes_the_manifest_into_the_run_dir(tmp_path, monkeypatch): + conn = _FakeConnection(["__GYM_JOB:bench_a:0:111 "]) + _install(monkeypatch, conn) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"])) + + manifest = Path(record.run_dir) / MANIFEST_NAME + assert SubmissionRecord.load(json.loads(manifest.read_text())) == record + + +def test_run_writes_the_local_index(tmp_path, monkeypatch): + conn = _FakeConnection(["__GYM_JOB:bench_a:0:111 "]) + _install(monkeypatch, conn) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"])) + + index = tmp_path / "cache" / "nemo-gym" / "jobs" / f"{record.gym_job_id}.json" + assert SubmissionRecord.load(json.loads(index.read_text())) == record + + +def test_run_records_a_failed_benchmark_without_disturbing_the_others(tmp_path, monkeypatch): + conn = _FakeConnection( + ["__GYM_JOB:bench_a:0:111 \n__GYM_JOB:bench_b:1:sbatch: error: bad account \n__GYM_JOB:bench_c:0:333 "] + ) + _install(monkeypatch, conn) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a", "bench_b", "bench_c"])) + + by_name = {b.benchmark: b for b in record.benchmarks} + assert by_name["bench_a"].job_id == "111" + assert by_name["bench_b"].job_id is None + assert "bad account" in by_name["bench_b"].error + assert by_name["bench_c"].job_id == "333" + assert [b.benchmark for b in record.failed] == ["bench_b"] + + +def test_run_records_a_benchmark_the_scheduler_never_answered_for(tmp_path, monkeypatch): + conn = _FakeConnection(["__GYM_JOB:bench_a:0:111 "]) + _install(monkeypatch, conn) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a", "bench_b"])) + + by_name = {b.benchmark: b for b in record.benchmarks} + assert by_name["bench_b"].job_id is None + assert "no result" in by_name["bench_b"].error + + +def test_two_runs_in_the_same_second_get_different_run_dirs(tmp_path, monkeypatch): + frozen = datetime(2026, 9, 9, 10, 2, 3, tzinfo=timezone.utc) + monkeypatch.setattr(slurm_module, "utc_now", lambda: frozen) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + + run_dirs = [] + for _ in range(2): + conn = _FakeConnection(["__GYM_JOB:bench_a:0:111 "]) + _install(monkeypatch, conn) + run_dirs.append(SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"])).run_dir) + + assert run_dirs[0] != run_dirs[1] + + +def test_a_failed_manifest_write_fails_the_submit_and_names_queued_jobs(tmp_path, monkeypatch): + class _NoWrite(_FakeConnection): + def write_text(self, remote, content): + raise RuntimeError("permission denied") + + conn = _NoWrite(["__GYM_JOB:bench_a:0:111 \n__GYM_JOB:bench_b:0:222 "]) + _install(monkeypatch, conn) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + + with pytest.raises(RuntimeError) as error: + SlurmExecutor().run(_submit_config(tmp_path, ["bench_a", "bench_b"])) + + # The jobs are already queued; an error that does not say so strands them. + message = str(error.value) + assert "111" in message and "222" in message + assert "permission denied" in message + + +def test_submitting_locally_runs_the_real_sbatch_command(tmp_path, monkeypatch): + # The one test that exercises the generated bash end to end, over a real + # LocalConnection. A command string that only a mocked `run` accepts would + # pass every other test in this file and still break every local submit. + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + sbatch = fake_bin / "sbatch" + sbatch.write_text("#!/bin/bash\necho 4242\n") + sbatch.chmod(0o755) + monkeypatch.setenv("PATH", f"{fake_bin}:{os.environ['PATH']}") + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + monkeypatch.setattr(slurm_module, "_validate_mounts", lambda config, connection: None) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"])) + + assert record.benchmarks[0].job_id == "4242" + assert (Path(record.run_dir) / MANIFEST_NAME).exists() + + +def test_submitting_locally_records_a_failing_sbatch(tmp_path, monkeypatch): + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + sbatch = fake_bin / "sbatch" + sbatch.write_text("#!/bin/bash\necho 'sbatch: error: Invalid account' >&2\nexit 1\n") + sbatch.chmod(0o755) + monkeypatch.setenv("PATH", f"{fake_bin}:{os.environ['PATH']}") + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + monkeypatch.setattr(slurm_module, "_validate_mounts", lambda config, connection: None) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"])) + + assert record.benchmarks[0].job_id is None + assert "Invalid account" in record.benchmarks[0].error + + +def test_submitting_locally_does_not_let_a_warning_become_the_job_id(tmp_path, monkeypatch): + # The end-to-end regression test for C1: a real bash pipeline, not a mocked + # `conn.run`, proves the generated command's own `2>&1` merge can't leak a + # warning on a successful sbatch into job_id. + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + sbatch = fake_bin / "sbatch" + sbatch.write_text('#!/bin/bash\necho "sbatch: Warning: can\'t honor --ntasks-per-node" >&2\necho 12345\n') + sbatch.chmod(0o755) + monkeypatch.setenv("PATH", f"{fake_bin}:{os.environ['PATH']}") + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + monkeypatch.setattr(slurm_module, "_validate_mounts", lambda config, connection: None) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"])) + + assert record.benchmarks[0].job_id == "12345" + assert record.benchmarks[0].error is None + + +def test_submitting_locally_keeps_a_multiline_error_on_one_marker_line(tmp_path, monkeypatch): + # Real-bash coverage for the `tr` flattening: a two-line sbatch error must + # still arrive whole in `error`, not truncated to whichever line survives. + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + sbatch = fake_bin / "sbatch" + sbatch.write_text( + "#!/bin/bash\n" + "echo 'sbatch: error: Invalid account' >&2\n" + "echo 'sbatch: error: try again with a valid account' >&2\n" + "exit 1\n" + ) + sbatch.chmod(0o755) + monkeypatch.setenv("PATH", f"{fake_bin}:{os.environ['PATH']}") + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + monkeypatch.setattr(slurm_module, "_validate_mounts", lambda config, connection: None) + + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"])) + + assert record.benchmarks[0].job_id is None + assert "Invalid account" in record.benchmarks[0].error + assert "try again with a valid account" in record.benchmarks[0].error + + +def test_dry_run_rejects_a_benchmark_name_that_would_break_the_marker(tmp_path): + # Bad names must fail before staging/copying even happens, and a dry run + # must not silently print a script listing for a benchmark that could + # never actually be submitted. + with pytest.raises(ValueError, match="benchmark name"): + SlurmExecutor().run(_submit_config(tmp_path, ["bad name"]), dry_run=True) + + +def test_dry_run_returns_no_record(tmp_path, monkeypatch, capsys): + record = SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"]), dry_run=True) + + assert record is None + assert "[dry-run]" in capsys.readouterr().out + + +def _config_with_driver_mounts(tmp_path, mounts): + return SubmitConfig.model_validate( + { + "services": {}, + "compute": {"hsg": {"type": "slurm", "account": "my-account", "hostname": None}}, + "driver": {"container": "gym:latest", "benchmarks": {"bench_a": {}}, "mounts": mounts}, + "job": {"output_path": str(tmp_path / "jobs")}, + } + ) + + +class TestValidateMounts: + """Both connections pipe commands to bash, so mount validation runs one + shell program either way. These go through a real `LocalConnection`: with + the old `isinstance` branch the local path never touched the shell, so the + check that actually runs in production was untested.""" + + def test_a_present_mount_src_passes(self, tmp_path): + src = tmp_path / "data" + src.mkdir() + + _validate_mounts(_config_with_driver_mounts(tmp_path, [f"{src}:/data"]), LocalConnection()) + + def test_a_missing_mount_src_is_named(self, tmp_path): + src = tmp_path / "absent" + + with pytest.raises(ValueError) as error: + _validate_mounts(_config_with_driver_mounts(tmp_path, [f"{src}:/data"]), LocalConnection()) + + assert str(src) in str(error.value) + assert "driver" in str(error.value) + + def test_no_mounts_asks_the_connection_nothing(self, tmp_path): + class _Explodes(LocalConnection): + def run(self, commands): + raise AssertionError("no mounts, so there is nothing to check") + + _validate_mounts(_config_with_driver_mounts(tmp_path, []), _Explodes()) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 81673d2cc2..70bba8ea2b 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -30,6 +30,7 @@ _RAY_SERVE_GATEWAY_SOURCE_PATH, _build_service_command, _build_vllm_command, + _build_vllm_multi_instance_multi_node_command, _build_vllm_ray_command, _build_vllm_ray_serve_command, _node_totals, @@ -534,6 +535,36 @@ def test_render_gym_cmd_prepare(): # --------------------------------------------------------------------------- +def test_driver_policy_model_type_defaults_to_openai_model(submit_config, bench_dir): + submit_config.driver.policy_model = "vllm_model" + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "--model-type openai_model" in script + + +def test_driver_policy_model_type_is_configurable(submit_config, bench_dir): + # Every certified NEL run of these benchmarks serves the policy as + # vllm_model, and lmarena_v3 ships its own vllm_model policy that a second + # composed server would collide with. + submit_config.driver.policy_model = "vllm_model" + submit_config.driver.policy_model_type = "vllm_model" + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "--model-type vllm_model" in script + assert "--model-type openai_model" not in script + + +def test_driver_policy_model_type_empty_composes_nothing(submit_config, bench_dir): + submit_config.driver.policy_model = "vllm_model" + submit_config.driver.policy_model_type = "" + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "--model-type" not in script + + def test_render_driver_entrypoint_no_install_no_prepare(): out = render_driver_entrypoint(None, None, None) assert out == '"${GYM_CMD[@]}"' @@ -542,7 +573,9 @@ def test_render_driver_entrypoint_no_install_no_prepare(): def test_render_driver_entrypoint_with_gym_install(): out = render_driver_entrypoint("https://github.com/NVIDIA-NeMo/gym", "main", None) assert "git clone" in out - assert "git checkout main" in out + # `git -C "$GYM_SRC/gym" checkout`, not `git checkout`: the clone is + # out-of-tree. See test_gym_install_does_not_clone_into_the_job_directory. + assert "checkout main" in out assert "uv venv --seed .venv" in out assert "source .venv/bin/activate" in out assert "uv pip install -e ." in out @@ -568,11 +601,81 @@ def test_render_driver_entrypoint_with_prepare(): def test_render_driver_entrypoint_install_and_prepare(): out = render_driver_entrypoint("https://github.com/NVIDIA-NeMo/gym", "v1.0", "gym eval prepare") assert "git clone" in out - assert "git checkout v1.0" in out + assert "checkout v1.0" in out assert "gym eval prepare" in out assert 'exec "$@"' in out +def test_worker_command_drops_api_server_count(): + """vLLM exits on `--api-server-count` in headless mode before loading anything: + "no API servers are started in headless mode". The flag is valid on the head + node and arrives via the service's own extra_args, so only the worker branch + is stripped. A real mmlu-prox run lost all five workers to this. + """ + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="/checkpoint", + tensor_parallel_size=4, + number_of_instances=2, + extra_args="--api-server-count 1 --enable-prefix-caching", + ) + block = _build_vllm_multi_instance_multi_node_command(service, total_nodes=2) + head, worker = block.split("else") + + assert "--api-server-count 1" in head + assert "--headless" in worker + assert "--api-server-count" not in worker + # Stripping must not take the neighbouring flag with it. + assert "--enable-prefix-caching" in worker + + +def test_multi_instance_multi_node_command_survives_the_shell(): + """vLLM's JSON flags are single-quoted, and the DP branches are embedded in a + single-quoted `bash -lc '...'`. Unescaped they end the block early and the + whole invocation word-splits -- a real mmlu-prox submission died of this with + "/usr/bin/env: Argument list too long". Run the rendered block through bash + and check the JSON arrives as one argument. + """ + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="/checkpoint", + tensor_parallel_size=4, + number_of_instances=2, + extra_args='--hf-overrides \'{"architectures":["Custom"],"norm_mean":[0.5,0.5]}\'', + ) + block = _build_vllm_multi_instance_multi_node_command(service, total_nodes=2) + + # Replace `vllm serve` with a printf that dumps one argument per line, so the + # test observes what the shell actually passed rather than the rendered text. + # Double quotes, because this substitution happens AFTER rendering and so is + # not itself escaped -- single quotes here would break the block the test is + # checking. + script = "SLURM_NODEID=0 HEAD_NODE_IP=1.2.3.4 " + block.replace("vllm serve", 'printf "%s\\n"', 2) + argv = subprocess.run(["bash", "-c", script], capture_output=True, text=True, check=True).stdout.splitlines() + + assert '{"architectures":["Custom"],"norm_mean":[0.5,0.5]}' in argv + + +def test_render_driver_entrypoint_prepare_arg_with_spaces_survives_the_shell(): + """A prepare argument containing spaces must reach Hydra as ONE word. + + The entrypoint body is embedded in a single-quoted `bash -c '...'`, so an + inner single quote ends the outer string instead of nesting. Without + escaping, gdpval's real prepare argument word-split and Hydra failed with + "no viable alternative at input '[{num_tasks:'". Asserting on the rendered + string would not catch that -- only running it through a shell does. + """ + arg = "+multistage.stages=[{num_tasks: 45, waivable: [timeout, transient]}]" + out = render_driver_entrypoint(None, None, f"printf '%s\\n' {shlex.quote(arg)}") + + script = out.replace('exec "$@"', ":").replace('"${GYM_CMD[@]}"', "''") + printed = subprocess.run(["bash", "-c", script], capture_output=True, text=True, check=True).stdout.splitlines() + + assert printed == [arg] + + def test_render_driver_entrypoint_no_install_no_prepare_has_no_set_e(): # The trivial path isn't wrapped in bash -c at all, so there's no # preamble for a failure to silently fall through in the first place. @@ -702,7 +805,10 @@ def test_build_sbatch_script_output_jsonl_fpath(submit_config, bench_dir): benchmark = submit_config.driver.benchmarks["gsm8k"] compute = next(iter(submit_config.compute.values())) script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) - assert "+output_jsonl_fpath=artifacts/rollouts.jsonl" in script + # Absolute: a relative output path lands wherever the container's cwd + # happens to be, and cwd cannot be moved without breaking a benchmark's + # cwd-relative prepare_script. + assert f"+output_jsonl_fpath={bench_dir}/artifacts/rollouts.jsonl" in script def test_build_sbatch_script_policy_model_flags(submit_config_with_policy, bench_dir): @@ -831,8 +937,8 @@ def test_build_sbatch_script_service_env_before_driver_env(bench_dir): script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) svc_env_idx = script.index("SVC_KEY=svc_val") drv_env_idx = script.index("DRV_KEY=drv_val") - svc_srun_idx = script.index("srun --overlap --no-container-mount-home --container-image=vllm:latest") - drv_srun_idx = script.index("srun --overlap --no-container-mount-home --container-image=python:3.12") + svc_srun_idx = script.index("--container-image=vllm:latest") + drv_srun_idx = script.index("--container-image=python:3.12") # Service env prefix appears before service srun; driver env prefix appears before driver srun. assert svc_env_idx < svc_srun_idx assert drv_env_idx < drv_srun_idx @@ -1014,11 +1120,18 @@ def test_build_sbatch_script_driver_mounts(bench_dir): assert "--container-mounts=/lustre/checkpoints:/ckpts" in driver_srun_line -def test_build_sbatch_script_no_mounts_by_default(submit_config, bench_dir): +def test_build_sbatch_script_no_service_mounts_by_default(submit_config, bench_dir): + """Services get no mounts unless configured. The driver is the exception and + always mounts the job directory, because that is where its artifacts go -- + see test_driver_can_write_its_artifacts_into_the_job_directory.""" benchmark = submit_config.driver.benchmarks["gsm8k"] compute = next(iter(submit_config.compute.values())) script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) - assert "--container-mounts" not in script + + service_lines = [line for line in script.splitlines() if "srun" in line and "--output=logs/driver.log" not in line] + assert service_lines, "expected at least one service srun line" + for line in service_lines: + assert "--container-mounts" not in line # --------------------------------------------------------------------------- @@ -1365,3 +1478,78 @@ def submit_config_with_policy(): "job": {"output_path": "/remote/jobs"}, } ) + + +def test_driver_can_write_its_artifacts_into_the_job_directory(): + """A run that cannot reach the job directory completes cleanly and produces + nothing: `output_jsonl_fpath` is relative, `#SBATCH --chdir` only sets the + host-side cwd of the batch script, and a Pyxis container starts in whatever + directory its image declares. Both the mount and the workdir are required. + """ + config = SubmitConfig.model_validate( + { + "services": {}, + "compute": {"hsg": {"type": "slurm", "account": "acct"}}, + "driver": { + "container": "gym:latest", + "mounts": ["/host/cache:/cache"], + "benchmarks": {"gpqa": {}}, + }, + "job": {"output_path": "/jobs"}, + } + ) + bench_dir = Path("/jobs/gym-job-x/gpqa") + + script = build_sbatch_script(config, "gpqa", config.driver.benchmarks["gpqa"], config.compute["hsg"], bench_dir) + + driver_line = next(line for line in script.splitlines() if "--output=logs/driver.log" in line) + assert f"{bench_dir}:{bench_dir}" in driver_line + # the caller's own mounts must survive alongside the injected one + assert "/host/cache:/cache" in driver_line + # The output path is absolute rather than cwd-relative, and cwd is left + # alone: a benchmark's prepare_script is resolved against cwd, so moving it + # breaks `gym eval prepare` on a file that exists. + assert f"+output_jsonl_fpath={bench_dir}/artifacts/rollouts.jsonl" in script + assert "--container-workdir" not in script + + +def test_driver_job_dir_is_mounted_even_with_no_configured_mounts(): + config = SubmitConfig.model_validate( + { + "services": {}, + "compute": {"hsg": {"type": "slurm", "account": "acct"}}, + "driver": {"container": "gym:latest", "benchmarks": {"gpqa": {}}}, + "job": {"output_path": "/jobs"}, + } + ) + bench_dir = Path("/jobs/gym-job-x/gpqa") + + script = build_sbatch_script(config, "gpqa", config.driver.benchmarks["gpqa"], config.compute["hsg"], bench_dir) + + driver_line = next(line for line in script.splitlines() if "--output=logs/driver.log" in line) + assert f"--container-mounts={bench_dir}:{bench_dir}" in driver_line + + +def test_gym_install_does_not_clone_into_the_job_directory(): + """The driver's cwd is the job directory. A clone there gives Gym a second + copy of every built-in asset, and named lookups (`--model-type + openai_model`) then abort as ambiguous against the installed copy.""" + entrypoint = render_driver_entrypoint(repo="https://github.com/NVIDIA-NeMo/gym", ref="abc123", prepare_cmd=None) + + assert "mktemp -d /tmp/gym-install-" in entrypoint + assert 'git clone https://github.com/NVIDIA-NeMo/gym "$GYM_SRC/gym"' in entrypoint + + +def test_gym_install_runs_from_the_install_root(): + """A benchmark's prepare_script is relative to cwd, and a runtime image bakes no + Gym, so the driver has to run from the clone or `gym eval prepare` cannot find + its own script. Safe because the clone is outside the job directory and the + driver's output path is absolute.""" + entrypoint = render_driver_entrypoint(repo="https://github.com/NVIDIA-NeMo/gym", ref="abc123", prepare_cmd=None) + + # Assert the behaviour, not the line layout: the `cd` is chained onto the + # checkout with && rather than standing on its own line. + assert 'cd "$GYM_SRC/gym"' in entrypoint + assert entrypoint.index('cd "$GYM_SRC/gym"') < entrypoint.index('exec "$@"') + # The only `cd` is into the clone -- nothing else may move cwd. + assert entrypoint.count("cd ") == entrypoint.count('cd "$GYM_SRC/gym"')