#700 H100 mig partitions Pull request opened by katrogan See https://docs.nvidia.com/datacenter/tesla/mig-user-guide/supported-mig-profiles.html#h100-mig-profiles Unfortunately I couldn't test due to lack of H100 quota on internal clusters Ideally would use

import flyte
from flyte import GPU, Resources

h100_mig_env = flyte.TaskEnvironment(
    name="h100_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="H100", quantity=1, partition="1g.10gb"),
    ),
)


@h100_mig_env.task
async def hello_h100_partition() -> str:
    import asyncio
    import os

    visible = os.environ.get("NVIDIA_VISIBLE_DEVICES", "none")
    print(f"Hello from H100 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 H100 partition (1g.10gb)."
    print(msg)
    return msg

flyteorg/flyte-sdk