feat: add multi-node serving and researcher role alias

[Pengfei Ni (feiskyer)]

Summary

Allow a Ray Serve application to use a multi-worker GPU profile instead of rejecting it as a single-Pod workload. Also align workspace creation and connection verification on the logical researcher role and grant the RayService access required for serving.

• Render a CPU-only Ray head and a fixed GPU worker pool, with profile-bound node/GPU counts, shared memory, mounts, environment variables, and worker anti-affinity.
• Add --app-args for application builders and --workload-profile-snapshot for explicitly offline client rendering. Keep connected workspace/profile checks mandatory for server dry-run and apply.
• Default new workspaces to researcher, accept the legacy tau-researcher-v1 alias, and update the generated CRDs, tests, examples, and documentation.
• Grant RayService lifecycle permissions in both Helm and Kustomize researcher RBAC. Verify create/get/list/patch/delete permissions during workspace connection checks.

Related issue

N/A.

Validation

Passed locally:

# cli/
go test -race -count=1 ./internal/serve ./internal/cli ./internal/workspace ./internal/workspaceconnection ./internal/reposcaffold
go vet ./internal/serve ./internal/cli ./internal/workspace ./internal/workspaceconnection ./internal/reposcaffold
go run honnef.co/go/tools/cmd/staticcheck@v0.7.0 ./internal/serve ./internal/cli ./internal/workspace ./internal/workspaceconnection ./internal/reposcaffold
go build ./cmd/tau ./cmd/tau-gen

# controllers/tau-core/
go test -race -count=1 ./...
go vet ./...

# Repository root
python3 scripts/check-license-headers.py
python3 site/scripts/check_content.py site/content/en/docs
bash -n controllers/tau-core/scripts/kind-e2e.sh
cmp controllers/tau-core/config/crd/bases/tau.azure.com_workspaces.yaml charts/tau-core-controller/crds/tau.azure.com_workspaces.yaml
helm lint charts/tau-core-controller
helm unittest --strict charts/tau-core-controller
bash charts/tau-core-controller/tests/scripts_test.sh
kubectl kustomize charts/tau-core-controller
git diff --cached --check

Go checks used GOTOOLCHAIN=local, GOFLAGS=-mod=readonly, and AI_RUNTIME_E2E=0 for tests. Missing cached dependencies were downloaded without changing module versions.

The offline eight-worker render has one CPU head, eight single-GPU workers, and builder-owned application configuration. Regression tests cover rejected live use of snapshots, argument validation order, annotation preservation, and unchanged bindings across role aliases. Additional tests check RayService grants in both RBAC sources and reject each missing RayService permission for either role name. Twenty offline CLI scenarios also produced identical exit codes, stdout, and stderr before and after the simplification pass.

Not run for this PR: full repository make check, full documentation site build, Kind integration tests, or live cluster/GPU tests.

Compatibility and operational impact

• Existing single-worker RayService and plain Deployment paths remain available. --nodes and --gpus assert profile values; application replicas do not resize the GPU worker pool.
• --app-args requires an explicit application builder and rejects CLI replica/autoscaling overrides and legacy --args.
• Snapshot rendering requires an explicit namespace and --dry-run=client, rejects --context, and marks its output as snapshot-derived. It does not authorize deployment.
• Upgrade the workspace CRD before using role: researcher. Both role names bind the same tau-researcher-v1 ClusterRole. Older CLIs that compare role names literally still need matching descriptor/workspace values.
• This PR adds RayService create/get/list/watch/delete/patch/update permissions to that role, scoped by the existing workspace RoleBinding. Update the Helm or Kustomize RBAC before using the new CLI: workspace-RBAC connection verification now requires create/get/list/patch/delete on RayServices. Updating only the CRD is insufficient.
• Distributed worker anti-affinity requires Kubernetes support for matchLabelKeys. RayService cluster upgrades can temporarily need capacity for both worker pools.
• No new dependencies, inference engine code, model-specific images, or deployment credentials are included.

Checklist

