fix(performance): apply recipe env in container

Author: yaoyu-33

scripts/performance/perf_plugins.py

@@ -37,9 +37,9 @@
 
 
 try:
-    from utils.utils import WorkloadBaseConfig, get_workload_base_config
+    from utils.utils import WorkloadBaseConfig
 except (ImportError, ModuleNotFoundError):
-    from .utils.utils import WorkloadBaseConfig, get_workload_base_config
+    from .utils.utils import WorkloadBaseConfig
 
 logger: logging.Logger = logging.getLogger(__name__)
 NSYS_SQLITE_EXPORT_ARG = "--export=sqlite"
@@ -495,16 +495,13 @@ def get_lock_gpu_freq_srun_cmd(job_dir, freq_mhz):
                 else lock_freq_cmd
             )
 
-    def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executor"):
-        """Enable the performance environment settings"""
-        workload_base_config = get_workload_base_config(
-            self.model_family_name,
-            self.model_recipe_name,
-            self.gpu,
-            self.compute_dtype,
-            self.train_task,
-            self.config_variant,
-        )
+    def setup_recipe_environment(
+        self,
+        task: Union["run.Partial", "run.Script", None],
+        executor: "run.Executor",
+        workload_base_config: WorkloadBaseConfig,
+    ) -> None:
+        """Apply recipe-dependent environment settings inside the training container."""
         tp_size = self.tp_size if self.tp_size is not None else workload_base_config.tensor_model_parallel_size
         pp_size = self.pp_size if self.pp_size is not None else workload_base_config.pipeline_model_parallel_size
         cp_size = self.cp_size if self.cp_size is not None else workload_base_config.context_parallel_size
@@ -554,15 +551,6 @@ def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executo
             nccl_pp_comm_chunksize = None
         self._set_nccl_pp_comm_chunksize(task, executor, nccl_pp_comm_chunksize, pp_size)
 
-        # Configure manual garbage collection
-        self._set_manual_gc(task, executor, self.enable_manual_gc, self.manual_gc_interval)
-
-        # Improve perf by steering power to tensor cores, may not work on all systems
-        self._set_vboost(task, executor, self.enable_vboost)
-
-        # Lock GPU graphics clock frequency for stable performance measurements
-        self._set_lock_gpu_freq(task, executor, self.lock_gpu_freq)
-
         # Set model-specific environment variables
         self._set_model_specific_environment_variables(
             task,
@@ -582,6 +570,14 @@ def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executo
         if self.deterministic:
             self._set_determinism_env_vars(executor)
 
+    def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executor"):
+        """Apply launcher-only settings without importing Megatron-Bridge on the login node."""
+        # Recipe-dependent environment variables are applied by run_script_with_env.py
+        # inside the training container, where the flat recipe can be imported safely.
+        self._set_manual_gc(task, executor, self.enable_manual_gc, self.manual_gc_interval)
+        self._set_vboost(task, executor, self.enable_vboost)
+        self._set_lock_gpu_freq(task, executor, self.lock_gpu_freq)
+
 
 @dataclass
 class PyTorchProfilerPluginScriptArgs:

scripts/performance/run_script_with_env.py

@@ -0,0 +1,74 @@
+# 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.
+
+"""Resolve recipe-dependent perf environment variables, then exec run_script.py.
+
+This entrypoint runs inside the training container.  It keeps recipe imports off
+the dependency-light login node while ensuring environment variables are present
+before the training Python process imports Torch, Transformer Engine, or Megatron.
+"""
+
+import os
+import sys
+from pathlib import Path
+
+from argument_parser import parse_cli_args
+from perf_plugins import PerfEnvPlugin
+from utils.utils import _workload_base_config_from_recipe, get_perf_recipe_by_name
+
+
+class _EnvironmentExecutor:
+    """Minimal executor adapter used by PerfEnvPlugin's environment helpers."""
+
+    def __init__(self) -> None:
+        self.env_vars = os.environ
+
+
+def main() -> None:
+    """Apply the selected recipe's environment and replace this process with run_script.py."""
+    parser = parse_cli_args()
+    args, _ = parser.parse_known_args()
+
+    recipe = get_perf_recipe_by_name(
+        model_recipe_name=args.model_recipe_name,
+        task=args.task,
+        num_gpus=args.num_gpus,
+        gpu=args.gpu,
+        precision=args.compute_dtype,
+        config_variant=args.config_variant,
+    )
+    workload_base_config = _workload_base_config_from_recipe(recipe, num_gpus=args.num_gpus)
+
+    plugin = PerfEnvPlugin(
+        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,
+        ep_size=args.expert_model_parallel_size,
+        model_family_name=args.model_family_name,
+        model_recipe_name=args.model_recipe_name,
+        gpu=args.gpu,
+        compute_dtype=args.compute_dtype,
+        train_task=args.task,
+        config_variant=args.config_variant,
+        deterministic=args.deterministic,
+    )
+    plugin.setup_recipe_environment(None, _EnvironmentExecutor(), workload_base_config)
+
+    run_script = Path(__file__).with_name("run_script.py")
+    os.execvpe(sys.executable, [sys.executable, str(run_script), *sys.argv[1:]], dict(os.environ))
+
+
+if __name__ == "__main__":
+    main()

scripts/performance/setup_experiment.py

@@ -52,7 +52,7 @@
     from .perf_plugins import NsysPlugin, PerfEnvPlugin, PyTorchProfilerPlugin
 
 SCRIPT_DIR = Path(__file__).parent.resolve()
-ENTRYPOINT_PEFORMANCE = "run_script.py"
+ENTRYPOINT_PEFORMANCE = "run_script_with_env.py"
 ENTRYPOINT_RECIPE = "run_recipe.py"
 
 logging.basicConfig(level=logging.DEBUG)