From 03b9f9170cedf09cb9330eb1d9fcd2778950cb81 Mon Sep 17 00:00:00 2001 From: chokevin Date: Tue, 20 Oct 2026 12:53:50 -0700 Subject: [PATCH 01/13] fix(controller): clean retained workspace namespace metadata --- .../internal/controller/controller_test.go | 89 ++++++++++++++++++- .../internal/controller/workspace_cleanup.go | 9 +- 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index a688a007..e11777a2 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -1257,9 +1257,14 @@ func TestWorkspaceDeleteCleansWorkspaceAccess(t *testing.T) { now := metav1.Now() workspace.Finalizers = []string{workspaceFinalizer} workspace.DeletionTimestamp = &now + workspace.Spec.Target.Namespace = "new-aurora" workspace.Status.Target.ResolvedNamespace = "aurora" targetNamespace := testNamespace("aurora") - targetNamespace.Labels = map[string]string{labelWorkspace: "aurora"} + targetNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + targetNamespace.Labels[labelTeam] = "vision" + targetNamespace.Labels["preserved"] = "true" + targetNamespace.Annotations[annotationResultScope] = workspace.Spec.Defaults.OutputRoot + targetNamespace.Annotations["preserved"] = "true" targetRole := testRole("aurora", defaultRoleName, "aurora") targetBinding := testRoleBinding("aurora", defaultRoleName, "aurora") targetServiceAccount := testServiceAccount("aurora", "tau-workload", "aurora") @@ -1298,6 +1303,35 @@ func TestWorkspaceDeleteCleansWorkspaceAccess(t *testing.T) { if err := c.Get(ctx, client.ObjectKey{Name: "aurora"}, targetNamespace); err != nil { t.Fatalf("target namespace was deleted: %v", err) } + for _, key := range []string{ + labelManagedBy, + labelTeam, + labelWorkspace, + labelWorkspaceLocalQueue, + labelKueueDefaultLocalQueue, + } { + if _, ok := targetNamespace.Labels[key]; ok { + t.Fatalf("target namespace retained workspace label %q: %#v", key, targetNamespace.Labels) + } + } + for _, key := range []string{annotationOwnerUID, annotationResultScope} { + if _, ok := targetNamespace.Annotations[key]; ok { + t.Fatalf("target namespace retained workspace annotation %q: %#v", key, targetNamespace.Annotations) + } + } + for key, want := range map[string]string{ + labelPSAEnforce: "baseline", + labelPSAAudit: "restricted", + labelPSAWarn: "restricted", + "preserved": "true", + } { + if got := targetNamespace.Labels[key]; got != want { + t.Fatalf("target namespace label %q = %q, want %q", key, got, want) + } + } + if got := targetNamespace.Annotations["preserved"]; got != "true" { + t.Fatalf("target namespace preserved annotation = %q, want true", got) + } if err := c.Get(ctx, client.ObjectKey{Name: clusterQueueReaderBindingName("aurora")}, &rbacv1.ClusterRoleBinding{}); !apierrors.IsNotFound(err) { t.Fatalf("expected ClusterQueue reader ClusterRoleBinding deleted, got err=%v", err) } @@ -1308,6 +1342,59 @@ func TestWorkspaceDeleteCleansWorkspaceAccess(t *testing.T) { } } +func TestWorkspaceDeletePreservesNamespaceMetadataForDifferentOwnerUID(t *testing.T) { + ctx := context.Background() + scheme := testScheme(t) + workspace := testWorkspace("aurora") + now := metav1.Now() + workspace.Finalizers = []string{workspaceFinalizer} + workspace.DeletionTimestamp = &now + workspace.Status.Target.ResolvedNamespace = "aurora" + targetNamespace := testNamespace("aurora") + targetNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + targetNamespace.Labels[labelTeam] = "vision" + targetNamespace.Labels["preserved"] = "true" + targetNamespace.Annotations[annotationOwnerUID] = "different-workspace-uid" + targetNamespace.Annotations[annotationResultScope] = workspace.Spec.Defaults.OutputRoot + targetNamespace.Annotations["preserved"] = "true" + c := fake.NewClientBuilder(). + WithScheme(scheme). + WithObjects(workspace, targetNamespace). + Build() + reconciler := newTestWorkspaceReconciler(c) + + if _, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: types.NamespacedName{Name: "aurora", Namespace: tauv1alpha1.SystemNamespace}}); err != nil { + t.Fatalf("Reconcile delete error = %v", err) + } + if err := c.Get(ctx, client.ObjectKey{Name: "aurora"}, targetNamespace); err != nil { + t.Fatalf("target namespace was deleted: %v", err) + } + for key, want := range map[string]string{ + labelManagedBy: labelManagedByValue, + labelTeam: "vision", + labelWorkspace: workspace.Name, + labelWorkspaceLocalQueue: workspace.Spec.Queue, + labelKueueDefaultLocalQueue: workspace.Spec.Queue, + labelPSAEnforce: "baseline", + labelPSAAudit: "restricted", + labelPSAWarn: "restricted", + "preserved": "true", + } { + if got := targetNamespace.Labels[key]; got != want { + t.Fatalf("target namespace label %q = %q, want %q", key, got, want) + } + } + for key, want := range map[string]string{ + annotationOwnerUID: "different-workspace-uid", + annotationResultScope: workspace.Spec.Defaults.OutputRoot, + "preserved": "true", + } { + if got := targetNamespace.Annotations[key]; got != want { + t.Fatalf("target namespace annotation %q = %q, want %q", key, got, want) + } + } +} + func TestWorkspaceDeleteWaitsForClusterQueueFinalizer(t *testing.T) { ctx := context.Background() scheme := testScheme(t) diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 6009f4ba..40672f53 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -121,7 +121,14 @@ func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, wor if err := r.cleanupWorkspaceClusterQueue(ctx, workspace); err != nil { return err } - return r.cleanupSystemReaderRBAC(ctx, workspace) + if err := r.cleanupSystemReaderRBAC(ctx, workspace); err != nil { + return err + } + targetNamespace := strings.TrimSpace(workspace.Status.Target.ResolvedNamespace) + if targetNamespace == "" { + targetNamespace = resolvedNamespace(workspace) + } + return r.cleanupStaleNamespaceMetadata(ctx, workspace, targetNamespace) } func (r *TauWorkspaceReconciler) cleanupWorkspaceClusterQueue(ctx context.Context, workspace *tauv1alpha1.TauWorkspace) error { From 46151fd69b954d9150ed996f3ecdd49ea6784f81 Mon Sep 17 00:00:00 2001 From: chokevin Date: Tue, 20 Oct 2026 13:51:06 -0700 Subject: [PATCH 02/13] test(controller): align retarget deletion expectation --- controllers/tau-core/internal/controller/controller_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index e11777a2..dc50ad80 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -2375,9 +2375,9 @@ func TestUnmigratedWorkspaceRetargetMigratesHistoryBeforeDeletion(t *testing.T) if err := deletionClient.Get(ctx, client.ObjectKeyFromObject(newNamespace), &gotNew); err != nil { t.Fatalf("Get new desired Namespace: %v", err) } - if gotNew.Annotations[annotationOwnerUID] != string(workspace.UID) || - gotNew.Labels[labelWorkspace] != workspace.Name { - t.Fatalf("normal retarget deletion did not preserve adopted Namespace ownership: labels=%#v annotations=%#v", gotNew.Labels, gotNew.Annotations) + if gotNew.Annotations[annotationOwnerUID] != "" || + gotNew.Labels[labelWorkspace] != "" { + t.Fatalf("normal retarget deletion retained adopted Namespace ownership: labels=%#v annotations=%#v", gotNew.Labels, gotNew.Annotations) } } From 0a391ac5a3519fbc165466908f0618fe9838d05c Mon Sep 17 00:00:00 2001 From: chokevin Date: Tue, 20 Oct 2026 16:45:59 -0700 Subject: [PATCH 03/13] fix(controller): preserve active queue routing metadata --- .../internal/controller/controller_test.go | 273 ++++++++++++++++++ .../internal/controller/workspace_cleanup.go | 124 +++++++- 2 files changed, 396 insertions(+), 1 deletion(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index dc50ad80..aba944f3 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -1395,6 +1395,260 @@ func TestWorkspaceDeletePreservesNamespaceMetadataForDifferentOwnerUID(t *testin } } +func TestWorkspaceDeletePreservesNamespaceRoutingMetadataForExternalQueueWorkloads(t *testing.T) { + for _, test := range []struct { + name string + status map[string]any + wantBlock bool + }{ + {name: "pending", wantBlock: true}, + { + name: "admitted", + status: map[string]any{ + "admission": map[string]any{"clusterQueue": "jobqueue"}, + }, + wantBlock: true, + }, + { + name: "finished", + status: map[string]any{ + "conditions": []any{map[string]any{"type": "Finished", "status": "True"}}, + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("training") + workload.SetNamespace(namespace.Name) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + if test.status != nil { + workload.Object["status"] = test.status + } + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, workload). + Build() + reconciler := newTestWorkspaceReconciler(c) + + _, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)}) + if test.wantBlock { + if err == nil || !strings.Contains(err.Error(), "before removing namespace routing metadata") { + t.Fatalf("Reconcile() error = %v, want routing metadata retention", err) + } + } else if err != nil { + t.Fatalf("Reconcile() error = %v", err) + } + + var gotNamespace corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if test.wantBlock { + if gotNamespace.Labels[labelWorkspace] != workspace.Name || + gotNamespace.Annotations[annotationOwnerUID] != string(workspace.UID) { + t.Fatalf("namespace routing metadata was removed while workload remained: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + var remaining tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &remaining); err != nil { + t.Fatal(err) + } + if !controllerutil.ContainsFinalizer(&remaining, workspaceFinalizer) { + t.Fatal("workspace finalizer was removed while routing metadata remained in use") + } + } else { + if gotNamespace.Labels[labelWorkspace] != "" || + gotNamespace.Annotations[annotationOwnerUID] != "" { + t.Fatalf("drained namespace retained stale metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &tauv1alpha1.TauWorkspace{}); !apierrors.IsNotFound(err) { + t.Fatalf("workspace finalizer was not removed after workload drain: %v", err) + } + } + for _, queue := range []*unstructured.Unstructured{localQueue, clusterQueue} { + got := newQueueObject(queue.GroupVersionKind()) + if err := c.Get(ctx, client.ObjectKeyFromObject(queue), got); err != nil { + t.Fatalf("external %s was mutated or deleted: %v", queue.GetKind(), err) + } + } + }) + } +} + +func TestWorkspaceDeleteIgnoresUnrelatedWorkloadsWhenCleaningNamespaceMetadata(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + namespace.Labels["queue-scope"] = "shared" + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := testClusterQueue("jobqueue") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchLabels": map[string]any{"queue-scope": "shared"}, + }, + } + selectorIndependent := newQueueObject(workloadGVK) + selectorIndependent.SetName("selector-independent") + selectorIndependent.SetNamespace(namespace.Name) + selectorIndependent.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + unrelated := newQueueObject(workloadGVK) + unrelated.SetName("unrelated") + unrelated.SetNamespace(namespace.Name) + unrelated.Object["spec"] = map[string]any{"queueName": "missing-queue"} + foreign := newQueueObject(workloadGVK) + foreign.SetName("foreign") + foreign.SetNamespace("other") + foreign.Object["spec"] = map[string]any{"queueName": "jobqueue"} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, selectorIndependent, unrelated, foreign). + Build() + + if _, err := newTestWorkspaceReconciler(c).Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)}); err != nil { + t.Fatalf("Reconcile() error = %v", err) + } + var gotNamespace corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if gotNamespace.Labels[labelWorkspace] != "" || + gotNamespace.Annotations[annotationOwnerUID] != "" { + t.Fatalf("unrelated workloads retained stale namespace metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } +} + +func TestWorkspaceDeleteFailsClosedWhenExternalQueueRoutingCannotBeVerified(t *testing.T) { + for _, test := range []struct { + name string + failGVK schema.GroupVersionKind + failKey client.ObjectKey + mutateQueue func(*unstructured.Unstructured) + wantError string + }{ + { + name: "local queue read", + failGVK: localQueueGVK, + failKey: client.ObjectKey{Name: "jobqueue", Namespace: "aurora"}, + wantError: "read LocalQueue aurora/jobqueue", + }, + { + name: "cluster queue read", + failGVK: clusterQueueGVK, + failKey: client.ObjectKey{Name: "jobqueue"}, + wantError: "read ClusterQueue jobqueue", + }, + { + name: "malformed selector", + mutateQueue: func(queue *unstructured.Unstructured) { + queue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": "invalid", + }, + } + }, + wantError: "read ClusterQueue \"jobqueue\" namespaceSelector", + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := testClusterQueue("jobqueue") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchLabels": map[string]any{labelWorkspace: workspace.Name}, + }, + } + if test.mutateQueue != nil { + test.mutateQueue(clusterQueue) + } + workload := newQueueObject(workloadGVK) + workload.SetName("training") + workload.SetNamespace(namespace.Name) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + base := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, workload). + Build() + reconciler := newTestWorkspaceReconciler(base) + if !test.failGVK.Empty() { + reconciler.APIReader = &objectReadFailingClient{ + Client: base, + gvk: test.failGVK, + key: test.failKey, + err: errors.New("access denied"), + } + } + + _, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)}) + if err == nil || !strings.Contains(err.Error(), test.wantError) { + t.Fatalf("Reconcile() error = %v, want %q", err, test.wantError) + } + var gotNamespace corev1.Namespace + if err := base.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if gotNamespace.Labels[labelWorkspace] != workspace.Name || + gotNamespace.Annotations[annotationOwnerUID] != string(workspace.UID) { + t.Fatalf("failed routing verification removed namespace metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + }) + } +} + func TestWorkspaceDeleteWaitsForClusterQueueFinalizer(t *testing.T) { ctx := context.Background() scheme := testScheme(t) @@ -4700,6 +4954,25 @@ type statusUpdateFailingClient struct { updateErr error } +type objectReadFailingClient struct { + client.Client + gvk schema.GroupVersionKind + key client.ObjectKey + err error +} + +func (c *objectReadFailingClient) Get( + ctx context.Context, + key client.ObjectKey, + obj client.Object, + opts ...client.GetOption, +) error { + if obj.GetObjectKind().GroupVersionKind() == c.gvk && key == c.key { + return c.err + } + return c.Client.Get(ctx, key, obj, opts...) +} + func (c *statusUpdateFailingClient) Status() client.SubResourceWriter { return &statusUpdateFailingWriter{SubResourceWriter: c.Client.Status(), client: c} } diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 40672f53..900f1e54 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -5,6 +5,7 @@ package controller import ( "context" + "errors" "fmt" "strings" @@ -15,6 +16,8 @@ import ( "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/runtime/schema" "sigs.k8s.io/controller-runtime/pkg/client" ) @@ -81,7 +84,7 @@ func (r *TauWorkspaceReconciler) cleanupStaleNamespaceMetadata( return nil } var namespace corev1.Namespace - if err := r.Get(ctx, client.ObjectKey{Name: namespaceName}, &namespace); err != nil { + if err := r.apiReader().Get(ctx, client.ObjectKey{Name: namespaceName}, &namespace); err != nil { if apierrors.IsNotFound(err) { return nil } @@ -92,6 +95,13 @@ func (r *TauWorkspaceReconciler) cleanupStaleNamespaceMetadata( workspace.UID == "" { return nil } + inUse, message, err := r.namespaceRoutingMetadataInUse(ctx, workspace, &namespace) + if err != nil { + return err + } + if inUse { + return errors.New(message) + } for _, key := range []string{ labelManagedBy, labelTeam, @@ -108,6 +118,118 @@ func (r *TauWorkspaceReconciler) cleanupStaleNamespaceMetadata( return r.Update(ctx, &namespace) } +func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespace *corev1.Namespace, +) (bool, string, error) { + establishedLocalQueue := strings.TrimSpace(workspace.Status.Queue.LocalQueue) + if establishedLocalQueue == "" { + return false, "", nil + } + workloads := &unstructured.UnstructuredList{} + workloads.SetGroupVersionKind(schema.GroupVersionKind{ + Group: workloadGVK.Group, Version: workloadGVK.Version, Kind: workloadGVK.Kind + "List", + }) + if err := r.apiReader().List(ctx, workloads, client.InNamespace(namespace.Name)); err != nil { + return false, "", fmt.Errorf("list Workloads in namespace %q before routing metadata cleanup: %w", namespace.Name, err) + } + cleanedLabels := make(map[string]string, len(namespace.Labels)) + for key, value := range namespace.Labels { + cleanedLabels[key] = value + } + for _, key := range []string{ + labelManagedBy, + labelTeam, + labelWorkspace, + labelWorkspaceLocalQueue, + labelKueueDefaultLocalQueue, + } { + delete(cleanedLabels, key) + } + for i := range workloads.Items { + workload := &workloads.Items[i] + if workloadFinished(workload) { + continue + } + queueName, found, err := unstructured.NestedString(workload.Object, "spec", "queueName") + if err != nil { + return false, "", fmt.Errorf("read Workload %s queueName: %w", client.ObjectKeyFromObject(workload), err) + } + if !found || strings.TrimSpace(queueName) != establishedLocalQueue { + continue + } + localQueue := newQueueObject(localQueueGVK) + localQueueKey := client.ObjectKey{Name: establishedLocalQueue, Namespace: namespace.Name} + if err := r.apiReader().Get(ctx, localQueueKey, localQueue); err != nil { + if apierrors.IsNotFound(err) { + continue + } + return false, "", fmt.Errorf("read LocalQueue %s before routing metadata cleanup: %w", localQueueKey, err) + } + clusterQueueName, found, err := unstructured.NestedString(localQueue.Object, "spec", "clusterQueue") + if err != nil { + return false, "", fmt.Errorf("read LocalQueue %s backing ClusterQueue: %w", localQueueKey, err) + } + if !found || strings.TrimSpace(clusterQueueName) == "" { + return false, "", fmt.Errorf("LocalQueue %s has no backing ClusterQueue", localQueueKey) + } + admittedClusterQueue, admitted, err := unstructured.NestedString(workload.Object, "status", "admission", "clusterQueue") + if err != nil { + return false, "", fmt.Errorf("read Workload %s admitted ClusterQueue: %w", client.ObjectKeyFromObject(workload), err) + } + if admitted && strings.TrimSpace(admittedClusterQueue) != "" { + clusterQueueName = admittedClusterQueue + } + clusterQueue := newQueueObject(clusterQueueGVK) + if err := r.apiReader().Get(ctx, client.ObjectKey{Name: clusterQueueName}, clusterQueue); err != nil { + return false, "", fmt.Errorf("read ClusterQueue %s before routing metadata cleanup: %w", clusterQueueName, err) + } + if expectedUID := strings.TrimSpace(workspace.Status.Queue.ClusterQueueUID); expectedUID != "" && + clusterQueueName == strings.TrimSpace(workspace.Status.Queue.ClusterQueue) && + string(clusterQueue.GetUID()) != expectedUID { + return false, "", fmt.Errorf( + "ClusterQueue %q UID changed from %q to %q before routing metadata cleanup", + clusterQueueName, + expectedUID, + clusterQueue.GetUID(), + ) + } + selectorMap, found, err := unstructured.NestedMap(clusterQueue.Object, "spec", "namespaceSelector") + if err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + labelSelector := metav1.LabelSelector{} + if found { + if err := runtime.DefaultUnstructuredConverter.FromUnstructured(selectorMap, &labelSelector); err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + } + selector, err := metav1.LabelSelectorAsSelector(&labelSelector) + if err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + if !selector.Matches(labels.Set(namespace.Labels)) { + return false, "", fmt.Errorf( + "ClusterQueue %q namespaceSelector does not match namespace %q for unfinished Workload %s", + clusterQueueName, + namespace.Name, + client.ObjectKeyFromObject(workload), + ) + } + if selector.Matches(labels.Set(cleanedLabels)) { + continue + } + return true, fmt.Sprintf( + "waiting for Workload %s using retained LocalQueue %q and ClusterQueue %q before removing namespace routing metadata", + client.ObjectKeyFromObject(workload), + establishedLocalQueue, + clusterQueueName, + ), nil + } + return false, "", nil +} + func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, workspace *tauv1alpha1.TauWorkspace) error { if err := r.cleanupStaleTargetRBAC(ctx, workspace, "", "", false); err != nil { return err From a1fbc7a0dcb01f046df8caf698b444832e94a369 Mon Sep 17 00:00:00 2001 From: chokevin Date: Thu, 22 Oct 2026 10:22:17 -0700 Subject: [PATCH 04/13] fix(controller): discover retained workload routes --- .../internal/controller/controller_test.go | 118 ++++++++++++++++-- .../internal/controller/workspace_cleanup.go | 11 +- 2 files changed, 109 insertions(+), 20 deletions(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index aba944f3..2abd164a 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -1397,23 +1397,53 @@ func TestWorkspaceDeletePreservesNamespaceMetadataForDifferentOwnerUID(t *testin func TestWorkspaceDeletePreservesNamespaceRoutingMetadataForExternalQueueWorkloads(t *testing.T) { for _, test := range []struct { - name string - status map[string]any - wantBlock bool + name string + workloadStatus map[string]any + queueStatus tauv1alpha1.WorkspaceQueueStatus + teamRef *tauv1alpha1.TauClusterObjectReference + wantBlock bool }{ - {name: "pending", wantBlock: true}, + { + name: "pending", + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + }, + wantBlock: true, + }, { name: "admitted", - status: map[string]any{ + workloadStatus: map[string]any{ "admission": map[string]any{"clusterQueue": "jobqueue"}, }, + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + }, + wantBlock: true, + }, + {name: "pending with empty queue status", wantBlock: true}, + { + name: "pending on compatibility alias", + teamRef: &tauv1alpha1.TauClusterObjectReference{Name: "vision"}, + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: workspaceClusterQueueName("aurora"), + ClusterQueue: workspaceClusterQueueName("aurora"), + }, wantBlock: true, }, { name: "finished", - status: map[string]any{ + workloadStatus: map[string]any{ "conditions": []any{map[string]any{"type": "Finished", "status": "True"}}, }, + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + }, }, } { t.Run(test.name, func(t *testing.T) { @@ -1425,12 +1455,9 @@ func TestWorkspaceDeletePreservesNamespaceRoutingMetadataForExternalQueueWorkloa workspace.DeletionTimestamp = &now workspace.Spec.Target.CreateNamespace = false workspace.Spec.Queue = "jobqueue" + workspace.Spec.TeamRef = test.teamRef workspace.Status.Target.ResolvedNamespace = workspace.Name - workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ - LocalQueue: "jobqueue", - ClusterQueue: "jobqueue", - ClusterQueueUID: "jobqueue-uid", - } + workspace.Status.Queue = test.queueStatus namespace := testNamespace(workspace.Name) namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") @@ -1448,8 +1475,8 @@ func TestWorkspaceDeletePreservesNamespaceRoutingMetadataForExternalQueueWorkloa workload.SetName("training") workload.SetNamespace(namespace.Name) workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} - if test.status != nil { - workload.Object["status"] = test.status + if test.workloadStatus != nil { + workload.Object["status"] = test.workloadStatus } c := fake.NewClientBuilder(). WithScheme(testScheme(t)). @@ -1501,6 +1528,71 @@ func TestWorkspaceDeletePreservesNamespaceRoutingMetadataForExternalQueueWorkloa } } +func TestLegacyWorkspaceDeletePreservesRoutingMetadataWithEmptyQueueStatus(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationV0Primary: "true"} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + namespace.Annotations = nil + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("training") + workload.SetNamespace(namespace.Name) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, workload). + WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). + Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + + var cleanupErr error + for range 4 { + if _, err := reconciler.Reconcile(ctx, request); err != nil { + cleanupErr = err + break + } + } + if cleanupErr == nil || !strings.Contains(cleanupErr.Error(), "before removing namespace routing metadata") { + t.Fatalf("Reconcile() error = %v, want routing metadata retention after ownership migration", cleanupErr) + } + var gotNamespace corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if gotNamespace.Labels[labelWorkspace] != workspace.Name || + gotNamespace.Annotations[annotationOwnerUID] != string(workspace.UID) { + t.Fatalf("legacy deletion removed migrated routing metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + var remaining tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &remaining); err != nil { + t.Fatal(err) + } + if remaining.Status.Queue != (tauv1alpha1.WorkspaceQueueStatus{}) { + t.Fatalf("legacy deletion unexpectedly populated queue status: %#v", remaining.Status.Queue) + } + if !controllerutil.ContainsFinalizer(&remaining, workspaceFinalizer) { + t.Fatal("workspace finalizer was removed while migrated routing metadata remained in use") + } +} + func TestWorkspaceDeleteIgnoresUnrelatedWorkloadsWhenCleaningNamespaceMetadata(t *testing.T) { ctx := context.Background() workspace := testWorkspace("aurora") diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 900f1e54..9304d1de 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -123,10 +123,6 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( workspace *tauv1alpha1.TauWorkspace, namespace *corev1.Namespace, ) (bool, string, error) { - establishedLocalQueue := strings.TrimSpace(workspace.Status.Queue.LocalQueue) - if establishedLocalQueue == "" { - return false, "", nil - } workloads := &unstructured.UnstructuredList{} workloads.SetGroupVersionKind(schema.GroupVersionKind{ Group: workloadGVK.Group, Version: workloadGVK.Version, Kind: workloadGVK.Kind + "List", @@ -156,11 +152,12 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( if err != nil { return false, "", fmt.Errorf("read Workload %s queueName: %w", client.ObjectKeyFromObject(workload), err) } - if !found || strings.TrimSpace(queueName) != establishedLocalQueue { + queueName = strings.TrimSpace(queueName) + if !found || queueName == "" { continue } localQueue := newQueueObject(localQueueGVK) - localQueueKey := client.ObjectKey{Name: establishedLocalQueue, Namespace: namespace.Name} + localQueueKey := client.ObjectKey{Name: queueName, Namespace: namespace.Name} if err := r.apiReader().Get(ctx, localQueueKey, localQueue); err != nil { if apierrors.IsNotFound(err) { continue @@ -223,7 +220,7 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( return true, fmt.Sprintf( "waiting for Workload %s using retained LocalQueue %q and ClusterQueue %q before removing namespace routing metadata", client.ObjectKeyFromObject(workload), - establishedLocalQueue, + queueName, clusterQueueName, ), nil } From fe6f70c3ae0bcdc9374446660b20eb6479f06e2f Mon Sep 17 00:00:00 2001 From: chokevin Date: Thu, 22 Oct 2026 13:13:51 -0700 Subject: [PATCH 05/13] fix(controller): retain workload identity during drain --- .../internal/controller/controller_test.go | 91 +++++++++++++++++++ .../internal/controller/workspace_cleanup.go | 72 +++++++++++---- 2 files changed, 144 insertions(+), 19 deletions(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index 2abd164a..e06601d1 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -1528,6 +1528,97 @@ func TestWorkspaceDeletePreservesNamespaceRoutingMetadataForExternalQueueWorkloa } } +func TestWorkspaceDeletePreservesWorkloadIdentityWhileExternalQueueWorkloadDrains(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + account := testServiceAccount(namespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: "training", + Namespace: namespace.Name, + UID: types.UID("training-job-uid"), + }, + Spec: batchv1.JobSpec{ + Template: corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{ + ServiceAccountName: account.Name, + RestartPolicy: corev1.RestartPolicyNever, + Containers: []corev1.Container{{ + Name: "trainer", + Image: "example.invalid/trainer", + Command: []string{"train"}, + }}, + }, + }, + }, + } + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(namespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: batchv1.SchemeGroupVersion.String(), + Kind: "Job", + Name: job.Name, + UID: job.UID, + }}) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, account, job, localQueue, clusterQueue, workload). + Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + + if _, err := reconciler.Reconcile(ctx, request); err == nil || + !strings.Contains(err.Error(), "before removing namespace routing metadata") { + t.Fatalf("Reconcile() error = %v, want workload drain wait", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { + t.Fatalf("workspace workload identity was removed before Workload drain: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(job), &batchv1.Job{}); err != nil { + t.Fatalf("pending Job was removed during Workload drain: %v", err) + } + + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + if err := c.Update(ctx, workload); err != nil { + t.Fatal(err) + } + if _, err := reconciler.Reconcile(ctx, request); err != nil { + t.Fatalf("Reconcile() after Workload drain = %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("workspace workload identity survived after Workload drain: %v", err) + } +} + func TestLegacyWorkspaceDeletePreservesRoutingMetadataWithEmptyQueueStatus(t *testing.T) { ctx := context.Background() workspace := testWorkspace("aurora") diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 9304d1de..43837f8a 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -80,22 +80,11 @@ func (r *TauWorkspaceReconciler) cleanupStaleNamespaceMetadata( workspace *tauv1alpha1.TauWorkspace, namespaceName string, ) error { - if reservedNamespaceReason(namespaceName, systemNamespace(r.SystemNamespace)) != "" { - return nil - } - var namespace corev1.Namespace - if err := r.apiReader().Get(ctx, client.ObjectKey{Name: namespaceName}, &namespace); err != nil { - if apierrors.IsNotFound(err) { - return nil - } + namespace, err := r.workspaceNamespaceForMetadataCleanup(ctx, workspace, namespaceName) + if err != nil || namespace == nil { return err } - if namespace.Labels[labelWorkspace] != workspace.Name || - namespace.Annotations[annotationOwnerUID] != string(workspace.UID) || - workspace.UID == "" { - return nil - } - inUse, message, err := r.namespaceRoutingMetadataInUse(ctx, workspace, &namespace) + inUse, message, err := r.namespaceRoutingMetadataInUse(ctx, workspace, namespace) if err != nil { return err } @@ -115,7 +104,49 @@ func (r *TauWorkspaceReconciler) cleanupStaleNamespaceMetadata( delete(namespace.Annotations, annotationResultScope) delete(namespace.Annotations, annotationOwnerUID) } - return r.Update(ctx, &namespace) + return r.Update(ctx, namespace) +} + +func (r *TauWorkspaceReconciler) workspaceNamespaceForMetadataCleanup( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespaceName string, +) (*corev1.Namespace, error) { + if reservedNamespaceReason(namespaceName, systemNamespace(r.SystemNamespace)) != "" { + return nil, nil + } + var namespace corev1.Namespace + if err := r.apiReader().Get(ctx, client.ObjectKey{Name: namespaceName}, &namespace); err != nil { + if apierrors.IsNotFound(err) { + return nil, nil + } + return nil, err + } + if namespace.Labels[labelWorkspace] != workspace.Name || + namespace.Annotations[annotationOwnerUID] != string(workspace.UID) || + workspace.UID == "" { + return nil, nil + } + return &namespace, nil +} + +func (r *TauWorkspaceReconciler) waitForNamespaceRoutingMetadataDrain( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespaceName string, +) error { + namespace, err := r.workspaceNamespaceForMetadataCleanup(ctx, workspace, namespaceName) + if err != nil || namespace == nil { + return err + } + inUse, message, err := r.namespaceRoutingMetadataInUse(ctx, workspace, namespace) + if err != nil { + return err + } + if inUse { + return errors.New(message) + } + return nil } func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( @@ -228,6 +259,13 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( } func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, workspace *tauv1alpha1.TauWorkspace) error { + targetNamespace := strings.TrimSpace(workspace.Status.Target.ResolvedNamespace) + if targetNamespace == "" { + targetNamespace = resolvedNamespace(workspace) + } + if err := r.waitForNamespaceRoutingMetadataDrain(ctx, workspace, targetNamespace); err != nil { + return err + } if err := r.cleanupStaleTargetRBAC(ctx, workspace, "", "", false); err != nil { return err } @@ -243,10 +281,6 @@ func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, wor if err := r.cleanupSystemReaderRBAC(ctx, workspace); err != nil { return err } - targetNamespace := strings.TrimSpace(workspace.Status.Target.ResolvedNamespace) - if targetNamespace == "" { - targetNamespace = resolvedNamespace(workspace) - } return r.cleanupStaleNamespaceMetadata(ctx, workspace, targetNamespace) } From 41b1b721ce8e0629fbbf5db12613028fc136073b Mon Sep 17 00:00:00 2001 From: chokevin Date: Thu, 22 Oct 2026 14:09:30 -0700 Subject: [PATCH 06/13] fix(controller): retain identity during retarget drain --- .../internal/controller/controller_test.go | 103 ++++++++++++++++++ .../internal/controller/workspace_sync.go | 14 ++- 2 files changed, 114 insertions(+), 3 deletions(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index e06601d1..e6243142 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -2818,6 +2818,109 @@ func TestUnmigratedWorkspaceRetargetMigratesHistoryBeforeDeletion(t *testing.T) } } +func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkloadDrains(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + workspace.Spec.Target.Namespace = "new-target" + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Phase = tauv1alpha1.WorkspacePhaseReady + workspace.Status.Target.ResolvedNamespace = "old-target" + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + } + oldNamespace := testNamespace("old-target") + oldNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + oldNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + newNamespace := testNamespace("new-target") + newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} + newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + account := testServiceAccount(oldNamespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: "training", + Namespace: oldNamespace.Name, + UID: types.UID("training-job-uid"), + }, + Spec: batchv1.JobSpec{ + Template: corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{ + ServiceAccountName: account.Name, + RestartPolicy: corev1.RestartPolicyNever, + Containers: []corev1.Container{{ + Name: "trainer", + Image: "example.invalid/trainer", + Command: []string{"train"}, + }}, + }, + }, + }, + } + localQueue := testLocalQueue(oldNamespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(oldNamespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: batchv1.SchemeGroupVersion.String(), + Kind: "Job", + Name: job.Name, + UID: job.UID, + }}) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, oldNamespace, newNamespace, account, job, localQueue, clusterQueue, workload). + WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). + Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + + if _, err := reconciler.Reconcile(ctx, request); err != nil { + t.Fatalf("Reconcile() during retarget drain = %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { + t.Fatalf("old workload identity was removed before Workload drain: %v", err) + } + var waiting tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != oldNamespace.Name { + t.Fatalf("resolved namespace = %q, want old target retained during drain", waiting.Status.Target.ResolvedNamespace) + } + + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + if err := c.Update(ctx, workload); err != nil { + t.Fatal(err) + } + reconcileWorkspaceTimes(t, reconciler, ctx, workspace.Name, 3) + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("old workload identity survived completed retarget: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != newNamespace.Name { + t.Fatalf("resolved namespace = %q, want completed retarget %q", waiting.Status.Target.ResolvedNamespace, newNamespace.Name) + } +} + func TestUnmigratedWorkspaceWithoutDefaultMigratesHistoryBeforeDeletion(t *testing.T) { ctx := context.Background() scheme := testScheme(t) diff --git a/controllers/tau-core/internal/controller/workspace_sync.go b/controllers/tau-core/internal/controller/workspace_sync.go index d8e1a405..c3b93c0f 100644 --- a/controllers/tau-core/internal/controller/workspace_sync.go +++ b/controllers/tau-core/internal/controller/workspace_sync.go @@ -103,9 +103,17 @@ func (r *TauWorkspaceReconciler) syncWorkspace(ctx context.Context, workspace *t } previousNamespace := workspace.Status.Target.ResolvedNamespace previousAccessReady := true - if err := r.cleanupOwnedNamespaceAccess(ctx, workspace, targetNamespace, keepServiceAccount); err != nil { - previousAccessReady = false - conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) + if previousNamespace != "" && previousNamespace != targetNamespace { + if err := r.waitForNamespaceRoutingMetadataDrain(ctx, workspace, previousNamespace); err != nil { + previousAccessReady = false + conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "NamespaceCleanupFailed", err.Error(), workspace.Generation)) + } + } + if previousAccessReady { + if err := r.cleanupOwnedNamespaceAccess(ctx, workspace, targetNamespace, keepServiceAccount); err != nil { + previousAccessReady = false + conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) + } } currentTargetReady := namespaceReady && rbacReady && rbacErr == nil && queueReady && previousAccessReady From ecfa4b616e195987918e70e4c51558d5422a0354 Mon Sep 17 00:00:00 2001 From: chokevin Date: Fri, 23 Oct 2026 16:02:39 -0700 Subject: [PATCH 07/13] fix(controller): reconcile access during retarget drain --- .../internal/controller/controller_test.go | 16 +++++++++++++++- .../internal/controller/workspace_cleanup.go | 6 +++--- .../internal/controller/workspace_sync.go | 6 +++++- 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index e6243142..214546b6 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -2840,6 +2840,12 @@ func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkload newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} account := testServiceAccount(oldNamespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + binding := testRoleBinding(oldNamespace.Name, defaultRoleName, workspace.Name) + binding.Subjects = []rbacv1.Subject{{ + Kind: rbacv1.GroupKind, + APIGroup: rbacv1.GroupName, + Name: "former-researchers", + }} job := &batchv1.Job{ ObjectMeta: metav1.ObjectMeta{ Name: "training", @@ -2883,7 +2889,7 @@ func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkload workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} c := fake.NewClientBuilder(). WithScheme(testScheme(t)). - WithObjects(workspace, oldNamespace, newNamespace, account, job, localQueue, clusterQueue, workload). + WithObjects(workspace, oldNamespace, newNamespace, account, binding, job, localQueue, clusterQueue, workload). WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). Build() reconciler := newTestWorkspaceReconciler(c) @@ -2895,6 +2901,14 @@ func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkload if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { t.Fatalf("old workload identity was removed before Workload drain: %v", err) } + var gotBinding rbacv1.RoleBinding + if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &gotBinding); err != nil { + t.Fatalf("Get old researcher RoleBinding: %v", err) + } + wantSubjects := []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, oldNamespace.Name)} + if !reflect.DeepEqual(gotBinding.Subjects, wantSubjects) { + t.Fatalf("old researcher subjects = %#v, want current subject %#v", gotBinding.Subjects, wantSubjects) + } var waiting tauv1alpha1.TauWorkspace if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { t.Fatal(err) diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 43837f8a..466e08c0 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -394,10 +394,10 @@ func (r *TauWorkspaceReconciler) cleanupStaleTargetRBAC( return nil } -func (r *TauWorkspaceReconciler) cleanupOwnedNamespaceAccess( +func (r *TauWorkspaceReconciler) cleanupOwnedNamespaceWorkloadIdentity( ctx context.Context, workspace *tauv1alpha1.TauWorkspace, - targetNamespace, keepServiceAccount string, + keepServiceAccount string, ) error { var namespaces corev1.NamespaceList if err := r.List(ctx, &namespaces, client.MatchingLabels{ @@ -420,7 +420,7 @@ func (r *TauWorkspaceReconciler) cleanupOwnedNamespaceAccess( return err } } - return r.reconcileOwnedNamespaceResearcherAccess(ctx, workspace, targetNamespace) + return nil } func canDeleteLegacyWorkspaceRole(role *rbacv1.Role, workspace *tauv1alpha1.TauWorkspace) bool { diff --git a/controllers/tau-core/internal/controller/workspace_sync.go b/controllers/tau-core/internal/controller/workspace_sync.go index c3b93c0f..c14d58ae 100644 --- a/controllers/tau-core/internal/controller/workspace_sync.go +++ b/controllers/tau-core/internal/controller/workspace_sync.go @@ -109,8 +109,12 @@ func (r *TauWorkspaceReconciler) syncWorkspace(ctx context.Context, workspace *t conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "NamespaceCleanupFailed", err.Error(), workspace.Generation)) } } + if err := r.reconcileOwnedNamespaceResearcherAccess(ctx, workspace, targetNamespace); err != nil { + previousAccessReady = false + conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) + } if previousAccessReady { - if err := r.cleanupOwnedNamespaceAccess(ctx, workspace, targetNamespace, keepServiceAccount); err != nil { + if err := r.cleanupOwnedNamespaceWorkloadIdentity(ctx, workspace, keepServiceAccount); err != nil { previousAccessReady = false conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) } From 992f1b6c39c268b1fb2052dc740bb67cc4ebfa57 Mon Sep 17 00:00:00 2001 From: chokevin Date: Fri, 23 Oct 2026 19:28:32 -0700 Subject: [PATCH 08/13] fix(controller): revoke researcher access during delete drain --- .../internal/controller/controller_test.go | 6 ++++- .../internal/controller/workspace_cleanup.go | 3 +++ .../internal/controller/workspace_rbac.go | 23 +++++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index 214546b6..3154318f 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -1546,6 +1546,7 @@ func TestWorkspaceDeletePreservesWorkloadIdentityWhileExternalQueueWorkloadDrain namespace := testNamespace(workspace.Name) namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) account := testServiceAccount(namespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + binding := testRoleBinding(namespace.Name, defaultRoleName, workspace.Name) job := &batchv1.Job{ ObjectMeta: metav1.ObjectMeta{ Name: "training", @@ -1589,7 +1590,7 @@ func TestWorkspaceDeletePreservesWorkloadIdentityWhileExternalQueueWorkloadDrain workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} c := fake.NewClientBuilder(). WithScheme(testScheme(t)). - WithObjects(workspace, namespace, account, job, localQueue, clusterQueue, workload). + WithObjects(workspace, namespace, account, binding, job, localQueue, clusterQueue, workload). Build() reconciler := newTestWorkspaceReconciler(c) request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} @@ -1601,6 +1602,9 @@ func TestWorkspaceDeletePreservesWorkloadIdentityWhileExternalQueueWorkloadDrain if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { t.Fatalf("workspace workload identity was removed before Workload drain: %v", err) } + if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &rbacv1.RoleBinding{}); !apierrors.IsNotFound(err) { + t.Fatalf("researcher RoleBinding survived deletion drain: %v", err) + } if err := c.Get(ctx, client.ObjectKeyFromObject(job), &batchv1.Job{}); err != nil { t.Fatalf("pending Job was removed during Workload drain: %v", err) } diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 466e08c0..6fa12599 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -263,6 +263,9 @@ func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, wor if targetNamespace == "" { targetNamespace = resolvedNamespace(workspace) } + if err := r.cleanupOwnedResearcherRoleBindings(ctx, workspace); err != nil { + return err + } if err := r.waitForNamespaceRoutingMetadataDrain(ctx, workspace, targetNamespace); err != nil { return err } diff --git a/controllers/tau-core/internal/controller/workspace_rbac.go b/controllers/tau-core/internal/controller/workspace_rbac.go index 75ce902f..14c42cda 100644 --- a/controllers/tau-core/internal/controller/workspace_rbac.go +++ b/controllers/tau-core/internal/controller/workspace_rbac.go @@ -126,6 +126,29 @@ func (r *TauWorkspaceReconciler) cleanupResearcherRBAC( return r.deleteOwnedObject(ctx, binding, workspace) } +func (r *TauWorkspaceReconciler) cleanupOwnedResearcherRoleBindings( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, +) error { + var bindings rbacv1.RoleBindingList + if err := r.List(ctx, &bindings, client.MatchingLabels{ + labelManagedBy: labelManagedByValue, + labelWorkspace: workspace.Name, + }); err != nil { + return err + } + for i := range bindings.Items { + binding := &bindings.Items[i] + if binding.Name != defaultRoleName { + continue + } + if err := r.deleteOwnedObject(ctx, binding, workspace); err != nil { + return err + } + } + return nil +} + func clusterQueueReaderBindingName(workspaceName string) string { return "tau-clusterqueue-reader-" + workspaceName } From a9236f8340733581808c535d68259a99dbc09799 Mon Sep 17 00:00:00 2001 From: chokevin Date: Mon, 26 Oct 2026 10:28:00 -0700 Subject: [PATCH 09/13] fix(controller): revoke obsolete identity before retarget drain --- .../internal/controller/controller_test.go | 299 +++++++++++------- .../internal/controller/workspace_cleanup.go | 12 + .../internal/controller/workspace_sync.go | 19 +- site/content/en/docs/reference/workspace.md | 11 + 4 files changed, 224 insertions(+), 117 deletions(-) diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index 3154318f..56e483ca 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -28,6 +28,7 @@ import ( ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/client/interceptor" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" "sigs.k8s.io/controller-runtime/pkg/event" "sigs.k8s.io/controller-runtime/pkg/reconcile" @@ -2823,119 +2824,199 @@ func TestUnmigratedWorkspaceRetargetMigratesHistoryBeforeDeletion(t *testing.T) } func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkloadDrains(t *testing.T) { - ctx := context.Background() - workspace := testWorkspace("aurora") - workspace.Finalizers = []string{workspaceFinalizer} - workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} - workspace.Spec.Target.Namespace = "new-target" - workspace.Spec.Target.CreateNamespace = false - workspace.Spec.Queue = "jobqueue" - workspace.Status.Phase = tauv1alpha1.WorkspacePhaseReady - workspace.Status.Target.ResolvedNamespace = "old-target" - workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ - LocalQueue: "jobqueue", - ClusterQueue: "jobqueue", - ClusterQueueUID: "jobqueue-uid", - } - oldNamespace := testNamespace("old-target") - oldNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) - oldNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} - newNamespace := testNamespace("new-target") - newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} - newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} - account := testServiceAccount(oldNamespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) - binding := testRoleBinding(oldNamespace.Name, defaultRoleName, workspace.Name) - binding.Subjects = []rbacv1.Subject{{ - Kind: rbacv1.GroupKind, - APIGroup: rbacv1.GroupName, - Name: "former-researchers", - }} - job := &batchv1.Job{ - ObjectMeta: metav1.ObjectMeta{ - Name: "training", - Namespace: oldNamespace.Name, - UID: types.UID("training-job-uid"), - }, - Spec: batchv1.JobSpec{ - Template: corev1.PodTemplateSpec{ - Spec: corev1.PodSpec{ - ServiceAccountName: account.Name, - RestartPolicy: corev1.RestartPolicyNever, - Containers: []corev1.Container{{ - Name: "trainer", - Image: "example.invalid/trainer", - Command: []string{"train"}, + for _, identityChange := range []string{"unchanged", "renamed", "client-id-changed", "removed", "deletion-fails", "deletion-pending"} { + t.Run(identityChange, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + workspace.Spec.Target.Namespace = "new-target" + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Phase = tauv1alpha1.WorkspacePhaseReady + workspace.Status.Target.ResolvedNamespace = "old-target" + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + } + oldNamespace := testNamespace("old-target") + oldNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + oldNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + newNamespace := testNamespace("new-target") + newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} + newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + account := testServiceAccount(oldNamespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + account.Labels[labelAzureWIUse] = "true" + account.Annotations[annotationAzureWIClientID] = workspace.Spec.WorkloadIdentity.ClientID + switch identityChange { + case "renamed", "deletion-fails", "deletion-pending": + workspace.Spec.WorkloadIdentity.ServiceAccountName = "replacement-workload" + case "client-id-changed": + workspace.Spec.WorkloadIdentity.ClientID = "replacement-client-id" + case "removed": + workspace.Spec.WorkloadIdentity = nil + } + if identityChange == "deletion-pending" { + account.Finalizers = []string{"example.invalid/retain"} + } + deletionBlocked := identityChange == "deletion-fails" || identityChange == "deletion-pending" + binding := testRoleBinding(oldNamespace.Name, defaultRoleName, workspace.Name) + binding.Subjects = []rbacv1.Subject{{ + Kind: rbacv1.GroupKind, + APIGroup: rbacv1.GroupName, + Name: "former-researchers", + }} + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: "training", + Namespace: oldNamespace.Name, + UID: types.UID("training-job-uid"), + }, + Spec: batchv1.JobSpec{ + Template: corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{ + ServiceAccountName: account.Name, + RestartPolicy: corev1.RestartPolicyNever, + Containers: []corev1.Container{{ + Name: "trainer", + Image: "example.invalid/trainer", + Command: []string{"train"}, + }}, + }, + }, + }, + } + localQueue := testLocalQueue(oldNamespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", }}, }, - }, - }, - } - localQueue := testLocalQueue(oldNamespace.Name, "jobqueue", "jobqueue") - localQueue.SetAnnotations(nil) - clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") - clusterQueue.Object["spec"] = map[string]any{ - "namespaceSelector": map[string]any{ - "matchExpressions": []any{map[string]any{ - "key": labelWorkspace, - "operator": "Exists", - }}, - }, - } - workload := newQueueObject(workloadGVK) - workload.SetName("job-training") - workload.SetNamespace(oldNamespace.Name) - workload.SetOwnerReferences([]metav1.OwnerReference{{ - APIVersion: batchv1.SchemeGroupVersion.String(), - Kind: "Job", - Name: job.Name, - UID: job.UID, - }}) - workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} - c := fake.NewClientBuilder(). - WithScheme(testScheme(t)). - WithObjects(workspace, oldNamespace, newNamespace, account, binding, job, localQueue, clusterQueue, workload). - WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). - Build() - reconciler := newTestWorkspaceReconciler(c) - request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + } + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(oldNamespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: batchv1.SchemeGroupVersion.String(), + Kind: "Job", + Name: job.Name, + UID: job.UID, + }}) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, oldNamespace, newNamespace, account, binding, job, localQueue, clusterQueue, workload). + WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). + WithInterceptorFuncs(interceptor.Funcs{ + Update: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.UpdateOption) error { + if _, ok := obj.(*rbacv1.RoleBinding); ok && obj.GetNamespace() == oldNamespace.Name && identityChange != "unchanged" { + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatal("researcher access was reassigned before obsolete identity revocation") + } + } + return c.Update(ctx, obj, opts...) + }, + Delete: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.DeleteOption) error { + if deletionBlocked && identityChange == "deletion-fails" { + if _, ok := obj.(*corev1.ServiceAccount); ok && client.ObjectKeyFromObject(obj) == client.ObjectKeyFromObject(account) { + return errors.New("obsolete identity deletion denied") + } + } + return c.Delete(ctx, obj, opts...) + }, + }). + Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} - if _, err := reconciler.Reconcile(ctx, request); err != nil { - t.Fatalf("Reconcile() during retarget drain = %v", err) - } - if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { - t.Fatalf("old workload identity was removed before Workload drain: %v", err) - } - var gotBinding rbacv1.RoleBinding - if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &gotBinding); err != nil { - t.Fatalf("Get old researcher RoleBinding: %v", err) - } - wantSubjects := []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, oldNamespace.Name)} - if !reflect.DeepEqual(gotBinding.Subjects, wantSubjects) { - t.Fatalf("old researcher subjects = %#v, want current subject %#v", gotBinding.Subjects, wantSubjects) - } - var waiting tauv1alpha1.TauWorkspace - if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { - t.Fatal(err) - } - if waiting.Status.Target.ResolvedNamespace != oldNamespace.Name { - t.Fatalf("resolved namespace = %q, want old target retained during drain", waiting.Status.Target.ResolvedNamespace) - } + for attempt := 0; attempt < 3; attempt++ { + if _, err := reconciler.Reconcile(ctx, request); err != nil { + t.Fatalf("Reconcile() during retarget drain = %v", err) + } + err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}) + if identityChange == "unchanged" || deletionBlocked { + if err != nil { + t.Fatalf("current workload identity was removed before Workload drain: %v", err) + } + } else if !apierrors.IsNotFound(err) { + t.Fatalf("obsolete workload identity survived retarget drain: %v", err) + } + } + var retainedJob batchv1.Job + if err := c.Get(ctx, client.ObjectKeyFromObject(job), &retainedJob); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(retainedJob.Spec, job.Spec) { + t.Fatalf("retarget changed the submitted Job: %#v", retainedJob.Spec) + } + if workspace.Spec.WorkloadIdentity != nil { + var currentAccount corev1.ServiceAccount + key := client.ObjectKey{Namespace: newNamespace.Name, Name: workspace.Spec.WorkloadIdentity.ServiceAccountName} + if err := c.Get(ctx, key, ¤tAccount); err != nil { + t.Fatal(err) + } + if currentAccount.Annotations[annotationAzureWIClientID] != workspace.Spec.WorkloadIdentity.ClientID { + t.Fatalf("new target identity = %#v, want current client ID", currentAccount.Annotations) + } + } + var gotBinding rbacv1.RoleBinding + if deletionBlocked { + if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &gotBinding); !apierrors.IsNotFound(err) { + t.Fatalf("old researcher binding survived incomplete identity revocation: %v", err) + } + } else { + if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &gotBinding); err != nil { + t.Fatalf("Get old researcher RoleBinding: %v", err) + } + wantSubjects := []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, oldNamespace.Name)} + if !reflect.DeepEqual(gotBinding.Subjects, wantSubjects) { + t.Fatalf("old researcher subjects = %#v, want current subject %#v", gotBinding.Subjects, wantSubjects) + } + } + var waiting tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != oldNamespace.Name { + t.Fatalf("resolved namespace = %q, want old target retained during drain", waiting.Status.Target.ResolvedNamespace) + } + if deletionBlocked { + assertCondition(t, waiting.Status.Conditions, tauv1alpha1.ConditionDriftDetected, metav1.ConditionTrue) + deletionBlocked = false + } + if identityChange == "deletion-pending" { + var retainedAccount corev1.ServiceAccount + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &retainedAccount); err != nil { + t.Fatal(err) + } + retainedAccount.Finalizers = nil + if err := c.Update(ctx, &retainedAccount); err != nil { + t.Fatal(err) + } + } - workload.Object["status"] = map[string]any{"conditions": []any{ - map[string]any{"type": "Finished", "status": "True"}, - }} - if err := c.Update(ctx, workload); err != nil { - t.Fatal(err) - } - reconcileWorkspaceTimes(t, reconciler, ctx, workspace.Name, 3) - if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { - t.Fatalf("old workload identity survived completed retarget: %v", err) - } - if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { - t.Fatal(err) - } - if waiting.Status.Target.ResolvedNamespace != newNamespace.Name { - t.Fatalf("resolved namespace = %q, want completed retarget %q", waiting.Status.Target.ResolvedNamespace, newNamespace.Name) + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + if err := c.Update(ctx, workload); err != nil { + t.Fatal(err) + } + reconcileWorkspaceTimes(t, reconciler, ctx, workspace.Name, 3) + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("old workload identity survived completed retarget: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != newNamespace.Name { + t.Fatalf("resolved namespace = %q, want completed retarget %q", waiting.Status.Target.ResolvedNamespace, newNamespace.Name) + } + }) } } diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 6fa12599..c6974079 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -576,6 +576,18 @@ func (r *TauWorkspaceReconciler) cleanupStaleWorkloadIdentityServiceAccounts( if err := r.deleteWithPreconditions(ctx, account); err != nil { return err } + var remaining corev1.ServiceAccount + if err := r.apiReader().Get(ctx, client.ObjectKeyFromObject(account), &remaining); err != nil { + if apierrors.IsNotFound(err) { + continue + } + return err + } + return fmt.Errorf( + "waiting for obsolete workload identity ServiceAccount %s deletion to complete; finalizers=%v", + client.ObjectKeyFromObject(&remaining), + remaining.Finalizers, + ) } return nil } diff --git a/controllers/tau-core/internal/controller/workspace_sync.go b/controllers/tau-core/internal/controller/workspace_sync.go index c14d58ae..1fe1c8f5 100644 --- a/controllers/tau-core/internal/controller/workspace_sync.go +++ b/controllers/tau-core/internal/controller/workspace_sync.go @@ -103,22 +103,25 @@ func (r *TauWorkspaceReconciler) syncWorkspace(ctx context.Context, workspace *t } previousNamespace := workspace.Status.Target.ResolvedNamespace previousAccessReady := true - if previousNamespace != "" && previousNamespace != targetNamespace { - if err := r.waitForNamespaceRoutingMetadataDrain(ctx, workspace, previousNamespace); err != nil { - previousAccessReady = false - conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "NamespaceCleanupFailed", err.Error(), workspace.Generation)) - } - } - if err := r.reconcileOwnedNamespaceResearcherAccess(ctx, workspace, targetNamespace); err != nil { + if err := r.cleanupOwnedNamespaceWorkloadIdentity(ctx, workspace, keepServiceAccount); err != nil { previousAccessReady = false conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) + if err := r.cleanupOwnedResearcherRoleBindings(ctx, workspace); err != nil { + conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) + } } if previousAccessReady { - if err := r.cleanupOwnedNamespaceWorkloadIdentity(ctx, workspace, keepServiceAccount); err != nil { + if err := r.reconcileOwnedNamespaceResearcherAccess(ctx, workspace, targetNamespace); err != nil { previousAccessReady = false conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) } } + if previousNamespace != "" && previousNamespace != targetNamespace { + if err := r.waitForNamespaceRoutingMetadataDrain(ctx, workspace, previousNamespace); err != nil { + previousAccessReady = false + conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "NamespaceCleanupFailed", err.Error(), workspace.Generation)) + } + } currentTargetReady := namespaceReady && rbacReady && rbacErr == nil && queueReady && previousAccessReady if workspace.Spec.WorkloadIdentity != nil { diff --git a/site/content/en/docs/reference/workspace.md b/site/content/en/docs/reference/workspace.md index a865f741..05aabb4c 100644 --- a/site/content/en/docs/reference/workspace.md +++ b/site/content/en/docs/reference/workspace.md @@ -49,6 +49,17 @@ The overall phase currently requires `RBACReady` and `QueueReady`; `DriftDetecte | `QueueReady=False` | Restore the named LocalQueue and its accessible backing ClusterQueue; fix that queue directly rather than creating a second queue with a different name to bypass the workspace spec. | | `DriftDetected=True` | Fix the dependency or workspace spec named in the message, then let the controller restore its owned namespace/RBAC objects. | +During namespace retargeting, unfinished Workloads on retained external queues +can delay cleanup of the old namespace's routing metadata. An unchanged workload +identity ServiceAccount is retained until those Workloads finish, while researcher +authorization is reconciled independently. Changing the ServiceAccount name or +client ID, or removing `workloadIdentity`, revokes the obsolete ServiceAccount +before granting the current researcher access to the old namespace, even during +the drain. Pending Jobs that use the revoked account can no longer create Pods; +the operator must cancel or replace affected work before the drain can complete. +Keep the identity unchanged until draining finishes when uninterrupted execution +is required. + ### Storage readiness is platform-owned `Ready` proves RBAC and queue reconciliation only, not storage. TauGrid 0.1 has no `StorageReady` condition on `TauWorkspace` or `TauCluster`, so a workspace can reach `Ready` while a workload's configured PVC is missing or unbound. Verify the platform-managed claim directly before handing the workspace to a researcher: From 1be0f575d6680eb61f07c3b0a0b6eb3efb6640c7 Mon Sep 17 00:00:00 2001 From: chokevin Date: Mon, 26 Oct 2026 12:52:58 -0700 Subject: [PATCH 10/13] fix(controller): recover drain access and retain source Job routes --- .../tau-core-controller/kustomize/rbac.yaml | 3 + .../tau-core-controller/templates/rbac.yaml | 3 + .../tau-core-controller/tests/rbac_test.yaml | 10 + .../internal/controller/controller_test.go | 27 ++ .../internal/controller/workspace_cleanup.go | 167 +++++++---- .../controller/workspace_cleanup_test.go | 266 ++++++++++++++++++ .../internal/controller/workspace_rbac.go | 39 ++- .../controller/workspace_rbac_test.go | 67 +++++ site/content/en/docs/reference/workspace.md | 12 +- 9 files changed, 525 insertions(+), 69 deletions(-) create mode 100644 controllers/tau-core/internal/controller/workspace_cleanup_test.go diff --git a/charts/tau-core-controller/kustomize/rbac.yaml b/charts/tau-core-controller/kustomize/rbac.yaml index 8010381f..a0e174e5 100644 --- a/charts/tau-core-controller/kustomize/rbac.yaml +++ b/charts/tau-core-controller/kustomize/rbac.yaml @@ -110,6 +110,9 @@ rules: - apiGroups: ["scheduling.k8s.io"] resources: ["priorityclasses"] verbs: ["get", "list", "watch"] + - apiGroups: ["batch"] + resources: ["jobs"] + verbs: ["get", "list", "watch"] - apiGroups: ["kueue.x-k8s.io"] resources: ["workloads"] verbs: ["get", "list", "watch"] diff --git a/charts/tau-core-controller/templates/rbac.yaml b/charts/tau-core-controller/templates/rbac.yaml index f043e848..d95c8a24 100644 --- a/charts/tau-core-controller/templates/rbac.yaml +++ b/charts/tau-core-controller/templates/rbac.yaml @@ -107,6 +107,9 @@ rules: - apiGroups: ["scheduling.k8s.io"] resources: ["priorityclasses"] verbs: ["get", "list", "watch"] + - apiGroups: ["batch"] + resources: ["jobs"] + verbs: ["get", "list", "watch"] - apiGroups: ["kueue.x-k8s.io"] resources: ["workloads"] verbs: ["get", "list", "watch"] diff --git a/charts/tau-core-controller/tests/rbac_test.yaml b/charts/tau-core-controller/tests/rbac_test.yaml index 18a0a85e..542d9086 100644 --- a/charts/tau-core-controller/tests/rbac_test.yaml +++ b/charts/tau-core-controller/tests/rbac_test.yaml @@ -2,6 +2,16 @@ suite: Researcher RayService permissions templates: - templates/rbac.yaml tests: + - it: grants the controller read-only source Job access for workspace drain checks + asserts: + - contains: + path: rules + content: + apiGroups: ["batch"] + resources: ["jobs"] + verbs: ["get", "list", "watch"] + documentIndex: 2 + - it: grants RayJob and RayService lifecycle permissions to the researcher role documentSelector: path: metadata.name diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index 56e483ca..f51e4fc9 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -2913,6 +2913,14 @@ func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkload WithObjects(workspace, oldNamespace, newNamespace, account, binding, job, localQueue, clusterQueue, workload). WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). WithInterceptorFuncs(interceptor.Funcs{ + Create: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.CreateOption) error { + if _, ok := obj.(*rbacv1.RoleBinding); ok && obj.GetNamespace() == oldNamespace.Name && identityChange != "unchanged" { + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatal("researcher access was restored before obsolete identity revocation") + } + } + return c.Create(ctx, obj, opts...) + }, Update: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.UpdateOption) error { if _, ok := obj.(*rbacv1.RoleBinding); ok && obj.GetNamespace() == oldNamespace.Name && identityChange != "unchanged" { if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { @@ -2999,6 +3007,25 @@ func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkload t.Fatal(err) } } + if identityChange == "deletion-fails" || identityChange == "deletion-pending" { + reconcileWorkspaceTimes(t, reconciler, ctx, workspace.Name, 3) + if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &gotBinding); err != nil { + t.Fatalf("researcher access was not restored after identity revocation recovered: %v", err) + } + wantSubjects := []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, oldNamespace.Name)} + if !reflect.DeepEqual(gotBinding.Subjects, wantSubjects) { + t.Fatalf("recovered researcher subjects = %#v, want %#v", gotBinding.Subjects, wantSubjects) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("obsolete identity survived successful recovery: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != oldNamespace.Name { + t.Fatal("recovery advanced the target before the Workload finished") + } + } workload.Object["status"] = map[string]any{"conditions": []any{ map[string]any{"type": "Finished", "status": "True"}, diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index c6974079..f03b49c5 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -10,6 +10,7 @@ import ( "strings" tauv1alpha1 "github.com/Azure/taugrid/controllers/tau-core/api/v1alpha1" + batchv1 "k8s.io/api/batch/v1" corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" @@ -19,6 +20,7 @@ import ( "k8s.io/apimachinery/pkg/labels" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" ) @@ -187,75 +189,128 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( if !found || queueName == "" { continue } - localQueue := newQueueObject(localQueueGVK) - localQueueKey := client.ObjectKey{Name: queueName, Namespace: namespace.Name} - if err := r.apiReader().Get(ctx, localQueueKey, localQueue); err != nil { - if apierrors.IsNotFound(err) { - continue - } - return false, "", fmt.Errorf("read LocalQueue %s before routing metadata cleanup: %w", localQueueKey, err) - } - clusterQueueName, found, err := unstructured.NestedString(localQueue.Object, "spec", "clusterQueue") - if err != nil { - return false, "", fmt.Errorf("read LocalQueue %s backing ClusterQueue: %w", localQueueKey, err) - } - if !found || strings.TrimSpace(clusterQueueName) == "" { - return false, "", fmt.Errorf("LocalQueue %s has no backing ClusterQueue", localQueueKey) - } - admittedClusterQueue, admitted, err := unstructured.NestedString(workload.Object, "status", "admission", "clusterQueue") + admittedClusterQueue, _, err := unstructured.NestedString(workload.Object, "status", "admission", "clusterQueue") if err != nil { return false, "", fmt.Errorf("read Workload %s admitted ClusterQueue: %w", client.ObjectKeyFromObject(workload), err) } - if admitted && strings.TrimSpace(admittedClusterQueue) != "" { - clusterQueueName = admittedClusterQueue - } - clusterQueue := newQueueObject(clusterQueueGVK) - if err := r.apiReader().Get(ctx, client.ObjectKey{Name: clusterQueueName}, clusterQueue); err != nil { - return false, "", fmt.Errorf("read ClusterQueue %s before routing metadata cleanup: %w", clusterQueueName, err) + inUse, message, err := r.queueRoutingMetadataInUse( + ctx, workspace, namespace, cleanedLabels, queueName, admittedClusterQueue, + fmt.Sprintf("Workload %s", client.ObjectKeyFromObject(workload)), + ) + if err != nil || inUse { + return inUse, message, err } - if expectedUID := strings.TrimSpace(workspace.Status.Queue.ClusterQueueUID); expectedUID != "" && - clusterQueueName == strings.TrimSpace(workspace.Status.Queue.ClusterQueue) && - string(clusterQueue.GetUID()) != expectedUID { - return false, "", fmt.Errorf( - "ClusterQueue %q UID changed from %q to %q before routing metadata cleanup", - clusterQueueName, - expectedUID, - clusterQueue.GetUID(), - ) + } + var jobs batchv1.JobList + if err := r.apiReader().List(ctx, &jobs, client.InNamespace(namespace.Name)); err != nil { + return false, "", fmt.Errorf("list Jobs in namespace %q before routing metadata cleanup: %w", namespace.Name, err) + } + workloadJobs := make(map[types.UID]string) + for i := range workloads.Items { + for _, owner := range workloads.Items[i].GetOwnerReferences() { + if owner.APIVersion == batchv1.SchemeGroupVersion.String() && owner.Kind == "Job" && owner.UID != "" { + workloadJobs[owner.UID] = owner.Name + } } - selectorMap, found, err := unstructured.NestedMap(clusterQueue.Object, "spec", "namespaceSelector") - if err != nil { - return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + for i := range jobs.Items { + job := &jobs.Items[i] + if job.UID != "" && workloadJobs[job.UID] == job.Name { + continue } - labelSelector := metav1.LabelSelector{} - if found { - if err := runtime.DefaultUnstructuredConverter.FromUnstructured(selectorMap, &labelSelector); err != nil { - return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + finished := false + for _, condition := range job.Status.Conditions { + if condition.Status == corev1.ConditionTrue && + (condition.Type == batchv1.JobComplete || condition.Type == batchv1.JobFailed) { + finished = true + break } } - selector, err := metav1.LabelSelectorAsSelector(&labelSelector) - if err != nil { - return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + queueName := strings.TrimSpace(job.Labels["kueue.x-k8s.io/queue-name"]) + if finished || queueName == "" { + continue } - if !selector.Matches(labels.Set(namespace.Labels)) { - return false, "", fmt.Errorf( - "ClusterQueue %q namespaceSelector does not match namespace %q for unfinished Workload %s", - clusterQueueName, - namespace.Name, - client.ObjectKeyFromObject(workload), - ) + inUse, message, err := r.queueRoutingMetadataInUse( + ctx, workspace, namespace, cleanedLabels, queueName, "", + fmt.Sprintf("Job %s", client.ObjectKeyFromObject(job)), + ) + if err != nil || inUse { + return inUse, message, err } - if selector.Matches(labels.Set(cleanedLabels)) { - continue + } + return false, "", nil +} + +func (r *TauWorkspaceReconciler) queueRoutingMetadataInUse( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespace *corev1.Namespace, + cleanedLabels map[string]string, + queueName, admittedClusterQueue, source string, +) (bool, string, error) { + localQueue := newQueueObject(localQueueGVK) + localQueueKey := client.ObjectKey{Name: queueName, Namespace: namespace.Name} + if err := r.apiReader().Get(ctx, localQueueKey, localQueue); err != nil { + if apierrors.IsNotFound(err) { + return false, "", nil + } + return false, "", fmt.Errorf("read LocalQueue %s before routing metadata cleanup: %w", localQueueKey, err) + } + clusterQueueName, found, err := unstructured.NestedString(localQueue.Object, "spec", "clusterQueue") + if err != nil { + return false, "", fmt.Errorf("read LocalQueue %s backing ClusterQueue: %w", localQueueKey, err) + } + if !found || strings.TrimSpace(clusterQueueName) == "" { + return false, "", fmt.Errorf("LocalQueue %s has no backing ClusterQueue", localQueueKey) + } + if strings.TrimSpace(admittedClusterQueue) != "" { + clusterQueueName = admittedClusterQueue + } + clusterQueue := newQueueObject(clusterQueueGVK) + if err := r.apiReader().Get(ctx, client.ObjectKey{Name: clusterQueueName}, clusterQueue); err != nil { + return false, "", fmt.Errorf("read ClusterQueue %s before routing metadata cleanup: %w", clusterQueueName, err) + } + if expectedUID := strings.TrimSpace(workspace.Status.Queue.ClusterQueueUID); expectedUID != "" && + clusterQueueName == strings.TrimSpace(workspace.Status.Queue.ClusterQueue) && + string(clusterQueue.GetUID()) != expectedUID { + return false, "", fmt.Errorf( + "ClusterQueue %q UID changed from %q to %q before routing metadata cleanup", + clusterQueueName, + expectedUID, + clusterQueue.GetUID(), + ) + } + selectorMap, found, err := unstructured.NestedMap(clusterQueue.Object, "spec", "namespaceSelector") + if err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + labelSelector := metav1.LabelSelector{} + if found { + if err := runtime.DefaultUnstructuredConverter.FromUnstructured(selectorMap, &labelSelector); err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) } - return true, fmt.Sprintf( - "waiting for Workload %s using retained LocalQueue %q and ClusterQueue %q before removing namespace routing metadata", - client.ObjectKeyFromObject(workload), - queueName, + } + selector, err := metav1.LabelSelectorAsSelector(&labelSelector) + if err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + if !selector.Matches(labels.Set(namespace.Labels)) { + return false, "", fmt.Errorf( + "ClusterQueue %q namespaceSelector does not match namespace %q for unfinished %s", clusterQueueName, - ), nil + namespace.Name, + source, + ) } - return false, "", nil + if selector.Matches(labels.Set(cleanedLabels)) { + return false, "", nil + } + return true, fmt.Sprintf( + "waiting for %s using retained LocalQueue %q and ClusterQueue %q before removing namespace routing metadata", + source, + queueName, + clusterQueueName, + ), nil } func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, workspace *tauv1alpha1.TauWorkspace) error { diff --git a/controllers/tau-core/internal/controller/workspace_cleanup_test.go b/controllers/tau-core/internal/controller/workspace_cleanup_test.go new file mode 100644 index 00000000..a536302d --- /dev/null +++ b/controllers/tau-core/internal/controller/workspace_cleanup_test.go @@ -0,0 +1,266 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +package controller + +import ( + "context" + "errors" + "reflect" + "strings" + "testing" + + tauv1alpha1 "github.com/Azure/taugrid/controllers/tau-core/api/v1alpha1" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/types" + "k8s.io/utils/ptr" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/client/interceptor" +) + +func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { + for _, lifecycle := range []string{"delete", "retarget"} { + t.Run(lifecycle, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + if lifecycle == "delete" { + workspace.DeletionTimestamp = &now + } + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + account := testServiceAccount(namespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + account.Annotations[annotationAzureWIClientID] = workspace.Spec.WorkloadIdentity.ClientID + newNamespace := testNamespace("new-target") + newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} + newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + if lifecycle == "retarget" { + workspace.Spec.Target.Namespace = newNamespace.Name + } + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: "training", Namespace: namespace.Name, UID: types.UID("training-job-uid"), + Labels: map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}, + }, + Spec: batchv1.JobSpec{ + Suspend: ptr.To(true), + Template: corev1.PodTemplateSpec{Spec: corev1.PodSpec{ + ServiceAccountName: account.Name, + RestartPolicy: corev1.RestartPolicyNever, + Containers: []corev1.Container{{Name: "trainer", Image: "example.invalid/trainer"}}, + }}, + }, + } + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchLabels": map[string]any{labelWorkspace: workspace.Name}, + }} + c := fake.NewClientBuilder().WithScheme(testScheme(t)). + WithObjects(workspace, namespace, newNamespace, account, job, localQueue, clusterQueue). + WithStatusSubresource(&tauv1alpha1.TauWorkspace{}).Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + for range 3 { + _, err := reconciler.Reconcile(ctx, request) + if lifecycle == "delete" && (err == nil || + !strings.Contains(err.Error(), "before removing namespace routing metadata")) { + t.Fatalf("Reconcile() = %v, want source Job drain wait", err) + } + if lifecycle == "retarget" && err != nil { + t.Fatal(err) + } + var waiting tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != namespace.Name { + t.Fatal("resolved target advanced before source Job drain") + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { + t.Fatalf("execution identity removed before source Job drain: %v", err) + } + var retained corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &retained); err != nil { + t.Fatal(err) + } + if retained.Labels[labelWorkspace] != workspace.Name { + t.Fatal("source Job lost namespace routing metadata") + } + } + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(namespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: batchv1.SchemeGroupVersion.String(), Kind: "Job", Name: job.Name, UID: job.UID, + }}) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + if err := c.Create(ctx, workload); err != nil { + t.Fatal(err) + } + if _, err := reconciler.Reconcile(ctx, request); lifecycle == "delete" && err == nil { + t.Fatal("workspace stopped draining after Kueue created the Workload") + } + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + if err := c.Update(ctx, workload); err != nil { + t.Fatal(err) + } + for range 3 { + if _, err := reconciler.Reconcile(ctx, request); err != nil { + t.Fatalf("Reconcile() after drain = %v", err) + } + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("execution account survived drain: %v", err) + } + var finished tauv1alpha1.TauWorkspace + err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &finished) + if lifecycle == "delete" && !apierrors.IsNotFound(err) { + t.Fatalf("workspace survived completed source Job drain: %v", err) + } + if lifecycle == "retarget" && (err != nil || finished.Status.Target.ResolvedNamespace != newNamespace.Name) { + t.Fatalf("retarget did not complete after source Job drain: %v, target=%q", err, finished.Status.Target.ResolvedNamespace) + } + }) + } +} + +func TestNamespaceRoutingMetadataSourceJobs(t *testing.T) { + for _, scenario := range []string{ + "suspended", "running", "complete", "failed", "false-completion", + "unlabeled", "other-namespace", "missing-queue", "selector-independent", + "workload-finished", "workload-wrong-uid", "workload-wrong-name", + "job-list-error", "queue-read-error", "cluster-queue-read-error", + "invalid-selector", "cluster-queue-uid-changed", + } { + t.Run(scenario, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + ClusterQueue: "jobqueue", ClusterQueueUID: "jobqueue-uid", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, "jobqueue") + job := &batchv1.Job{ObjectMeta: metav1.ObjectMeta{ + Name: "training", Namespace: namespace.Name, UID: types.UID("training-job-uid"), + Labels: map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}, + }, Spec: batchv1.JobSpec{Suspend: ptr.To(true)}} + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchLabels": map[string]any{labelWorkspace: workspace.Name}, + }} + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(namespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: batchv1.SchemeGroupVersion.String(), Kind: "Job", Name: job.Name, UID: job.UID, + }}) + workload.Object["spec"] = map[string]any{"queueName": "jobqueue"} + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + wantInUse, wantError := true, false + objects := []client.Object{namespace, job, localQueue, clusterQueue} + switch scenario { + case "running": + job.Spec.Suspend = ptr.To(false) + case "complete", "failed", "false-completion": + conditionType, status := batchv1.JobComplete, corev1.ConditionTrue + if scenario == "failed" { + conditionType = batchv1.JobFailed + } + if scenario == "false-completion" { + status = corev1.ConditionFalse + } else { + wantInUse = false + } + job.Status.Conditions = []batchv1.JobCondition{{Type: conditionType, Status: status}} + case "unlabeled": + job.Labels = nil + wantInUse = false + case "other-namespace": + job.Namespace = "unrelated" + wantInUse = false + case "missing-queue": + job.Labels["kueue.x-k8s.io/queue-name"] = "missing" + wantInUse = false + case "selector-independent": + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{}} + wantInUse = false + case "workload-finished", "workload-wrong-uid", "workload-wrong-name": + if scenario == "workload-finished" { + wantInUse = false + } else { + owner := workload.GetOwnerReferences()[0] + if scenario == "workload-wrong-uid" { + owner.UID = types.UID("different-job-uid") + } else { + owner.Name = "different-job" + } + workload.SetOwnerReferences([]metav1.OwnerReference{owner}) + } + objects = append(objects, workload) + case "job-list-error", "queue-read-error", "cluster-queue-read-error": + wantInUse, wantError = false, true + case "invalid-selector": + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{"key": labelWorkspace, "operator": "Invalid"}}, + }} + wantInUse, wantError = false, true + case "cluster-queue-uid-changed": + clusterQueue.SetUID(types.UID("replacement-uid")) + wantInUse, wantError = false, true + } + baseClient := fake.NewClientBuilder().WithScheme(testScheme(t)).WithObjects(objects...).Build() + c := interceptor.NewClient(baseClient, interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*batchv1.JobList); ok && scenario == "job-list-error" { + return errors.New("source Jobs unavailable") + } + return c.List(ctx, list, opts...) + }, + Get: func(ctx context.Context, c client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error { + if queue, ok := obj.(*unstructured.Unstructured); ok && + (scenario == "queue-read-error" && queue.GroupVersionKind() == localQueueGVK || + scenario == "cluster-queue-read-error" && queue.GroupVersionKind() == clusterQueueGVK) { + return errors.New("queue unavailable") + } + return c.Get(ctx, key, obj, opts...) + }, + }) + inUse, message, err := newTestWorkspaceReconciler(c).namespaceRoutingMetadataInUse(ctx, workspace, namespace) + if inUse != wantInUse || (err != nil) != wantError { + t.Fatalf("namespaceRoutingMetadataInUse() = %v, %q, %v; want inUse=%v error=%v", inUse, message, err, wantInUse, wantError) + } + if inUse && !strings.Contains(message, "Job "+client.ObjectKeyFromObject(job).String()) { + t.Fatalf("drain message does not identify the source Job: %q", message) + } + for _, queue := range []*unstructured.Unstructured{localQueue, clusterQueue} { + var got unstructured.Unstructured + got.SetGroupVersionKind(queue.GroupVersionKind()) + if err := baseClient.Get(ctx, client.ObjectKeyFromObject(queue), &got); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(got.Object, queue.Object) { + t.Fatalf("external queue was mutated: %#v", got.Object) + } + } + }) + } +} diff --git a/controllers/tau-core/internal/controller/workspace_rbac.go b/controllers/tau-core/internal/controller/workspace_rbac.go index 14c42cda..dcc04f0a 100644 --- a/controllers/tau-core/internal/controller/workspace_rbac.go +++ b/controllers/tau-core/internal/controller/workspace_rbac.go @@ -11,6 +11,7 @@ import ( corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" apiequality "k8s.io/apimachinery/pkg/api/equality" + apierrors "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" @@ -32,18 +33,7 @@ func (r *TauWorkspaceReconciler) reconcileRBAC(ctx context.Context, workspace *t if workspace.Spec.PrincipalRef == nil || workspace.Spec.KubernetesSubject == nil { return false, "workspace-rbac authorization requires principalRef and kubernetesSubject", nil } - binding := &rbacv1.RoleBinding{ObjectMeta: metav1.ObjectMeta{Name: defaultRoleName, Namespace: targetNamespace}} - if err := r.getAndValidateWorkspaceOwnership(ctx, binding, workspace); err != nil { - return false, "refusing to adopt existing researcher RoleBinding", err - } - _, err := controllerutil.CreateOrUpdate(ctx, r.Client, binding, func() error { - binding.Labels = workspaceLabels(workspace.Name) - setOwnerUIDAnnotation(binding, workspace.UID) - binding.RoleRef = rbacv1.RoleRef{APIGroup: rbacv1.GroupName, Kind: "ClusterRole", Name: defaultRoleName} - binding.Subjects = []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, targetNamespace)} - return nil - }) - if err != nil { + if err := r.reconcileResearcherRoleBinding(ctx, workspace, targetNamespace); err != nil { return false, "failed to reconcile researcher RoleBinding", err } if err := r.reconcileClusterQueueReaderRBAC(ctx, workspace, targetNamespace); err != nil { @@ -55,6 +45,25 @@ func (r *TauWorkspaceReconciler) reconcileRBAC(ctx context.Context, workspace *t return true, "researcher namespace and ClusterQueue reader bindings are reconciled", nil } +func (r *TauWorkspaceReconciler) reconcileResearcherRoleBinding( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespace string, +) error { + binding := &rbacv1.RoleBinding{ObjectMeta: metav1.ObjectMeta{Name: defaultRoleName, Namespace: namespace}} + if err := r.getAndValidateWorkspaceOwnership(ctx, binding, workspace); err != nil { + return fmt.Errorf("refusing to adopt existing researcher RoleBinding: %w", err) + } + _, err := controllerutil.CreateOrUpdate(ctx, r.Client, binding, func() error { + binding.Labels = workspaceLabels(workspace.Name) + setOwnerUIDAnnotation(binding, workspace.UID) + binding.RoleRef = rbacv1.RoleRef{APIGroup: rbacv1.GroupName, Kind: "ClusterRole", Name: defaultRoleName} + binding.Subjects = []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, namespace)} + return nil + }) + return err +} + func authorizationMode(workspace *tauv1alpha1.TauWorkspace) string { if workspace.Spec.Authorization == nil || workspace.Spec.Authorization.Mode == "" { return tauv1alpha1.AuthorizationModeWorkspaceRBAC @@ -94,6 +103,12 @@ func (r *TauWorkspaceReconciler) reconcileExistingResearcherAccess( ) error { binding := &rbacv1.RoleBinding{ObjectMeta: metav1.ObjectMeta{Name: defaultRoleName, Namespace: namespace}} if err := r.Get(ctx, client.ObjectKeyFromObject(binding), binding); err != nil { + if apierrors.IsNotFound(err) && + namespace == workspace.Status.Target.ResolvedNamespace && + authorizationMode(workspace) != tauv1alpha1.AuthorizationModeClusterWide && + workspace.Spec.PrincipalRef != nil && workspace.Spec.KubernetesSubject != nil { + return r.reconcileResearcherRoleBinding(ctx, workspace, namespace) + } return client.IgnoreNotFound(err) } if !ownedByWorkspace(binding.Labels, workspace.Name) { diff --git a/controllers/tau-core/internal/controller/workspace_rbac_test.go b/controllers/tau-core/internal/controller/workspace_rbac_test.go index 0cfdf68f..b95497a8 100644 --- a/controllers/tau-core/internal/controller/workspace_rbac_test.go +++ b/controllers/tau-core/internal/controller/workspace_rbac_test.go @@ -10,7 +10,9 @@ import ( "reflect" "testing" + tauv1alpha1 "github.com/Azure/taugrid/controllers/tau-core/api/v1alpha1" rbacv1 "k8s.io/api/rbac/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/util/yaml" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" @@ -96,6 +98,71 @@ func TestKustomizeControllerRoleAllowsLegacyRoleDeletionWithoutMutation(t *testi t.Fatalf("Kustomize controller role must permit identity-checked Role deletion without Role mutation: %#v", role.Rules) } +func TestKustomizeControllerRoleAllowsReadOnlySourceJobDiscovery(t *testing.T) { + role := kustomizeClusterRole(t, "tau-core-controller") + want := rbacv1.PolicyRule{ + APIGroups: []string{"batch"}, + Resources: []string{"jobs"}, + Verbs: []string{"get", "list", "watch"}, + } + for _, rule := range role.Rules { + if reflect.DeepEqual(rule, want) { + return + } + } + t.Fatalf("controller role must allow read-only source Job discovery: %#v", role.Rules) +} + +func TestRetainedNamespaceRestoresOnlyCurrentAuthorizedResearcherBinding(t *testing.T) { + for _, scenario := range []string{"group", "service-account", "cluster-wide", "missing-subject", "missing-principal", "abandoned-namespace", "foreign-namespace"} { + t.Run(scenario, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("research") + workspace.Status.Target.ResolvedNamespace = "retained" + namespace := testNamespace("retained") + namespace.Labels = workspaceLabels(workspace.Name) + namespace.Annotations[annotationOwnerUID] = string(workspace.UID) + wantBinding := scenario == "group" || scenario == "service-account" + switch scenario { + case "service-account": + workspace.Spec.KubernetesSubject = &tauv1alpha1.KubernetesSubject{Kind: rbacv1.ServiceAccountKind, Name: "current-researcher"} + case "cluster-wide": + workspace.Spec.Authorization = &tauv1alpha1.WorkspaceAuthorization{Mode: tauv1alpha1.AuthorizationModeClusterWide} + case "missing-subject": + workspace.Spec.KubernetesSubject = nil + case "missing-principal": + workspace.Spec.PrincipalRef = nil + case "abandoned-namespace": + workspace.Status.Target.ResolvedNamespace = "another-namespace" + case "foreign-namespace": + namespace.Annotations[annotationOwnerUID] = "another-workspace-uid" + } + c := fake.NewClientBuilder().WithScheme(testScheme(t)).WithObjects(namespace).Build() + if err := newTestWorkspaceReconciler(c).reconcileOwnedNamespaceResearcherAccess(ctx, workspace, "new-target"); err != nil { + t.Fatal(err) + } + var binding rbacv1.RoleBinding + err := c.Get(ctx, client.ObjectKey{Namespace: namespace.Name, Name: defaultRoleName}, &binding) + if !wantBinding { + if !apierrors.IsNotFound(err) { + t.Fatalf("unexpected researcher access was created: %v, %#v", err, binding) + } + return + } + if err != nil { + t.Fatal(err) + } + wantSubjects := []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, namespace.Name)} + if !reflect.DeepEqual(binding.Subjects, wantSubjects) || + binding.RoleRef != (rbacv1.RoleRef{APIGroup: rbacv1.GroupName, Kind: "ClusterRole", Name: defaultRoleName}) || + !ownedByWorkspace(binding.Labels, workspace.Name) || + binding.Annotations[annotationOwnerUID] != string(workspace.UID) { + t.Fatalf("restored researcher access does not match current workspace: %#v", binding) + } + }) + } +} + func kustomizeClusterRole(t *testing.T, name string) rbacv1.ClusterRole { t.Helper() file, err := os.Open("../../../../charts/tau-core-controller/kustomize/rbac.yaml") diff --git a/site/content/en/docs/reference/workspace.md b/site/content/en/docs/reference/workspace.md index 05aabb4c..58ad95e3 100644 --- a/site/content/en/docs/reference/workspace.md +++ b/site/content/en/docs/reference/workspace.md @@ -50,7 +50,10 @@ The overall phase currently requires `RBACReady` and `QueueReady`; `DriftDetecte | `DriftDetected=True` | Fix the dependency or workspace spec named in the message, then let the controller restore its owned namespace/RBAC objects. | During namespace retargeting, unfinished Workloads on retained external queues -can delay cleanup of the old namespace's routing metadata. An unchanged workload +can delay cleanup of the old namespace's routing metadata. Submitted, unfinished +Jobs that have not received a derived Workload yet also retain routing metadata +on deletion or retargeting, using their queue label and the live queue selector. +An unchanged workload identity ServiceAccount is retained until those Workloads finish, while researcher authorization is reconciled independently. Changing the ServiceAccount name or client ID, or removing `workloadIdentity`, revokes the obsolete ServiceAccount @@ -60,6 +63,13 @@ the operator must cancel or replace affected work before the drain can complete. Keep the identity unchanged until draining finishes when uninterrupted execution is required. +If identity cleanup fails or deletion is delayed by a finalizer, researcher +bindings are removed until obsolete identity revocation is confirmed. After +cleanup recovers, the current researcher's binding is restored in the retained +resolved namespace even while its Workloads are unfinished, so the researcher +can inspect or cancel the remaining work. Cluster-wide authorization does not +create a researcher binding. + ### Storage readiness is platform-owned `Ready` proves RBAC and queue reconciliation only, not storage. TauGrid 0.1 has no `StorageReady` condition on `TauWorkspace` or `TauCluster`, so a workspace can reach `Ready` while a workload's configured PVC is missing or unbound. Verify the platform-managed claim directly before handing the workspace to a researcher: From f6c59da54f2e2c30b55be595986e711609be5dc2 Mon Sep 17 00:00:00 2001 From: chokevin Date: Mon, 26 Oct 2026 13:58:41 -0700 Subject: [PATCH 11/13] fix(controller): retain routes for RayJobs awaiting Workloads --- .../tau-core-controller/kustomize/rbac.yaml | 3 + .../tau-core-controller/templates/rbac.yaml | 3 + .../tau-core-controller/tests/rbac_test.yaml | 7 + .../internal/controller/workspace_cleanup.go | 44 +++ .../controller/workspace_cleanup_test.go | 338 +++++++++++++----- .../controller/workspace_rbac_test.go | 26 +- site/content/en/docs/reference/workspace.md | 13 +- 7 files changed, 330 insertions(+), 104 deletions(-) diff --git a/charts/tau-core-controller/kustomize/rbac.yaml b/charts/tau-core-controller/kustomize/rbac.yaml index a0e174e5..230f1a71 100644 --- a/charts/tau-core-controller/kustomize/rbac.yaml +++ b/charts/tau-core-controller/kustomize/rbac.yaml @@ -113,6 +113,9 @@ rules: - apiGroups: ["batch"] resources: ["jobs"] verbs: ["get", "list", "watch"] + - apiGroups: ["ray.io"] + resources: ["rayjobs"] + verbs: ["get", "list", "watch"] - apiGroups: ["kueue.x-k8s.io"] resources: ["workloads"] verbs: ["get", "list", "watch"] diff --git a/charts/tau-core-controller/templates/rbac.yaml b/charts/tau-core-controller/templates/rbac.yaml index d95c8a24..0a13fb34 100644 --- a/charts/tau-core-controller/templates/rbac.yaml +++ b/charts/tau-core-controller/templates/rbac.yaml @@ -110,6 +110,9 @@ rules: - apiGroups: ["batch"] resources: ["jobs"] verbs: ["get", "list", "watch"] + - apiGroups: ["ray.io"] + resources: ["rayjobs"] + verbs: ["get", "list", "watch"] - apiGroups: ["kueue.x-k8s.io"] resources: ["workloads"] verbs: ["get", "list", "watch"] diff --git a/charts/tau-core-controller/tests/rbac_test.yaml b/charts/tau-core-controller/tests/rbac_test.yaml index 542d9086..8eb7f22c 100644 --- a/charts/tau-core-controller/tests/rbac_test.yaml +++ b/charts/tau-core-controller/tests/rbac_test.yaml @@ -11,6 +11,13 @@ tests: resources: ["jobs"] verbs: ["get", "list", "watch"] documentIndex: 2 + - contains: + path: rules + content: + apiGroups: ["ray.io"] + resources: ["rayjobs"] + verbs: ["get", "list", "watch"] + documentIndex: 2 - it: grants RayJob and RayService lifecycle permissions to the researcher role documentSelector: diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index f03b49c5..098be79e 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -206,11 +206,15 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( return false, "", fmt.Errorf("list Jobs in namespace %q before routing metadata cleanup: %w", namespace.Name, err) } workloadJobs := make(map[types.UID]string) + workloadRayJobs := make(map[types.UID]string) for i := range workloads.Items { for _, owner := range workloads.Items[i].GetOwnerReferences() { if owner.APIVersion == batchv1.SchemeGroupVersion.String() && owner.Kind == "Job" && owner.UID != "" { workloadJobs[owner.UID] = owner.Name } + if owner.APIVersion == "ray.io/v1" && owner.Kind == "RayJob" && owner.UID != "" { + workloadRayJobs[owner.UID] = owner.Name + } } } for i := range jobs.Items { @@ -238,9 +242,49 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( return inUse, message, err } } + rayJobs := &unstructured.UnstructuredList{} + rayJobs.SetGroupVersionKind(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayJobList"}) + if err := r.apiReader().List(ctx, rayJobs, client.InNamespace(namespace.Name)); err != nil { + if meta.IsNoMatchError(err) || apierrors.IsNotFound(err) { + return false, "", nil + } + return false, "", fmt.Errorf("list RayJobs in namespace %q before routing metadata cleanup: %w", namespace.Name, err) + } + for i := range rayJobs.Items { + rayJob := &rayJobs.Items[i] + if rayJob.GetUID() != "" && workloadRayJobs[rayJob.GetUID()] == rayJob.GetName() { + continue + } + queueName := strings.TrimSpace(rayJob.GetLabels()["kueue.x-k8s.io/queue-name"]) + if queueName == "" { + continue + } + finished, err := sourceRayJobFinished(rayJob) + if err != nil { + return false, "", err + } + if finished { + continue + } + inUse, message, err := r.queueRoutingMetadataInUse( + ctx, workspace, namespace, cleanedLabels, queueName, "", + fmt.Sprintf("RayJob %s", client.ObjectKeyFromObject(rayJob)), + ) + if err != nil || inUse { + return inUse, message, err + } + } return false, "", nil } +func sourceRayJobFinished(rayJob *unstructured.Unstructured) (bool, error) { + value, _, err := unstructured.NestedString(rayJob.Object, "status", "jobDeploymentStatus") + if err != nil { + return false, fmt.Errorf("read RayJob %s jobDeploymentStatus before routing metadata cleanup: %w", client.ObjectKeyFromObject(rayJob), err) + } + return value == "Complete" || value == "Failed", nil +} + func (r *TauWorkspaceReconciler) queueRoutingMetadataInUse( ctx context.Context, workspace *tauv1alpha1.TauWorkspace, diff --git a/controllers/tau-core/internal/controller/workspace_cleanup_test.go b/controllers/tau-core/internal/controller/workspace_cleanup_test.go index a536302d..3d8c6bab 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup_test.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup_test.go @@ -14,8 +14,10 @@ import ( batchv1 "k8s.io/api/batch/v1" corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/types" "k8s.io/utils/ptr" ctrl "sigs.k8s.io/controller-runtime" @@ -25,115 +27,273 @@ import ( ) func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { - for _, lifecycle := range []string{"delete", "retarget"} { - t.Run(lifecycle, func(t *testing.T) { + for _, kind := range []string{"Job", "RayJob"} { + t.Run(kind, func(t *testing.T) { + for _, lifecycle := range []string{"delete", "retarget"} { + t.Run(lifecycle, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + if lifecycle == "delete" { + workspace.DeletionTimestamp = &now + } + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + account := testServiceAccount(namespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + account.Annotations[annotationAzureWIClientID] = workspace.Spec.WorkloadIdentity.ClientID + newNamespace := testNamespace("new-target") + newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} + newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + if lifecycle == "retarget" { + workspace.Spec.Target.Namespace = newNamespace.Name + } + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: "training", Namespace: namespace.Name, UID: types.UID("training-job-uid"), + Labels: map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}, + }, + Spec: batchv1.JobSpec{ + Suspend: ptr.To(true), + Template: corev1.PodTemplateSpec{Spec: corev1.PodSpec{ + ServiceAccountName: account.Name, + RestartPolicy: corev1.RestartPolicyNever, + Containers: []corev1.Container{{Name: "trainer", Image: "example.invalid/trainer"}}, + }}, + }, + } + var source client.Object = job + sourceAPIVersion := batchv1.SchemeGroupVersion.String() + if kind == "RayJob" { + sourceAPIVersion = "ray.io/v1" + rayJob := newQueueObject(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayJob"}) + rayJob.SetName(job.Name) + rayJob.SetNamespace(job.Namespace) + rayJob.SetUID(job.UID) + rayJob.SetLabels(job.Labels) + rayJob.Object["spec"] = map[string]any{ + "suspend": true, "entrypoint": "python train.py", + "shutdownAfterJobFinishes": true, + "rayClusterSpec": map[string]any{"headGroupSpec": map[string]any{ + "template": map[string]any{"spec": map[string]any{"serviceAccountName": account.Name}}, + }}, + } + source = rayJob + } + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchLabels": map[string]any{labelWorkspace: workspace.Name}, + }} + c := fake.NewClientBuilder().WithScheme(testScheme(t)). + WithObjects(workspace, namespace, newNamespace, account, source, localQueue, clusterQueue). + WithStatusSubresource(&tauv1alpha1.TauWorkspace{}).Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + for range 3 { + _, err := reconciler.Reconcile(ctx, request) + if lifecycle == "delete" && (err == nil || + !strings.Contains(err.Error(), "before removing namespace routing metadata")) { + t.Fatalf("Reconcile() = %v, want source Job drain wait", err) + } + if lifecycle == "retarget" && err != nil { + t.Fatal(err) + } + var waiting tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != namespace.Name { + t.Fatal("resolved target advanced before source Job drain") + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { + t.Fatalf("execution identity removed before source Job drain: %v", err) + } + var retained corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &retained); err != nil { + t.Fatal(err) + } + if retained.Labels[labelWorkspace] != workspace.Name { + t.Fatal("source Job lost namespace routing metadata") + } + } + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(namespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: sourceAPIVersion, Kind: kind, Name: source.GetName(), UID: source.GetUID(), + }}) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + if err := c.Create(ctx, workload); err != nil { + t.Fatal(err) + } + if _, err := reconciler.Reconcile(ctx, request); lifecycle == "delete" && err == nil { + t.Fatal("workspace stopped draining after Kueue created the Workload") + } + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + if err := c.Update(ctx, workload); err != nil { + t.Fatal(err) + } + for range 3 { + if _, err := reconciler.Reconcile(ctx, request); err != nil { + t.Fatalf("Reconcile() after drain = %v", err) + } + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("execution account survived drain: %v", err) + } + var finished tauv1alpha1.TauWorkspace + err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &finished) + if lifecycle == "delete" && !apierrors.IsNotFound(err) { + t.Fatalf("workspace survived completed source Job drain: %v", err) + } + if lifecycle == "retarget" && (err != nil || finished.Status.Target.ResolvedNamespace != newNamespace.Name) { + t.Fatalf("retarget did not complete after source Job drain: %v, target=%q", err, finished.Status.Target.ResolvedNamespace) + } + }) + } + }) + } +} + +func TestNamespaceRoutingMetadataSourceRayJobs(t *testing.T) { + for _, scenario := range []string{ + "suspended", "running", "complete", "failed", "retrying", "succeeded-driver", + "unlabeled", "other-namespace", "missing-queue", "selector-independent", + "workload-finished", "workload-wrong-uid", "workload-wrong-name", "workload-wrong-kind", + "api-not-installed", "api-not-found", "list-forbidden", "list-error", + "malformed-status", "queue-read-error", "invalid-selector", + } { + t.Run(scenario, func(t *testing.T) { ctx := context.Background() workspace := testWorkspace("aurora") - workspace.Finalizers = []string{workspaceFinalizer} - workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} - now := metav1.Now() - if lifecycle == "delete" { - workspace.DeletionTimestamp = &now - } - workspace.Spec.Target.CreateNamespace = false - workspace.Spec.Queue = "jobqueue" - workspace.Status.Target.ResolvedNamespace = workspace.Name namespace := testNamespace(workspace.Name) - namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) - account := testServiceAccount(namespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) - account.Annotations[annotationAzureWIClientID] = workspace.Spec.WorkloadIdentity.ClientID - newNamespace := testNamespace("new-target") - newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} - newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} - if lifecycle == "retarget" { - workspace.Spec.Target.Namespace = newNamespace.Name - } - job := &batchv1.Job{ - ObjectMeta: metav1.ObjectMeta{ - Name: "training", Namespace: namespace.Name, UID: types.UID("training-job-uid"), - Labels: map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}, - }, - Spec: batchv1.JobSpec{ - Suspend: ptr.To(true), - Template: corev1.PodTemplateSpec{Spec: corev1.PodSpec{ - ServiceAccountName: account.Name, - RestartPolicy: corev1.RestartPolicyNever, - Containers: []corev1.Container{{Name: "trainer", Image: "example.invalid/trainer"}}, - }}, - }, - } + namespace.Labels = workspaceNamespaceLabels(workspace.Name, "jobqueue") + rayJob := newQueueObject(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayJob"}) + rayJob.SetName("training") + rayJob.SetNamespace(namespace.Name) + rayJob.SetUID(types.UID("rayjob-uid")) + rayJob.SetLabels(map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}) + rayJob.Object["spec"] = map[string]any{"suspend": true} localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") localQueue.SetAnnotations(nil) clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ "matchLabels": map[string]any{labelWorkspace: workspace.Name}, }} - c := fake.NewClientBuilder().WithScheme(testScheme(t)). - WithObjects(workspace, namespace, newNamespace, account, job, localQueue, clusterQueue). - WithStatusSubresource(&tauv1alpha1.TauWorkspace{}).Build() - reconciler := newTestWorkspaceReconciler(c) - request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} - for range 3 { - _, err := reconciler.Reconcile(ctx, request) - if lifecycle == "delete" && (err == nil || - !strings.Contains(err.Error(), "before removing namespace routing metadata")) { - t.Fatalf("Reconcile() = %v, want source Job drain wait", err) - } - if lifecycle == "retarget" && err != nil { - t.Fatal(err) - } - var waiting tauv1alpha1.TauWorkspace - if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { - t.Fatal(err) - } - if waiting.Status.Target.ResolvedNamespace != namespace.Name { - t.Fatal("resolved target advanced before source Job drain") - } - if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { - t.Fatalf("execution identity removed before source Job drain: %v", err) - } - var retained corev1.Namespace - if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &retained); err != nil { - t.Fatal(err) - } - if retained.Labels[labelWorkspace] != workspace.Name { - t.Fatal("source Job lost namespace routing metadata") - } - } workload := newQueueObject(workloadGVK) - workload.SetName("job-training") + workload.SetName("rayjob-training") workload.SetNamespace(namespace.Name) workload.SetOwnerReferences([]metav1.OwnerReference{{ - APIVersion: batchv1.SchemeGroupVersion.String(), Kind: "Job", Name: job.Name, UID: job.UID, + APIVersion: "ray.io/v1", Kind: "RayJob", Name: rayJob.GetName(), UID: rayJob.GetUID(), }}) - workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} - if err := c.Create(ctx, workload); err != nil { - t.Fatal(err) - } - if _, err := reconciler.Reconcile(ctx, request); lifecycle == "delete" && err == nil { - t.Fatal("workspace stopped draining after Kueue created the Workload") - } + workload.Object["spec"] = map[string]any{"queueName": "jobqueue"} workload.Object["status"] = map[string]any{"conditions": []any{ map[string]any{"type": "Finished", "status": "True"}, }} - if err := c.Update(ctx, workload); err != nil { - t.Fatal(err) - } - for range 3 { - if _, err := reconciler.Reconcile(ctx, request); err != nil { - t.Fatalf("Reconcile() after drain = %v", err) + objects := []client.Object{namespace, rayJob, localQueue, clusterQueue} + wantInUse, wantError := true, false + switch scenario { + case "running": + rayJob.Object["spec"] = map[string]any{"suspend": false} + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Running", "jobStatus": "RUNNING"} + case "complete", "failed": + state := "Complete" + if scenario == "failed" { + state = "Failed" } + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": state} + wantInUse = false + case "retrying": + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Retrying", "jobStatus": "FAILED"} + case "succeeded-driver": + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Running", "jobStatus": "SUCCEEDED"} + case "unlabeled": + rayJob.SetLabels(nil) + wantInUse = false + case "other-namespace": + rayJob.SetNamespace("unrelated") + wantInUse = false + case "missing-queue": + rayJob.SetLabels(map[string]string{"kueue.x-k8s.io/queue-name": "missing"}) + wantInUse = false + case "selector-independent": + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{}} + wantInUse = false + case "workload-finished", "workload-wrong-uid", "workload-wrong-name", "workload-wrong-kind": + owner := workload.GetOwnerReferences()[0] + switch scenario { + case "workload-finished": + wantInUse = false + case "workload-wrong-uid": + owner.UID = types.UID("another-uid") + case "workload-wrong-name": + owner.Name = "another-job" + case "workload-wrong-kind": + owner.APIVersion, owner.Kind = "batch/v1", "Job" + } + workload.SetOwnerReferences([]metav1.OwnerReference{owner}) + objects = append(objects, workload) + case "api-not-installed", "api-not-found": + wantInUse = false + case "list-forbidden", "list-error", "queue-read-error": + wantInUse, wantError = false, true + case "malformed-status": + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": int64(42)} + wantInUse, wantError = false, true + case "invalid-selector": + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{"key": labelWorkspace, "operator": "Invalid"}}, + }} + wantInUse, wantError = false, true } - if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { - t.Fatalf("execution account survived drain: %v", err) + baseClient := fake.NewClientBuilder().WithScheme(testScheme(t)).WithObjects(objects...).Build() + c := interceptor.NewClient(baseClient, interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if list.GetObjectKind().GroupVersionKind().Kind == "RayJobList" { + resource := schema.GroupResource{Group: "ray.io", Resource: "rayjobs"} + switch scenario { + case "api-not-installed": + return &meta.NoKindMatchError{GroupKind: schema.GroupKind{Group: "ray.io", Kind: "RayJob"}} + case "api-not-found": + return apierrors.NewNotFound(resource, "") + case "list-forbidden": + return apierrors.NewForbidden(resource, "", errors.New("RayJob discovery denied")) + case "list-error": + return errors.New("RayJob discovery unavailable") + } + } + return c.List(ctx, list, opts...) + }, + Get: func(ctx context.Context, c client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error { + if scenario == "queue-read-error" && obj.GetObjectKind().GroupVersionKind() == localQueueGVK { + return errors.New("queue unavailable") + } + return c.Get(ctx, key, obj, opts...) + }, + }) + inUse, message, err := newTestWorkspaceReconciler(c).namespaceRoutingMetadataInUse(ctx, workspace, namespace) + if inUse != wantInUse || (err != nil) != wantError { + t.Fatalf("namespaceRoutingMetadataInUse() = %v, %q, %v; want inUse=%v error=%v", inUse, message, err, wantInUse, wantError) } - var finished tauv1alpha1.TauWorkspace - err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &finished) - if lifecycle == "delete" && !apierrors.IsNotFound(err) { - t.Fatalf("workspace survived completed source Job drain: %v", err) + if inUse && !strings.Contains(message, "RayJob "+client.ObjectKeyFromObject(rayJob).String()) { + t.Fatalf("drain message does not identify the source RayJob: %q", message) } - if lifecycle == "retarget" && (err != nil || finished.Status.Target.ResolvedNamespace != newNamespace.Name) { - t.Fatalf("retarget did not complete after source Job drain: %v, target=%q", err, finished.Status.Target.ResolvedNamespace) + for _, queue := range []*unstructured.Unstructured{localQueue, clusterQueue} { + got := newQueueObject(queue.GroupVersionKind()) + if err := baseClient.Get(ctx, client.ObjectKeyFromObject(queue), got); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(got.Object, queue.Object) { + t.Fatalf("external queue was mutated: %#v", got.Object) + } } }) } diff --git a/controllers/tau-core/internal/controller/workspace_rbac_test.go b/controllers/tau-core/internal/controller/workspace_rbac_test.go index b95497a8..53847932 100644 --- a/controllers/tau-core/internal/controller/workspace_rbac_test.go +++ b/controllers/tau-core/internal/controller/workspace_rbac_test.go @@ -99,18 +99,22 @@ func TestKustomizeControllerRoleAllowsLegacyRoleDeletionWithoutMutation(t *testi } func TestKustomizeControllerRoleAllowsReadOnlySourceJobDiscovery(t *testing.T) { - role := kustomizeClusterRole(t, "tau-core-controller") - want := rbacv1.PolicyRule{ - APIGroups: []string{"batch"}, - Resources: []string{"jobs"}, - Verbs: []string{"get", "list", "watch"}, - } - for _, rule := range role.Rules { - if reflect.DeepEqual(rule, want) { - return - } + for _, source := range []struct{ group, resource string }{{"batch", "jobs"}, {"ray.io", "rayjobs"}} { + t.Run(source.resource, func(t *testing.T) { + role := kustomizeClusterRole(t, "tau-core-controller") + want := rbacv1.PolicyRule{ + APIGroups: []string{source.group}, + Resources: []string{source.resource}, + Verbs: []string{"get", "list", "watch"}, + } + for _, rule := range role.Rules { + if reflect.DeepEqual(rule, want) { + return + } + } + t.Fatalf("controller role must allow read-only source Job discovery: %#v", role.Rules) + }) } - t.Fatalf("controller role must allow read-only source Job discovery: %#v", role.Rules) } func TestRetainedNamespaceRestoresOnlyCurrentAuthorizedResearcherBinding(t *testing.T) { diff --git a/site/content/en/docs/reference/workspace.md b/site/content/en/docs/reference/workspace.md index 58ad95e3..9b0ad31f 100644 --- a/site/content/en/docs/reference/workspace.md +++ b/site/content/en/docs/reference/workspace.md @@ -51,10 +51,15 @@ The overall phase currently requires `RBACReady` and `QueueReady`; `DriftDetecte During namespace retargeting, unfinished Workloads on retained external queues can delay cleanup of the old namespace's routing metadata. Submitted, unfinished -Jobs that have not received a derived Workload yet also retain routing metadata -on deletion or retargeting, using their queue label and the live queue selector. -An unchanged workload -identity ServiceAccount is retained until those Workloads finish, while researcher +Jobs and RayJobs that have not received a derived Workload yet also retain routing +metadata on deletion or retargeting, using their queue label and the live queue +selector. Completed or failed sources do not block cleanup. A RayJob is finished +only when `status.jobDeploymentStatus` is `Complete` or `Failed`, matching Kueue; +a failed driver with a retry still pending retains its route. RayJob discovery +is skipped when the optional RayJob API is not installed; other discovery errors +block cleanup until the source state can be read. Both Helm and Kustomize grant +the controller read-only discovery of Jobs and RayJobs. +An unchanged workload identity ServiceAccount is retained until those Workloads finish, while researcher authorization is reconciled independently. Changing the ServiceAccount name or client ID, or removing `workloadIdentity`, revokes the obsolete ServiceAccount before granting the current researcher access to the old namespace, even during From 325aa539c0c7cd086ac67f4cf8b35211b8911c69 Mon Sep 17 00:00:00 2001 From: chokevin Date: Mon, 26 Oct 2026 17:48:39 -0700 Subject: [PATCH 12/13] fix(controller): retain workspace routes for RayServices --- .../tau-core-controller/kustomize/rbac.yaml | 2 +- .../tau-core-controller/templates/rbac.yaml | 2 +- .../tau-core-controller/tests/rbac_test.yaml | 2 +- .../internal/controller/workspace_cleanup.go | 62 ++-- .../controller/workspace_cleanup_test.go | 308 ++++++++++-------- .../controller/workspace_rbac_test.go | 12 +- site/content/en/docs/reference/workspace.md | 15 +- 7 files changed, 233 insertions(+), 170 deletions(-) diff --git a/charts/tau-core-controller/kustomize/rbac.yaml b/charts/tau-core-controller/kustomize/rbac.yaml index 230f1a71..c1ce7b1c 100644 --- a/charts/tau-core-controller/kustomize/rbac.yaml +++ b/charts/tau-core-controller/kustomize/rbac.yaml @@ -114,7 +114,7 @@ rules: resources: ["jobs"] verbs: ["get", "list", "watch"] - apiGroups: ["ray.io"] - resources: ["rayjobs"] + resources: ["rayjobs", "rayservices"] verbs: ["get", "list", "watch"] - apiGroups: ["kueue.x-k8s.io"] resources: ["workloads"] diff --git a/charts/tau-core-controller/templates/rbac.yaml b/charts/tau-core-controller/templates/rbac.yaml index 0a13fb34..b2c0cbc3 100644 --- a/charts/tau-core-controller/templates/rbac.yaml +++ b/charts/tau-core-controller/templates/rbac.yaml @@ -111,7 +111,7 @@ rules: resources: ["jobs"] verbs: ["get", "list", "watch"] - apiGroups: ["ray.io"] - resources: ["rayjobs"] + resources: ["rayjobs", "rayservices"] verbs: ["get", "list", "watch"] - apiGroups: ["kueue.x-k8s.io"] resources: ["workloads"] diff --git a/charts/tau-core-controller/tests/rbac_test.yaml b/charts/tau-core-controller/tests/rbac_test.yaml index 8eb7f22c..a13a8dec 100644 --- a/charts/tau-core-controller/tests/rbac_test.yaml +++ b/charts/tau-core-controller/tests/rbac_test.yaml @@ -15,7 +15,7 @@ tests: path: rules content: apiGroups: ["ray.io"] - resources: ["rayjobs"] + resources: ["rayjobs", "rayservices"] verbs: ["get", "list", "watch"] documentIndex: 2 diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 098be79e..18fc840d 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -242,36 +242,40 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( return inUse, message, err } } - rayJobs := &unstructured.UnstructuredList{} - rayJobs.SetGroupVersionKind(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayJobList"}) - if err := r.apiReader().List(ctx, rayJobs, client.InNamespace(namespace.Name)); err != nil { - if meta.IsNoMatchError(err) || apierrors.IsNotFound(err) { - return false, "", nil - } - return false, "", fmt.Errorf("list RayJobs in namespace %q before routing metadata cleanup: %w", namespace.Name, err) - } - for i := range rayJobs.Items { - rayJob := &rayJobs.Items[i] - if rayJob.GetUID() != "" && workloadRayJobs[rayJob.GetUID()] == rayJob.GetName() { - continue - } - queueName := strings.TrimSpace(rayJob.GetLabels()["kueue.x-k8s.io/queue-name"]) - if queueName == "" { - continue - } - finished, err := sourceRayJobFinished(rayJob) - if err != nil { - return false, "", err - } - if finished { - continue + for _, kind := range []string{"RayJob", "RayService"} { + sources := &unstructured.UnstructuredList{} + sources.SetGroupVersionKind(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: kind + "List"}) + if err := r.apiReader().List(ctx, sources, client.InNamespace(namespace.Name)); err != nil { + if meta.IsNoMatchError(err) || apierrors.IsNotFound(err) { + continue + } + return false, "", fmt.Errorf("list %ss in namespace %q before routing metadata cleanup: %w", kind, namespace.Name, err) } - inUse, message, err := r.queueRoutingMetadataInUse( - ctx, workspace, namespace, cleanedLabels, queueName, "", - fmt.Sprintf("RayJob %s", client.ObjectKeyFromObject(rayJob)), - ) - if err != nil || inUse { - return inUse, message, err + for i := range sources.Items { + source := &sources.Items[i] + queueName := strings.TrimSpace(source.GetLabels()["kueue.x-k8s.io/queue-name"]) + if queueName == "" { + continue + } + if kind == "RayJob" { + if source.GetUID() != "" && workloadRayJobs[source.GetUID()] == source.GetName() { + continue + } + finished, err := sourceRayJobFinished(source) + if err != nil { + return false, "", err + } + if finished { + continue + } + } + inUse, message, err := r.queueRoutingMetadataInUse( + ctx, workspace, namespace, cleanedLabels, queueName, "", + fmt.Sprintf("%s %s", kind, client.ObjectKeyFromObject(source)), + ) + if err != nil || inUse { + return inUse, message, err + } } } return false, "", nil diff --git a/controllers/tau-core/internal/controller/workspace_cleanup_test.go b/controllers/tau-core/internal/controller/workspace_cleanup_test.go index 3d8c6bab..76f18260 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup_test.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup_test.go @@ -27,7 +27,7 @@ import ( ) func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { - for _, kind := range []string{"Job", "RayJob"} { + for _, kind := range []string{"Job", "RayJob", "RayService"} { t.Run(kind, func(t *testing.T) { for _, lifecycle := range []string{"delete", "retarget"} { t.Run(lifecycle, func(t *testing.T) { @@ -68,9 +68,9 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { } var source client.Object = job sourceAPIVersion := batchv1.SchemeGroupVersion.String() - if kind == "RayJob" { + if kind != "Job" { sourceAPIVersion = "ray.io/v1" - rayJob := newQueueObject(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayJob"}) + rayJob := newQueueObject(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: kind}) rayJob.SetName(job.Name) rayJob.SetNamespace(job.Namespace) rayJob.SetUID(job.UID) @@ -82,6 +82,16 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { "template": map[string]any{"spec": map[string]any{"serviceAccountName": account.Name}}, }}, } + if kind == "RayService" { + rayJob.Object["spec"] = map[string]any{ + "serveConfigV2": "applications:\n - name: default\n import_path: serve:app\n", + "rayClusterConfig": map[string]any{ + "headGroupSpec": map[string]any{ + "template": map[string]any{"spec": map[string]any{"serviceAccountName": account.Name}}, + }, + }, + } + } source = rayJob } localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") @@ -141,6 +151,21 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { if err := c.Update(ctx, workload); err != nil { t.Fatal(err) } + if kind == "RayService" { + _, err := reconciler.Reconcile(ctx, request) + if lifecycle == "delete" && err == nil { + t.Fatal("live RayService stopped retaining routing after its Workload finished") + } + if lifecycle == "retarget" && err != nil { + t.Fatal(err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { + t.Fatalf("live RayService lost execution identity: %v", err) + } + if err := c.Delete(ctx, source); err != nil { + t.Fatal(err) + } + } for range 3 { if _, err := reconciler.Reconcile(ctx, request); err != nil { t.Fatalf("Reconcile() after drain = %v", err) @@ -163,137 +188,160 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { } } -func TestNamespaceRoutingMetadataSourceRayJobs(t *testing.T) { - for _, scenario := range []string{ - "suspended", "running", "complete", "failed", "retrying", "succeeded-driver", - "unlabeled", "other-namespace", "missing-queue", "selector-independent", - "workload-finished", "workload-wrong-uid", "workload-wrong-name", "workload-wrong-kind", - "api-not-installed", "api-not-found", "list-forbidden", "list-error", - "malformed-status", "queue-read-error", "invalid-selector", - } { - t.Run(scenario, func(t *testing.T) { - ctx := context.Background() - workspace := testWorkspace("aurora") - namespace := testNamespace(workspace.Name) - namespace.Labels = workspaceNamespaceLabels(workspace.Name, "jobqueue") - rayJob := newQueueObject(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayJob"}) - rayJob.SetName("training") - rayJob.SetNamespace(namespace.Name) - rayJob.SetUID(types.UID("rayjob-uid")) - rayJob.SetLabels(map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}) - rayJob.Object["spec"] = map[string]any{"suspend": true} - localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") - localQueue.SetAnnotations(nil) - clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") - clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ - "matchLabels": map[string]any{labelWorkspace: workspace.Name}, - }} - workload := newQueueObject(workloadGVK) - workload.SetName("rayjob-training") - workload.SetNamespace(namespace.Name) - workload.SetOwnerReferences([]metav1.OwnerReference{{ - APIVersion: "ray.io/v1", Kind: "RayJob", Name: rayJob.GetName(), UID: rayJob.GetUID(), - }}) - workload.Object["spec"] = map[string]any{"queueName": "jobqueue"} - workload.Object["status"] = map[string]any{"conditions": []any{ - map[string]any{"type": "Finished", "status": "True"}, - }} - objects := []client.Object{namespace, rayJob, localQueue, clusterQueue} - wantInUse, wantError := true, false - switch scenario { - case "running": - rayJob.Object["spec"] = map[string]any{"suspend": false} - rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Running", "jobStatus": "RUNNING"} - case "complete", "failed": - state := "Complete" - if scenario == "failed" { - state = "Failed" - } - rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": state} - wantInUse = false - case "retrying": - rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Retrying", "jobStatus": "FAILED"} - case "succeeded-driver": - rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Running", "jobStatus": "SUCCEEDED"} - case "unlabeled": - rayJob.SetLabels(nil) - wantInUse = false - case "other-namespace": - rayJob.SetNamespace("unrelated") - wantInUse = false - case "missing-queue": - rayJob.SetLabels(map[string]string{"kueue.x-k8s.io/queue-name": "missing"}) - wantInUse = false - case "selector-independent": - clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{}} - wantInUse = false - case "workload-finished", "workload-wrong-uid", "workload-wrong-name", "workload-wrong-kind": - owner := workload.GetOwnerReferences()[0] - switch scenario { - case "workload-finished": - wantInUse = false - case "workload-wrong-uid": - owner.UID = types.UID("another-uid") - case "workload-wrong-name": - owner.Name = "another-job" - case "workload-wrong-kind": - owner.APIVersion, owner.Kind = "batch/v1", "Job" - } - workload.SetOwnerReferences([]metav1.OwnerReference{owner}) - objects = append(objects, workload) - case "api-not-installed", "api-not-found": - wantInUse = false - case "list-forbidden", "list-error", "queue-read-error": - wantInUse, wantError = false, true - case "malformed-status": - rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": int64(42)} - wantInUse, wantError = false, true - case "invalid-selector": - clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ - "matchExpressions": []any{map[string]any{"key": labelWorkspace, "operator": "Invalid"}}, - }} - wantInUse, wantError = false, true - } - baseClient := fake.NewClientBuilder().WithScheme(testScheme(t)).WithObjects(objects...).Build() - c := interceptor.NewClient(baseClient, interceptor.Funcs{ - List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { - if list.GetObjectKind().GroupVersionKind().Kind == "RayJobList" { - resource := schema.GroupResource{Group: "ray.io", Resource: "rayjobs"} +func TestNamespaceRoutingMetadataSourceRayResources(t *testing.T) { + for _, kind := range []string{"RayJob", "RayService"} { + t.Run(kind, func(t *testing.T) { + for _, scenario := range []string{ + "suspended", "running", "complete", "failed", "retrying", "succeeded-driver", + "unlabeled", "other-namespace", "missing-queue", "selector-independent", + "workload-finished", "workload-wrong-uid", "workload-wrong-name", "workload-wrong-kind", + "api-not-installed", "api-not-found", "list-forbidden", "list-error", + "malformed-status", "queue-read-error", "invalid-selector", + "rayjob-api-not-installed", "rayjob-api-not-found", + } { + t.Run(scenario, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, "jobqueue") + rayJob := newQueueObject(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: kind}) + rayJob.SetName("training") + rayJob.SetNamespace(namespace.Name) + rayJob.SetUID(types.UID("rayjob-uid")) + rayJob.SetLabels(map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}) + if kind == "RayService" { + rayJob.Object["spec"] = map[string]any{"rayClusterConfig": map[string]any{}} + } else { + rayJob.Object["spec"] = map[string]any{"suspend": true} + } + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchLabels": map[string]any{labelWorkspace: workspace.Name}, + }} + workload := newQueueObject(workloadGVK) + workload.SetName("rayjob-training") + workload.SetNamespace(namespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: "ray.io/v1", Kind: kind, Name: rayJob.GetName(), UID: rayJob.GetUID(), + }}) + workload.Object["spec"] = map[string]any{"queueName": "jobqueue"} + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + objects := []client.Object{namespace, rayJob, localQueue, clusterQueue} + wantInUse, wantError := true, false + switch scenario { + case "running": + rayJob.Object["spec"] = map[string]any{"suspend": false} + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Running", "jobStatus": "RUNNING"} + case "complete", "failed": + state := "Complete" + if scenario == "failed" { + state = "Failed" + } + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": state} + wantInUse = kind == "RayService" + case "retrying": + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Retrying", "jobStatus": "FAILED"} + case "succeeded-driver": + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": "Running", "jobStatus": "SUCCEEDED"} + case "unlabeled": + rayJob.SetLabels(nil) + wantInUse = false + case "other-namespace": + rayJob.SetNamespace("unrelated") + wantInUse = false + case "missing-queue": + rayJob.SetLabels(map[string]string{"kueue.x-k8s.io/queue-name": "missing"}) + wantInUse = false + case "selector-independent": + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{}} + wantInUse = false + case "workload-finished", "workload-wrong-uid", "workload-wrong-name", "workload-wrong-kind": + owner := workload.GetOwnerReferences()[0] switch scenario { - case "api-not-installed": - return &meta.NoKindMatchError{GroupKind: schema.GroupKind{Group: "ray.io", Kind: "RayJob"}} - case "api-not-found": - return apierrors.NewNotFound(resource, "") - case "list-forbidden": - return apierrors.NewForbidden(resource, "", errors.New("RayJob discovery denied")) - case "list-error": - return errors.New("RayJob discovery unavailable") + case "workload-finished": + wantInUse = kind == "RayService" + case "workload-wrong-uid": + owner.UID = types.UID("another-uid") + case "workload-wrong-name": + owner.Name = "another-job" + case "workload-wrong-kind": + owner.Kind = "RayService" + if kind == "RayService" { + owner.Kind = "RayJob" + } } + workload.SetOwnerReferences([]metav1.OwnerReference{owner}) + objects = append(objects, workload) + case "api-not-installed", "api-not-found": + wantInUse = false + case "list-forbidden", "list-error", "queue-read-error": + wantInUse, wantError = false, true + case "malformed-status": + rayJob.Object["status"] = map[string]any{"jobDeploymentStatus": int64(42)} + wantInUse, wantError = kind == "RayService", kind == "RayJob" + case "rayjob-api-not-installed", "rayjob-api-not-found": + wantInUse = kind == "RayService" + case "invalid-selector": + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{"key": labelWorkspace, "operator": "Invalid"}}, + }} + wantInUse, wantError = false, true } - return c.List(ctx, list, opts...) - }, - Get: func(ctx context.Context, c client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error { - if scenario == "queue-read-error" && obj.GetObjectKind().GroupVersionKind() == localQueueGVK { - return errors.New("queue unavailable") + baseClient := fake.NewClientBuilder().WithScheme(testScheme(t)).WithObjects(objects...).Build() + c := interceptor.NewClient(baseClient, interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + listKind := list.GetObjectKind().GroupVersionKind().Kind + if listKind == "RayJobList" { + if scenario == "rayjob-api-not-installed" { + return &meta.NoKindMatchError{GroupKind: schema.GroupKind{Group: "ray.io", Kind: "RayJob"}} + } + if scenario == "rayjob-api-not-found" { + return apierrors.NewNotFound(schema.GroupResource{Group: "ray.io", Resource: "rayjobs"}, "") + } + } + if listKind == kind+"List" { + resource := schema.GroupResource{Group: "ray.io", Resource: strings.ToLower(kind) + "s"} + switch scenario { + case "api-not-installed": + return &meta.NoKindMatchError{GroupKind: schema.GroupKind{Group: "ray.io", Kind: kind}} + case "api-not-found": + return apierrors.NewNotFound(resource, "") + case "list-forbidden": + return apierrors.NewForbidden(resource, "", errors.New("RayJob discovery denied")) + case "list-error": + return errors.New("RayJob discovery unavailable") + } + } + return c.List(ctx, list, opts...) + }, + Get: func(ctx context.Context, c client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error { + if scenario == "queue-read-error" && obj.GetObjectKind().GroupVersionKind() == localQueueGVK { + return errors.New("queue unavailable") + } + return c.Get(ctx, key, obj, opts...) + }, + }) + inUse, message, err := newTestWorkspaceReconciler(c).namespaceRoutingMetadataInUse(ctx, workspace, namespace) + if inUse != wantInUse || (err != nil) != wantError { + t.Fatalf("namespaceRoutingMetadataInUse() = %v, %q, %v; want inUse=%v error=%v", inUse, message, err, wantInUse, wantError) } - return c.Get(ctx, key, obj, opts...) - }, - }) - inUse, message, err := newTestWorkspaceReconciler(c).namespaceRoutingMetadataInUse(ctx, workspace, namespace) - if inUse != wantInUse || (err != nil) != wantError { - t.Fatalf("namespaceRoutingMetadataInUse() = %v, %q, %v; want inUse=%v error=%v", inUse, message, err, wantInUse, wantError) - } - if inUse && !strings.Contains(message, "RayJob "+client.ObjectKeyFromObject(rayJob).String()) { - t.Fatalf("drain message does not identify the source RayJob: %q", message) - } - for _, queue := range []*unstructured.Unstructured{localQueue, clusterQueue} { - got := newQueueObject(queue.GroupVersionKind()) - if err := baseClient.Get(ctx, client.ObjectKeyFromObject(queue), got); err != nil { - t.Fatal(err) - } - if !reflect.DeepEqual(got.Object, queue.Object) { - t.Fatalf("external queue was mutated: %#v", got.Object) - } + if inUse && !strings.Contains(message, kind+" "+client.ObjectKeyFromObject(rayJob).String()) { + t.Fatalf("drain message does not identify the source %s: %q", kind, message) + } + for _, queue := range []*unstructured.Unstructured{localQueue, clusterQueue} { + got := newQueueObject(queue.GroupVersionKind()) + if err := baseClient.Get(ctx, client.ObjectKeyFromObject(queue), got); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(got.Object, queue.Object) { + t.Fatalf("external queue was mutated: %#v", got.Object) + } + } + }) } }) } diff --git a/controllers/tau-core/internal/controller/workspace_rbac_test.go b/controllers/tau-core/internal/controller/workspace_rbac_test.go index 53847932..6a663482 100644 --- a/controllers/tau-core/internal/controller/workspace_rbac_test.go +++ b/controllers/tau-core/internal/controller/workspace_rbac_test.go @@ -99,12 +99,18 @@ func TestKustomizeControllerRoleAllowsLegacyRoleDeletionWithoutMutation(t *testi } func TestKustomizeControllerRoleAllowsReadOnlySourceJobDiscovery(t *testing.T) { - for _, source := range []struct{ group, resource string }{{"batch", "jobs"}, {"ray.io", "rayjobs"}} { - t.Run(source.resource, func(t *testing.T) { + for _, source := range []struct { + group string + resources []string + }{ + {"batch", []string{"jobs"}}, + {"ray.io", []string{"rayjobs", "rayservices"}}, + } { + t.Run(source.group, func(t *testing.T) { role := kustomizeClusterRole(t, "tau-core-controller") want := rbacv1.PolicyRule{ APIGroups: []string{source.group}, - Resources: []string{source.resource}, + Resources: source.resources, Verbs: []string{"get", "list", "watch"}, } for _, rule := range role.Rules { diff --git a/site/content/en/docs/reference/workspace.md b/site/content/en/docs/reference/workspace.md index 9b0ad31f..10c8d183 100644 --- a/site/content/en/docs/reference/workspace.md +++ b/site/content/en/docs/reference/workspace.md @@ -51,14 +51,19 @@ The overall phase currently requires `RBACReady` and `QueueReady`; `DriftDetecte During namespace retargeting, unfinished Workloads on retained external queues can delay cleanup of the old namespace's routing metadata. Submitted, unfinished -Jobs and RayJobs that have not received a derived Workload yet also retain routing +Jobs, RayJobs, and RayServices that have not received a derived Workload yet retain routing metadata on deletion or retargeting, using their queue label and the live queue -selector. Completed or failed sources do not block cleanup. A RayJob is finished +selector. Completed or failed batch sources do not block cleanup. A RayJob is finished only when `status.jobDeploymentStatus` is `Complete` or `Failed`, matching Kueue; a failed driver with a retry still pending retains its route. RayJob discovery -is skipped when the optional RayJob API is not installed; other discovery errors -block cleanup until the source state can be read. Both Helm and Kustomize grant -the controller read-only discovery of Jobs and RayJobs. +is skipped when the optional RayJob API is not installed; RayService discovery +is checked independently and also tolerates an absent API. Other discovery errors +block cleanup until the source state can be read. RayServices have no terminal +job state: a queued service retains required routing metadata and unchanged +execution identity until the service is deleted, even if a derived Workload +reports `Finished`. Delete the service before completing workspace +deletion or retargeting. Both Helm and Kustomize grant the controller read-only +discovery of Jobs, RayJobs, and RayServices. An unchanged workload identity ServiceAccount is retained until those Workloads finish, while researcher authorization is reconciled independently. Changing the ServiceAccount name or client ID, or removing `workloadIdentity`, revokes the obsolete ServiceAccount From 9283d7d07218cac379cd9e896193226f30d2f528 Mon Sep 17 00:00:00 2001 From: chokevin Date: Mon, 26 Oct 2026 19:14:46 -0700 Subject: [PATCH 13/13] fix(controller): retain routes for queued Deployments --- .../tau-core-controller/kustomize/rbac.yaml | 3 + .../tau-core-controller/templates/rbac.yaml | 3 + .../tau-core-controller/tests/rbac_test.yaml | 7 + .../internal/controller/workspace_cleanup.go | 19 +++ .../controller/workspace_cleanup_test.go | 128 +++++++++++++++++- .../controller/workspace_rbac_test.go | 1 + site/content/en/docs/reference/workspace.md | 8 ++ 7 files changed, 164 insertions(+), 5 deletions(-) diff --git a/charts/tau-core-controller/kustomize/rbac.yaml b/charts/tau-core-controller/kustomize/rbac.yaml index c1ce7b1c..8b4a05d2 100644 --- a/charts/tau-core-controller/kustomize/rbac.yaml +++ b/charts/tau-core-controller/kustomize/rbac.yaml @@ -113,6 +113,9 @@ rules: - apiGroups: ["batch"] resources: ["jobs"] verbs: ["get", "list", "watch"] + - apiGroups: ["apps"] + resources: ["deployments"] + verbs: ["get", "list", "watch"] - apiGroups: ["ray.io"] resources: ["rayjobs", "rayservices"] verbs: ["get", "list", "watch"] diff --git a/charts/tau-core-controller/templates/rbac.yaml b/charts/tau-core-controller/templates/rbac.yaml index b2c0cbc3..5998a570 100644 --- a/charts/tau-core-controller/templates/rbac.yaml +++ b/charts/tau-core-controller/templates/rbac.yaml @@ -110,6 +110,9 @@ rules: - apiGroups: ["batch"] resources: ["jobs"] verbs: ["get", "list", "watch"] + - apiGroups: ["apps"] + resources: ["deployments"] + verbs: ["get", "list", "watch"] - apiGroups: ["ray.io"] resources: ["rayjobs", "rayservices"] verbs: ["get", "list", "watch"] diff --git a/charts/tau-core-controller/tests/rbac_test.yaml b/charts/tau-core-controller/tests/rbac_test.yaml index a13a8dec..6d0dec00 100644 --- a/charts/tau-core-controller/tests/rbac_test.yaml +++ b/charts/tau-core-controller/tests/rbac_test.yaml @@ -4,6 +4,13 @@ templates: tests: - it: grants the controller read-only source Job access for workspace drain checks asserts: + - contains: + path: rules + content: + apiGroups: ["apps"] + resources: ["deployments"] + verbs: ["get", "list", "watch"] + documentIndex: 2 - contains: path: rules content: diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 18fc840d..8cb321be 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -10,6 +10,7 @@ import ( "strings" tauv1alpha1 "github.com/Azure/taugrid/controllers/tau-core/api/v1alpha1" + appsv1 "k8s.io/api/apps/v1" batchv1 "k8s.io/api/batch/v1" corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" @@ -242,6 +243,24 @@ func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( return inUse, message, err } } + var deployments appsv1.DeploymentList + if err := r.apiReader().List(ctx, &deployments, client.InNamespace(namespace.Name)); err != nil { + return false, "", fmt.Errorf("list Deployments in namespace %q before routing metadata cleanup: %w", namespace.Name, err) + } + for i := range deployments.Items { + deployment := &deployments.Items[i] + queueName := strings.TrimSpace(deployment.Spec.Template.Labels["kueue.x-k8s.io/queue-name"]) + if deployment.Spec.Template.Labels["kueue.x-k8s.io/managed"] != "true" || queueName == "" { + continue + } + inUse, message, err := r.queueRoutingMetadataInUse( + ctx, workspace, namespace, cleanedLabels, queueName, "", + fmt.Sprintf("Deployment %s", client.ObjectKeyFromObject(deployment)), + ) + if err != nil || inUse { + return inUse, message, err + } + } for _, kind := range []string{"RayJob", "RayService"} { sources := &unstructured.UnstructuredList{} sources.SetGroupVersionKind(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: kind + "List"}) diff --git a/controllers/tau-core/internal/controller/workspace_cleanup_test.go b/controllers/tau-core/internal/controller/workspace_cleanup_test.go index 76f18260..7baa3b17 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup_test.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup_test.go @@ -11,6 +11,7 @@ import ( "testing" tauv1alpha1 "github.com/Azure/taugrid/controllers/tau-core/api/v1alpha1" + appsv1 "k8s.io/api/apps/v1" batchv1 "k8s.io/api/batch/v1" corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" @@ -27,7 +28,7 @@ import ( ) func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { - for _, kind := range []string{"Job", "RayJob", "RayService"} { + for _, kind := range []string{"Job", "RayJob", "RayService", "Deployment"} { t.Run(kind, func(t *testing.T) { for _, lifecycle := range []string{"delete", "retarget"} { t.Run(lifecycle, func(t *testing.T) { @@ -68,7 +69,26 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { } var source client.Object = job sourceAPIVersion := batchv1.SchemeGroupVersion.String() - if kind != "Job" { + if kind == "Deployment" { + sourceAPIVersion = appsv1.SchemeGroupVersion.String() + source = &appsv1.Deployment{ + ObjectMeta: job.ObjectMeta, + Spec: appsv1.DeploymentSpec{ + Replicas: ptr.To(int32(1)), + Selector: &metav1.LabelSelector{MatchLabels: map[string]string{"app": job.Name}}, + Template: corev1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + "app": job.Name, "kueue.x-k8s.io/queue-name": "jobqueue", + "kueue.x-k8s.io/managed": "true", + }, + Annotations: map[string]string{"kueue.x-k8s.io/pod-suspending-parent": "deployment"}, + }, + Spec: job.Spec.Template.Spec, + }, + }, + } + } else if kind != "Job" { sourceAPIVersion = "ray.io/v1" rayJob := newQueueObject(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: kind}) rayJob.SetName(job.Name) @@ -138,6 +158,11 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { workload.SetOwnerReferences([]metav1.OwnerReference{{ APIVersion: sourceAPIVersion, Kind: kind, Name: source.GetName(), UID: source.GetUID(), }}) + if kind == "Deployment" { + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: "v1", Kind: "Pod", Name: "service-pod", UID: types.UID("service-pod-uid"), + }}) + } workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} if err := c.Create(ctx, workload); err != nil { t.Fatal(err) @@ -151,16 +176,16 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { if err := c.Update(ctx, workload); err != nil { t.Fatal(err) } - if kind == "RayService" { + if kind == "RayService" || kind == "Deployment" { _, err := reconciler.Reconcile(ctx, request) if lifecycle == "delete" && err == nil { - t.Fatal("live RayService stopped retaining routing after its Workload finished") + t.Fatalf("live %s stopped retaining routing after its Workload finished", kind) } if lifecycle == "retarget" && err != nil { t.Fatal(err) } if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { - t.Fatalf("live RayService lost execution identity: %v", err) + t.Fatalf("live %s lost execution identity: %v", kind, err) } if err := c.Delete(ctx, source); err != nil { t.Fatal(err) @@ -188,6 +213,99 @@ func TestWorkspacePreservesRoutingForJobBeforeWorkloadCreation(t *testing.T) { } } +func TestNamespaceRoutingMetadataSourceDeployments(t *testing.T) { + for _, scenario := range []string{ + "managed", "unmanaged", "missing-managed-label", "root-label-only", "template-label-only", + "template-queue-differs", "zero-replicas", "paused", "deleting", "other-namespace", + "missing-queue", "selector-independent", "list-forbidden", "queue-read-error", + } { + t.Run(scenario, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, "jobqueue") + deployment := &appsv1.Deployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "service", Namespace: namespace.Name, + Labels: map[string]string{"kueue.x-k8s.io/queue-name": "jobqueue"}, + }, + Spec: appsv1.DeploymentSpec{ + Replicas: ptr.To(int32(1)), + Template: corev1.PodTemplateSpec{ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + "kueue.x-k8s.io/managed": "true", "kueue.x-k8s.io/queue-name": "jobqueue", + }, + }}, + }, + } + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{ + "matchLabels": map[string]any{labelWorkspace: workspace.Name}, + }} + wantInUse, wantError := true, false + switch scenario { + case "unmanaged": + deployment.Spec.Template.Labels["kueue.x-k8s.io/managed"] = "false" + wantInUse = false + case "missing-managed-label": + delete(deployment.Spec.Template.Labels, "kueue.x-k8s.io/managed") + wantInUse = false + case "root-label-only": + delete(deployment.Spec.Template.Labels, "kueue.x-k8s.io/queue-name") + wantInUse = false + case "template-label-only": + deployment.Labels = nil + case "template-queue-differs": + deployment.Labels["kueue.x-k8s.io/queue-name"] = "unrelated" + case "zero-replicas": + deployment.Spec.Replicas = ptr.To(int32(0)) + case "paused": + deployment.Spec.Paused = true + case "deleting": + now := metav1.Now() + deployment.DeletionTimestamp = &now + deployment.Finalizers = []string{"example.invalid/drain"} + case "other-namespace": + deployment.Namespace = "unrelated" + wantInUse = false + case "missing-queue": + deployment.Spec.Template.Labels["kueue.x-k8s.io/queue-name"] = "missing" + wantInUse = false + case "selector-independent": + clusterQueue.Object["spec"] = map[string]any{"namespaceSelector": map[string]any{}} + wantInUse = false + case "list-forbidden", "queue-read-error": + wantInUse, wantError = false, true + } + baseClient := fake.NewClientBuilder().WithScheme(testScheme(t)). + WithObjects(namespace, deployment, localQueue, clusterQueue).Build() + c := interceptor.NewClient(baseClient, interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*appsv1.DeploymentList); ok && scenario == "list-forbidden" { + return apierrors.NewForbidden(schema.GroupResource{Group: "apps", Resource: "deployments"}, "", errors.New("discovery denied")) + } + return c.List(ctx, list, opts...) + }, + Get: func(ctx context.Context, c client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error { + if scenario == "queue-read-error" && obj.GetObjectKind().GroupVersionKind() == localQueueGVK { + return errors.New("queue unavailable") + } + return c.Get(ctx, key, obj, opts...) + }, + }) + inUse, message, err := newTestWorkspaceReconciler(c).namespaceRoutingMetadataInUse(ctx, workspace, namespace) + if inUse != wantInUse || (err != nil) != wantError { + t.Fatalf("namespaceRoutingMetadataInUse() = %v, %q, %v; want inUse=%v error=%v", inUse, message, err, wantInUse, wantError) + } + if inUse && !strings.Contains(message, "Deployment "+client.ObjectKeyFromObject(deployment).String()) { + t.Fatalf("drain message does not identify source Deployment: %q", message) + } + }) + } +} + func TestNamespaceRoutingMetadataSourceRayResources(t *testing.T) { for _, kind := range []string{"RayJob", "RayService"} { t.Run(kind, func(t *testing.T) { diff --git a/controllers/tau-core/internal/controller/workspace_rbac_test.go b/controllers/tau-core/internal/controller/workspace_rbac_test.go index 6a663482..4566be1c 100644 --- a/controllers/tau-core/internal/controller/workspace_rbac_test.go +++ b/controllers/tau-core/internal/controller/workspace_rbac_test.go @@ -104,6 +104,7 @@ func TestKustomizeControllerRoleAllowsReadOnlySourceJobDiscovery(t *testing.T) { resources []string }{ {"batch", []string{"jobs"}}, + {"apps", []string{"deployments"}}, {"ray.io", []string{"rayjobs", "rayservices"}}, } { t.Run(source.group, func(t *testing.T) { diff --git a/site/content/en/docs/reference/workspace.md b/site/content/en/docs/reference/workspace.md index 10c8d183..28ac0fa6 100644 --- a/site/content/en/docs/reference/workspace.md +++ b/site/content/en/docs/reference/workspace.md @@ -64,6 +64,14 @@ execution identity until the service is deleted, even if a derived Workload reports `Finished`. Delete the service before completing workspace deletion or retargeting. Both Helm and Kustomize grant the controller read-only discovery of Jobs, RayJobs, and RayServices. +Kueue-managed Deployments submitted with `tau serve deploy --kind=deployment` +also retain routing before their Pod Workloads exist. Discovery uses the Pod +template's managed and queue labels, not just the Deployment's own labels. +A Deployment has no terminal job state: a present managed Deployment retains +its required route even at zero replicas or after a Pod Workload finishes. +Delete the Deployment before completing workspace deletion or retargeting. +Both controller RBAC manifests include read-only Deployment discovery; +unreadable Deployment lists block cleanup. An unchanged workload identity ServiceAccount is retained until those Workloads finish, while researcher authorization is reconciled independently. Changing the ServiceAccount name or client ID, or removing `workloadIdentity`, revokes the obsolete ServiceAccount