Merge pull request #9 from TaekyungHeo/allow-override

Allow overriding of runner mappings in Runner class

src/cloudai/runner/core/runner.py

@@ -61,8 +61,6 @@ def register(cls, system_type: str) -> Callable:
         """
 
         def decorator(runner_class: Type[BaseRunner]) -> Type[BaseRunner]:
-            if system_type in cls._runners:
-                raise KeyError(f"Runner for {system_type} already registered.")
             cls._runners[system_type] = runner_class
             return runner_class
 

tests/runner/core/test_runner.py

@@ -0,0 +1,21 @@
+from unittest.mock import patch
+
+from cloudai.runner.core import BaseRunner, Runner
+
+
+def test_register_multiple_runners():
+    """
+    Test registering multiple different runners for the same type to ensure that
+    only the last registered runner is kept.
+    """
+    with patch.dict("cloudai.runner.Runner._runners", clear=True):
+
+        @Runner.register("slurm")
+        class FirstSlurmRunner(BaseRunner):
+            pass
+
+        @Runner.register("slurm")
+        class SecondSlurmRunner(BaseRunner):
+            pass
+
+        assert Runner._runners["slurm"] is SecondSlurmRunner