Skip to content

Use managed KubeRay by default in Kueue and Ray examples - #6046

Open
Li Chen (AaronLiChen) wants to merge 4 commits into
Azure:masterfrom
AaronLiChen:lichen5/managed-kuberay-example-4044
Open

Li Chen (AaronLiChen) wants to merge 4 commits into
Azure:masterfrom
AaronLiChen:lichen5/managed-kuberay-example-4044

Conversation

@AaronLiChen

@AaronLiChen Li Chen (AaronLiChen) commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Make managed KubeRay and Kueue the recommended executable path for examples/kueue-and-ray-on-aks, aligned with docs PR4044, while preserving the separate self-managed Terraform alternative.

The checked-in examples remain GPU-based. Per the author's request, all four Ray applications were functionally executed on CPU on a real AKS cluster, using temporary test-only adaptations outside Git. The actual Aurora/Qwen2.5-7B models, training/adapter/storage flows, RayService endpoint, and Ray Data/vLLM consumer were exercised. These bounded runs are not full-dataset training, model-quality, GPU-kernel, NCCL, or GPU-performance acceptance.

Merge/publication order

  1. Merge docs(examples): use managed Kueue for queue and autoscaler demos #6040 first. This branch is based on its unmerged head 473df1acfefa883d21cfb43bdf3f5977ab23df83; those changes currently appear in the aggregate diff. docs(examples): use managed Kueue for queue and autoscaler demos #6040's branch/manifests remain unchanged.
  2. Merge this producer before publishing consumer docs PR4044. Its separate [BUG] nodepool update with --enable-artifact-streaming fails #4043 Kueue filename/TOC integration is still outstanding.

Changes

  • A real managed-first Module1 CLI installer: public preview discovery, compatible auto-upgrade, Azure/controller/API readiness, explicit subscription, private kubeconfig.
  • Read-only ownership/orphan-CRD checks. Reuse healthy managed instances; no Helm adoption, protected-setting overrides, forced finalizers, automatic state migration, or uninstall/reinstall recovery.
  • Default fixed-quota Ray path needs no CPU autoscaler. Optional --with-cpu-autoscaling checks the CAS-owned established/served ProvisioningRequest v1 API before Kueue startup; never creates/replaces the shared CRD.
  • Explicit managed private-storage, account-scoped Workload Identity/federation, and WeatherBench2/ViGGO preparation with Entra-authenticated uploads and shared-key access disabled. Update all module/workload guides and storage comments.
  • Retain the full Terraform alternative and existing cluster/operator resource addresses/Helm ownership. The Ray image output default follows the corrected Module3 GPU pin; no resource migration or Terraform apply is performed.
  • Pod-local Python3 compatibility wrapper for KubeRay1.7's generated submitter health command in the three RayJobs and reusable template, without replacing generated commands or changing operator settings.
  • Correct batch submitters returning1 after successful client dry runs. Add 25 secretless offline tests and GitHub Actions coverage, including GPU-default and rank-local checkpoint regressions.

Bugs found by real execution

  • 41201b00 — Ray Serve/protobuf compatibility: the published Ray2.55.1 image has protobuf7.36.2 but Serve calls removed FieldDescriptor.label, failing deployment before model loading. Consistently pin published py3.12-ray2.56.1-cuda13.0, with the upstream compatibility helper and matching RAY_VERSION. The actual new-image protobuf conversion and full CPU RayService GET/POST passed. An application runtime-env pin alone wouldn't fix the base-image controller.
  • a9160984 — complete LoRA artifact selection: a real four-rank training step succeeded, but the old retrieval returned a nonzero rank's RNG-only checkpoint instead of rank zero's80MB model. Now skip incomplete ranks, accept nested/direct adapter layouts, and require nonempty adapter configuration + weights before retrieval/upload. The full real training/upload/inference flow passed after the fix. Five new unit regressions cover RNG-only-first ordering and fail-closed upload behavior.

