fix(perf): keep recipe imports off login nodes

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

scripts/performance/README.md

@@ -260,6 +260,7 @@ Mounting cached files is not enough by itself. If `HF_HUB_OFFLINE` remains `0`,
 
 ##### Logging arguments
 
+- `--experiment-name`: NeMo-Run experiment name used for scheduler bookkeeping. This is independent of W&B naming.
 - `-l/--log_dir`: Directory for logging experiment results. Defaults to `NEMORUN_HOME`.
   - Make sure the environment variable `NEMORUN_HOME=<log_dir>` is accessible and set correctly in your virtual environment.
   - You can run `export NEMORUN_HOME=<log_dir>` in your terminal. You can add it your bashrc file (or equivalent for your OS/Linux distro) for setting it permanently.

scripts/performance/argument_parser.py

@@ -874,6 +874,14 @@ def parse_cli_args():
 
     # Logging
     logging_args = parser.add_argument_group("Logging arguments")
+    logging_args.add_argument(
+        "--experiment-name",
+        "--experiment_name",
+        dest="experiment_name",
+        type=str,
+        help="NeMo-Run experiment name. Independent of the Weights & Biases run name.",
+        required=False,
+    )
     logging_args.add_argument(
         "-wdk",
         "--wandb_key",

scripts/performance/setup_experiment.py

@@ -23,7 +23,6 @@
 import tempfile
 import time
 from pathlib import Path
-from types import SimpleNamespace
 from typing import Any, Dict, List, Optional
 
 import nemo_run as run
@@ -38,7 +37,7 @@
         kubeflow_executor,
         slurm_executor,
     )
-    from utils.utils import configure_slurm_gpu_tuning, get_exp_name_config, select_config_variant_interactive
+    from utils.utils import configure_slurm_gpu_tuning, select_config_variant_interactive
 except (ImportError, ModuleNotFoundError):
     from .argument_parser import NUM_GPUS_PER_NODE_MAP, parse_cli_args
     from .utils.executors import (
@@ -47,7 +46,7 @@
         kubeflow_executor,
         slurm_executor,
     )
-    from .utils.utils import configure_slurm_gpu_tuning, get_exp_name_config, select_config_variant_interactive
+    from .utils.utils import configure_slurm_gpu_tuning, select_config_variant_interactive
 
 try:
     import wandb
@@ -79,6 +78,7 @@ def _filter_run_script_args(argv: List[str]) -> List[str]:
     * ``--additional_slurm_params`` — Slurm orchestration only.
     * ``--enable_vboost`` / ``--lock_gpu_freq`` — applied directly to the
       Slurm executor before submission.
+    * ``--experiment-name`` — selects the NeMo-Run experiment name.
     * ``--csp`` — launcher-only; selects the CSP fabric plugin. The rank-local
       script forwards unrecognized args to Hydra, which rejects ``--csp``.
     * ``--kubeflow_*`` — consumed here to build the Kubeflow TrainJob. Several
@@ -97,6 +97,8 @@ def _is_launcher_only(flag: str) -> bool:
             "--additional_slurm_params",
             "--csp",
             "--enable_vboost",
+            "--experiment-name",
+            "--experiment_name",
             "--lock_gpu_freq",
         ) or flag.startswith("--kubeflow_")
 
@@ -115,6 +117,26 @@ def _is_launcher_only(flag: str) -> bool:
     return filtered_args
 
 
