Merge pull request #603 from NVIDIA/am/reg-bash

Register BashCmd workload

Author: amaslenn

src/cloudai/registration.py

@@ -40,6 +40,7 @@ def register_all():
     )
 
     # Import all workloads and their strategies
+    from cloudai.workloads.bash_cmd.bash_cmd import BashCmdCommandGenStrategy, BashCmdTestDefinition
     from cloudai.workloads.chakra_replay import (
         ChakraReplayGradingStrategy,
         ChakraReplayReportGenerationStrategy,
@@ -156,6 +157,7 @@ def register_all():
     )
 
     Registry().add_command_gen_strategy(SlurmSystem, AIDynamoTestDefinition, AIDynamoSlurmCommandGenStrategy)
+    Registry().add_command_gen_strategy(SlurmSystem, BashCmdTestDefinition, BashCmdCommandGenStrategy)
 
     Registry().add_installer("slurm", SlurmInstaller)
     Registry().add_installer("standalone", StandaloneInstaller)
@@ -183,6 +185,7 @@ def register_all():
     Registry().add_test_definition("TritonInference", TritonInferenceTestDefinition)
     Registry().add_test_definition("NIXLBench", NIXLBenchTestDefinition)
     Registry().add_test_definition("AIDynamo", AIDynamoTestDefinition)
+    Registry().add_test_definition("BashCmd", BashCmdTestDefinition)
 
     Registry().add_agent("grid_search", GridSearchAgent)
 

tests/test_init.py

@@ -26,6 +26,7 @@
     AIDynamoSlurmCommandGenStrategy,
     AIDynamoTestDefinition,
 )
+from cloudai.workloads.bash_cmd import BashCmdCommandGenStrategy, BashCmdTestDefinition
 from cloudai.workloads.chakra_replay import (
     ChakraReplayGradingStrategy,
     ChakraReplaySlurmCommandGenStrategy,
@@ -117,6 +118,7 @@ def test_runners():
     (SlurmSystem, TritonInferenceTestDefinition): TritonInferenceSlurmCommandGenStrategy,
     (SlurmSystem, NIXLBenchTestDefinition): NIXLBenchSlurmCommandGenStrategy,
     (SlurmSystem, AIDynamoTestDefinition): AIDynamoSlurmCommandGenStrategy,
+    (SlurmSystem, BashCmdTestDefinition): BashCmdCommandGenStrategy,
 }
 ALL_STRATEGIES = {
     (GradingStrategy, SlurmSystem, ChakraReplayTestDefinition): ChakraReplayGradingStrategy,
@@ -175,7 +177,7 @@ def test_installers():
 
 def test_definitions():
     test_defs = Registry().test_definitions_map
-    assert len(test_defs) == 14
+    assert len(test_defs) == 15
     for tdef in [
         ("UCCTest", UCCTestDefinition),
         ("NcclTest", NCCLTestDefinition),
@@ -191,6 +193,7 @@ def test_definitions():
         ("TritonInference", TritonInferenceTestDefinition),
         ("NIXLBench", NIXLBenchTestDefinition),
         ("AIDynamo", AIDynamoTestDefinition),
+        ("BashCmd", BashCmdTestDefinition),
     ]:
         assert test_defs[tdef[0]] == tdef[1]