GPU resource requests, CUDA guards, actual model choices, and default worker counts are retained. The only application logic change is device-independent adapter validation/retrieval; no CPU resource/device/library profile is committed.

Real AKS validation — 2026-10-24

Dedicated normal ARM AKS1.35.9, Central US EUAP, with the actual public-preview managed extensions. No mock cluster, self-hosted replacement, shared-cluster mutation, or dry-run-only claim.

Exact #6040 — independently passed twice

Fresh real-AKS rerun result at immutable 473df1ac:

  • CPU user pool 1→3 Ready nodes.
  • Actual ProvisioningRequest Provisioned=True, best-effort-atomic class, then Workload Admitted=True.
  • Exact sample Job 3/3 completions in205.5s. Earlier independent real-cluster run also passed in217s.
  • Demo cleanup passed; all three managed-default UIDs retained.

Managed installation/reuse, controller/API checks, account-scoped identity/federation, real WeatherBench2/ViGGO uploads, and remote-worker identity Blob read/write also passed.

All four CPU functional workflows — PASS

Workflow Actual acceptance evidence
Aurora fine-tune Real pretrained Aurora, one LoRA optimizer step; finite loss69340.55→68923.04; uploaded/downloaded adapter with192finite tensors and96nonzero trained LoRA-B matrices.
Online serving Actual RayService loaded that fresh adapter; GET health + POST6h forecast passed; four finite surface-variable summaries, shape1×1×52×100, correct forecast timestamp.
LLM training Actual four-worker Ray Train/Gloo/BF16 Qwen2.5-7B step, loss1.8241358; valid80,792,096byte rank8 adapter/config/latest pointer;392finite tensors/196nonzero trained LoRA-B matrices.
Batch inference Actual Ray Data + official vLLM0.18.0 CPU build, same Qwen7B and newly trained Blob adapter; generated2nonempty predictions and uploaded predictions/metrics. Adapter download and device_config=cpu observed.

CPU test scope (not production defaults)

  • Same actual Aurora and Qwen2.5-7B models, not smaller substitutes. CPU-only copied payloads/manifests remove GPU resource/device requirements for the test only.
  • Aurora1step. LLM4CPU workers,1step/8samples/cutoff512, warmup0, Gloo/BF16. LLaMA-Factory's hardcoded GPU ScalingConfig is overridden in memory only in the disposable CPU process.
  • Batch2samples/8outputtokens with official CPU vLLM and matched CPU torch/audio/vision wheels. Native CPU libnuma/tcmalloc/OpenMP dependencies were supplied from checksum-verified Microsoft Azure Linux packages/the official wheel into the owned test cache without root, added capabilities, or system installation. CPU affinity uses supported nobind.
  • This small inference smoke had 0/2 exact matches; its pass verifies execution, adapter consumption, and outputs, not accuracy. No GPU/NCCL/performance pass is claimed.

The earlier isolated A100 allocation returned Azure Compute InternalExecutionError before any GPU node registered. That does not establish a quota or physical-capacity root cause; that environment was deleted. The CPU-only environment's owned demo resources were cleaned up with managed defaults/shared API retained. Its CPU VMs, disks, and node resource group are verified deleted; the parent AKS resource group is still Deleting, so full group deletion is not claimed yet. Shared GPU/NPD labs/operators were untouched.

Checks and remaining limits

  • 25 offline tests and all six GitHub checks pass at a9160984; shell syntax/whitespace checks pass.
  • Four server schema/admission dry runs passed, independently of the actual application runs above.
  • Documentation also removes the protected accelerator label write rejected by real AKS; no node-label validation bypass is used.
  • Full GPU execution, NCCL, performance/accuracy, publication, merge, production deployment, and SME approval are not implied. AI-assisted implementation and validation.

@AaronLiChen

Li Chen (AaronLiChen) commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor Author

Real AKS CPU functional validation — all four workflows PASS

