feat(server): resize pools over the admin API, and two things it needed to work

deploy/operator/internal/controller/builders.go

@@ -317,11 +317,16 @@ func envFor(idep *inferav1alpha1.InferaDeployment, svc inferav1alpha1.ServiceSpe
 	}
 	if useK8sDiscovery(idep) {
 		// Pod identity for self-registration (worker) + selector for the server.
+		// POD_IP is what a worker advertises: it binds 0.0.0.0, and the address
+		// it registers is the one the router dials, so without this it
+		// registers its bind host and every request to it is unreachable.
 		env = append(env,
 			corev1.EnvVar{Name: "POD_NAME", ValueFrom: &corev1.EnvVarSource{
 				FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.name"}}},
 			corev1.EnvVar{Name: "POD_NAMESPACE", ValueFrom: &corev1.EnvVarSource{
 				FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.namespace"}}},
+			corev1.EnvVar{Name: "POD_IP", ValueFrom: &corev1.EnvVarSource{
+				FieldRef: &corev1.ObjectFieldSelector{FieldPath: "status.podIP"}}},
 		)
 		if svc.ComponentType == inferav1alpha1.ComponentTypeServer {
 			env = append(env, corev1.EnvVar{
@@ -424,6 +429,24 @@ func injectWorkerRolloutDefaults(
 // exposes the service port (so buildServerService has a target). Used when
 // ServiceSpec.ExtraPodSpec is set (an external orchestrator renders the full
 // pod template).
+// appendEnvIfAbsent adds each variable the container does not already declare.
+//
+// A template supplied by an external orchestrator may well set these itself,
+// and a duplicate name in a container's env is not an error -- the last one
+// wins, silently overriding what the author wrote.
+func appendEnvIfAbsent(env []corev1.EnvVar, add ...corev1.EnvVar) []corev1.EnvVar {
+	present := make(map[string]bool, len(env))
+	for _, e := range env {
+		present[e.Name] = true
+	}
+	for _, e := range add {
+		if !present[e.Name] {
+			env = append(env, e)
+		}
+	}
+	return env
+}
+
 func podTemplateFromExtra(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec) corev1.PodTemplateSpec {
 	spec := *svc.ExtraPodSpec.DeepCopy()
 	port := servicePort(svc)
@@ -448,11 +471,29 @@ func podTemplateFromExtra(idep *inferav1alpha1.InferaDeployment, svcName string,
 			spec.Containers[idx].Ports = append(spec.Containers[idx].Ports,
 				corev1.ContainerPort{ContainerPort: port})
 		}
-		// k8s discovery: the server reads its watch scope from an env var so we
-		// don't have to rewrite the externally-supplied entrypoint command.
-		if useK8sDiscovery(idep) && svc.ComponentType == inferav1alpha1.ComponentTypeServer {
-			spec.Containers[idx].Env = append(spec.Containers[idx].Env, corev1.EnvVar{
-				Name: "INFERA_K8S_LABEL_SELECTOR", Value: discoveryLabelSelector(idep.Name)})
+		if useK8sDiscovery(idep) {
+			// Pod identity, for both component types: a worker registers by
+			// patching its own Pod annotation, and the server finds the
+			// deployment it belongs to from its own Pod labels. Rendering the
+			// template elsewhere does not change that either needs to know
+			// which Pod it is.
+			spec.Containers[idx].Env = appendEnvIfAbsent(spec.Containers[idx].Env,
+				corev1.EnvVar{Name: "POD_NAME", ValueFrom: &corev1.EnvVarSource{
+					FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.name"}}},
+				corev1.EnvVar{Name: "POD_NAMESPACE", ValueFrom: &corev1.EnvVarSource{
+					FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.namespace"}}},
+				corev1.EnvVar{Name: "POD_IP", ValueFrom: &corev1.EnvVarSource{
+					FieldRef: &corev1.ObjectFieldSelector{FieldPath: "status.podIP"}}},
+			)
+			// The server reads its watch scope from an env var so we don't have
+			// to rewrite the externally-supplied entrypoint command.
+			if svc.ComponentType == inferav1alpha1.ComponentTypeServer {
+				spec.Containers[idx].Env = appendEnvIfAbsent(spec.Containers[idx].Env,
+					corev1.EnvVar{
+						Name:  "INFERA_K8S_LABEL_SELECTOR",
+						Value: discoveryLabelSelector(idep.Name),
+					})
+			}
 		}
 	}
 	// Bind the discovery ServiceAccount (workers patch their own Pod; the
@@ -615,11 +656,30 @@ func buildDiscoveryRole(idep *inferav1alpha1.InferaDeployment) *rbacv1.Role {
 			Namespace: idep.Namespace,
 			Labels:    labelsFor(idep.Name, "disc"),
 		},
-		Rules: []rbacv1.PolicyRule{{
-			APIGroups: []string{""},
-			Resources: []string{"pods"},
-			Verbs:     []string{"get", "list", "watch", "patch"},
-		}},
+		Rules: []rbacv1.PolicyRule{
+			{
+				APIGroups: []string{""},
+				Resources: []string{"pods"},
+				Verbs:     []string{"get", "list", "watch", "patch"},
+			},
+			{
+				// The server's scaling API writes replica counts back to the CR
+				// it belongs to. Named via ResourceNames so the grant reaches
+				// exactly this deployment: every Pod here shares one identity,
+				// so an unrestricted grant would also let any worker resize the
+				// fleet, and a worker has no business doing that.
+				//
+				// Present whether or not --enable-scaling-api is set. The flag
+				// lives on the server's command line and the operator does not
+				// parse it; a permission nothing exercises costs nothing, while
+				// discovering it is absent only after enabling the feature
+				// costs a redeploy.
+				APIGroups:     []string{inferav1alpha1.GroupVersion.Group},
+				Resources:     []string{"inferadeployments"},
+				ResourceNames: []string{idep.Name},
+				Verbs:         []string{"get", "patch"},
+			},
+		},
 	}
 }
 

deploy/operator/internal/controller/builders_test.go

@@ -10,6 +10,8 @@ import (
 	"testing"
 
 	corev1 "k8s.io/api/core/v1"
+
+	inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1"
 )
 
 // The grace period is the only thing standing between a graceful drain and a
@@ -248,3 +250,106 @@ func TestDrainTimeoutStillAcceptsOrdinaryValues(t *testing.T) {
 		}
 	}
 }
+
+// Pod identity is what k8s discovery is built on: a worker patches its own Pod
+// annotation to register, and the server reads its own labels to find the
+// deployment it belongs to. Both need POD_NAME, which the operator injects --
+// on the path that renders the pod itself. A template supplied through
+// extraPodSpec took a different path and got the watch selector but not the
+// identity, so registration and the scaling API both failed on exactly the
+// deployments the PD example tells people to write.
+func TestExtraPodSpecStillGetsPodIdentity(t *testing.T) {
+	for _, ct := range []inferav1alpha1.ComponentType{
+		inferav1alpha1.ComponentTypeServer,
+		inferav1alpha1.ComponentTypeWorker,
+	} {
+		idep := idepWith(1)
+		idep.Spec.DiscoveryBackend = "kubernetes"
+		svc := inferav1alpha1.ServiceSpec{
+			ComponentType: ct,
+			ExtraPodSpec: &corev1.PodSpec{
+				Containers: []corev1.Container{{Name: "main", Image: "x"}},
+			},
+		}
+		tmpl := podTemplateFromExtra(idep, "svc", svc)
+		got := map[string]bool{}
+		for _, e := range tmpl.Spec.Containers[0].Env {
+			got[e.Name] = true
+		}
+		for _, want := range []string{"POD_NAME", "POD_NAMESPACE"} {
+			if !got[want] {
+				t.Errorf("%s: extraPodSpec container has no %s; "+
+					"self-registration and the scaling API both need it", ct, want)
+			}
+		}
+	}
+}
+
+// A template that sets these itself keeps its own values: a duplicate env name
+// is not an error, the last one wins, and appending ours would silently
+// override whatever the author had in mind.
+func TestExtraPodSpecKeepsItsOwnPodIdentity(t *testing.T) {
+	idep := idepWith(1)
+	idep.Spec.DiscoveryBackend = "kubernetes"
+	svc := inferav1alpha1.ServiceSpec{
+		ComponentType: inferav1alpha1.ComponentTypeWorker,
+		ExtraPodSpec: &corev1.PodSpec{Containers: []corev1.Container{{
+			Name:  "main",
+			Image: "x",
+			Env:   []corev1.EnvVar{{Name: "POD_NAME", Value: "chosen-by-the-author"}},
+		}}},
+	}
+	tmpl := podTemplateFromExtra(idep, "svc", svc)
+
+	seen := 0
+	for _, e := range tmpl.Spec.Containers[0].Env {
+		if e.Name != "POD_NAME" {
+			continue
+		}
+		seen++
+		if e.Value != "chosen-by-the-author" {
+			t.Errorf("POD_NAME = %q, want the template's own value", e.Value)
+		}
+	}
+	if seen != 1 {
+		t.Errorf("POD_NAME appears %d times, want 1", seen)
+	}
+}
+
+// A worker binds 0.0.0.0 and advertises something else, because the address it
+// registers is the one the router dials. Under k8s discovery it resolves that
+// from POD_IP -- the logic is already there and reads the downward API -- so
+// leaving the variable out makes the worker register 0.0.0.0 and every request
+// to it fail with "worker unreachable".
+//
+// Measured before this was injected: the worker came up healthy, registered,
+// and the router returned {"error":"worker 0.0.0.0:8080 unreachable"} for the
+// first inference request.
+func TestWorkersLearnTheirOwnAddress(t *testing.T) {
+	idep := idepWith(1)
+	idep.Spec.DiscoveryBackend = "kubernetes"
+
+	check := func(t *testing.T, env []corev1.EnvVar, where string) {
+		t.Helper()
+		for _, e := range env {
+			if e.Name != "POD_IP" {
+				continue
+			}
+			if e.ValueFrom == nil || e.ValueFrom.FieldRef == nil ||
+				e.ValueFrom.FieldRef.FieldPath != "status.podIP" {
+				t.Errorf("%s: POD_IP is not read from the downward API", where)
+			}
+			return
+		}
+		t.Errorf("%s: no POD_IP; the worker would advertise its bind address", where)
+	}
+
+	svc := inferav1alpha1.ServiceSpec{ComponentType: inferav1alpha1.ComponentTypeWorker}
+	check(t, envFor(idep, svc), "rendered pod")
+
+	svc.ExtraPodSpec = &corev1.PodSpec{
+		Containers: []corev1.Container{{Name: "main", Image: "x"}},
+	}
+	tmpl := podTemplateFromExtra(idep, "worker", svc)
+	check(t, tmpl.Spec.Containers[0].Env, "extraPodSpec")
+}

deploy/operator/internal/controller/scale_paths_test.go

@@ -314,3 +314,53 @@ func TestEditingTheChildLWSIsAlsoReverted(t *testing.T) {
 		t.Fatalf("LWS scale survived reconciliation: %d groups, want it reverted to 2", got)
 	}
 }
+
+// The server's scaling API writes replica counts back to the CR, which needs a
+// grant the discovery identity did not previously carry. Every Pod in a
+// deployment shares that identity, so the grant has to name the one CR it may
+// touch: without ResourceNames a worker could resize the fleet it belongs to,
+// or any other deployment in the namespace.
+func TestTheDiscoveryRoleCanWriteOnlyItsOwnDeployment(t *testing.T) {
+	idep := idepWith(2)
+	role := buildDiscoveryRole(idep)
+
+	var found *rbacv1.PolicyRule
+	for i := range role.Rules {
+		for _, res := range role.Rules[i].Resources {
+			if res == "inferadeployments" {
+				found = &role.Rules[i]
+			}
+		}
+	}
+	if found == nil {
+		t.Fatal("no grant for inferadeployments: the scaling API would 403")
+	}
+	if len(found.ResourceNames) != 1 || found.ResourceNames[0] != idep.Name {
+		t.Fatalf("ResourceNames = %v, want exactly [%s]: an unscoped grant lets "+
+			"any Pod here resize any deployment in the namespace",
+			found.ResourceNames, idep.Name)
+	}
+	for _, verb := range found.Verbs {
+		switch verb {
+		case "get", "patch":
+		default:
+			t.Errorf("verb %q is more than the scaling API needs", verb)
+		}
+	}
+}
+
+// Pods are a separate rule and must stay unscoped: the server lists and watches
+// every worker Pod, which ResourceNames cannot express.
+func TestThePodGrantIsUnchanged(t *testing.T) {
+	role := buildDiscoveryRole(idepWith(1))
+	for _, rule := range role.Rules {
+		for _, res := range rule.Resources {
+			if res != "pods" {
+				continue
+			}
+			if len(rule.ResourceNames) != 0 {
+				t.Fatal("the Pod grant must not be scoped by name; discovery lists all of them")
+			}
+		}
+	}
+}

infera/server/__main__.py

@@ -26,6 +26,7 @@
 from infera.router.policy.factory import build_policy
 from infera.server.app import init_app
 from infera.server.args import parse_server_args
+from infera.server.scaling import DeploymentScaler
 
 logging.basicConfig(level=logging.INFO)
 # httpx logs every outbound request (etcd keepalives every ~10s, and one
@@ -260,12 +261,24 @@ def on_worker_removed(worker_id: str) -> None:
             request_max_retries=args.request_max_retries,
             breaker=breaker,
         )
+    scaler = None
+    if args.enable_scaling_api:
+        # Resolved lazily on first call: the deployment is read from this Pod's
+        # labels, and failing at startup over a permission the operator grants
+        # would take down a server whose main job does not need it.
+        scaler = DeploymentScaler(namespace=args.k8s_namespace or None)
+        logger.info(
+            "scaling API enabled: GET/POST /v1/admin/scale resizes this "
+            "deployment's pools (needs RBAC to patch inferadeployments)"
+        )
+
     app = init_app(
         registry,
         router,
         kv=policy.kv_client,
         kvd_socket_path=args.kvd_socket_path,
         enable_profiling=args.enable_profiling,
+        scaler=scaler,
     )
     app.include_router(
         make_stats_router(