controlapi: Add worker pool metadata to workers (#403)

Decouple the actor resumption scheduler from Kubernetes by avoiding
direct queries to WorkerPool CRDs. Instead, cache the worker pool's
labels and sandbox class on the Worker object in Redis during pod
synchronization, and update the scheduler to evaluate worker eligibility
using these cached fields.

Author: sfunkenhauser

cmd/ateapi/internal/controlapi/functional_test.go

@@ -283,7 +283,7 @@ func setupTest(t *testing.T, ns string) *testContext {
 
 	ctx, cancel := context.WithCancel(context.Background())
 
-	syncer := NewWorkerPoolSyncer(persistence, workerInformer)
+	syncer := NewWorkerPoolSyncer(persistence, workerInformer, workerPoolLister)
 	syncer.Start(ctx)
 
 	workerFactory.Start(ctx.Done())
@@ -1044,15 +1044,17 @@ func TestListActors_PageSizeValidation(t *testing.T) {
 
 // TestListWorkers tests that workers mirrored to Redis are listed.
 // Workflow:
-// 1. Creates a mock worker Pod in Kubernetes.
-// 2. Waits for the background WorkerPoolSyncer to mirror it to Redis.
-// 3. Calls ListWorkers RPC.
-// 4. Verifies that the worker appears in the response.
+// 1. Creates a mock WorkerPool in Kubernetes.
+// 2. Creates a mock worker Pod in Kubernetes belonging to that pool.
+// 3. Waits for the background WorkerPoolSyncer to mirror it to Redis.
+// 4. Calls ListWorkers RPC.
+// 5. Verifies that the worker appears in the response.
 func TestListWorkers(t *testing.T) {
 	ns := namespaceForTest("ns-list-workers")
 	tc := setupTest(t, ns)
 	defer tc.cleanup()
 
+	createWorkerPool(t, tc, ns, "pool1", map[string]string{"foo": "bar"})
 	createWorkerPod(t, tc, ns, "worker-1", "node1", "pool1")
 
 	listResp, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{})
@@ -1075,6 +1077,8 @@ func TestListWorkers(t *testing.T) {
 			NodeName:        "node1",
 			Ip:              "127.0.0.1",
 			Version:         1,
+			SandboxClass:    "gvisor",
+			Labels:          map[string]string{"foo": "bar"},
 		},
 	}
 
@@ -1176,8 +1180,10 @@ func TestResumeActor(t *testing.T) {
 				Atespace: testAtespace,
 			},
 		},
-		Ip:       "127.0.0.1",
-		NodeName: "node1",
+		Ip:           "127.0.0.1",
+		NodeName:     "node1",
+		SandboxClass: "gvisor",
+		Labels:       map[string]string{poolLabelKey: ns},
 	}
 
 	if diff := cmp.Diff(wantWorker, actorWorker, protocmp.Transform(), protocmp.IgnoreFields(&ateapipb.Worker{}, "version"), protocmp.IgnoreFields(&ateapipb.Worker{}, "worker_pod_uid")); diff != "" {
@@ -1298,34 +1304,6 @@ func TestResumeActor_NoWorkers(t *testing.T) {
 	assertGrpcError(t, err, codes.FailedPrecondition, "no free workers available")
 }
 
-// TestResumeActor_NoEligiblePool tests the distinct failure mode from
-// TestResumeActor_NoWorkers: here no WorkerPool's labels satisfy the
-// template's WorkerSelector at all, so there isn't even a pool to look for
-// free workers in.
-func TestResumeActor_NoEligiblePool(t *testing.T) {
-	ns := namespaceForTest("ns-resume-no-eligible-pool")
-	tc := setupTest(t, ns)
-	defer tc.cleanup()
-
-	createTemplateWithSelector(t, tc, ns, "tmpl1", &metav1.LabelSelector{
-		MatchLabels: map[string]string{"nonexistent": ns},
-	})
-
-	createResp, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
-		ActorRef:               &ateapipb.ActorRef{Atespace: testAtespace, Name: "id1"},
-		ActorTemplateNamespace: ns,
-		ActorTemplateName:      "tmpl1",
-	})
-	if err != nil {
-		t.Fatalf("CreateActor failed: %v", err)
-	}
-
-	_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
-		ActorRef: &ateapipb.ActorRef{Atespace: testAtespace, Name: createResp.GetActor().GetActorId()},
-	})
-	assertGrpcError(t, err, codes.FailedPrecondition, "no worker pool matches the template's sandboxClass and the template/actor selectors")
-}
-
 // TestResumeActor_MultiPoolSelector exercises the AND-of-two-selectors path
 // end to end: a template's WorkerSelector gates two pools, and the actor's
 // worker_selector narrows to just one of them.