Executed on a newly created real ARM AKS1.35.9 target, using public managed Kueue 0.19.6-aks.0.6.0 and KubeRay 1.7.0-aks.0.1.3. No GPUs were allocated, and no shared GPU/NPD cluster or operator was modified.

The published examples remain GPU-based. CPU copies/overlays were temporary and outside Git, as requested by the author. Actual pretrained Aurora and Qwen2.5-7B models were retained—not downsized substitutes.

CPU functional check Observed evidence
Aurora fine-tune Real1step LoRA; loss69340.55→68923.04; Blob adapter/metrics downloaded and verified.192finite tensors,96nonzero trained LoRA-B matrices.
Online serving RayService loaded that fresh adapter; actual GET health and POST6h forecast succeeded. Four finite output-variable summaries, shape1×1×52×100, prediction timestamp2021-01-01T06:00.
LLM training Actual 4ranks, CPU/Gloo/BF16, Qwen7B one-step training; loss1.8241358. Valid80,792,096byte rank8 adapter/config/latest pointer;392finite tensors/196nonzero trained LoRA-B matrices.
Batch inference Actual Ray Data + official vLLM0.18.0 CPU build, same Qwen7B and the newly trained Azure Blob LoRA. Engine logs confirm device_config=cpu and adapter download. Two nonempty prediction records plus metrics uploaded and downloaded.

What this does—and doesn't—validate

  • CPU profile: Aurora1step; LLM4CPU workers/1step/8samples/cutoff512/warmup0; inference2samples/8outputtokens. These are bounded integration checks, not full training, accuracy, GPU/NCCL, or GPU-performance acceptance. The tiny inference smoke had0/2 exact matches; no accuracy pass is claimed.
  • CPU-only device/resource edits, matched CPU torch/audio/vision wheels, and in-memory CPU ScalingConfig adaptation stayed outside the repository. Required native CPU libraries came from checksum-verified Microsoft Azure Linux packages/the official Intel wheel in an owned test cache, without root, extra capabilities, system installation, or operator-setting changes.
  • Checked-in GPU requests, CUDA guards, models, and default worker counts remain intact. Source fixes are device-independent: published GPU-capable Ray2.56.1 pin (41201b00) and complete LoRA checkpoint validation/selection (a9160984).

The tests caught real bugs rather than stopping at admission/dry-run: Ray2.55/protobuf7 prevented Serve deployment; then rank-local RNG state was being returned/uploaded instead of rank zero's actual adapter. The complete corrected training→Blob→inference chain passed after those fixes. All25 offline tests and all six PR checks pass at a9160984.

Kueue #6040 was also rerun on real AKS

Recorded independently on #6040: exact immutable 473df1ac manifests, pool1→3Ready nodes, actual ProvisioningRequest Provisioned=True, Workload Admitted=True, exact Job3/3 completed in205.5s. Demo cleanup retained all three managed-default UIDs.

Only run-owned model jobs/queues/flavors/namespace were cleaned up. Managed-default UIDs and the shared ProvisioningRequest API were preserved before full teardown. The completed CPU environment's VMs, disks, and node resource group are verified deleted. Its parent AKS resource group remains Deleting; full group deletion is not claimed yet. No shared environment was touched. PR is ready for review, not merged, published, or SME-approved.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Subscription-scoping fixes remain, and the preview-operator integration needs human review because full GPU execution remains unvalidated.

1 open issue
What changed in this PR

Makes managed KubeRay and Kueue the recommended path for the AKS examples while preserving the separate self-managed Terraform lab. Builds on #6040 and prepares the examples for the companion Learn documentation.

Changes:

  • Adds a managed-extension installer with ownership checks, readiness verification, and storage, identity, and dataset guidance.
  • Updates GPU runtime defaults and submitter compatibility; validates LoRA artifacts before retrieval and upload.
  • Adds offline regression tests and GitHub Actions coverage.

Paths below are relative to examples/kueue-and-ray-on-aks/, except the GitHub Actions workflow.

