From d87214b08f044bfe4fdfbd643ba911d181880d95 Mon Sep 17 00:00:00 2001 From: Srijan Upadhyay Date: Wed, 19 Aug 2026 13:23:17 +0530 Subject: [PATCH] feat(run): let downstream executors opt in to JobGroup JobGroup.__post_init__ gated on `executor_type in SUPPORTED_EXECUTORS`, which is an exact class-identity check, so any Executor defined outside nemo_run was rejected and so was any subclass of a supported executor. Downstream packages had no extension point short of mutating the class attribute or sniffing type names. Add Executor.supports_job_group(), False on the base and True on SlurmExecutor, DockerExecutor and LocalExecutor, and gate JobGroup on it. The SUPPORTED_EXECUTORS membership check stays as a fallback so anything that already appends to that list keeps working. The merge dispatch now uses issubclass, so a SlurmExecutor or DockerExecutor subclass keeps the group merge semantics instead of silently taking the no-merge path. Closes #537 Signed-off-by: Srijan Upadhyay --- nemo_run/core/execution/base.py | 8 +++++ nemo_run/core/execution/docker.py | 4 +++ nemo_run/core/execution/local.py | 4 +++ nemo_run/core/execution/slurm.py | 4 +++ nemo_run/run/job.py | 9 +++-- test/run/test_job.py | 57 +++++++++++++++++++++++++++++++ 6 files changed, 83 insertions(+), 3 deletions(-) diff --git a/nemo_run/core/execution/base.py b/nemo_run/core/execution/base.py index cfc02a2c..2a87ed50 100644 --- a/nemo_run/core/execution/base.py +++ b/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()) diff --git a/nemo_run/core/execution/docker.py b/nemo_run/core/execution/docker.py index 14d9774c..06608a26 100644 --- a/nemo_run/core/execution/docker.py +++ b/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 diff --git a/nemo_run/core/execution/local.py b/nemo_run/core/execution/local.py index e8954bae..d79ebe6d 100644 --- a/nemo_run/core/execution/local.py +++ b/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, diff --git a/nemo_run/core/execution/slurm.py b/nemo_run/core/execution/slurm.py index 35fa9ee3..61e364e6 100644 --- a/nemo_run/core/execution/slurm.py +++ b/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 diff --git a/nemo_run/run/job.py b/nemo_run/run/job.py index 35a6cf4e..0219e16f 100644 --- a/nemo_run/run/job.py +++ b/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) diff --git a/test/run/test_job.py b/test/run/test_job.py index e4684b0c..7ac273a2 100644 --- a/test/run/test_job.py +++ b/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(