Fix success condition

Author: allkoow

conf/common/test/osu_test.toml

@@ -21,6 +21,6 @@ description = "OSU Benchmark example"
 [cmd_args]
 "docker_image_url" = "artifactory.nvidia.com/sw-nbu-swx-hpcx-docker-local/dlfw/pytorch:25.06-hpcx-v2.26-gcc-ubuntu24.04-cuda12-x86_64-latest-nightly"
 "benchmarks_dir" = "/opt/hpcx/ompi/tests/osu-micro-benchmarks"
-"benchmark" = ["osu_allreduce", "osu_allgather"]
+"benchmark" = "osu_allreduce"
 "iterations" = 10
 "message_size" = "1024"

conf/common/test_scenario/osu_test.toml

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 name = "osu_test_scenario"
+job_status_check = true
 
 [[Tests]]
 id = "Tests.1"

conf/osu/funk.toml

@@ -104,7 +104,7 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult:
                 ),
             )
 
-        if "# Size       Avg Latency(us)   Min Latency(us)   Max Latency(us)  Iterations" not in content:
+        if "# Size" not in content:
             return JobStatusResult(
                 is_successful=False,
                 error_message=(

doc/workloads/index.md

@@ -21,6 +21,16 @@
 from .osu_bench import OSUBenchCmdArgs, OSUBenchTestDefinition
 
 
+FULL_FLAG_UNSUPPORTED = [
+    "osu_latency", "osu_latency_mt", "osu_latency_mp", "osu_bw",
+    "osu_bibw", "osu_latency_persistent", "osu_bw_persistent",
+    "osu_bibw_persistent", "osu_multi_lat", "osu_mbw_mr",
+    "osu_put_latency", "osu_get_latency", "osu_acc_latency",
+    "osu_get_acc_latency", "osu_cas_latency", "osu_fop_latency",
+    "osu_put_bw", "osu_get_bw", "osu_put_bibw", "osu_init", "osu_hello"
+]
+
+
 class OSUBenchSlurmCommandGenStrategy(SlurmCommandGenStrategy):
     """Command generation strategy for OSU Benchmark test on Slurm systems."""
 
@@ -37,14 +47,25 @@ def generate_test_command(self) -> List[str]:
         binary = f"{args.benchmarks_dir}/{args.benchmark}"
         srun_command_parts = [binary]
 
-        general = {"docker_image_url", "location", "benchmark"}
+        general = {"docker_image_url", "benchmarks_dir", "benchmark"}
 
         for name, value in args.model_dump(exclude=general).items():
             if value is None:
                 continue
 
             flag = f"--{name.replace('_', '-')}"
-            srun_command_parts.append(f"{flag} {value}")
+
+            if isinstance(value, bool) and value:
+                argument = flag
+            else:
+                argument = f"{flag} {value}"
+
+            # Some benchmarks don't support the full flag; suppress it
+            # to avoid errors.
+            if name == "full" and args.benchmark in FULL_FLAG_UNSUPPORTED:
+                continue
+
+            srun_command_parts.append(argument)
 
         if self.test_run.test.extra_cmd_args:
             srun_command_parts.append(self.test_run.test.extra_args_str)

src/cloudai/workloads/osu_bench/osu_bench.py

@@ -10,8 +10,8 @@
 
 export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1)
 
-srun --export=ALL --mpi=pmix -N1 --container-image=artifactory.nvidia.com/sw-nbu-swx-hpcx-docker-local/dlfw/pytorch:25.06-hpcx-v2.26-gcc-ubuntu24.04-cuda12-x86_64-latest-nightly --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
+srun --export=ALL --mpi=pmix --container-image=nvcr.io#nvidia/pytorch:24.02-py3 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
 
-srun --export=ALL --mpi=pmix -N1 --container-image=artifactory.nvidia.com/sw-nbu-swx-hpcx-docker-local/dlfw/pytorch:25.06-hpcx-v2.26-gcc-ubuntu24.04-cuda12-x86_64-latest-nightly --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
+srun --export=ALL --mpi=pmix --container-image=nvcr.io#nvidia/pytorch:24.02-py3 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
 
-srun --export=ALL --mpi=pmix -N1 --container-image=artifactory.nvidia.com/sw-nbu-swx-hpcx-docker-local/dlfw/pytorch:25.06-hpcx-v2.26-gcc-ubuntu24.04-cuda12-x86_64-latest-nightly --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash -c "source __OUTPUT_DIR__/output/env_vars.sh; /opt/hpcx/ompi/tests/osu-micro-benchmarks/osu_allreduce -m 1024 -i 10 -f"
+srun --export=ALL --mpi=pmix --container-image=nvcr.io#nvidia/pytorch:24.02-py3 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash -c "source __OUTPUT_DIR__/output/env_vars.sh; /opt/hpcx/ompi/tests/osu-micro-benchmarks/osu_allreduce --message-size 1024 --iterations 10 --full"

src/cloudai/workloads/osu_bench/slurm_command_gen_strategy.py

@@ -314,8 +314,8 @@ def test_req(request, slurm_system: SlurmSystem, partial_tr: partial[TestRun]) -
                 description="osu-bench",
                 test_template_name="osu-bench",
                 cmd_args=OSUBenchCmdArgs(
-                    docker_image_url="artifactory.nvidia.com/sw-nbu-swx-hpcx-docker-local/dlfw/pytorch:25.06-hpcx-v2.26-gcc-ubuntu24.04-cuda12-x86_64-latest-nightly",
-                    location="/opt/hpcx/ompi/tests/osu-micro-benchmarks",
+                    docker_image_url="nvcr.io#nvidia/pytorch:24.02-py3",
+                    benchmarks_dir="/opt/hpcx/ompi/tests/osu-micro-benchmarks",
                     benchmark="osu_allreduce",
                     iterations=10,
                     message_size="1024",