cmd/ateapi/internal/controlapi/syncer.go

@@ -18,9 +18,11 @@ import (
 	"context"
 	"errors"
 	"log/slog"
+	"maps"
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/internal/resources"
+	listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -31,15 +33,17 @@ import (
 // WorkerPoolSyncer reconciles the state of worker pods from Kubernetes Informer
 // into the store.
 type WorkerPoolSyncer struct {
-	persistence    store.Interface
-	workerInformer cache.SharedIndexInformer
+	persistence      store.Interface
+	workerInformer   cache.SharedIndexInformer
+	workerPoolLister listersv1alpha1.WorkerPoolLister
 }
 
 // NewWorkerPoolSyncer creates a new WorkerPoolSyncer.
-func NewWorkerPoolSyncer(persistence store.Interface, workerInformer cache.SharedIndexInformer) *WorkerPoolSyncer {
+func NewWorkerPoolSyncer(persistence store.Interface, workerInformer cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister) *WorkerPoolSyncer {
 	return &WorkerPoolSyncer{
-		persistence:    persistence,
-		workerInformer: workerInformer,
+		persistence:      persistence,
+		workerInformer:   workerInformer,
+		workerPoolLister: workerPoolLister,
 	}
 }
 
@@ -113,17 +117,26 @@ func (s *WorkerPoolSyncer) syncWorkerToStore(ctx context.Context, pod *corev1.Po
 		return
 	}
 
-	w, err := s.persistence.GetWorker(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name)
+	poolName := pod.Labels[workerPodLabel]
+	pool, err := s.workerPoolLister.WorkerPools(pod.Namespace).Get(poolName)
+	if err != nil {
+		slog.ErrorContext(ctx, "Failed to get WorkerPool for worker pod", slog.String("worker", pod.Namespace+"/"+pod.Name), slog.String("pool", poolName), slog.Any("err", err))
+		return
+	}
+
+	w, err := s.persistence.GetWorker(ctx, pod.Namespace, poolName, pod.Name)
 	if err != nil {
 		if errors.Is(err, store.ErrNotFound) {
 			slog.InfoContext(ctx, "Syncer: creating worker in store", slog.String("worker", pod.Namespace+"/"+pod.Name))
 			worker := &ateapipb.Worker{
 				WorkerNamespace: pod.Namespace,
-				WorkerPool:      pod.Labels[workerPodLabel],
+				WorkerPool:      poolName,
 				WorkerPod:       pod.Name,
 				Ip:              pod.Status.PodIP,
 				WorkerPodUid:    string(pod.UID),
 				NodeName:        pod.Spec.NodeName,
+				SandboxClass:    string(pool.Spec.SandboxClass),
+				Labels:          pool.GetLabels(),
 			}
 			// TODO(thockin): for now this is the only place Workers are
 			// created.  If/when this becomes a regular API, validation should
@@ -143,10 +156,25 @@ func (s *WorkerPoolSyncer) syncWorkerToStore(ctx context.Context, pod *corev1.Po
 		return
 	}
 
+	changed := false
 	if w.Ip != pod.Status.PodIP {
 		// TODO: I don't think this is possible, but handling this case so we can log it just in case we can reproduce it.
 		slog.InfoContext(ctx, "Syncer: updating worker in store (IP changed)", slog.String("worker", pod.Namespace+"/"+pod.Name))
 		w.Ip = pod.Status.PodIP
+		changed = true
+	}
+	if w.SandboxClass != string(pool.Spec.SandboxClass) {
+		slog.InfoContext(ctx, "Syncer: updating worker in store (SandboxClass changed)", slog.String("worker", pod.Namespace+"/"+pod.Name))
+		w.SandboxClass = string(pool.Spec.SandboxClass)
+		changed = true
+	}
+	if !maps.Equal(w.Labels, pool.GetLabels()) {
+		slog.InfoContext(ctx, "Syncer: updating worker in store (labels changed)", slog.String("worker", pod.Namespace+"/"+pod.Name))
+		w.Labels = pool.GetLabels()
+		changed = true
+	}
+
+	if changed {
 		if err = s.persistence.UpdateWorker(ctx, w, w.Version); err != nil {
 			slog.ErrorContext(ctx, "Failed to update worker in store", slog.Any("err", err))
 		}

cmd/ateapi/internal/controlapi/syncer_test.go

@@ -18,47 +18,75 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"testing"
 	"time"
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest"
+	atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
+	atefake "github.com/agent-substrate/substrate/pkg/client/clientset/versioned/fake"
+	"github.com/agent-substrate/substrate/pkg/client/informers/externalversions"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/client-go/kubernetes/fake"
 )
 
 // setupSyncerTest sets up a real store with fake Redis and a fake K8s client with informer.
-func setupSyncerTest(t *testing.T, ctx context.Context) (store.Interface, *fake.Clientset, func()) {
+func setupSyncerTest(t *testing.T, ctx context.Context, initPools ...*atev1alpha1.WorkerPool) (store.Interface, *fake.Clientset, *atefake.Clientset, func()) {
 	t.Helper()
 
 	persistence, cleanup := storetest.SetupTestStore(t)
 
 	fakeK8s := fake.NewSimpleClientset()
 	workerFactory, workerInformer := WorkerPodInformer(fakeK8s)
 
-	syncer := NewWorkerPoolSyncer(persistence, workerInformer)
+	objects := make([]runtime.Object, len(initPools))
+	for i, pool := range initPools {
+		objects[i] = pool
+	}
+	//nolint:staticcheck // NewSimpleClientset is the only available fake clientset for versioned CRDs.
+	fakeAte := atefake.NewSimpleClientset(objects...)
+	ateInformerFactory := externalversions.NewSharedInformerFactory(fakeAte, 0)
+	workerPoolLister := ateInformerFactory.Api().V1alpha1().WorkerPools().Lister()
+
+	syncer := NewWorkerPoolSyncer(persistence, workerInformer, workerPoolLister)
 	syncer.Start(ctx)
 
 	workerFactory.Start(ctx.Done())
+	ateInformerFactory.Start(ctx.Done())
+
 	workerFactory.WaitForCacheSync(ctx.Done())
+	ateInformerFactory.WaitForCacheSync(ctx.Done())
 
-	return persistence, fakeK8s, cleanup
+	return persistence, fakeK8s, fakeAte, cleanup
 }
 
 func TestSyncer_Lifecycle(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	persistence, fakeK8s, cleanup := setupSyncerTest(t, ctx)
-	defer cleanup()
-
 	ns := "ns-syncer-lifecycle"
 	podName := "worker-unit-1"
 	poolName := "pool1"
 
+	pool := &atev1alpha1.WorkerPool{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      poolName,
+			Namespace: ns,
+			Labels:    map[string]string{"foo": "bar"},
+		},
+		Spec: atev1alpha1.WorkerPoolSpec{
+			SandboxClass: "gvisor",
+		},
+	}
+
+	persistence, fakeK8s, _, cleanup := setupSyncerTest(t, ctx, pool)
+	defer cleanup()
+
 	// 1. Verify no workers in Redis initially
 	workers, err := persistence.ListWorkers(context.Background())
 	if err != nil {
@@ -127,7 +155,16 @@ func TestSyncer_Lifecycle(t *testing.T) {
 			}
 			return false, err
 		}
-		return w.Ip == "127.0.0.1", nil
+		if w.Ip != "127.0.0.1" {
+			return false, nil
+		}
+		if w.SandboxClass != "gvisor" {
+			return false, fmt.Errorf("expected SandboxClass gvisor, got %q", w.SandboxClass)
+		}
+		if !maps.Equal(w.Labels, map[string]string{"foo": "bar"}) {
+			return false, fmt.Errorf("expected labels map[foo:bar], got %v", w.Labels)
+		}
+		return true, nil
 	})
 	if err != nil {
 		t.Fatalf("Worker not found in Redis after update: %v", err)
@@ -159,10 +196,20 @@ func TestSyncer_DeleteBoundWorker_ClearsActor(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	persistence, fakeK8s, cleanup := setupSyncerTest(t, ctx)
-	defer cleanup()
-
 	ns, pool, pod, ip := "ns-orphan", "pool1", "worker-orphan", "10.0.0.1"
+	workerPool := &atev1alpha1.WorkerPool{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      pool,
+			Namespace: ns,
+			Labels:    map[string]string{"foo": "bar"},
+		},
+		Spec: atev1alpha1.WorkerPoolSpec{
+			SandboxClass: "gvisor",
+		},
+	}
+
+	persistence, fakeK8s, _, cleanup := setupSyncerTest(t, ctx, workerPool)
+	defer cleanup()
 	if _, err := fakeK8s.CoreV1().Pods(ns).Create(ctx,
 		&corev1.Pod{
 			ObjectMeta: metav1.ObjectMeta{
@@ -240,3 +287,76 @@ func TestSyncer_DeleteBoundWorker_ClearsActor(t *testing.T) {
 		t.Errorf("External SnapshotUriPrefix must be preserved")
 	}
 }
+
+func TestSyncer_OmittedFields(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ns := "ns-syncer-omitted"
+	podName := "worker-unit-1"
+	poolName := "pool1"
+
+	// Create a pool with omitted sandbox class and no labels
+	pool := &atev1alpha1.WorkerPool{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      poolName,
+			Namespace: ns,
+		},
+		Spec: atev1alpha1.WorkerPoolSpec{
+			// Spec has no SandboxClass and no Labels
+		},
+	}
+
+	persistence, fakeK8s, _, cleanup := setupSyncerTest(t, ctx, pool)
+	defer cleanup()
+
+	// Create a pod
+	pod := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      podName,
+			Namespace: ns,
+			UID:       "08675309-4a65-6e6e-7973-6e756d626572",
+			Labels: map[string]string{
+				workerPodLabel: poolName,
+			},
+		},
+		Spec: corev1.PodSpec{
+			NodeName:   "node1",
+			Containers: []corev1.Container{{Name: "main", Image: "nginx"}},
+		},
+		Status: corev1.PodStatus{
+			Phase:  corev1.PodRunning,
+			PodIP:  "127.0.0.1",
+			PodIPs: []corev1.PodIP{{IP: "127.0.0.1"}},
+		},
+	}
+
+	_, err := fakeK8s.CoreV1().Pods(ns).Create(context.Background(), pod, metav1.CreateOptions{})
+	if err != nil {
+		t.Fatalf("failed to create pod: %v", err)
+	}
+
+	// Verify that it is created in Redis with empty SandboxClass and empty Labels
+	err = wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 2*time.Second, true, func(ctx context.Context) (bool, error) {
+		w, err := persistence.GetWorker(ctx, ns, poolName, podName)
+		if err != nil {
+			if errors.Is(err, store.ErrNotFound) {
+				return false, nil
+			}
+			return false, err
+		}
+		if w.Ip != "127.0.0.1" {
+			return false, nil
+		}
+		if w.SandboxClass != "" {
+			return false, fmt.Errorf("expected SandboxClass to be empty, got %q", w.SandboxClass)
+		}
+		if len(w.Labels) != 0 {
+			return false, fmt.Errorf("expected labels to be empty, got %v", w.Labels)
+		}
+		return true, nil
+	})
+	if err != nil {
+		t.Fatalf("Worker state check failed: %v", err)
+	}
+}

cmd/ateapi/internal/controlapi/workflow.go

@@ -166,7 +166,7 @@ func (w *ActorWorkflow) ResumeActor(ctx context.Context, atespace, id string, bo
 
 	steps := []WorkflowStep[*ResumeInput, *ResumeState]{
 		&LoadActorForResumeStep{store: w.store, actorTemplateLister: w.actorTemplateLister},
-		&AssignWorkerStep{store: w.store, workerCache: w.workerCache, workerPoolLister: w.workerPoolLister},
+		&AssignWorkerStep{store: w.store, workerCache: w.workerCache},
 		&CallAteletRestoreStep{store: w.store, dialer: w.dialer, kubeClient: w.kubeClient, secretCache: w.secretCache, workerPoolLister: w.workerPoolLister, sandboxConfigLister: w.sandboxConfigLister},
 		&FinalizeRunningStep{store: w.store},
 	}