Azureml K3s Compute Target Setup

Azure ML K3s Compute Target Setup

Set up K3s on an Ubuntu host with an NVIDIA GPU, connect the host and cluster to Azure Arc, and attach the cluster to Azure ML as a Kubernetes compute target. Then prove real CUDA execution with a bounded smoke job before submitting a long-running training job. A GPU reservation in the Azure ML job spec is not proof the container can see the device; validate at the container level every time the runtime stack changes.

Prerequisites

Choose Where Each Command Runs

Host setup and the kubectl checks in this skill run on the GPU host, because the K3s kubeconfig points at the host-local API server. The attach script reaches the cluster through Arc cluster connect, so it runs from a workstation, as do job submission commands.

Before running a host step, confirm you are on the target host. Run these checks and compare the hostname with the intended host:

hostname
nvidia-smi -L
systemctl is-active k3s

If nvidia-smi is missing or the hostname does not match, you are not on the GPU host. Connect to it before continuing:

• Before the host is Arc-connected, use regular SSH or console access to reach it, then clone this repository on the host.
• After Connect the Host to Arc completes, use az ssh arc. It needs no public IP address or inbound port, and requires Owner or Contributor on the Arc-enabled server and the Microsoft.HybridConnectivity resource provider.

Enable SSH on the Arc-enabled server once, from a workstation:

az extension add --name ssh
az provider register -n Microsoft.HybridConnectivity

machine_id=/subscriptions/<subscription-id>/resourceGroups/<arc-resource-group>/providers/Microsoft.HybridCompute/machines/<arc-server-name>
az rest --method put \
  --uri "https://management.azure.com${machine_id}/providers/Microsoft.HybridConnectivity/endpoints/default?api-version=2023-03-15" \
  --body '{"properties": {"type": "default"}}'
az rest --method put \
  --uri "https://management.azure.com${machine_id}/providers/Microsoft.HybridConnectivity/endpoints/default/serviceconfigurations/SSH?api-version=2023-03-15" \
  --body '{"properties": {"serviceName": "SSH", "port": 22}}'

Open a session on the host, then run the host steps inside it:

az ssh arc --resource-group <arc-resource-group> --name <arc-server-name> --local-user <host-user>

Inside the session, point kubectl at the protected kubeconfig written by the K3s installer:

export KUBECONFIG="$HOME/.local/share/physical-ai-toolchain/hil/kubeconfig.yaml"

Connect the Host to Arc

Run on the GPU host. Preview, then connect the Ubuntu server:

data-pipeline/setup/edge/03-connect-arc-server.sh \
  --subscription-id <subscription-id> --tenant-id <tenant-id> \
  --resource-group <arc-resource-group> --location <location> \
  --server-name <arc-server-name> --config-preview

data-pipeline/setup/edge/03-connect-arc-server.sh \
  --subscription-id <subscription-id> --tenant-id <tenant-id> \
  --resource-group <arc-resource-group> --location <location> \
  --server-name <arc-server-name>

Install K3s and Configure the GPU Default Runtime

Run on the GPU host. Install the pinned, owned K3s compute plane, passing --default-runtime nvidia on any host where GPU-backed Azure ML jobs will run:

data-pipeline/setup/hil/01-install-k3s.sh --default-runtime nvidia --config-preview
data-pipeline/setup/hil/01-install-k3s.sh --default-runtime nvidia

[!IMPORTANT] Azure ML InstanceType limits.nvidia.com/gpu only reserves the device through the Kubernetes device plugin. It does not set a pod runtimeClassName, so on a K3s host whose default runtime is not NVIDIA, the container starts with no /dev/nvidia* devices and torch.cuda.device_count() returns 0 even though the job reports a reserved GPU. Setting default-runtime: nvidia in the K3s config is required whenever the Azure ML Kubernetes extension does not expose per-pod runtime-class configuration.

Then run 05-enable-k3s-gpu.sh on the host. Its preview checks the driver, the NVIDIA Container Toolkit, the K3s service and API, and the nvidia RuntimeClass, and says whether the run will restart K3s:

data-pipeline/setup/hil/05-enable-k3s-gpu.sh --config-preview
data-pipeline/setup/hil/05-enable-k3s-gpu.sh

