Skip to main content
LabsThis feature is in Labs. Contact your account team to enable this feature.
Astro supports graphics processing unit (GPU) workloads in dedicated clusters on Amazon Web Services (AWS). You can run model training, inference, and compute-intensive data preparation in the same Deployments that orchestrate your data pipelines, so your extract, transform, and load (ETL), training, and inference steps all run on one platform. During Labs, GPU workloads run through the KubernetesPodOperator (KPO), and Astronomer enables GPU support by request for a limited number of Organizations. This document describes what the GPU support Labs program includes, its limitations, and how to request access.

​
What’s included

GPU support in Labs includes:
  • NVIDIA GPUs in dedicated clusters on AWS: Astronomer adds GPU node pools to your dedicated cluster.
  • Execution with the KPO: Any containerized GPU workload runs as a Kubernetes Pod that Astro schedules onto a GPU node.
  • Whole GPUs: Tasks run on dedicated, whole GPU instances. Tasks don’t share GPUs or use fractions of a GPU.
  • A choice of GPU instance types: You choose from the NVIDIA instance types available in your cluster’s AWS region, such as A100-class and H100-class instances. Astronomer confirms the instance types with you during onboarding.

​
Limitations

Labs is the earliest stage of the Astronomer feature lifecycle. GPU support in Labs has the following limitations:
  • Experimental behavior: Capabilities can change, and breaking changes can occur as Astronomer incorporates feedback.
  • Limited support: Support is best effort through your Astronomer account team, with no uptime service level agreement and no guarantee of continued development.
  • Astronomer-managed provisioning: Astronomer adds and changes the GPU node pools in your cluster. You can’t manage GPU node pools yourself.
  • Long-running tasks: Routine cluster maintenance and upgrades can interrupt very long tasks, such as large training runs. Share your expected task runtimes with Astronomer during onboarding. A maximum task duration might apply.
  • KPO only: GPU support doesn’t extend to tasks that run on Astro executor, Celery executor, or Kubernetes executor workers.
  • Dedicated clusters on AWS only: GPU support isn’t available in standard clusters or in clusters on other cloud providers.

​
Billing

GPU support doesn’t add an Astronomer software charge during Labs. You pay for the GPU compute that your cluster uses, and Astronomer confirms the rates for your selected instance types during onboarding.

​
Request access

To join the GPU support Labs program:
  1. Contact your Astronomer account team to confirm your interest and describe your target use case.
  2. Work with Astronomer to choose GPU instance types, an AWS region, and your expected workloads. Astronomer then adds GPU node pools to your dedicated cluster.
  3. Run GPU workloads with the KPO and share feedback with your account team.
Astronomer uses feedback from Labs participants to decide on pricing, supported GPU types, and future GPU capabilities.

​
Run a GPU workload

After Astronomer adds GPU node pools to your cluster, request GPUs in a KPO task by setting the nvidia.com/gpu resource in the task’s container resources. Astro detects the GPU request and schedules the task’s Pod onto a GPU node, so you don’t need to configure tolerations for GPU nodes. The following example runs nvidia-smi in an NVIDIA CUDA container to confirm that the task can access a GPU. The task logs show the GPU model, driver version, and CUDA version.
Set nvidia.com/gpu to the number of whole GPUs that the task needs. To also set CPU and memory for the task, add them to the same V1ResourceRequirements object. See Configure task-level Pod resources. A task that requests GPUs fails in the following cases:
  • Your cluster has no GPU node pools. The task fails with the error Pod requests nvidia.com/gpu but the cluster has no GPU node pools.
  • The task requests more CPU, memory, or GPUs than your GPU node pools support. The error names the maximum for each resource that the task exceeds.