Rediger

Run batch inference with Ray on AKS

In this article, you submit a RayJob that runs vLLM offline batch inference using the LoRA adapter produced by the LLM training example. The job reads the viggo test split and LoRA adapter from Azure Blob Storage, generates predictions on a single GPU, and uploads the results.

Important

Open-source software is mentioned throughout AKS documentation and samples. Software that you deploy is excluded from AKS service-level agreements, limited warranty, and Azure support. As you use open-source technology alongside AKS, consult the support options available from the respective communities and project maintainers to develop a plan.

Microsoft takes responsibility for building the open-source packages that we deploy on AKS. That responsibility includes having complete ownership of the build, scan, sign, validate, and hotfix process, along with control over the binaries in container images. For more information, see Vulnerability management for AKS and AKS support coverage.

Prerequisites

Set environment variables

Navigate to the batch inference example in the cloned repository and configure the required environment variables:

cd <path-to-cloned-repo>/AKS/examples/kueue-and-ray-on-aks/3-workloads/batch-inference
export AZURE_STORAGE_ACCOUNT_NAME=$(terraform -chdir=../../1-infrastructure/terraform output -raw storage_account_name)
source env.example

Note

If you configured team queues (Option B), set export QUEUE_NAME=team-a or export QUEUE_NAME=team-b before running source env.example. The default QUEUE_NAME=default only works with the single-queue configuration (Option A).

Submit the workload

Submit the batch inference RayJob:

./submit.sh

The script creates a ConfigMap from the inference script, renders the manifest template via envsubst, and applies it. Kueue admits the job when 1 GPU is available in the configured queue.

Tip

Run ./submit.sh --dry-run to validate the rendered manifest without applying it to the cluster.

Monitor progress

Find and export the job name if you're in a new shell:

export JOB_NAME=$(kubectl -n ray get rayjob --no-headers -o custom-columns=":metadata.name" | grep batch-inference)

Watch the RayJob status and Kueue admission:

kubectl -n ray get rayjob ${JOB_NAME} -w
kubectl -n ray get workload -w

Expected output when the job completes:

NAME                         JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME               AGE
batch-inference-xxxxxxxxxx   SUCCEEDED    Complete            2026-01-01T00:00:00Z   2026-01-01T00:09:00Z   9m

NAME                                      QUEUE     RESERVED IN     ADMITTED   FINISHED   AGE
rayjob-batch-inference-xxxxxxxxxx-xxxxx   default   cluster-queue   True       True       9m

Tail the worker logs:

kubectl -n ray logs -l ray.io/cluster=$(kubectl -n ray get rayjob ${JOB_NAME} -o jsonpath='{.status.rayClusterName}') -f --tail=100

Verify results

List prediction files in blob storage:

az storage blob list -c llm-pipeline --prefix "inference/${JOB_NAME}/" \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login -o table

Expected output:

Name                                                 Blob Type    Blob Tier    Length    Content Type
---------------------------------------------------  -----------  -----------  --------  ------------------------
inference/<job-name>/metrics.json                    BlockBlob    Hot          128       application/octet-stream
inference/<job-name>/predictions.jsonl               BlockBlob    Hot          411777    application/octet-stream

Download and inspect the predictions:

az storage blob download -c llm-pipeline \
  -n "inference/${JOB_NAME}/predictions.jsonl" \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login \
  --file /tmp/predictions.jsonl
head -1 /tmp/predictions.jsonl | python3 -m json.tool

Expected output:

{
    "input": "I'm wondering, have you played any games you found genuinely shocking?",
    "expected": "request(specifier[shocking])",
    "generated": "request_explanation"
}

Configuration reference

Variable Default Description
AZURE_STORAGE_ACCOUNT_NAME (required) Storage account from Module 1
JOB_NAME batch-inference-<timestamp> Unique RayJob name
QUEUE_NAME default Kueue LocalQueue name
LLM_DATA_CONTAINER llm-pipeline Blob container with viggo test data
LLM_LORA_CONTAINER llm-pipeline Blob container with LoRA adapters
CONFIGMAP_NAME batch-inference-scripts Name for the ConfigMap holding the inference script

Clean up resources

Delete the RayJob and its ConfigMap:

kubectl -n ray delete rayjob ${JOB_NAME}
kubectl -n ray delete configmap ${CONFIGMAP_NAME}

To tear down all infrastructure, see Deploy infrastructure — Clean up resources.

Next steps

  • Last updated on