Merge pull request #590 from NVIDIA/am/ucc-upd

Update UCC configs

Author: amaslenn

conf/common/test/ucc_test_allgather.toml conf/common/test/ucc_test.tomlconf/common/test/ucc_test_allgather.toml renamed to conf/common/test/ucc_test.toml

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,9 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-name = "ucc_test_allgather"
-description = "allgather"
+name = "ucc_base_test"
+description = "Base config for UCC tests"
 test_template_name = "UCCTest"
 
 [cmd_args]
-"collective" = "allgather"
+docker_image_url = "nvcr.io/nvidia/pytorch:25.06-py3"

conf/common/test/ucc_test_allreduce.toml

@@ -18,33 +18,46 @@ name = "ucc_test"
 
 [[Tests]]
 id = "Tests.alltoall"
-test_name = "ucc_test_alltoall"
+test_name = "ucc_base_test"
+description = "UCC alltoall"
 time_limit = "00:20:00"
 num_nodes = 2
+  [Tests.cmd_args]
+  collective = "alltoall"
 
 [[Tests]]
 id = "Tests.allgather"
-test_name = "ucc_test_allgather"
+test_name = "ucc_base_test"
+description = "UCC allgather"
 time_limit = "00:20:00"
 num_nodes = 2
+  [Tests.cmd_args]
+  collective = "allgather"
   [[Tests.dependencies]]
   type = "start_post_comp"
   id = "Tests.alltoall"
 
 [[Tests]]
 id = "Tests.allreduce"
-test_name = "ucc_test_allreduce"
+test_name = "ucc_base_test"
+description = "UCC allreduce"
 time_limit = "00:20:00"
 num_nodes = 2
+  [Tests.cmd_args]
+  collective = "allreduce"
+  e = "4G"
   [[Tests.dependencies]]
   type = "start_post_comp"
   id = "Tests.allgather"
 
 [[Tests]]
 id = "Tests.reduce_scatter"
-test_name = "ucc_test_reduce_scatter"
+test_name = "ucc_base_test"
+description = "UCC reduce_scatter"
 time_limit = "00:20:00"
 num_nodes = 2
+  [Tests.cmd_args]
+  collective = "reduce_scatter"
   [[Tests.dependencies]]
   type = "start_post_comp"
   id = "Tests.allreduce"

conf/common/test/ucc_test_alltoall.toml

@@ -23,7 +23,7 @@
 class UCCCmdArgs(CmdArgs):
     """UCC test command arguments."""
 
-    docker_image_url: str = "nvcr.io/nvidia/pytorch:24.02-py3"
+    docker_image_url: str
     collective: Union[
         Literal[
             "allgather",

conf/common/test/ucc_test_reduce_scatter.toml

@@ -72,7 +72,7 @@ def ucc_tr(slurm_system: SlurmSystem) -> TestRun:
                 name="ucc_test",
                 description="ucc_test",
                 test_template_name="ucc_test",
-                cmd_args=UCCCmdArgs(),
+                cmd_args=UCCCmdArgs(docker_image_url="url://fake/ucc"),
             ),
             test_template=TestTemplate(system=slurm_system),
         ),

conf/common/test_scenario/ucc_test.toml

@@ -33,7 +33,7 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> UCCTestSlurmCommandGenS
         "cmd_args_data, extra_cmd_args, expected_command",
         [
             (
-                {"collective": "allgather", "b": 8, "e": "256M"},
+                {"collective": "allgather", "b": 8, "e": "256M", "docker_image_url": "url://fake/ucc"},
                 {"--max-steps": "100"},
                 [
                     "/opt/hpcx/ucc/bin/ucc_perftest",
@@ -46,7 +46,7 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> UCCTestSlurmCommandGenS
                 ],
             ),
             (
-                {"collective": "allreduce", "b": 4, "e": "8M"},
+                {"collective": "allreduce", "b": 4, "e": "8M", "docker_image_url": "url://fake/ucc"},
                 {},
                 [
                     "/opt/hpcx/ucc/bin/ucc_perftest",
@@ -68,6 +68,7 @@ def test_generate_test_command(
         expected_command: list[str],
     ) -> None:
         ucc_cmd_args = UCCCmdArgs(
+            docker_image_url=cmd_args_data["docker_image_url"],
             collective=cmd_args_data["collective"],
             b=cmd_args_data["b"],
             e=cmd_args_data.get("e", "8M"),

src/cloudai/workloads/ucc_test/ucc.py

@@ -278,7 +278,12 @@ def test_req(request, slurm_system: SlurmSystem, partial_tr: partial[TestRun]) -
             partial_tr,
             slurm_system,
             "ucc",
-            UCCTestDefinition(name="ucc", description="ucc", test_template_name="ucc", cmd_args=UCCCmdArgs()),
+            UCCTestDefinition(
+                name="ucc",
+                description="ucc",
+                test_template_name="ucc",
+                cmd_args=UCCCmdArgs(docker_image_url="nvcr.io/nvidia/pytorch:24.02-py3"),
+            ),
             UCCTestSlurmCommandGenStrategy,
         ),
         "nccl": lambda: create_test_run(

tests/report_generation_strategy/test_ucc_report_generation_strategy.py

@@ -104,7 +104,12 @@ def test_chakra_docker_image_is_required():
 @pytest.mark.parametrize(
     "test",
     [
-        UCCTestDefinition(name="ucc", description="desc", test_template_name="ucc", cmd_args=UCCCmdArgs()),
+        UCCTestDefinition(
+            name="ucc",
+            description="desc",
+            test_template_name="ucc",
+            cmd_args=UCCCmdArgs(docker_image_url="fake://url/ucc"),
+        ),
         NCCLTestDefinition(name="nccl", description="desc", test_template_name="nccl", cmd_args=NCCLCmdArgs()),
         GPTTestDefinition(
             name="gpt",