• The change is focused and does not include unrelated edits.
• Tests cover the changed behavior, or I explained why tests are not needed.
• Documentation and examples are updated when public behavior changes.
• Generated files are updated from their source and included in this pull request.
• New dependencies are publicly available and their license and provenance were reviewed. No new dependencies.
• The diff contains no secrets, private endpoints, customer data, or Microsoft-internal information.
• Fork-based CI can validate the change without repository secrets, or maintainer-only validation is clearly identified.

[yliu382]

Summary

Adds multi-node Ray Serve (CPU head + fixed GPU worker pool) driven by workload-profile workerCount/gpusPerWorker/placement, plus two auxiliary flags (--app-args for builder-owned Ray Serve apps, --workload-profile-snapshot for offline --dry-run=client rendering), and renames the default researcher role to researcher while preserving tau-researcher-v1 as a backward-compatible alias across CRD enum, CLI, workspace connection descriptor, verifier, adoption compare, and reconciler. Coverage is unusually thorough for a change of this size and I did not find a functional or security regression.

Test Plan & Verification Results

Executed against the PR HEAD (51642e5) checked out under /home/vibe/taugrid:

• Build: ✅ go build ./... on both cli/ and controllers/tau-core/ clean.
• Unit tests (controllers/tau-core): ✅ go test ./api/... ./internal/... all pass, including the new TestWorkspaceRoleAliasIsBackwardCompatible-style loop that verifies switching spec.role between researcher and tau-researcher-v1 produces byte-identical RoleBinding/ClusterRoleBinding/Role lists.
• Unit tests (cli): ✅ internal/serve, internal/workspace, internal/workspaceconnection, internal/reposcaffold. ⚠️ 1 failure in internal/cli: TestLoadServeAppArgsRejectsInvalidInput/non-string_nested_key — the reviewer sandbox ships Go 1.27.1 and GOTOOLCHAIN=auto did not fetch the pinned go1.26.7, so encoding/json now stringifies non-string interface keys instead of failing. Not a PR bug (author's CI runs on the pinned toolchain), but see the note in Findings.
• Lint/vet: ✅ go vet ./... clean on both modules.
• CRD sync check: ✅ md5sum matches between charts/tau-core-controller/crds/tau.azure.com_workspaces.yaml and controllers/tau-core/config/crd/bases/tau.azure.com_workspaces.yaml.
• Cluster deployment / e2e: Skipped for this PR — it is a CLI + CRD schema change, not a chart/controller behavior change, and there is no reachable cluster surface that would exercise the new --nodes/--shm-size/--app-args render paths without also standing up a full KubeRay + Kueue stack. Behavioral coverage is instead exercised by the new TestRenderDistributedRayService, TestServeSnapshotProvenanceOnDeploymentChildren, and the eight-rank snapshot fixture tests — all pass.
• Regression check: ✅ Full unit test suites of cli and controllers/tau-core pass (the single failure above excepted).

Security Scan

• Secrets/tokens: ✅ none introduced. HF_TOKEN in tests uses envspec.Secret(...) valueFrom.
• Privileged/hostNetwork/hostPID/capabilities: ✅ none added.
• RBAC changes: ✅ no new verbs or resources. The researcher alias binds the existing tau-researcher-v1 ClusterRole — no permission surface change (asserted by TestKubectlVerifierResearcherRoleCompatibility/alias cannot bypass permissions and by the RBAC-parity assertions in controller_test.go).
• Container images: ✅ n/a — no image references changed.
• Network exposure: ✅ no new Services / Ingress / removed NetworkPolicies. Ray Serve http_options.host: 0.0.0.0 was already the shape; only the port line is now emitted explicitly.
• Snapshot boundary: ✅ good defence — resolveSnapshotServeProfile rejects --dry-run=server, empty apply, or --context overrides; snapshot mode never touches the cluster (Restore: func() {}), and TestServeSnapshotCannotAuthorizeLiveOperations locks that down.

Breaking Changes

• CRD enum broadened, not narrowed (enum: [tau-researcher-v1] → [researcher, tau-researcher-v1]): existing TauWorkspace objects with role: tau-researcher-v1 remain valid. New CRD must be installed before CLI-generated objects with role: researcher are applied — this is explicitly called out in cli/README.md L102-107. Good.
• Default researcher role in newly-generated artefacts changes (create, reposcaffold, workspace connection descriptor) from tau-researcher-v1 → researcher. Adoption comparison in sameAdoptionIntent treats both as equal, and KubectlVerifier accepts either side, so re-running tau workspace create against a live workspace that still says tau-researcher-v1 is a no-op. Verified in TestResearcherRoleCompatibility.
• selectServeWorkloadProfile signature changed (added explicitNodes *int, kind string). Only internal callers; no exported break.
• Ray Serve config now emits http_options.host/port, and RenderDeployment/Ray render swapped a custom yamlWriter for strings.Builder (identical output). Verified by the surrounding regression tests.
• No proto, no gRPC, no downstream ConfigMap/Secret shape change.

Findings

Nit — serve.go L163-164, --nodes upper bound. The check int64(nodes) > 2147483647 compares an int (a value that Go already caps at math.MaxInt64 on 64-bit) after casting to int64. On 64-bit hosts this bound only triggers when someone passes an explicit value > 2^31-1, which is fine, but the intent (Workers is int32 downstream, so we need to fit in an int32) reads more clearly as nodes > math.MaxInt32. Not blocking.

Observation — render.go L168-171. copyResources(p.Resources.Requests) shallow-copies the top-level map; nested values (e.g. quantities as resource.Quantity) still share references. The tests confirm the caller's profile isn't mutated for the current render paths, but if a future callsite mutates a nested map[string]any from a resource entry the aliasing would surface. Consider using a deep copy helper (or runtime.DeepCopyJSON-style) if that ever becomes possible. Not blocking.

Observation — Go toolchain drift. The single failing non-string_nested_key case (see Test Plan) is a Go 1.27 stdlib behavior change (encoding/json now marshals map[interface{}]interface{}{1: "value"} as {"1":"value"} instead of erroring). This test asserts on the pre-1.27 error path. Author's CI on the pinned 1.26.7 will pass, but this test will start failing once the repo bumps .go-version to any release ≥ 1.27. Cheap follow-up: assert against actual invalid inputs (e.g. func(){}, chan int) that encoding/json still rejects. Not blocking this PR.

Positive callouts:

• The failure-mode matrix in TestServeSnapshotFailsClosed (namespace required, context conflict, scope mismatch, GPU/node conflict, deployment cardinality, unavailable profile, startup args, hash tamper) is a great template for future offline-render features.
• TestRenderDistributedRayService explicitly asserts head has no GPU tolerations, no GPU resource requests, no worker TAS annotations, and workers carry matchLabelKeys: [ray.io/cluster] for anti-affinity across rollouts. Exactly the right invariants.
• Snapshot provenance annotation tau.azure.com/workload-profile-source: snapshot propagates to child Deployment/Service/HPA (TestServeSnapshotProvenanceOnDeploymentChildren) — makes offline-rendered artefacts trivially auditable.

Verdict

APPROVE. No hard-block conditions met. Build, vet, and unit tests pass under both modules; RBAC-parity, snapshot-safety, and CRD-sync invariants are covered by new tests; description matches code; alias migration is documented in cli/README.md.

[gossion]

Review found one blocking permission issue.

[yliu382]

Summary

Re-review of the delta since my previous APPROVE (51642e5 → 9fd79e4). The two new commits fill a natural gap left by the earlier multi-node-serving work: the researcher role and CLI connection verifier now also cover RayService alongside RayJob, plus new tests and kind-e2e coverage. Delta is small (10 files, +142/-5), focused, and well tested. Approving.

Test Plan & Verification Results

Focused on the delta commits (0be8c24 grant researcher access, 9fd79e4 controller smoke fixtures) rather than re-running everything on the previously approved base.

• Build: ✅ go build ./... in both cli/ and controllers/tau-core/
• Unit tests: ✅
  • cli/internal/workspaceconnection — passes (covers the new TestKubectlVerifierRejectsMissingRayServicePermission matrix across researcher / tau-researcher-v1 × create/get/list/patch/delete, and the existing role-compat suite now including rayservices.ray.io)
  • controllers/tau-core/internal/controller — passes (TestKustomizeResearcherRoleGrantsRayServicePermissions decodes the kustomize rbac.yaml, finds the tau-researcher-v1 ClusterRole, and asserts the exact RayJob+RayService rule)
• Lint/vet: ✅ go vet clean in touched packages
• Helm: ✅
  • helm lint charts/tau-core-controller clean
  • helm unittest --strict charts/tau-core-controller — all 10 tests pass, including new Researcher RayService permissions suite
  • kubectl kustomize charts/tau-core-controller renders cleanly
• Cluster deployment: skipped — taugrid is not deployed to the aks-ai-runtime reviewer cluster; kind-e2e.sh boundary checks are inspected statically
• Behavioral verification: verifier now issues kubectl auth can-i {create,get,list,patch,delete} rayservices.ray.io in the workspace namespace, in the same style as the existing RayJob checks. kind-e2e.sh installs a mock rayservices.ray.io CRD, waits for condition=Established (nice hardening — previously CRDs were applied without a readiness gate), and boundary-checks that all seven RayService verbs are yes in the target namespace and no in the system namespace.
• Regression: ✅ No changes outside the RayService-permission surface; previously validated multi-node-serving behavior, --app-args, --workload-profile-snapshot, and researcher/tau-researcher-v1 alias behavior are untouched.

Security Scan

• Secrets/tokens: ✅ None
• Privilege escalation: ✅ No new privileged / hostNetwork / hostPID / added capabilities
• RBAC changes: ⚠️ (justified) — tau-researcher-v1 ClusterRole extended to include rayservices in the same rule that already granted rayjobs (create,get,list,watch,delete,patch,update). This is bound via the existing per-workspace RoleBinding (namespace-scoped), not cluster-wide, so the effective grant is scoped to a workspace's target namespace and follows the established pattern. Justified in the PR description and required by the multi-node RayService serving feature merged earlier in this PR.
• Container images: ✅ No image changes
• Secret handling: ✅ No secret env/volume changes
• Network exposure: ✅ No Service/Ingress/NetworkPolicy changes

Breaking Changes

• Workspace RBAC contract change: KubectlVerifier.Verify now requires create/get/list/patch/delete on rayservices.ray.io. Any pre-existing workspace that hasn't upgraded the Helm/Kustomize RBAC will fail connection verification with the new CLI. This is clearly and repeatedly documented (PR body, cli/README.md, site/content/en/docs/developer-guide/serve-model.md, site/content/en/docs/platform-admin-guide/setup-guides/handoff.md), with explicit upgrade instructions ("Install the updated Helm or Kustomize RBAC before connecting with this CLI"). Migration guidance is sufficient — not a hard block.

Findings

Nothing blocking. A few small observations:

• controllers/tau-core/internal/controller/workspace_rbac_test.go uses a repo-relative path (../../../../charts/tau-core-controller/kustomize/rbac.yaml) to load the kustomize file. This is fine because Go tests run with pwd == package dir, but it does couple the controller test to the repo layout — noting for future refactors, not blocking.
• Nice touch: kind-e2e.sh now kubectl wait --for=condition=Established on the mock workload CRDs. This closes a real race window where the RESTMapper could reject subsequent auth can-i calls against rayservices.ray.io before the CRD was fully registered. Recommend keeping this pattern for any future mock CRDs added to the script.
• The new verifier subtest matrix (role × verb) explicitly proves that either role alias enforces every RayService permission — great coverage against future silent regressions in the alias-compat path.

Verdict

APPROVE — the delta is a clean, well-tested extension of the RayJob RBAC pattern to RayService, matching the multi-node RayService serving feature already approved in this PR. All local checks (build, unit, vet, helm lint/unittest, kustomize render) pass, the upgrade requirement is thoroughly documented, and the RBAC change stays inside the existing workspace-scoped RoleBinding boundary.

[gossion]

LGTM

[Li Chen (AaronLiChen)]

Reviewed the current head. The new unit/CLI tests cover rendering, snapshot validation, app args, role aliases, and RBAC. The Kind controller smoke covers the logical role and RayService auth can-i checks, but it uses a permissive mock RayService CRD and does not create a RayService. The runtime Kind smoke installs real KubeRay/Kueue but exercises only Job and RayJob. Thus there is still no end-to-end coverage for RayService admission, generated head/worker Pods, custom ports, shared memory, or readiness.

[yliu382]

Summary

Re-review of the delta since my previous APPROVE (9fd79e4 → d6c2c7f). The single new commit "Fix RayService Contracts and Upgrade Tau CRDs" is a focused, well-tested extension of the multi-node serving work: (a) the CLI now upgrades tau.azure.com CRDs from the selected chart before helm upgrade, (b) the distributed Ray renderer publishes the Serve port on workers, adds a CriticalAddonsOnly toleration on the CPU head, and propagates topology labels/pod-priority to the head, and (c) the kind smoke script now exercises RayService end-to-end and the legacy → aliased researcher CRD upgrade. Delta is 16 files, +564/-24. Approving.

Test Plan & Verification Results

Focused on the delta (d6c2c7f) — previous base already approved.

• Build: ✅ go build ./... in cli/ clean.
• Unit tests: ✅
  • cli/internal/cli — ok (0.32s). Covers new TestClusterInstallUpgradesCRDsBeforeExistingRelease (validates render → server-dry-run → apply → wait order, plus failure propagation at each step), TestTauGridCRDManifestExcludesWorkloadsAndThirdPartyCRDs (filters out rayservices.ray.io and user CRs), TestKindSmokeExercisesRayServiceAndLegacyCRDUpgrade, and TestServeRayWorkloadPriority (workspace + pod priority separation).
  • cli/internal/serve — ok (0.02s). Covers TestRenderKindRayServiceFixture (offline fixture generation) and TestDistributedRayServicePorts (Serve port matches on head, workers, and serveConfigV2.http_options.port for both 8000 and 9000; workers do not advertise dashboard/gcs-server).
• Lint/vet: ✅ go vet ./... clean on cli/.
• Cluster deployment: skipped — CRD upgrade path is a pre-helm upgrade step that runs against a live cluster with tau-core-controller Helm state, which isn't installed on the aks-ai-runtime reviewer cluster. The kind-smoke script is inspected statically and asserted via TestKindSmokeExercisesRayServiceAndLegacyCRDUpgrade.
• Behavioral verification (from tests + code reading):
  • upgradeTauGridCRDs renders the exact chart+version+values used for install (clusterInstallRenderArgs), filters manifest to apiextensions.k8s.io/v1 CRDs in group tau.azure.com, runs kubectl apply --dry-run=server first, then kubectl apply --field-manager=taugrid-crds, then kubectl wait --for=condition=Established -f - with spec.Timeout. Existing CRs are retained (no --prune). Third-party CRDs (rayservices.ray.io, kueue, etc.) are explicitly excluded, and non-CRD tau resources like TauWorkspace are filtered by kind check.
  • Head pod tolerations now [{key: CriticalAddonsOnly, operator: Exists, effect: NoSchedule}], matching the AKS system pool taint pattern already used elsewhere. GPU-only tolerations remain worker-only.
  • Options.ServePort is now honored on both head and workers, matching serveConfigV2.http_options.port. Worker containers keep only the serve port; dashboard (8265) and gcs-server (6379) remain head-only.
  • render.go now merges topoPlan.Labels into RayService labels and sets priorityClassName on the head pod spec when the topology plan supplies one — the two-level (workload vs pod) priority split is validated by TestServeRayWorkloadPriority.
• Regression: ✅ Previously validated behavior — multi-node-serving contract, --app-args, --workload-profile-snapshot, researcher/tau-researcher-v1 alias, and RayService RBAC — is untouched. The head-toleration change is intentional (previous test asserted no tolerations; it is updated in the same commit to assert the new CriticalAddonsOnly toleration).

Security Scan

• Secrets/tokens: ✅ None.
• Privilege escalation: ✅ No new privileged/hostNetwork/hostPID/capabilities. The head toleration for CriticalAddonsOnly:NoSchedule is a scheduling hint, not a privilege escalation.
• RBAC changes: ✅ None in this delta. taugrid-crds is used as a Server-Side-Apply --field-manager, not a Kubernetes principal.
• Container images: ✅ No image changes. Fixture references mcr.microsoft.com/aks/ai-runtime/ray:....
• Secret handling: ✅ No secret env/volume changes.
• Network exposure: ✅ No new Services/Ingress/NetworkPolicies. Ports advertised on worker containers are containerPorts only; no cluster-scoped exposure.

Breaking Changes

• CRD upgrade behavior change on existing clusters (documented): tau cluster install now updates tau.azure.com CRDs before every helm upgrade. Existing CRs are retained (no --prune), and this is called out in both the plan output ("Tau CRDs: update from the selected chart before an existing-release upgrade; not rolled back by Helm") and the command long-form help. This is the intended fix for the CRD-vs-controller version drift and is safe under Kubernetes' additive-CRD semantics.
• Head pod scheduling change: CPU head now tolerates CriticalAddonsOnly:NoSchedule. This can cause the head to land on system nodes it previously would have been repelled from, which is desirable on AKS but worth noting for downstream operators who tightly constrain system-pool placement. Non-blocking; consistent with system-node affinity already present.
• Ray worker container port surface change: Worker containers now declare the Serve container port. This changes the pod spec shape but has no runtime impact for RayJob (which doesn't create a Service) and enables per-worker Serve routing for RayService. Documented indirectly through the new tests and kind-smoke coverage.

Findings

Nothing blocking. Small observations only:

• cli/internal/cli/cluster_install_crds.go:69 — the errors.As(err, &typeErr) branch handles the Updating chart dependencies progress lines that helm template emits before the first YAML document; the test TestTauGridCRDManifestExcludesWorkloadsAndThirdPartyCRDs exercises this exact case. Nice defensive parse.
• cli/internal/cli/cluster_install_crds.go:40 — after runner.Raw(... apply ...) returns, out is printed before the error check. This is correct (we want configured/unchanged output visible even when the apply fails, so operators can see partial progress), but a brief comment noting the intent would help future readers. Non-blocking.
• cli/scripts/kind-smoke-e2e.sh:454 — the legacy-workspace regression flow (patch CRD → create with legacy role → reinstall → assert UID preserved and role updated) is a strong migration test. The --dry-run=server role-update verification is a nice touch: it confirms the new schema accepts the alias without mutating the smoke resource.
• cli/internal/serve/ray_cluster.go:124 — head port list is built by appending the head-only ports to the worker port slice; since ports is reassigned each pod iteration this is safe. If future changes ever hoist ports out of the loop, watch for aliasing.

Verdict

APPROVE. Delta is focused, well tested, and matches the documented behavior. All local checks (build, vet, unit tests for cli/internal/cli and cli/internal/serve) pass. The new CRD-upgrade path is guarded by server dry-run + Established wait + explicit chart-version parity, and the Ray renderer changes are covered by both unit fixtures and the kind smoke script.

[Pengfei Ni (feiskyer)]

Li Chen (@AaronLiChen) Added real Kueue/KubeRay RayService smoke coverage in d6c2c7f, with fixture memory sizing corrected in c560465: admission priority, two ready workers on port 9000, memory-backed shared memory, and an HTTP response. The fixture is CPU-adapted; GPU allocation and cross-host placement remain unit-tested, not GPU E2E claims. All CI is now green, including both Kind smoke jobs.