revert plugin

Signed-off-by: Malay Nagda <malayn@nvidia.com>

Author: malay-nagda

scripts/performance/perf_plugins.py

@@ -35,12 +35,6 @@
 from nemo_run import Plugin, Script, SlurmExecutor
 
 
-try:
-    from utils.utils import configure_slurm_gpu_tuning
-except (ImportError, ModuleNotFoundError):
-    from .utils.utils import configure_slurm_gpu_tuning
-
-
 logger: logging.Logger = logging.getLogger(__name__)
 NSYS_SQLITE_EXPORT_ARG = "--export=sqlite"
 
@@ -70,26 +64,6 @@ def _ensure_sqlite_nsys_export(nsys_extra_args: list[str]) -> list[str]:
     return nsys_extra_args + [NSYS_SQLITE_EXPORT_ARG]
 
 
-@dataclass(kw_only=True)
-class SlurmGpuTuningPlugin(Plugin):
-    """Apply GPU tuning after NeMo Run assigns the executor's job directory."""
-
-    enable_vboost: bool = False
-    lock_gpu_freq: int | None = None
-    peak_mem_clk: int | None = None
-
-    def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executor") -> None:
-        """Add Slurm setup commands using the assigned per-task job directory."""
-        del task
-        if isinstance(executor, SlurmExecutor):
-            configure_slurm_gpu_tuning(
-                executor,
-                enable_vboost=self.enable_vboost,
-                lock_gpu_freq=self.lock_gpu_freq,
-                peak_mem_clk=self.peak_mem_clk,
-            )
-
-
 @dataclass
 class NsysPluginScriptArgs:
     """Arguments for NsysPlugin to pass to run.Script."""

scripts/performance/setup_experiment.py

@@ -38,7 +38,7 @@
         kubeflow_executor,
         slurm_executor,
     )
-    from utils.utils import get_exp_name_config, select_config_variant_interactive
+    from utils.utils import configure_slurm_gpu_tuning, get_exp_name_config, 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 +47,7 @@
         kubeflow_executor,
         slurm_executor,
     )
-    from .utils.utils import get_exp_name_config, select_config_variant_interactive
+    from .utils.utils import configure_slurm_gpu_tuning, get_exp_name_config, select_config_variant_interactive
 
 try:
     import wandb
@@ -57,9 +57,9 @@
     HAVE_WANDB = False
 
 try:
-    from perf_plugins import NsysPlugin, PreemptionPlugin, PyTorchProfilerPlugin, SlurmGpuTuningPlugin
+    from perf_plugins import NsysPlugin, PreemptionPlugin, PyTorchProfilerPlugin
 except (ImportError, ModuleNotFoundError):
-    from .perf_plugins import NsysPlugin, PreemptionPlugin, PyTorchProfilerPlugin, SlurmGpuTuningPlugin
+    from .perf_plugins import NsysPlugin, PreemptionPlugin, PyTorchProfilerPlugin
 
 SCRIPT_DIR = Path(__file__).parent.resolve()
 ENTRYPOINT_BOOTSTRAP = "bootstrap.py"
@@ -704,17 +704,15 @@ def main(
             packager=packager,
             enable_pct_binding=enable_pct_binding,
         )
-    plugins = []
-
-    if not kubeflow_namespace and (enable_vboost or lock_gpu_freq is not None or peak_mem_clk is not None):
-        plugins.append(
-            SlurmGpuTuningPlugin(
-                enable_vboost=enable_vboost,
-                lock_gpu_freq=lock_gpu_freq,
-                peak_mem_clk=peak_mem_clk,
-            )
+        configure_slurm_gpu_tuning(
+            executor,
+            enable_vboost=enable_vboost,
+            lock_gpu_freq=lock_gpu_freq,
+            peak_mem_clk=peak_mem_clk,
         )
 
+    plugins = []
+
     # Long-convergence runs are split across walltime slices and resume from the last
     # checkpoint each slice. Without a preemption signal, Slurm hard-kills the slice at
     # the time limit before a checkpoint is written, so no progress persists and the

scripts/performance/utils/utils.py

@@ -79,14 +79,9 @@ def configure_slurm_gpu_tuning(
     lock_gpu_freq: int | None,
     peak_mem_clk: int | None,
 ) -> None:
-    """Add optional GPU tuning commands to an assigned Slurm executor."""
-    if not (enable_vboost or lock_gpu_freq is not None or peak_mem_clk is not None):
-        return
-
+    """Add optional GPU tuning commands to a Slurm executor before submission."""
     commands = []
-    job_dir = executor.job_dir
-    if not job_dir:
-        raise ValueError("Slurm executor must be assigned to a job before configuring GPU tuning.")
+    job_dir = executor.tunnel.job_dir
 
     if enable_vboost:
         commands.append(

tests/unit_tests/scripts/performance/test_run_script_entrypoint.py

@@ -25,7 +25,6 @@
     sys.path.insert(0, str(_PERF_SCRIPTS_DIR))
 
 import bootstrap
-import perf_plugins
 import run_recipe
 import run_script
 import setup_experiment
@@ -163,35 +162,26 @@ def test_compatibility_overrides_preserve_legacy_manual_gc_defaults():
     assert recipe.train.manual_gc_interval == 100
 
 
-def test_gpu_tuning_plugin_uses_assigned_slurm_job_dir():
-    executor = perf_plugins.SlurmExecutor(
-        account="test",
+def test_gpu_tuning_options_are_applied_directly_to_slurm_executor():
+    executor = SimpleNamespace(
         nodes=2,
-        tunnel=perf_plugins.run.LocalTunnel(job_dir="/experiments"),
+        tunnel=SimpleNamespace(job_dir="/job/dir"),
         setup_lines="existing setup\n",
     )
-    executor.assign(
-        "name_123",
-        "/experiments/name/name_123",
-        task_id="task",
-        task_dir="task",
-    )
-    plugin = perf_plugins.SlurmGpuTuningPlugin(
+
+    utils.configure_slurm_gpu_tuning(
+        executor,
         enable_vboost=True,
         lock_gpu_freq=1200,
         peak_mem_clk=2600,
     )
-    plugin.setup(SimpleNamespace(), executor)
 
     assert executor.setup_lines.startswith("existing setup\n")
     assert "sudo nvidia-smi boost-slider --vboost 1" in executor.setup_lines
     assert "--ntasks=2" in executor.setup_lines
     assert "sudo nvidia-smi -lgc 1200" in executor.setup_lines
-    assert (
-        "srun --ntasks=2 --ntasks-per-node=1 "
-        "--output /experiments/name/name_123/task/peak_mem_clock.out" in executor.setup_lines
-    )
-    assert "--error /experiments/name/name_123/task/peak_mem_clock.err" in executor.setup_lines
+    assert "srun --ntasks=2 --ntasks-per-node=1 --output /job/dir/peak_mem_clock.out" in executor.setup_lines
+    assert "--error /job/dir/peak_mem_clock.err" in executor.setup_lines
     assert "sudo nvidia-smi -lmc 2600,2600" in executor.setup_lines