Let downstream executors opt in to JobGroup (#537)

nemo_run/core/execution/base.py

@@ -102,6 +102,14 @@ class Executor(ConfigurableMixin):
     def info(self) -> str:
         return self.__class__.__qualname__
 
+    @classmethod
+    def supports_job_group(cls) -> bool:
+        """Whether this executor can back a :class:`~nemo_run.run.job.JobGroup`.
+
+        Executors defined outside nemo_run opt in by overriding this to return True.
+        """
+        return False
+
     def clone(self) -> Self:
         return fdl.build(self.to_config())
 

nemo_run/core/execution/docker.py

@@ -138,6 +138,10 @@ class DockerExecutor(Executor):
     run_as_group: bool = field(init=False, default=False)
     resource_group: list["DockerExecutor"] = field(init=False, default_factory=list)
 
+    @classmethod
+    def supports_job_group(cls) -> bool:
+        return True
+
     @classmethod
     def merge(
         cls: Type["DockerExecutor"], executors: list["DockerExecutor"], num_tasks: int

nemo_run/core/execution/local.py

@@ -39,6 +39,10 @@ class LocalExecutor(Executor):
     ntasks_per_node: int = 1
     nodes: int = 1
 
+    @classmethod
+    def supports_job_group(cls) -> bool:
+        return True
+
     def assign(
         self,
         exp_id: str,

nemo_run/core/execution/slurm.py

@@ -356,6 +356,10 @@ class ResourceRequest:
     resource_group: list[ResourceRequest] = field(init=False, default_factory=list)
     run_as_group: bool = field(init=False, default=False)
 
+    @classmethod
+    def supports_job_group(cls) -> bool:
+        return True
+
     @classmethod
     def merge(
         cls: Type["SlurmExecutor"], executors: list["SlurmExecutor"], num_tasks: int

nemo_run/run/job.py

@@ -266,13 +266,16 @@ def __post_init__(self):
 
         assert len(executor_types) == 1, "All executors must be of the same type."
         executor_type = list(executor_types)[0]
-        assert executor_type in self.SUPPORTED_EXECUTORS, "Unsupported executor type."
-        if executor_type == SlurmExecutor:
+        assert executor_type.supports_job_group() or executor_type in self.SUPPORTED_EXECUTORS, (
+            f"Unsupported executor type {executor_type.__name__}. Executors opt in to JobGroup "
+            "by overriding Executor.supports_job_group()."
+        )
+        if issubclass(executor_type, SlurmExecutor):
             self._merge = True
             self.executors = SlurmExecutor.merge(
                 cast(list[SlurmExecutor], executors), num_tasks=len(self.tasks)
             )
-        elif executor_type == DockerExecutor:
+        elif issubclass(executor_type, DockerExecutor):
             self._merge = True
             self.executors = DockerExecutor.merge(
                 cast(list[DockerExecutor], executors), num_tasks=len(self.tasks)

test/run/test_job.py

@@ -13,12 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import dataclass
 from unittest.mock import MagicMock, patch
 
 import pytest
 from torchx.specs.api import AppState
 
 from nemo_run.config import Partial, Script
+from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.docker import DockerExecutor
 from nemo_run.core.execution.slurm import SlurmExecutor
 from nemo_run.run.job import Job, JobGroup
@@ -384,6 +386,61 @@ def test_job_group_init_mixed_executor_types(simple_task):
         )
 
 
+def test_job_group_builtin_executors_opt_in():
+    # SUPPORTED_EXECUTORS and supports_job_group() must not drift apart.
+    for executor_type in JobGroup.SUPPORTED_EXECUTORS:
+        assert executor_type.supports_job_group(), executor_type.__name__
+
+
+def test_job_group_accepts_downstream_executor(simple_task):
+    @dataclass(kw_only=True)
+    class CustomExecutor(Executor):
+        @classmethod
+        def supports_job_group(cls) -> bool:
+            return True
+
+    executor = CustomExecutor(job_dir="/tmp/custom")
+    job_group = JobGroup(
+        id="test-group",
+        tasks=[simple_task, simple_task],
+        executors=executor,
+    )
+
+    assert not job_group._merge
+    assert job_group.executors == [executor, executor]
+
+
+def test_job_group_rejects_executor_without_opt_in(simple_task):
+    @dataclass(kw_only=True)
+    class UnsupportedExecutor(Executor):
+        pass
+
+    with pytest.raises(AssertionError, match="Unsupported executor type"):
+        JobGroup(
+            id="test-group",
+            tasks=[simple_task],
+            executors=UnsupportedExecutor(job_dir="/tmp/unsupported"),
+        )
+
+
+def test_job_group_slurm_subclass_keeps_group_semantics(simple_task):
+    @dataclass(kw_only=True)
+    class CustomSlurmExecutor(SlurmExecutor):
+        pass
+
+    job_group = JobGroup(
+        id="test-group",
+        tasks=[simple_task, simple_task],
+        executors=CustomSlurmExecutor(
+            account="test_account", partition="test", job_dir="/tmp/test"
+        ),
+    )
+
+    assert job_group._merge
+    assert isinstance(job_group.executors, CustomSlurmExecutor)
+    assert job_group.executors.run_as_group
+
+
 def test_job_group_properties(simple_task, docker_executor):
     # Mock the property behavior directly
     job_group = JobGroup(