Skip to content

Fix #1267: [torchx/schedulers] Add support for kueue_scheduler#1276

Open
JiwaniZakir wants to merge 1 commit intometa-pytorch:mainfrom
JiwaniZakir:fix/1267-torchx-schedulers-add-support-for-kueue
Open

Fix #1267: [torchx/schedulers] Add support for kueue_scheduler#1276
JiwaniZakir wants to merge 1 commit intometa-pytorch:mainfrom
JiwaniZakir:fix/1267-torchx-schedulers-add-support-for-kueue

Conversation

@JiwaniZakir
Copy link
Copy Markdown

Closes #1267

Adds KueueScheduler as a namespace-package plugin under torchx_plugins/schedulers/kueue.py, implementing the plugin architecture described in the issue. KueueScheduler subclasses KubernetesScheduler and overrides _submit_dryrun to inject the kueue.x-k8s.io/queue-name label into the job resource metadata before submission. The queue name is read from cfg.get("queue"), defaulting to "default" if unset. The factory function kueue_scheduler is registered via @register.scheduler(name="kueue"), making it auto-discoverable at runtime when the wheel containing torchx_plugins/ is installed.

Test plan:
Two unit tests added to torchx/schedulers/test/kubernetes_scheduler_test.py under KueueSchedulerTest:

  • test_factory_returns_kueue_scheduler: calls kueue_scheduler(session_name="test") and asserts the returned object is a KueueScheduler instance with session_name == "test".
  • test_queue_label_added_to_metadata: patches KubernetesScheduler._submit_dryrun with a MagicMock, invokes KueueScheduler(session_name="test")._submit_dryrun with Opts(queue="my-kueue-queue"), and asserts that resource["metadata"]["labels"]["kueue.x-k8s.io/queue-name"] == "my-kueue-queue" and that the original dryrun_info object is returned unmodified.

This PR was created with AI assistance (Claude). The changes were reviewed by quality gates and a critic model before submission.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[torchx/schedulers] Add support for kueue_scheduler

2 participants