Bring your towel and join one of the fastest growing data communities. Welcome to our second-generation open source orchestration platform, a completely rethought approach to dataflow automation.

Prefect Community

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 <@URG3ZTNH0>, 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"```


We also use this same approach for what we call "High Memory Workers." We have certain parts of our data science pipeline that need a large amount of RAM on Dask workers, e.g. some code that manipulates a large amount of data in a pandas dataframe, etc. (We're migrating to Dask dataframes to avoid this, but some of our legacy code isn't converted yet.)

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

Hi <@URRN7598X>, 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.

<@URRN7598X> 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... :slightly_smiling_face:) 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.