Hello! Wondering if there’s a timeframe for task affinity — would love to be able to specify my model training task on a k8s resource with GPU.

Hi @Jake Schmidt, we already take advantage of this and it works great! Docs are here: https://docs.prefect.io/api/unreleased/engine/executors.html#daskexecutor and some code snippets from how we use it:

@task(tags=["dask-resource:GPU=1"])
def task_that_uses_gpu():

and then the relevant YAML section for our k8s GPU workers:

containers:
        - args:
            - dask-worker
            - dask-scheduler:8786
            - --resources
            - "GPU=1"

@Joe Schmid this is intriguing, sorry for my lack of knowledge where is that YAML placed? / More info on that file?

Hi @Jackson Maxfield Brown, I should have explained more -- in this scenario, we are creating our own long-running Dask cluster using dask-kubernetes and running in AWS. The YAML snippet that I showed is from a Kubernetes "Deployment" specification for Dask workers running on machines with a GPU and the snippet shows starting the Dask workers with a parameter called

resources

and passing that parameter the value

GPU=1

Prefect can then use task tagging (the other snippet I showed) to route tasks only to Dask workers that have appropriate resources. It's really powerful and has been very successful for us.

Ahhhhh I see. Makes sense though and this is useful info. We have an internal SLURM cluster that splits our CPU and GPU nodes so to get this to work I think we would have to set up a dask scheduler on the head node which is probably not ideal but a Dask Kubernetes configuration like that would be great.

@Jackson Maxfield Brown Yeah, the Dask resources / task affinity in Prefect is really cool. I haven't done it on SLURM (though I have a friend who might be able to get me access to a supercomputer to try this out... 🙂) but it looks like from these Dask docs it should be doable: https://jobqueue.dask.org/en/latest/examples.html#slurm-deployment-providing-additional-arguments-to-the-dask-workers (See the bottom example)

Hmmm I think you're right that this is possible. We would just be passing the "queue" to spawn the worker with. Huh. Will have to give this a try tomorrow.