Let downstream executors opt in to JobGroup (#537) - #587
Open
CodersAcademy006 wants to merge 1 commit into
Open
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #537.
Problem
JobGroup.__post_init__gated onexecutor_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:Executordefined outside nemo_run, such as nemo-skills'RayExecutor, could never construct aJobGroup. The only workarounds available downstream were mutating the class attribute or sniffingtype(executor).__name__.class MySlurmExecutor(SlurmExecutor)was rejected even though it is aSlurmExecutorin 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()returnsFalseon the base class.SlurmExecutor,DockerExecutorandLocalExecutoroverride it toTrue.JobGroup.__post_init__gates onexecutor_type.supports_job_group() or executor_type in self.SUPPORTED_EXECUTORS. The membership clause is kept deliberately, so anyone who already extendedSUPPORTED_EXECUTORSin the wild keeps working, and the assertion message now names the rejected class and how to opt in.executor_type == SlurmExecutortoissubclass(...). Without that, relaxing the assertion would let aSlurmExecutorsubclass through and then silently drop it into the_merge = Falsepath, 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
JobGroupto launch end to end:JobGroup.launchroutes throughEXECUTOR_MAPPINGintorchx_backend/schedulers/api.py, which has no Ray entry, soget_executor_strstill raisesKeyError. 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_inasserts every class inSUPPORTED_EXECUTORSreportssupports_job_group(), so the two lists cannot drift apart.test_job_group_accepts_downstream_executorbuilds aJobGroupfrom an out-of-treeExecutorsubclass that opts in, and checks it takes the no-merge path with one executor per task.test_job_group_rejects_executor_without_opt_inkeeps the rejection for executors that do not opt in.test_job_group_slurm_subclass_keeps_group_semanticspins the second half of the fix: aSlurmExecutorsubclass merges and comes back withrun_as_groupset.Verified locally:
test/run/test_job.py47 passed. Reverting only thenemo_run/half fails 3 of the 4 new tests, so they do catch the regression.test/runplustest/core/executionmatch themainbaseline exactly (24 pre-existing failures either side, from skypilot and kubeflow extras missing in my environment).ruff format --diffandruff checkclean.@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.