The script installs the digest-pinned NVIDIA device plugin so the node advertises nvidia.com/gpu. When nvidia is already the default runtime, as after --default-runtime nvidia, it leaves K3s running. On hosts installed without that option, it writes the K3s default-runtime: nvidia drop-in and restarts K3s once, so run it when no job containers are running.

Verify the NVIDIA device plugin is healthy and the node advertises allocatable GPUs before attaching Azure ML:

kubectl get pods -n kube-system -l name=nvidia-device-plugin-ds
kubectl get node <node-name> -o jsonpath='{.status.allocatable.nvidia\.com/gpu}'

Connect the K3s Cluster to Arc

Run on the GPU host. Connect the cluster with the OIDC issuer and workload identity enabled, and grant yourself cluster-admin so Arc cluster connect accepts your identity:

data-pipeline/setup/edge/05-connect-arc-kubernetes.sh \
  --subscription-id <subscription-id> --tenant-id <tenant-id> \
  --resource-group <arc-resource-group> --location <location> \
  --cluster-name <arc-cluster-name> --kubeconfig <kubeconfig-path> \
  --enable-workload-identity --cluster-admin-signed-in-user --config-preview

data-pipeline/setup/edge/05-connect-arc-kubernetes.sh \
  --subscription-id <subscription-id> --tenant-id <tenant-id> \
  --resource-group <arc-resource-group> --location <location> \
  --cluster-name <arc-cluster-name> --kubeconfig <kubeconfig-path> \
  --enable-workload-identity --cluster-admin-signed-in-user

The attach script applies InstanceTypes through Arc cluster connect, which checks K3s RBAC for the identity that runs it. When someone else runs the attach, grant them with --cluster-admin-object-id <object-id> instead.

Attach the Cluster to Azure ML

infrastructure/setup/02-deploy-azureml-extension.sh targets the Terraform-managed AKS cluster. For an Arc-connected K3s cluster, run 05-attach-hil-azureml-compute.sh from a workstation that has this repository's Terraform outputs. Preview, then attach:

infrastructure/setup/05-attach-hil-azureml-compute.sh \
  --arc-cluster-resource-id /subscriptions/<subscription-id>/resourceGroups/<arc-resource-group>/providers/Microsoft.Kubernetes/connectedClusters/<arc-cluster-name> \
  --compute-name <compute-name> --config-preview

infrastructure/setup/05-attach-hil-azureml-compute.sh \
  --arc-cluster-resource-id /subscriptions/<subscription-id>/resourceGroups/<arc-resource-group>/providers/Microsoft.Kubernetes/connectedClusters/<arc-cluster-name> \
  --compute-name <compute-name> --require-gpu

The script:

1. Installs a training-only Azure ML extension from azureml-arc-config.template.json, with the extension's own device plugin and DCGM exporter off, or updates only the settings that differ on an existing extension.
2. Waits for the InstanceType CRD and applies azureml-instance-types-hil.yaml through Arc cluster connect: defaultinstancetype for CPU jobs and gpu for one GPU. The gpu type has no node selector, so the node needs no label.
3. Attaches the cluster as a Kubernetes compute with a system-assigned identity in the azureml namespace.
4. Grants that identity AzureML Data Scientist on the workspace, Storage Blob Data Contributor on its storage account, and AcrPull on its container registry, so jobs can pull images from it.

Compute names are 16 characters at most, and the default, k8s-<cluster>, is truncated, so pass --compute-name. --require-gpu stops before any change unless a node reports allocatable nvidia.com/gpu. For InstanceTypes that request more GPUs, pass your own manifest with --instance-types-manifest, and request only what the node advertises.

The extension creates an Azure Relay namespace and hybrid connection in the Arc cluster's resource group. Don't modify them, because the compute depends on them.

Azure ML Credentials and Dataset Inputs

Store HF_TOKEN in the untracked repository-root .env.local, never as a CLI argument or in chat. The submission script loads .env.local and forwards HF_TOKEN to the job, so --hf-token is not needed. Tokens passed as CLI arguments are visible to any process inspecting the host, so rotate a token immediately if it was ever exposed that way. .amlignore already excludes .env and .env.* from the Azure ML code snapshot.

Data asset mount failures during job start often mean the registered asset version does not resolve against its backing datastore. Register a new datastore-backed version and reference that explicit version rather than reusing a broken one:

az ml data create --name <dataset-name> --version <next-version> \
  --type uri_folder --path azureml://datastores/<datastore>/paths/<path>