Hi, I am trying to assign GPU (Nvidia MIG - Multi-Instance GPU) on flyte. However, I cannot see how I assign it on each tasks in flyte (there were no mentions in flyte docs). Does anyone have a solution to it?

Hi, @gifted-house-14547! You can assign GPU to a task as follows:

@task(requests=Resources(gpu=1, mem="3Gi", storage="1Gi"))
def pytorch_mnist_task(hp: Hyperparameters) -> TrainingOutputs:
   ...

@tall-lock-23197 Ah, the point I want to ask is flyte recognize "nvidia MIG (splitted instances into one gpu)" when I assign the way you mentioned. the solution you said recieves integer parameter, not gpu instance name. I want to try to use MIG on flyte.

You can configure it here by following this guide.

@tall-lock-23197 I have also read

values.yaml

file & guide docs, but these does not provide specific answers of MIG support. I think I have waited answers from those who you tagged

@gifted-house-14547 As per nvidia’s documentation, You will get different GPU in MIG with different name Like this, You can add them in config as mention here

kubectl describe node
...
Capacity:
<http://nvidia.com/mig-1g.5gb|nvidia.com/mig-1g.5gb>:   1
<http://nvidia.com/mig-2g.10gb|nvidia.com/mig-2g.10gb>:  1
<http://nvidia.com/mig-3g.20gb|nvidia.com/mig-3g.20gb>:  1

This is the global config for all task, Flyte will add taints and tolerations to all the task if they require GPU. Currently only Pod Task has option to configure taints and tolerations, Flyte core task doesn’t have option to configure the taints and tolerations, it will always use the default GPU. cc: @hallowed-mouse-14616 @high-park-82026 Please correct me if i am wrong.

@tall-lock-23197 I think I got a solution to assign MIG on flyte. flyte provides Integrations with k8s plugin, which can define pod spec in the flyte codes. https://docs.flyte.org/projects/cookbook/en/latest/auto/integrations/kubernetes/pod/pod.html

Nice! Is that working? If so, could you share with us the pod spec you used?

@tall-lock-23197 sure I will share the result soon! I think flyte docs should mention about MIG strategy if I succeeded.

Oh yes! Thank you.

Yes that would be awesome