|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import re |
| 4 | +from pathlib import Path |
| 5 | + |
| 6 | +REPO_ROOT = Path(__file__).resolve().parents[2] |
| 7 | + |
| 8 | +MEGATRON_BRIDGE_REPO = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git" |
| 9 | +NEMO_RL_REPO = "https://github.com/NVIDIA-NeMo/RL.git" |
| 10 | +AUTOMODEL_REPO = "https://github.com/NVIDIA-NeMo/Automodel.git" |
| 11 | + |
| 12 | +SUPER_V3_SHA = "f570c0529c81b57cb2ae909bd31a19408c7f4583" |
| 13 | +NANO_V3_SHA = "1cedb0a9c5f79d2cd2b5226a86b794b9f0e048a8" |
| 14 | +NEMO_RL_SUPER_V3_SHA = "bb0a7d43931950a74522e159f7117543a87b580b" |
| 15 | +AUTOMODEL_NEMOTRON_OMNI_SHA = "7dfec6130ddf675cc9721d1619945dcc743f0095" |
| 16 | + |
| 17 | +SUPER3_MEGATRON_BRIDGE_DOCS = ( |
| 18 | + REPO_ROOT / "docs/nemotron/super3/pretrain.md", |
| 19 | + REPO_ROOT / "docs/nemotron/super3/sft.md", |
| 20 | + REPO_ROOT / "docs/nemotron/super3/quantization.md", |
| 21 | +) |
| 22 | +NANO3_MEGATRON_BRIDGE_DOCS = ( |
| 23 | + REPO_ROOT / "docs/nemotron/nano3/pretrain.md", |
| 24 | + REPO_ROOT / "docs/nemotron/nano3/sft.md", |
| 25 | +) |
| 26 | +SUPER3_RL_DOC = REPO_ROOT / "docs/nemotron/super3/rl/index.md" |
| 27 | +AUTOMODEL_COOKBOOK = ( |
| 28 | + REPO_ROOT |
| 29 | + / "usage-cookbook/Nemotron-3-Nano-Omni/automodel/automodel_training_cookbook.md" |
| 30 | +) |
| 31 | + |
| 32 | + |
| 33 | +def _read(path: Path) -> str: |
| 34 | + return path.read_text(encoding="utf-8") |
| 35 | + |
| 36 | + |
| 37 | +def _assert_lower_sha(value: str) -> None: |
| 38 | + assert re.fullmatch(r"[0-9a-f]{40}", value) |
| 39 | + |
| 40 | + |
| 41 | +def _assert_checkout_pinned(path: Path, *, repo: str, branch: str, sha: str) -> None: |
| 42 | + text = _read(path) |
| 43 | + |
| 44 | + assert repo in text, f"{path} missing repo {repo}" |
| 45 | + assert branch in text, f"{path} missing branch context {branch}" |
| 46 | + assert sha in text, f"{path} missing pinned SHA {sha}" |
| 47 | + _assert_lower_sha(sha) |
| 48 | + assert re.search(rf"^git checkout {re.escape(sha)}(?:\s+#\s*{re.escape(branch)})?$", text, re.MULTILINE), ( |
| 49 | + f"{path} does not checkout pinned {branch} SHA {sha}" |
| 50 | + ) |
| 51 | + |
| 52 | + |
| 53 | +def test_super3_megatron_bridge_docs_pin_super_v3_revision() -> None: |
| 54 | + for path in SUPER3_MEGATRON_BRIDGE_DOCS: |
| 55 | + _assert_checkout_pinned( |
| 56 | + path, |
| 57 | + repo=MEGATRON_BRIDGE_REPO, |
| 58 | + branch="super-v3", |
| 59 | + sha=SUPER_V3_SHA, |
| 60 | + ) |
| 61 | + assert "git checkout super-v3" not in _read(path) |
| 62 | + |
| 63 | + |
| 64 | +def test_nano3_megatron_bridge_docs_pin_nano_v3_revision() -> None: |
| 65 | + for path in NANO3_MEGATRON_BRIDGE_DOCS: |
| 66 | + _assert_checkout_pinned( |
| 67 | + path, |
| 68 | + repo=MEGATRON_BRIDGE_REPO, |
| 69 | + branch="nano-v3", |
| 70 | + sha=NANO_V3_SHA, |
| 71 | + ) |
| 72 | + assert "git checkout nano-v3" not in _read(path) |
| 73 | + |
| 74 | + |
| 75 | +def test_super3_rl_doc_pins_nemo_rl_super_v3_revision() -> None: |
| 76 | + text = _read(SUPER3_RL_DOC) |
| 77 | + |
| 78 | + _assert_checkout_pinned( |
| 79 | + SUPER3_RL_DOC, |
| 80 | + repo=NEMO_RL_REPO, |
| 81 | + branch="super-v3", |
| 82 | + sha=NEMO_RL_SUPER_V3_SHA, |
| 83 | + ) |
| 84 | + assert "git clone --recursive -b super-v3" in text |
| 85 | + assert "git submodule update --init --recursive" in text |
| 86 | + |
| 87 | + |
| 88 | +def test_automodel_cookbook_pins_nemotron_omni_revision() -> None: |
| 89 | + text = _read(AUTOMODEL_COOKBOOK) |
| 90 | + |
| 91 | + _assert_checkout_pinned( |
| 92 | + AUTOMODEL_COOKBOOK, |
| 93 | + repo=AUTOMODEL_REPO, |
| 94 | + branch="nemotron-omni", |
| 95 | + sha=AUTOMODEL_NEMOTRON_OMNI_SHA, |
| 96 | + ) |
| 97 | + assert "git clone -b nemotron-omni" in text |
| 98 | + assert "gitlab-master.nvidia.com" not in text |
0 commit comments