<#6921 Add label for gpu partitions used> Pull req...
# flyte-github
a
#6921 Add label for gpu partitions used Pull request opened by katrogan Tracking issue MIG partitions allow us to slice specific GPU types into 7 slices. This updates allows us record the respective fraction of gpu usage based on number of slices requested as a pod attribute Why are the changes needed? What changes were proposed in this pull request? How was this patch tested? Deployed changes to a cluster where we have partitioned A100s Ran
Copy code
a100_mig_env = flyte.TaskEnvironment(
    name="a100_mig",
    image="<http://ghcr.io/flyteorg/flyte:py3.11-v2.0.0b43|ghcr.io/flyteorg/flyte:py3.11-v2.0.0b43>",
    resources=Resources(
        cpu="1",
        memory="4Gi",
        gpu=GPU(device="A100 80G", quantity=1, partition="1g.10gb"),
    ),
)


@a100_mig_env.task
async def hello_a100_partition() -> str:
    import asyncio
    import os

    visible = os.environ.get("NVIDIA_VISIBLE_DEVICES", "none")
    print(f"Hello from A100 80G MIG partition (1g.10gb)! NVIDIA_VISIBLE_DEVICES={visible}")

    total_seconds = 10 * 60
    interval = 60
    for elapsed in range(0, total_seconds, interval):
        remaining = (total_seconds - elapsed) // 60
        print(f"Running... {elapsed // 60} min elapsed, {remaining} min remaining")
        await asyncio.sleep(interval)

    msg = "Done! Ran for 10 minutes on A100 80G partition (1g.10gb)."
    print(msg)
    return msg
Labels Please add one or more of the following labels to categorize your PR: • added: For new features. • changed: For changes in existing functionality. • deprecated: For soon-to-be-removed features. • removed: For features being removed. • fixed: For any bug fixed. • security: In case of vulnerabilities This is important to improve the readability of release notes. Setup process Screenshots Check all the applicable boxes • I updated the documentation accordingly. • All new and existing tests passed. • All commits are signed-off. Related PRs Docs link •
main
• #6583 • #6921 👈 flyteorg/flyte