I want to use Kueue as the batch scheduler for Fly...
# flyte-support
f
I want to use Kueue as the batch scheduler for Flyte with PyTorch. The documentation mentions Yunikorn and K8s scheduler plugins with co-scheduling. I don't see any reason why Kueue won't work. https://kueue.sigs.k8s.io/docs/tasks/run/kubeflow/pytorchjobs/ There is a slight difference though. In order to send a job to Kueue, one needs a PyTorchJob metadata.label that specifies the queue. However, the PyTorch plugin does not expose a way to customize the labels. This capability should be relatively simple to add. I am willing to contribute this to the PyTorch plugin backend and flytekit. Has anyone used Kueue before with Flyte?
f
Yes at union we are experimenting with it and I know cradle Is using it
f
@freezing-airport-6809 Looks like in Flyte there is no easy way to modify the metadata.labels of the PyTorchJob submitted by the plugin. This capability needs to be added. Am I missing something?
I've created https://github.com/flyteorg/flyte/pull/5677 for Flyte backend side changes I wonder if anyone can take a look, cc: @high-accountant-32689, @cool-lifeguard-49380, @freezing-airport-6809
👀 1
c
When specifying
--labels
for
pyflyte run
, they will only be set on the pod spec but not the pytorchjob?
f
Not on the PyTorch job from what I can tell reading the code, but will try it out.
It actually does!
Thanks @cool-lifeguard-49380. From my experiment, it does set it on both the Pod and the PyTorchJob.
👍 1
c
Perfect. When adding general logic like this to a plugin, I think it’s good to add it to all plugins so that the UX doesn’t diverge too much between different ones.