Polyaxon & AMD GPUs

How to use Polyaxon and AMD GPUs together

+

The AMD device plugin makes GPUs schedulable as amd.com/gpu; ROCm and the workload image determine whether an allocated device is actually usable. Polyaxon keeps the hardware request, placement, image, and run evidence together.

See the AMD Kubernetes device plugin for the upstream configuration and requirements.

Prerequisites

• ROCm-capable AMD GPUs with supported host drivers and firmware.
• The AMD GPU device plugin exposing amd.com/gpu on the Kubernetes nodes.
• A tested ROCm framework image and an AMD-specific Polyaxon queue or preset.

Configuration

Request an AMD GPU

Target the AMD pool and request the resource name published by the official Kubernetes device plugin.

version: 1.1
kind: component
name: amd-gpu-readiness
run:
  kind: job
  container:
    image: rocm/pytorch:rocm7.14.1_ubuntu24.04_py3.12_pytorch_release_2.12.0
    command: ["python", "-c"]
    args: ["import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name())"]
    resources:
      limits:
        amd.com/gpu: "1"

ROCm exposes accelerators through PyTorch's torch.cuda API for compatibility; that name does not mean the workload is using CUDA. Select the hardware through your configured queue and preset; no provider-specific node label is required by this example.

Submit the operation

Run the component through the target Polyaxon project, queue, preset, and approval path.

polyaxon run -f amd-readiness.yaml

Verify the ROCm execution path

Confirm the pod lands on the AMD pool, receives one device, loads the expected framework build, and emits logs and artifacts normally.

polyaxon ops dashboard

Deployment checks

• Pin a compatible driver, device integration, container runtime, framework image, and firmware combination.
• Publish accelerator selection through centrally managed Polyaxon presets instead of user-specific labels.
• Track allocatable devices, hardware health, utilization, memory, thermals, and workload failures.
• Validate one device, one node, and multi-node communication separately before production adoption.
• Run RCCL tests before enabling multi-node AMD workloads and keep CUDA and ROCm images in separate presets.

Troubleshooting

The resource does not appear

Inspect the node driver, firmware, device integration, kubelet registration, node labels, and allocatable resources.

The container cannot access the device

Verify the requested custom resource, runtime class, permissions, device mounts, and image compatibility.

Performance is unexpectedly low

Check topology, NUMA placement, clocks, memory pressure, host-device transfer, communication libraries, and framework build.

References

• AMD Kubernetes device plugin — Official installation, amd.com/gpu resource, node labels, and verification guidance.
• AMD ROCm PyTorch images — AMD-published image tags for matching ROCm, Ubuntu, Python, and PyTorch versions.
• Polyaxon job runtime — Job containers, resources, connections, initialization, and lifecycle settings.
• Polyaxon resource scheduling — Kubernetes requests, limits, custom resources, queues, and placement controls.