Merge pull request #188 from TaekyungHeo/k8s-nccl-report-bugfix

Fix bug in generating NCCL test reports on K8s
NVIDIA · Sep 10, 2024 · 03c23d6 · 03c23d6
2 parents d96c0b6 + 9b57912
commit 03c23d6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/cloudai/__init__.py b/src/cloudai/__init__.py
@@ -103,7 +103,9 @@
 
 Registry().add_strategy(InstallStrategy, [SlurmSystem], [NcclTest], NcclTestSlurmInstallStrategy)
 Registry().add_strategy(InstallStrategy, [SlurmSystem], [NeMoLauncher], NeMoLauncherSlurmInstallStrategy)
-Registry().add_strategy(ReportGenerationStrategy, [SlurmSystem], [NcclTest], NcclTestReportGenerationStrategy)
+Registry().add_strategy(
+    ReportGenerationStrategy, [SlurmSystem, KubernetesSystem], [NcclTest], NcclTestReportGenerationStrategy
+)
 Registry().add_strategy(CommandGenStrategy, [StandaloneSystem], [Sleep], SleepStandaloneCommandGenStrategy)
 Registry().add_strategy(CommandGenStrategy, [SlurmSystem], [Sleep], SleepSlurmCommandGenStrategy)
 Registry().add_strategy(JsonGenStrategy, [KubernetesSystem], [Sleep], SleepKubernetesJsonGenStrategy)

diff --git a/tests/test_init.py b/tests/test_init.py
@@ -120,6 +120,7 @@ def test_runners():
         ((ReportGenerationStrategy, SlurmSystem, ChakraReplay), ChakraReplayReportGenerationStrategy),
         ((ReportGenerationStrategy, SlurmSystem, JaxToolbox), JaxToolboxReportGenerationStrategy),
         ((ReportGenerationStrategy, SlurmSystem, NcclTest), NcclTestReportGenerationStrategy),
+        ((ReportGenerationStrategy, KubernetesSystem, NcclTest), NcclTestReportGenerationStrategy),
         ((ReportGenerationStrategy, SlurmSystem, NeMoLauncher), NeMoLauncherReportGenerationStrategy),
         ((ReportGenerationStrategy, SlurmSystem, Sleep), SleepReportGenerationStrategy),
         ((ReportGenerationStrategy, SlurmSystem, UCCTest), UCCTestReportGenerationStrategy),