Hello everyone, I was wondering if/how we can specify the type of GPU or the memory for the GPU in metaflow? It currently selects GPUs with 15GB of memory but I would need one with 24B. Is there a way to choose one rather than the other? Thanks!

are you using

@batch

or

@kubernetes

?

None, I am using

--with kubernetes

and

@ressource

@stale-cartoon-78113 how are GPUs currently configured in your cluster?

You can specify the gpu count and gpu vendor via the cli or decorators.

$ python3 myflow.py run --with kubernetes:gpu=1,gpu_vendor=nvidia

gpu

is the count of gpus and

gpu_vendor

has to be either

nvidia

or

amd

Internally, metaflow takes these options and adds them to the

resources

section in the Kubernetes pod yaml. Do you know if Kubernetes offers a way to specify the

memory for the GPU

in it's pod spec?

I was able to do what I wanted specifying the EC2 instance I want, thank you for your help!

Just out of curiosity, what did you end up using? This might help someone else trying to solve the same problem.

I am also interested in how this is possible when you have different GPUs in a cluster and want to select a certain kind for a job

We can wait for @stale-cartoon-78113 to tell us what worked for him specifically. Generally speaking, there could be two ways of doing this: • If there are gpu of different kinds (

nvidia

vs

amd

), you can use the appropriate

gpu_vendor

option as mentioned earlier in the thread. • If the gpus are connected to nodes with specific labels, you could also use the

node_selector

option. And FWIW, if you want the nodes to be automatically labelled, there are some projects like this (which you'd have to run on your own). If these don't seem to fit the specific needs, let me know and we can brainstorm about other options.

@User Ended up specifying the EC2 instance with the target GPU and seems to work