Skip to content

Fix omni3 CLI test collection and restore super3/omni3 command registration #294

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
temujinkz wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix omni3 CLI test collection and restore super3/omni3 command registration #294

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions src/nemotron/cli/commands/super3/data/prep/_typer_group.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,12 @@

from __future__ import annotations

from nemotron.cli.commands.super3.data.prep.pretrain import META as PRETRAIN_META, pretrain
from nemotron.cli.commands.super3.data.prep.rl import META as RL_META, rl
from nemotron.cli.commands.super3.data.prep.sft import META as SFT_META, sft
from nemo_runspec.recipe_typer import RecipeTyper
from nemotron.cli.commands.super3.data.prep.pretrain import META as PRETRAIN_META
from nemotron.cli.commands.super3.data.prep.pretrain import pretrain
from nemotron.cli.commands.super3.data.prep.rl import rl_app
from nemotron.cli.commands.super3.data.prep.sft import META as SFT_META
from nemotron.cli.commands.super3.data.prep.sft import sft

# Create prep app using RecipeTyper
prep_app = RecipeTyper(
Expand All @@ -43,4 +45,4 @@

prep_app.add_recipe_command(pretrain, meta=PRETRAIN_META)
prep_app.add_recipe_command(sft, meta=SFT_META)
prep_app.add_recipe_command(rl, meta=RL_META)
prep_app.add_typer(rl_app, name="rl")
114 changes: 1 addition & 113 deletions tests/recipes/omni3/test_cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,106 +4,15 @@

import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from typer.testing import CliRunner

from nemotron.cli.bin.nemotron import app
from nemotron.cli.commands.omni3.build import (
BUILD_IMAGE,
REMOTE_CODE_ROOT,
STAGES,
STAGE_ALIASES,
_make_build_script,
)
from nemotron.cli.commands.omni3.data.prep.rl import _make_job_name

runner = CliRunner()


class TestOmni3BuildScript:
"""Unit tests for the inline bash script the build dispatcher
submits to slurm via ``run.Script(inline=...)``.

These pin the build pipeline's three phases (enroot install →
podman build → enroot import → squashfs) and the path conventions
that downstream stages depend on.
"""

def test_renders_three_phases(self):
script = _make_build_script(
stage_dir=Path("src/nemotron/recipes/omni3/stage0_sft"),
image_tag="nemotron/omni3-sft:latest",
sqsh_name="omni3-sft.sqsh",
enroot_version="3.5.0",
)
# Phase 1: enroot installs from GitHub release RPMs.
assert "enroot-3.5.0-1.el8.x86_64.rpm" in script
assert "enroot+caps-3.5.0-1.el8.x86_64.rpm" in script
# Phase 2: podman build against the in-container Dockerfile path.
assert (
f"{REMOTE_CODE_ROOT}/src/nemotron/recipes/omni3/stage0_sft/Dockerfile"
in script
)
assert "-t nemotron/omni3-sft:latest" in script
# Phase 3: enroot import from podman storage to squashfs.
assert "enroot import --output" in script
assert "podman://nemotron/omni3-sft:latest" in script
assert "/nemotron-cache/containers/omni3-sft.sqsh" in script

def test_strict_bash_mode(self):
script = _make_build_script(
stage_dir=Path("src/nemotron/recipes/omni3/stage0_sft"),
image_tag="t:latest",
sqsh_name="t.sqsh",
enroot_version="3.5.0",
)
assert "set -euo pipefail" in script

def test_extra_podman_args_are_quoted_safely(self):
script = _make_build_script(
stage_dir=Path("src/nemotron/recipes/omni3/stage0_sft"),
image_tag="t:latest",
sqsh_name="t.sqsh",
enroot_version="3.5.0",
extra_podman_args=["--build-arg", "FOO=bar baz", "--squash-all"],
)
# ``shlex.quote`` should wrap the multi-word arg in single quotes
# so the shell sees one token, not three.
assert "'FOO=bar baz'" in script
assert "--squash-all" in script

def test_enroot_version_pin_propagates(self):
script = _make_build_script(
stage_dir=Path("src/nemotron/recipes/omni3/stage0_sft"),
image_tag="t:latest",
sqsh_name="t.sqsh",
enroot_version="4.1.2",
)
assert "enroot-4.1.2-1.el8.x86_64.rpm" in script
assert "v4.1.2" in script


class TestOmni3BuildStageRegistry:
"""Adding/removing stages should be a single-row change in
``STAGES``; this test guards the public surface."""

def test_known_stages_have_image_and_sqsh_names(self):
for canonical, spec in STAGES.items():
assert spec["image_tag"], canonical
assert spec["sqsh_name"].endswith(".sqsh"), canonical

def test_aliases_resolve_to_known_stages(self):
for alias, canonical in STAGE_ALIASES.items():
assert canonical in STAGES, (alias, canonical)

def test_build_image_uses_enroot_uri_form(self):
# Pyxis on dlw misroutes bare ``quay.io/...`` strings to
# docker.io; the ``#`` separator is required.
assert BUILD_IMAGE.startswith("docker://"), BUILD_IMAGE
assert "#" in BUILD_IMAGE, BUILD_IMAGE


class TestRLJobNameUniqueness:
"""Regression tests for the parallel-batch race that produced
``Config file not found: config.yaml`` when sibling RL configs (mpo,
Expand All @@ -120,9 +29,7 @@ def test_parallel_invocations_unique(self):
# Simulate the exact failure mode: 4+ submissions in the same
# wall-clock second from the same dispatcher process.
with ThreadPoolExecutor(max_workers=8) as pool:
names = list(
pool.map(_make_job_name, ["omni3-data-prep-rl"] * 32)
)
names = list(pool.map(_make_job_name, ["omni3-data-prep-rl"] * 32))
assert len(set(names)) == len(names), names

def test_name_starts_with_prefix(self):
Expand All @@ -138,16 +45,11 @@ def test_help_succeeds(self):
def test_commit4_commands_exist(self):
result = runner.invoke(app, ["omni3", "--help"])
assert result.exit_code == 0
assert "build" in result.output
assert "sft" in result.output
assert "data" in result.output
assert "model" in result.output
assert "rl" in result.output

def test_build_help_succeeds(self):
result = runner.invoke(app, ["omni3", "build", "--help"])
assert result.exit_code == 0

def test_sft_help_succeeds(self):
result = runner.invoke(app, ["omni3", "sft", "--help"])
assert result.exit_code == 0
Expand All @@ -160,10 +62,6 @@ def test_data_prep_rl_help_succeeds(self):
result = runner.invoke(app, ["omni3", "data", "prep", "rl", "--help"])
assert result.exit_code == 0

def test_build_rl_help_succeeds(self):
result = runner.invoke(app, ["omni3", "build", "rl", "--help"])
assert result.exit_code == 0

def test_rl_group_help_succeeds(self):
result = runner.invoke(app, ["omni3", "rl", "--help"])
assert result.exit_code == 0
Expand Down Expand Up @@ -213,16 +111,6 @@ def test_pipe_help_succeeds(self):


class TestDryRun:
def test_build_dry_run_sft_stage(self, monkeypatch):
monkeypatch.setattr(sys, "argv", ["nemotron", "omni3", "build", "sft", "-d"])
result = runner.invoke(app, ["omni3", "build", "sft", "-d"])
assert result.exit_code == 0, f"dry-run failed: {result.output}\n{result.exception}"
# The dispatcher resolves the alias to the canonical stage name
# and shows the produced sqsh + image tag.
assert "stage0_sft" in result.output
assert "omni3-sft.sqsh" in result.output
assert "nemotron/omni3-sft:latest" in result.output

def test_pipe_dry_run_succeeds(self, monkeypatch):
# Vision RL launcher landed upstream — pipe runs all 4 stages by
# default and surfaces ``omni3-vision-rl-model:latest`` as the
Expand Down