[Extended Resources] GPU Accelerators

[jeevb]

TL;DR

Adds support for specifying extended resources (e.g. accelerator type - gpu device, partition size etc.) to allocate to a task.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

Added support for populating ExtendedResources from the @task decorator or node .with_overrides(). Currently only supports the ability to specify the GPU accelerator to use for a specific task node. The following syntaxes are valid when specifying a task:

No preference of GPU accelerator to use:

@task(limits=Resources(gpu="1"))
def my_task() -> None:
    ...

Schedule on a specific GPU accelerator:

from flytekit.extras.accelerators import T4


@task(
    limits=Resources(gpu="1"),
    accelerator=T4,
)
def my_task() -> None:
    ...

Schedule on a Multi-instance GPU (MIG) accelerator with no preference of partition size:

from flytekit.extras.accelerators import A100


@task(
    limits=Resources(gpu="1"),
    accelerator=A100,
)
def my_task() -> None:
    ...

Schedule on a Multi-instance GPU (MIG) accelerator with a specific partition size:

from flytekit.extras.accelerators import A100


@task(
    limits=Resources(gpu="1"),
    accelerator=A100.partition_1g_5gb,
)
def my_task() -> None:
    ...

Schedule on an unpartitioned Multi-instance GPU (MIG) accelerator:

from flytekit.extras.accelerators import A100


@task(
    limits=Resources(gpu="1"),
    accelerator=A100.unpartitioned,
)
def my_task() -> None:
    ...

An override can also be specified for the GPU accelerator to use, as follows:

from flytekit.extra.accelerators import A100, T4


@task(
    limits=Resources(gpu="1"),
    accelerator=T4,
)
def my_task() -> str:
    return "hello"

@workflow
def my_wf() -> str:
    return my_task().with_overrides(
        accelerator=A100.partition_1g_5gb
    )

Tracking Issue

https://github.com/flyteorg/flyte/issues/

Follow-up issue

NA
OR
https://github.com/flyteorg/flyte/issues/

[eapolinario]

one minor comment, otherwise LGTM.

[wild-endeavor]

let's not use the _ anymore? i'm actively trying to remove it from the code base.

[jeevb]

Will clean up. Are you using tasks_pb2 directly, or still aliasing as core_task?

[wild-endeavor]

ditto

[kumare3]

@jeevb can you remind me why can we not do

@task(resources=Resources(gpu= NvidiaTeslaA100("1"))
   ...

Is this because of requests and limits?

Also can we drop
NvidiaTesla -> Simply keep A100 I think this is well understood, lets be less verbose

Also this is too verbose NvidiaTeslaA100.with_partition_size(NvidiaTeslaA100.partition_sizes.PARTITION_1G_5GB)
Can we make it
A100.partition(A100.partitions.5GB) (what is the use of 1G?)

[jeevb]

@jeevb can you remind me why can we not do

@task(resources=Resources(gpu= NvidiaTeslaA100("1"))
   ...

Is this because of requests and limits?

In a way, yes. It introduces room for confusion when specifying required GPUs - would a request of T4(“1”) and a limit of A100(“1”) be valid? We could work around it by adding validations that GPU types match in requests and limits. In addition, we’ll have to break up the Resources object into count and GPU type at serialization time anyway since there are fundamental differences in the way we define resources in the Container and K8sPod models.

Also can we drop NvidiaTesla -> Simply keep A100 I think this is well understood, lets be less verbose

👍

Also this is too verbose NvidiaTeslaA100.with_partition_size(NvidiaTeslaA100.partition_sizes.PARTITION_1G_5GB) Can we make it A100.partition(A100.partitions.5GB) (what is the use of 1G?)

See: https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi

Each partition size definition combines the compute: 1G, and memory: 5GB. 1g.5gb is the format that Nvidia uses. While GCP doesn’t support it currently, 3g.20gb and 4g.20gb are both valid for the A100, so specifying memory alone will not be sufficient.

An argument can be made that it would be better to be consistent with existing semantics instead of creating new ones, even if they may be more verbose. Note GCP’s use of the gke-gpu-partition-size label instead of gke-gpu-partition, for instance. That isn’t necessarily stopping us from doing it differently (both seem plausible), but it might come with the risk of introducing additional cognitive burden to users.

Also, A100.partition and A100.partitions might be too close. I can see if A100(A100.partitions.X) is feasible perhaps. Thoughts?

[jeevb]

@kumare3: Went ahead and made the change, so the MIG definition is now:

• A100 (unspecified partition size)
• A100.partitioned(A100.partitions.PARTITION_1G_5GB) (explicit partition size)
• A100.unpartitioned (explicit unpartitioned)

Wdyt?

[kumare3]

One question, what if we don't support partitioning - what happens if a user selects it - or thinking about it another way - what is the reason a user chooses a partitioned (specific) gpu. The reason to ask is to see if we can arrive with something so
Simple that it is obvious

[jeevb]

One question, what if we don't support partitioning - what happens if a user selects it - or thinking about it another way - what is the reason a user chooses a partitioned (specific) gpu. The reason to ask is to see if we can arrive with something so
Simple that it is obvious

In a scenario where multiple A100 node pools are available, a user might choose to explicitly schedule on one with a larger partition size, or one that is unpartitioned. If the cluster only has a single GPU node group, however, the extra partition specification is not needed. For all other cases, users will likely need an escape hatch to schedule on specific GPU configs.

[jeevb]

Final defs (following offline conversation with @kumare3):

• A100 (unspecified partition size)
• A100.partition_1g_5gb (explicit partition size)
• A100.unpartitioned (explicit unpartitioned)