Skip to content

Let downstream executors opt in to JobGroup (#537) - #587

Open
CodersAcademy006 wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
CodersAcademy006:feat/537-jobgroup-executor-extension
Open

Let downstream executors opt in to JobGroup (#537)#587
CodersAcademy006 wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
CodersAcademy006:feat/537-jobgroup-executor-extension

Conversation

@CodersAcademy006

Copy link
Copy Markdown
Contributor

Closes #537.

Problem

JobGroup.__post_init__ gated on executor_type in self.SUPPORTED_EXECUTORS (nemo_run/run/job.py:269). That is an exact class-identity check against [SlurmExecutor, DockerExecutor, LocalExecutor], so two things failed:

  1. Any Executor defined outside nemo_run, such as nemo-skills' RayExecutor, could never construct a JobGroup. The only workarounds available downstream were mutating the class attribute or sniffing type(executor).__name__.
  2. Any subclass of a supported executor failed too. A plain class MySlurmExecutor(SlurmExecutor) was rejected even though it is a SlurmExecutor in every meaningful sense.

Change

@nicgupta-nvidia offered three designs on the issue and asked which would be accepted. This implements design 3, the classmethod on the base, since it is the one that needs no mutation of upstream state and shows up in IDEs and in the API docs:

  • Executor.supports_job_group() returns False on the base class. SlurmExecutor, DockerExecutor and LocalExecutor override it to True.
  • JobGroup.__post_init__ gates on executor_type.supports_job_group() or executor_type in self.SUPPORTED_EXECUTORS. The membership clause is kept deliberately, so anyone who already extended SUPPORTED_EXECUTORS in the wild keeps working, and the assertion message now names the rejected class and how to opt in.
  • The merge dispatch changed from executor_type == SlurmExecutor to issubclass(...). Without that, relaxing the assertion would let a SlurmExecutor subclass through and then silently drop it into the _merge = False path, which is a worse failure than the rejection it replaces.

What this does not do

As the issue itself notes, relaxing the assertion is necessary but not sufficient for a Ray-backed JobGroup to launch end to end: JobGroup.launch routes through EXECUTOR_MAPPING in torchx_backend/schedulers/api.py, which has no Ray entry, so get_executor_str still raises KeyError. That is a separate architectural question and is left alone here. This PR only makes the extension point real.

Tests

Four tests in test/run/test_job.py:

  • test_job_group_builtin_executors_opt_in asserts every class in SUPPORTED_EXECUTORS reports supports_job_group(), so the two lists cannot drift apart.
  • test_job_group_accepts_downstream_executor builds a JobGroup from an out-of-tree Executor subclass that opts in, and checks it takes the no-merge path with one executor per task.
  • test_job_group_rejects_executor_without_opt_in keeps the rejection for executors that do not opt in.
  • test_job_group_slurm_subclass_keeps_group_semantics pins the second half of the fix: a SlurmExecutor subclass merges and comes back with run_as_group set.

Verified locally: test/run/test_job.py 47 passed. Reverting only the nemo_run/ half fails 3 of the 4 new tests, so they do catch the regression. test/run plus test/core/execution match the main baseline exactly (24 pre-existing failures either side, from skypilot and kubeflow extras missing in my environment). ruff format --diff and ruff check clean.

@ko3n1g this one has been sitting since May 25 with no reviewer, and the reporter explicitly asked for a design verdict before sending code. If you would rather have design 1 or 2 from the issue, say so and I will swap it, the diff is small either way.

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 NVIDIA-NeMo#537

Signed-off-by: Srijan Upadhyay <srjnupadhyay@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

JobGroup.SUPPORTED_EXECUTORS is closed to downstream Executor subclasses

1 participant