File Description
README.md Presents the managed-first workflow and installation alternatives.
3-workloads/​tests/​test_managed_ray.py Tests runtime pins, submitter wrappers, and dry runs.
3-workloads/​tests/​test_lora_checkpoints.py Tests complete-adapter selection and upload rejection.
3-workloads/​README.md Documents managed prerequisites and runtime compatibility.
3-workloads/​online-serving/​README.md Uses managed storage and operator prerequisites.
3-workloads/​online-serving/​env.example Updates runtime defaults and storage guidance.
3-workloads/​llm-training/​submit.sh Corrects dry-run exit status and storage guidance.
3-workloads/​llm-training/​README.md Documents managed setup and adapter validation.
3-workloads/​llm-training/​manifests/​rayjob.yaml.tmpl Adds the submitter’s Python compatibility wrapper.
3-workloads/​llm-training/​llm_training.py Rejects incomplete rank-local adapter checkpoints.
3-workloads/​llm-training/​env.example Updates runtime defaults and storage guidance.
3-workloads/​cas-batch-job/​README.md Uses managed provisioning defaults and scoped demo cleanup.
3-workloads/​batch-inference/​submit.sh Corrects dry-run exit status and storage guidance.
3-workloads/​batch-inference/​README.md Switches prerequisites to managed setup.
3-workloads/​batch-inference/​manifests/​rayjob.yaml.tmpl Adds the submitter’s Python compatibility wrapper.
3-workloads/​batch-inference/​env.example Updates runtime defaults and storage guidance.
3-workloads/​aurora-finetune/​submit.sh Corrects dry-run exit status and storage guidance.
3-workloads/​aurora-finetune/​README.md Links managed infrastructure and dataset preparation.
3-workloads/​aurora-finetune/​manifests/​rayjob.yaml.tmpl Adds the submitter’s Python compatibility wrapper.
3-workloads/​aurora-finetune/​env.example Updates runtime defaults and input-file guidance.
3-workloads/​_template/​rayservice.yaml.tmpl Updates runtime and storage comments only.
3-workloads/​_template/​rayjob.yaml.tmpl Adds the reusable wrapper and updates comments.
2-kueue-queues/​tests/​test_managed_kueue.py Tests managed references, ownership, and queue routing.
2-kueue-queues/​tests/​requirements.txt Pins the YAML test dependency.
2-kueue-queues/​README.md Documents managed installation and queue selection.
2-kueue-queues/​manifests/​README.md Clarifies manifest ownership and prerequisites.
2-kueue-queues/​manifests/​40-autoscale-queue.yaml References managed defaults instead of creating them.
2-kueue-queues/​manifests/​10-resource-flavors.yaml Clarifies customer ownership in comments only.
1-infrastructure/​terraform/​variables.tf Aligns the Ray image output default.
1-infrastructure/​terraform/​README.md Labels Terraform as the self-managed alternative.
1-infrastructure/​terraform/​legacy-guide.md Preserves the full Terraform walkthrough.
1-infrastructure/​README.md Directs readers to managed setup.
1-infrastructure/​managed/​tests/​test_preflight.py Tests ownership and API safety checks.
1-infrastructure/​managed/​tests/​test_installer.py Tests installer branches with fake command-line tools.
1-infrastructure/​managed/​README.md Documents installation, storage, identity, and cleanup.
1-infrastructure/​managed/​preflight.py Adds read-only ownership and API validation.
1-infrastructure/​managed/​install-extensions.sh Installs or reuses managed extensions with readiness checks.
1-infrastructure/​managed/​data.md Documents authenticated dataset preparation and uploads.
1-infrastructure/​managed/​.gitignore Excludes local state and Python cache files.
.github/​workflows/​kueue-ray-examples.yml Runs shell syntax and offline contract checks.

Review effort: Balanced

kubectl get resourceflavor scalepool # NotFound
kubectl get namespace cas-kueue-demo # NotFound
az aks nodepool delete \
--resource-group "$RESOURCE_GROUP" --cluster-name "$CLUSTER_NAME" \

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants