Add GPU-capable demo cluster image

[pingsutw]

Summary

Adds a GPU variant of the demo-bundled image so users on an NVIDIA-enabled host can run:

flyte start demo --image ghcr.io/flyteorg/flyte-demo:gpu-latest

and submit Flyte tasks with Resources(gpu=1) — no PodTemplate, no runtimeClassName needed — and get a GPU.

What's in the image

• Dockerfile.gpu — stages NVIDIA Container Toolkit v1.19.x (nvidia-ctk, nvidia-container-runtime, libnvidia-container) into the rancher/k3s final image. Two subtle prereqs the OCI hook needs:
  • Libs under /usr/lib/<arch-triple>/ — the nvidia-ctk OCI hook runs without inheriting LD_LIBRARY_PATH.
  • A statically-linked /sbin/ldconfig from debian:bookworm-slim — rancher/k3s ships none, and the toolkit's update-ldcache hook bind-mounts it into workload pods.
• containerd-config.toml.tmpl — sets default_runtime_name = "nvidia". GPU pods get GPUs automatically; non-GPU pods are unaffected (nvidia-container-runtime is a passthrough).
• nvidia-device-plugin.yaml — RuntimeClass nvidia + DaemonSet (nvcr.io/nvidia/k8s-device-plugin:v0.17.0) auto-applied by k3s on startup.
• Makefile — new build-gpu target.
• CI — new build-and-push step publishing gpu-latest, gpu-nightly, and gpu-<sha> tags to both flyte-demo and flyte-sandbox-v2.

Companion change

A matching --gpu flag on flyte start demo (adds --gpus all to docker run) will land in flyteorg/flyte-sdk. Image is useful without it via manual docker run --gpus all ….

Test plan

import flyte

torch_image = flyte.Image.from_debian_base(
    python_version=(3, 12),
    registry="localhost:30000",
).with_pip_packages("torch")

gpu_env = flyte.TaskEnvironment(
    name="torch_gpu",
    image=torch_image,
    resources=flyte.Resources(cpu=1, memory="2Gi", gpu=1),
)


@gpu_env.task
def check_torch() -> str:
    import torch

    info = (
        f"torch={torch.__version__}\n"
        f"cuda_available={torch.cuda.is_available()}\n"
        f"device_count={torch.cuda.device_count()}\n"
        f"device_name={torch.cuda.get_device_name(0) if torch.cuda.is_available() else None}"
    )
    print(info)
    return info


if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(check_torch)
    print("name:", r.name)
    print("url:", r.url)
    r.wait()

• Layered test image (NVIDIA toolkit + device plugin on top of flyte-demo:nightly) verified on an A10G: k3s auto-registers nvidia runtime, node advertises nvidia.com/gpu: 1, Flyte task with Resources(gpu=1) ran torch 2.11.0+cu130 reporting cuda_available=True, device_name="NVIDIA A10G".
• Full multi-stage Dockerfile.gpu has NOT been built locally — this CI run is the first build. Expect potential fixup iterations.

Follow-ups

• GPU type filters (Resources(gpu="A10G:1")) require GPU Feature Discovery to label the node.
• MIG support (A100/H100 partitioning) needs a device-plugin ConfigMap with migStrategy.

• main
  • Flyte 2 #6583
    • Add GPU-capable demo cluster image 👈
      • Rename flyte-demo to flyte-devbox #7247