Merge pull request #638 from NVIDIA/am/slurm-container-upd

Updates for SlurmContainer workload

Author: amaslenn

conf/common/test_scenario/slurm_container.toml

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+name = "slurm-container"
+
+[[Tests]]
+id = "nccl.alltoall"
+num_nodes = 2
+time_limit = "00:20:00"
+
+name = "nccl-alltoall"
+description = "NCCL alltoall via SlurmContainer"
+test_template_name = "SlurmContainer"
+
+  [Tests.cmd_args]
+  docker_image_url = "nvcr.io#nvidia/pytorch:25.06-py3"
+  cmd = "alltoall_perf_mpi --nthreads 1 --ngpus 1 --minbytes 128 --maxbytes 4G --stepbytes 1M --op sum --datatype float --root 0 --iters 100 --warmup_iters 50 --agg_iters 1 --average 1 --parallel_init 0 --check 1 --blocking 0 --cudagraph 0 --stepfactor 2"

src/cloudai/registration.py

@@ -95,11 +95,7 @@ def register_all():
         SleepStandaloneCommandGenStrategy,
         SleepTestDefinition,
     )
-    from cloudai.workloads.slurm_container import (
-        SlurmContainerCommandGenStrategy,
-        SlurmContainerReportGenerationStrategy,
-        SlurmContainerTestDefinition,
-    )
+    from cloudai.workloads.slurm_container import SlurmContainerCommandGenStrategy, SlurmContainerTestDefinition
     from cloudai.workloads.triton_inference import (
         TritonInferenceReportGenerationStrategy,
         TritonInferenceSlurmCommandGenStrategy,
@@ -201,7 +197,6 @@ def register_all():
     Registry().add_report(NeMoRunTestDefinition, NeMoRunReportGenerationStrategy)
     Registry().add_report(NeMoRunTestDefinition, NeMoRunDataStoreReportGenerationStrategy)
     Registry().add_report(NemotronTestDefinition, JaxToolboxReportGenerationStrategy)
-    Registry().add_report(SlurmContainerTestDefinition, SlurmContainerReportGenerationStrategy)
     Registry().add_report(UCCTestDefinition, UCCTestReportGenerationStrategy)
     Registry().add_report(TritonInferenceTestDefinition, TritonInferenceReportGenerationStrategy)
     Registry().add_report(NIXLBenchTestDefinition, NIXLBenchReportGenerationStrategy)

src/cloudai/workloads/slurm_container/__init__.py

@@ -14,13 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .report_generation_strategy import SlurmContainerReportGenerationStrategy
 from .slurm_command_gen_strategy import SlurmContainerCommandGenStrategy
 from .slurm_container import SlurmContainerCmdArgs, SlurmContainerTestDefinition
 
 __all__ = [
     "SlurmContainerCmdArgs",
     "SlurmContainerCommandGenStrategy",
-    "SlurmContainerReportGenerationStrategy",
     "SlurmContainerTestDefinition",
 ]

src/cloudai/workloads/slurm_container/report_generation_strategy.py

@@ -65,7 +65,6 @@
     NeMoRunTestDefinition,
 )
 from cloudai.workloads.nixl_bench import NIXLBenchReportGenerationStrategy, NIXLBenchTestDefinition
-from cloudai.workloads.slurm_container import SlurmContainerReportGenerationStrategy, SlurmContainerTestDefinition
 from cloudai.workloads.triton_inference import TritonInferenceReportGenerationStrategy, TritonInferenceTestDefinition
 from cloudai.workloads.ucc_test import UCCTestDefinition, UCCTestReportGenerationStrategy
 
@@ -481,7 +480,7 @@ def test_default(self):
         assert len(reporters) == 0
 
     def test_default_reporters_size(self):
-        assert len(Registry().reports_map) == 13
+        assert len(Registry().reports_map) == 12
 
     @pytest.mark.parametrize(
         "tdef,expected_reporters",
@@ -494,7 +493,6 @@ def test_default_reporters_size(self):
             (NeMoLauncherTestDefinition, {NeMoLauncherReportGenerationStrategy}),
             (NeMoRunTestDefinition, {NeMoRunReportGenerationStrategy, NeMoRunDataStoreReportGenerationStrategy}),
             (NemotronTestDefinition, {JaxToolboxReportGenerationStrategy}),
-            (SlurmContainerTestDefinition, {SlurmContainerReportGenerationStrategy}),
             (UCCTestDefinition, {UCCTestReportGenerationStrategy}),
             (TritonInferenceTestDefinition, {TritonInferenceReportGenerationStrategy}),
             (NIXLBenchTestDefinition, {NIXLBenchReportGenerationStrategy}),