Add DDLB workload

conf/common/test/ddlb_test.toml

@@ -0,0 +1,22 @@
+# 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 = "ddlb_test"
+description = "DDLB test configuration"
+test_template_name = "DDLBTest"
+
+[cmd_args]
+docker_image_url = "gitlab-master.nvidia.com/nsarkauskas/ddlb:latest"

conf/common/test_scenario/ddlb_test.toml

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# 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");
+# 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 = "ddlb-test"
+
+[[Tests]]
+id = "Tests.ddlb"
+test_name = "ddlb_test"
+num_nodes = 1
+time_limit = "00:30:00"

src/cloudai/registration.py

@@ -92,6 +92,10 @@ def register_all():
         NcclTestRunAIJsonGenStrategy,
         NcclTestSlurmCommandGenStrategy,
     )
+    from cloudai.workloads.ddlb import (
+        DDLBTestDefinition,
+        DDLBTestSlurmCommandGenStrategy,
+    )
     from cloudai.workloads.nemo_launcher import (
         NeMoLauncherGradingStrategy,
         NeMoLauncherReportGenerationStrategy,
@@ -155,6 +159,7 @@ def register_all():
 
     Registry().add_command_gen_strategy(SlurmSystem, MegatronRunTestDefinition, MegatronRunSlurmCommandGenStrategy)
     Registry().add_command_gen_strategy(SlurmSystem, NCCLTestDefinition, NcclTestSlurmCommandGenStrategy)
+    Registry().add_command_gen_strategy(SlurmSystem, DDLBTestDefinition, DDLBTestSlurmCommandGenStrategy)
     Registry().add_strategy(GradingStrategy, [SlurmSystem], [SleepTestDefinition], SleepGradingStrategy)
 
     Registry().add_command_gen_strategy(SlurmSystem, NeMoLauncherTestDefinition, NeMoLauncherSlurmCommandGenStrategy)
@@ -204,6 +209,7 @@ def register_all():
 
     Registry().add_test_definition("UCCTest", UCCTestDefinition)
     Registry().add_test_definition("NcclTest", NCCLTestDefinition)
+    Registry().add_test_definition("DDLBTest", DDLBTestDefinition)
     Registry().add_test_definition("ChakraReplay", ChakraReplayTestDefinition)
     Registry().add_test_definition("Sleep", SleepTestDefinition)
     Registry().add_test_definition("NeMoLauncher", NeMoLauncherTestDefinition)

src/cloudai/workloads/ddlb/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+
+from .ddlb import DDLBCmdArgs, DDLBTestDefinition
+from .slurm_command_gen_strategy import DDLBTestSlurmCommandGenStrategy
+
+__all__ = [
+    "DDLBCmdArgs",
+    "DDLBTestDefinition",
+    "DDLBTestSlurmCommandGenStrategy",
+]

src/cloudai/workloads/ddlb/ddlb.py

@@ -0,0 +1,94 @@
+# 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.
+
+from typing import Literal, Optional, Union
+
+from cloudai.core import DockerImage, Installable, JobStatusResult, TestRun
+from cloudai.models.workload import CmdArgs, TestDefinition
+
+
+class DDLBCmdArgs(CmdArgs):
+    """DDLB test command arguments."""
+
+    docker_image_url: str
+
+
+class DDLBTestDefinition(TestDefinition):
+    """Test object for DDLB."""
+
+    cmd_args: DDLBCmdArgs
+    _docker_image: Optional[DockerImage] = None
+
+    @property
+    def extra_args_str(self) -> str:
+        parts = []
+        for k, v in self.extra_cmd_args.items():
+            parts.append(f"{k} {v}" if v else k)
+        return " ".join(parts)
+
+    @property
+    def docker_image(self) -> DockerImage:
+        if not self._docker_image:
+            self._docker_image = DockerImage(url=self.cmd_args.docker_image_url)
+        return self._docker_image
+
+    @property
+    def installables(self) -> list[Installable]:
+        return [self.docker_image]
+
+    def was_run_successful(self, tr: TestRun) -> JobStatusResult:
+        stdout_path = tr.output_path / "stdout.txt"
+        if stdout_path.is_file():
+            with stdout_path.open("r") as file:
+                content = file.read()
+
+                # Check for specific error patterns
+                if "Error" in content:
+                    return JobStatusResult(
+                        is_successful=False,
+                        error_message=(
+                            f"DDLB test failure detected in {stdout_path}. "
+                            "Possible reasons include network errors or remote process exits. "
+                            "Please review the DDLB test output and errors in the file first. "
+                            "If the issue persists, contact the system administrator."
+                        ),
+                    )
+
+                # Identify missing success indicators
+                if "Benchmark Results" not in content:
+                    error_message = (
+                        f"Missing success indicators in {stdout_path}: 'Benchmark Results'. "
+                        "These keywords are expected to be present in stdout.txt, usually towards the end of the file. "
+                        "Please review the DDLB test output and errors in the file. "
+                        "Ensure the DDLB test ran to completion. You can run the generated sbatch script manually "
+                        f"and check if {stdout_path} is created and contains the expected keywords. "
+                        "If the issue persists, contact the system administrator."
+                    )
+
+                    return JobStatusResult(is_successful=False, error_message=error_message)
+
+                return JobStatusResult(is_successful=True)
+
+        return JobStatusResult(
+            is_successful=False,
+            error_message=(
+                f"stdout.txt file not found in the specified output directory {tr.output_path}. "
+                "This file is expected to be created as a result of the DDLB test run. "
+                "Please ensure the DDLB test was executed properly and that stdout.txt is generated. "
+                f"You can run the generated DDLB test command manually and verify the creation of {stdout_path}. "
+                "If the issue persists, contact the system administrator."
+            ),
+        )

src/cloudai/workloads/ddlb/slurm_command_gen_strategy.py

@@ -0,0 +1,40 @@
+# 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.
+
+from typing import List, cast
+
+from cloudai.systems.slurm import SlurmCommandGenStrategy
+
+from .ddlb import DDLBTestDefinition
+
+
+class DDLBTestSlurmCommandGenStrategy(SlurmCommandGenStrategy):
+    """Command generation strategy for DDLB tests on Slurm systems."""
+
+    def _container_mounts(self) -> List[str]:
+        return []
+
+    def image_path(self) -> str | None:
+        tdef: DDLBTestDefinition = cast(DDLBTestDefinition, self.test_run.test.test_definition)
+        return str(tdef.docker_image.installed_path)
+
+    def generate_test_command(self) -> List[str]:
+        tdef: DDLBTestDefinition = cast(DDLBTestDefinition, self.test_run.test.test_definition)
+        return ["python scripts/run_benchmark.py"]
+
+    def gen_srun_success_check(self) -> str:
+        output_file = self.test_run.output_path / "stdout.txt"
+        return f'grep -q "Benchmark Results" {output_file} && echo 1 || echo 0'