+def _default_experiment_name(
+    *,
+    use_recipes: bool,
+    model_recipe_name: str,
+    task: str,
+    compute_dtype: str,
+    num_gpus: int,
+    gpu: str,
+    config_variant: str | None,
+) -> str:
+    """Build a stable experiment name without importing a performance recipe."""
+    if use_recipes:
+        return f"{model_recipe_name}_{task}_{num_gpus}gpu_{gpu}"
+
+    fields = [task, model_recipe_name, compute_dtype, f"gpus{num_gpus}", gpu]
+    if config_variant and config_variant.lower() not in {"v1", "v2"}:
+        fields.append(config_variant.lower())
+    return "_".join(fields)
+
+
 def _build_nemorun_script(
     *,
     script_path: str,
@@ -456,18 +478,11 @@ def main(
     export_nsys_sqlite: bool,
     pytorch_profiler: bool,
     moe_a2a_overlap: bool,
-    tp_size: Optional[int],
-    pp_size: Optional[int],
-    cp_size: Optional[int],
-    vp_size: Optional[int],
-    ep_size: Optional[int],
-    etp_size: Optional[int],
-    micro_batch_size: Optional[int],
-    global_batch_size: Optional[int],
     wandb_key: str,
     wandb_project_name: str,
     wandb_experiment_name: str,
     wandb_entity_name: str,
+    experiment_name: Optional[str],
     profiling_start_step: int,
     profiling_stop_step: int,
     record_memory_history: bool,
@@ -564,35 +579,23 @@ def main(
         logger.warning("--export_nsys_sqlite was set without --enable_nsys; no Nsys SQLite export will be generated.")
 
     script_name = ENTRYPOINT_BOOTSTRAP
-    if use_recipes:
-        exp_name = (
-            wandb_experiment_name
-            if wandb_experiment_name is not None
-            else f"{model_recipe_name}_{task}_{num_gpus}gpu_{gpu}"
+    # Keep the historical W&B-name fallback for callers that relied on it, but
+    # let scheduling use a neutral name. The lightweight default deliberately
+    # avoids resolving a recipe: effective parallelism, batches, and process
+    # environment are finalized by bootstrap.py inside the submitted container.
+    exp_name = (
+        experiment_name
+        or wandb_experiment_name
+        or _default_experiment_name(
+            use_recipes=use_recipes,
+            model_recipe_name=model_recipe_name,
+            task=task,
+            compute_dtype=compute_dtype,
+            num_gpus=num_gpus,
+            gpu=gpu,
+            config_variant=config_variant,
         )
-
-    else:
-        if wandb_experiment_name is not None:
-            # CI supplies the complete experiment name. Avoid resolving a perf recipe on the
-            # login node in this path: recipe imports belong in the training container.
-            exp_name = wandb_experiment_name
-        else:
-            # Create a simple namespace with the args needed by get_exp_name_config
-            args_for_config = SimpleNamespace(
-                num_gpus=num_gpus,
-                tensor_model_parallel_size=tp_size,
-                pipeline_model_parallel_size=pp_size,
-                context_parallel_size=cp_size,
-                virtual_pipeline_model_parallel_size=vp_size,
-                expert_model_parallel_size=ep_size,
-                expert_tensor_parallel_size=etp_size,
-                micro_batch_size=micro_batch_size,
-                global_batch_size=global_batch_size,
-            )
-            exp_config = get_exp_name_config(
-                args_for_config, model_family_name, model_recipe_name, gpu, compute_dtype, task, config_variant
-            )
-            exp_name = f"{task}_{model_recipe_name}_{compute_dtype}_{exp_config}"
+    )
 
     if pretrained_checkpoint is not None:
         custom_mounts.append(f"{pretrained_checkpoint}:{pretrained_checkpoint}")
@@ -1002,18 +1005,11 @@ def main(
         export_nsys_sqlite=args.export_nsys_sqlite,
         pytorch_profiler=args.pytorch_profiler,
         moe_a2a_overlap=args.moe_a2a_overlap,
-        tp_size=args.tensor_model_parallel_size,
-        pp_size=args.pipeline_model_parallel_size,
-        cp_size=args.context_parallel_size,
-        vp_size=args.virtual_pipeline_model_parallel_size,
-        ep_size=args.expert_model_parallel_size,
-        etp_size=args.expert_tensor_parallel_size,
-        micro_batch_size=args.micro_batch_size,
-        global_batch_size=args.global_batch_size,
         wandb_key=args.wandb_key,
         wandb_project_name=args.wandb_project_name,
         wandb_experiment_name=args.wandb_experiment_name,
         wandb_entity_name=args.wandb_entity_name,
+        experiment_name=args.experiment_name,
         profiling_start_step=args.profiling_start_step,
         profiling_stop_step=args.profiling_stop_step,
         record_memory_history=args.record_memory_history,

scripts/training/recipe_metadata.py

@@ -56,6 +56,7 @@
         "llama3_70b_peft_8gpu_h100_bf16_config",
         "llama3_70b_sft_32gpu_h100_bf16_config",
         "qwen3_235b_a22b_pretrain_256gpu_h100_bf16_config",
+        "qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config",
     }
 )
 

tests/unit_tests/scripts/performance/test_setup_experiment.py

@@ -0,0 +1,171 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""Tests for lightweight performance submission from a Slurm login node."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+
+pytestmark = pytest.mark.unit
+
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+_PERF_SCRIPTS_DIR = _REPO_ROOT / "scripts" / "performance"
+if str(_PERF_SCRIPTS_DIR) not in sys.path:
+    sys.path.insert(0, str(_PERF_SCRIPTS_DIR))
+
+import setup_experiment
+from argument_parser import parse_cli_args
+
+
+def test_default_experiment_name_uses_only_cli_metadata() -> None:
+    assert (
+        setup_experiment._default_experiment_name(
+            use_recipes=False,
+            model_recipe_name="deepseek_v3",
+            task="pretrain",
+            compute_dtype="fp8_mx",
+            num_gpus=256,
+            gpu="vr200",
+            config_variant="large_scale",
+        )
+        == "pretrain_deepseek_v3_fp8_mx_gpus256_vr200_large_scale"
+    )
+    assert (
+        setup_experiment._default_experiment_name(
+            use_recipes=True,
+            model_recipe_name="llama3_8b",
+            task="pretrain",
+            compute_dtype="bf16",
+            num_gpus=8,
+            gpu="h100",
+            config_variant=None,
+        )
+        == "llama3_8b_pretrain_8gpu_h100"
+    )
+
+
+def test_experiment_name_is_launcher_only_and_recipe_arguments_are_unchanged() -> None:
+    recipe_args = [
+        "--tensor_model_parallel_size",
+        "2",
+        "--global_batch_size",
+        "64",
+        "++model.num_layers=2",
+        "++env_vars.TEST_SENTINEL=1",
+    ]
+    argv = [
+        "--experiment-name",
+        "scheduler-name",
+        "--wandb_experiment_name",
+        "wandb-name",
+        *recipe_args,
+    ]
+
+    assert setup_experiment._filter_run_script_args(argv) == [
+        "--wandb_experiment_name",
+        "wandb-name",
+        *recipe_args,
+    ]
+
+
+@pytest.mark.parametrize("flag", ["--experiment-name", "--experiment_name"])
+def test_experiment_name_cli_aliases(flag: str) -> None:
+    parser = parse_cli_args()
+    args, unknown = parser.parse_known_args(
+        [
+            "--model_family_name",
+            "llama",
+            "--model_recipe_name",
+            "llama3_8b",
+            "--num_gpus",
+            "8",
+            "--gpu",
+            "h100",
+            flag,
+            "scheduler-name",
+        ]
+    )
+
+    assert unknown == []
+    assert args.experiment_name == "scheduler-name"
+
+
+def test_submission_dry_run_does_not_import_bridge_or_mcore(tmp_path: Path) -> None:
+    blocker_dir = tmp_path / "login_node"
+    blocker_dir.mkdir()
+    (blocker_dir / "sitecustomize.py").write_text(
+        """
+import importlib.abc
+import sys
+
+class _RejectMegatron(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname == "megatron" or fullname.startswith("megatron."):
+            raise RuntimeError(f"login-node import forbidden: {fullname}")
+        return None
+
+sys.meta_path.insert(0, _RejectMegatron())
+"""
+    )
+    environment = os.environ.copy()
+    environment["NEMORUN_HOME"] = str(tmp_path / "nemorun")
+    environment["PYTHONPATH"] = os.pathsep.join(
+        value for value in (str(blocker_dir), environment.get("PYTHONPATH")) if value
+    )
+    command = [
+        sys.executable,
+        str(_PERF_SCRIPTS_DIR / "setup_experiment.py"),
+        "--model_family_name",
+        "llama",
+        "--model_recipe_name",
+        "llama3_8b",
+        "--num_gpus",
+        "8",
+        "--gpus_per_node",
+        "8",
+        "--gpu",
+        "h100",
+        "--account",
+        "test-account",
+        "--partition",
+        "batch",
+        "--container_image",
+        "example.invalid/nemo:test",
+        "--packager",
+        "none",
+        "--dryrun",
+    ]
+
+    result = subprocess.run(
+        command,
+        cwd=_REPO_ROOT,
+        env=environment,
+        capture_output=True,
+        text=True,
+        check=False,
+        timeout=30,
+    )
+
+    output = result.stdout + result.stderr
+    assert result.returncode == 0, output
+    assert "bootstrap.py" in output
+    assert "pretrain_llama3_8b_bf16_gpus8_h100" in output
+    assert "login-node import forbidden" not in output