diff --git a/.github/workflows/build-and-push.yml b/.github/workflows/build-and-push.yml index 004b5d18..e3a1b26a 100644 --- a/.github/workflows/build-and-push.yml +++ b/.github/workflows/build-and-push.yml @@ -21,25 +21,37 @@ jobs: - image_name: client context: . dockerfile: examples/autoresearch/Dockerfile + - image_name: placement-controller + context: controller + dockerfile: controller/Dockerfile permissions: contents: read packages: write steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: Free Disk Space + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /opt/ghc + sudo rm -rf "/usr/local/share/boost" + sudo rm -rf "$AGENT_TOOLSDIRECTORY" - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 + uses: docker/setup-buildx-action@b5ca514318bd6ebac0fb2aedd5d36ec1b5c232a2 # v3.10.0 - name: Login to GHCR - uses: docker/login-action@v3 + uses: docker/login-action@9780b0c442fbb1117ed29e0efdff1e18412f7567 # v3.3.0 with: registry: ghcr.io username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} - name: Build and push - uses: docker/build-push-action@v5 + uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0 with: context: ${{ matrix.context }} file: ${{ matrix.dockerfile }} diff --git a/.github/workflows/build-pr.yml b/.github/workflows/build-pr.yml index a63c4d5b..d78b0def 100644 --- a/.github/workflows/build-pr.yml +++ b/.github/workflows/build-pr.yml @@ -18,17 +18,29 @@ jobs: - image_name: client context: . dockerfile: examples/autoresearch/Dockerfile + - image_name: placement-controller + context: controller + dockerfile: controller/Dockerfile permissions: contents: read steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: Free Disk Space + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /opt/ghc + sudo rm -rf "/usr/local/share/boost" + sudo rm -rf "$AGENT_TOOLSDIRECTORY" - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 + uses: docker/setup-buildx-action@b5ca514318bd6ebac0fb2aedd5d36ec1b5c232a2 # v3.10.0 - name: Build - uses: docker/build-push-action@v5 + uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0 with: context: ${{ matrix.context }} file: ${{ matrix.dockerfile }} diff --git a/.github/workflows/controller-tests.yml b/.github/workflows/controller-tests.yml new file mode 100644 index 00000000..95f57219 --- /dev/null +++ b/.github/workflows/controller-tests.yml @@ -0,0 +1,42 @@ +name: controller tests + +on: + pull_request: + paths: + - "controller/**" + - ".github/workflows/controller-tests.yml" + push: + branches: + - main + paths: + - "controller/**" + - ".github/workflows/controller-tests.yml" + +jobs: + test: + runs-on: ubuntu-latest + permissions: + contents: read + defaults: + run: + working-directory: controller + steps: + - name: Checkout + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: Set up Go + uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0 + with: + go-version-file: controller/go.mod + cache-dependency-path: controller/go.sum + + - name: gofmt + run: test -z "$(gofmt -l .)" + + - name: vet + run: go vet ./... + + - name: test + run: go test ./... diff --git a/controller/.gitignore b/controller/.gitignore new file mode 100644 index 00000000..8ded3c35 --- /dev/null +++ b/controller/.gitignore @@ -0,0 +1,5 @@ +# Local tool binaries (controller-gen) and build output. The repo's top-level +# .gitignore has no Go section, so the kubebuilder scaffold's ignores live here. +bin/ +*.test +*.out diff --git a/controller/Dockerfile b/controller/Dockerfile new file mode 100644 index 00000000..ea559ece --- /dev/null +++ b/controller/Dockerfile @@ -0,0 +1,28 @@ +# Build the placement controller. +# +# Two stages so the shipped image holds a static binary and nothing else: the +# controller runs with a read-only root filesystem as a non-root user, and has +# no reason to carry a shell or a package manager into a cluster that hands it +# permission to create pods. +FROM golang:1.26 AS build + +WORKDIR /src + +# Dependencies first, so a source-only change does not re-download the module +# graph. +COPY go.mod go.sum ./ +RUN go mod download + +COPY cmd/ cmd/ +COPY api/ api/ +COPY internal/ internal/ + +ARG TARGETARCH +RUN CGO_ENABLED=0 GOOS=linux GOARCH=${TARGETARCH:-amd64} \ + go build -trimpath -ldflags="-s -w" -o /out/manager ./cmd/manager + +FROM gcr.io/distroless/static:nonroot +WORKDIR / +COPY --from=build /out/manager /manager +USER 65532:65532 +ENTRYPOINT ["/manager"] diff --git a/controller/Makefile b/controller/Makefile new file mode 100644 index 00000000..493a4fbf --- /dev/null +++ b/controller/Makefile @@ -0,0 +1,89 @@ +IMG ?= ghcr.io/gke-labs/open-rl/placement-controller:latest +CONTROLLER_TOOLS_VERSION ?= v0.19.0 + +CRD_MANIFEST := ../k8s/deploy/scheduler/00-openrlworker-crd.yaml +LOCALBIN := $(shell pwd)/bin +CONTROLLER_GEN := $(LOCALBIN)/controller-gen +GEN_DIR := $(LOCALBIN)/generated + +.PHONY: all +all: generate manifests fmt vet test build + +$(LOCALBIN): + mkdir -p $(LOCALBIN) + +$(CONTROLLER_GEN): $(LOCALBIN) + GOBIN=$(LOCALBIN) go install sigs.k8s.io/controller-tools/cmd/controller-gen@$(CONTROLLER_TOOLS_VERSION) + +# api/v1alpha1/zz_generated.deepcopy.go. Checked in, so a plain `go build` +# works without the toolchain. +.PHONY: generate +generate: $(CONTROLLER_GEN) + $(CONTROLLER_GEN) object paths="./api/..." + +# The CRD schema, from the markers on the api/ types. kubebuilder would write +# this to config/crd/bases; there is no config/ tree here (see PROJECT), so it +# lands directly on the manifest k8s/deploy applies. The types' doc comments +# become the schema descriptions, so the prose lives in Go, not in the YAML. +.PHONY: manifests +manifests: $(CONTROLLER_GEN) + $(CONTROLLER_GEN) crd paths="./api/..." output:crd:dir=$(GEN_DIR)/crd + cp $(GEN_DIR)/crd/openrl.io_openrlworkers.yaml $(CRD_MANIFEST) + +# Fails when the checked-in CRD no longer matches the types. For CI, and for +# the case where the generator cannot be run on the machine holding the diff. +.PHONY: manifests-check +manifests-check: $(CONTROLLER_GEN) + $(CONTROLLER_GEN) crd paths="./api/..." output:crd:dir=$(GEN_DIR)/crd + diff -u $(CRD_MANIFEST) $(GEN_DIR)/crd/openrl.io_openrlworkers.yaml + +# RBAC stays hand-written and this target only checks it. controller-gen emits +# one ClusterRole; the deployed policy is deliberately split, cluster-scoped for +# nodes and ResourceSlices and namespaced for everything else, so the controller +# cannot reach pods or claims outside its own namespace. Compare the verb sets. +.PHONY: rbac +rbac: $(CONTROLLER_GEN) + $(CONTROLLER_GEN) rbac:roleName=open-rl-scheduler paths="./internal/..." output:rbac:dir=$(GEN_DIR)/rbac + @echo "generated to $(GEN_DIR)/rbac; reconcile by hand with k8s/deploy/scheduler/01-scheduler.yaml" + +.PHONY: fmt +fmt: + go fmt ./... + +.PHONY: vet +vet: + go vet ./... + +.PHONY: test +test: + go test ./... -race -count=1 + +.PHONY: build +build: + go build -o bin/manager ./cmd/manager + +.PHONY: run +run: + go run ./cmd/manager --leader-elect=false --namespace=$${OPEN_RL_WORKER_NAMESPACE:-default} + +# The pipeline against a real API server and fake GPUs: kind plus the DRA +# example driver. No hardware needed. See hack/kind-smoke.sh for the knobs. +.PHONY: smoke +smoke: + ./hack/kind-smoke.sh + +.PHONY: docker-build +docker-build: + docker build -t $(IMG) . + +.PHONY: docker-push +docker-push: + docker push $(IMG) + +.PHONY: deploy +deploy: + kubectl apply -k ../k8s/deploy/scheduler + +.PHONY: clean +clean: + rm -rf bin diff --git a/controller/PROJECT b/controller/PROJECT new file mode 100644 index 00000000..25f4ca84 --- /dev/null +++ b/controller/PROJECT @@ -0,0 +1,27 @@ +# Kubebuilder project metadata. +# https://book.kubebuilder.io/reference/project-config.html +# +# Written by hand rather than by `kubebuilder init`: the tree already existed in +# the kubebuilder layout (api/, cmd/manager, internal/controller, controller-gen +# markers on the types) and this file is what makes the CLI and its plugins +# recognise it, so `kubebuilder create api` and `kubebuilder edit` work here. +# +# One deviation from a stock scaffold: there is no config/ kustomize tree. This +# repo deploys every component from k8s/deploy, and a second deploy path would +# be a second thing to keep correct. `make manifests` generates the CRD straight +# into k8s/deploy/dra-placement instead. +domain: openrl.io +layout: +- go.kubebuilder.io/v4 +projectName: open-rl-placement-controller +repo: github.com/gke-labs/open-rl/controller +resources: +- api: + crdVersion: v1 + namespaced: true + controller: true + domain: openrl.io + kind: OpenRLWorker + path: github.com/gke-labs/open-rl/controller/api/v1alpha1 + version: v1alpha1 +version: "3" diff --git a/controller/README.md b/controller/README.md new file mode 100644 index 00000000..a58ba448 --- /dev/null +++ b/controller/README.md @@ -0,0 +1,94 @@ +# The GPU scheduler + +A worker says how much accelerator memory it needs and which owner it belongs +to. The scheduler decides which bundle of accelerators it lands on and who it +takes turns with. That is the whole contract. + +```yaml +apiVersion: openrl.io/v1alpha1 +kind: OpenRLWorker +metadata: + name: adapter-a +spec: + role: trainer # which node pools may host it + modelId: adapter-a # its identity everywhere + memory: 6Gi # total accelerator memory, from the estimator + ownerId: Qwen/Qwen3-0.6B # optional: the unit of fairness it belongs to +``` + +Everything else — device count, per-device split, claim, node — is derived +and reported back in `status`. + +## The model, in one sentence + +**A claim is a bundle of accelerators; several workers may be assigned to it; +exactly one of them is resident at a time.** + +- There is no co-residency in V1. Whatever the workers share, at most one + process's state is loaded on the allocation; everyone else is suspended in + host RAM. So estimates are never summed — each worker only has to fit the + allocation *by itself* — and a handoff finishes suspending the outgoing + worker before the next one is restored. +- The owner ID is an opaque string, compared and never interpreted. It is the + unit of fairness: turns rotate between owners, so an owner never gets extra + turns for having more processes, requests, or adapters. Naming none makes + you an owner of one. Placement ignores it entirely. +- `role` selects nodes, never claims: a trainer and a sampler share one GPU + by turns. +- No sharding, so device count is plain ceiling division — derived, never + requested. + +## Layout + +| path | what it is | +| --- | --- | +| `api/v1alpha1` | the CRD: the request, and what was decided about it | +| `internal/placement` | the decision. Pure functions, no Kubernetes imports | +| `internal/controller` | the part that reads and writes Kubernetes objects | +| `docs/design.md` | the design | + +## Try it + +The behaviors live in `internal/placement/behavior_test.go`: workers arriving +and leaving, with the estimator's real tier figures on the hardware we run, +played through the same `Decide` the controller calls. + +``` +go test ./... +``` + +For the pipeline — real API server, real kube-scheduler, real DRA — there is +a kind smoke test that needs no hardware (the DRA example driver publishes +fake GPUs): + +``` +make smoke # kind + fake GPUs +USE_EXISTING_CLUSTER=1 DEVICE_CLASS=gpu.nvidia.com ./hack/kind-smoke.sh # real GPUs +``` + +The controller only ever reads ResourceSlices and node labels, so fake and +real devices exercise the identical path; only the two env values differ. + +## Deploy + +``` +kubectl apply -k ../k8s/deploy/scheduler +kubectl label node openrl.io/enabled=true openrl.io/trainer=true openrl.io/max-workers-per-claim=4 +``` + +Applying it changes nothing about a running cluster: the scheduler only acts +on OpenRLWorker objects. Node labels are policy, never hardware — the DRA +driver's ResourceSlices say what devices actually exist. + +Labeling a node opts its GPUs in **exclusively**: the scheduler counts a +device as free unless one of its own claims holds it, so other GPU workloads +on an enabled node are invisible to placement and will collide with it. Give +OpenRL whole nodes. + +## Everything else + +Assumptions and caveats, the estimator, worker identity, claim lifecycle, +and the future optimizations all live in +[`docs/design.md`](docs/design.md); a file-by-file tour with a suggested +reading order is [`docs/layout.md`](docs/layout.md). If the code and any +document disagree, the code is right. diff --git a/controller/api/v1alpha1/groupversion_info.go b/controller/api/v1alpha1/groupversion_info.go new file mode 100644 index 00000000..744db76f --- /dev/null +++ b/controller/api/v1alpha1/groupversion_info.go @@ -0,0 +1,23 @@ +// Package v1alpha1 contains the OpenRLWorker API, the placement request the +// gateway writes for every worker process it wants running. +// +// See docs/designs/012-dynamic-placement.md. +// +kubebuilder:object:generate=true +// +groupName=openrl.io +package v1alpha1 + +import ( + "k8s.io/apimachinery/pkg/runtime/schema" + "sigs.k8s.io/controller-runtime/pkg/scheme" +) + +var ( + // GroupVersion is the group and version this package's types belong to. + GroupVersion = schema.GroupVersion{Group: "openrl.io", Version: "v1alpha1"} + + // SchemeBuilder registers this package's types with a runtime.Scheme. + SchemeBuilder = &scheme.Builder{GroupVersion: GroupVersion} + + // AddToScheme adds this package's types to a runtime.Scheme. + AddToScheme = SchemeBuilder.AddToScheme +) diff --git a/controller/api/v1alpha1/openrlworker_types.go b/controller/api/v1alpha1/openrlworker_types.go new file mode 100644 index 00000000..1093c4a9 --- /dev/null +++ b/controller/api/v1alpha1/openrlworker_types.go @@ -0,0 +1,194 @@ +package v1alpha1 + +import ( + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +// WorkerRole is which half of the training loop a worker runs. It selects +// which node pools may host the worker, and nothing else. +// +kubebuilder:validation:Enum=trainer;sampler +type WorkerRole string + +const ( + RoleTrainer WorkerRole = "trainer" + RoleSampler WorkerRole = "sampler" +) + +// Phase is a coarse summary of where a worker is in scheduling. +// +kubebuilder:validation:Enum=Pending;Placing;Running;Failed +type Phase string + +const ( + // PhasePending means no claim has been assigned, usually for want of capacity. + PhasePending Phase = "Pending" + // PhasePlacing means a claim and pod exist but the allocation is not yet observed. + PhasePlacing Phase = "Placing" + // PhaseRunning means the pod is running on an observed allocation. + PhaseRunning Phase = "Running" + // PhaseFailed means the request cannot be satisfied as written. + PhaseFailed Phase = "Failed" +) + +// Condition types set on an OpenRLWorker. +const ( + // ConditionPlaced reports whether the worker holds a claim and a pod. + ConditionPlaced = "Placed" +) + +// PodTemplateRef names the ConfigMap holding the pod spec this worker is +// rendered from. The template carries what scheduling has no opinion about; +// the controller overwrites nodeSelector and resourceClaims, because those +// are the decision. +type PodTemplateRef struct { + // Name of the ConfigMap in the controller's namespace. + // +kubebuilder:validation:MinLength=1 + Name string `json:"name"` + + // Key within the ConfigMap holding the pod YAML. + // +kubebuilder:default=pod.yaml + Key string `json:"key,omitempty"` +} + +// ContainerOverlay is what the caller stamps onto the template's first +// container: everything the pod needs to know about the model arrives here. +type ContainerOverlay struct { + // Image replaces the template's image when set. + Image string `json:"image,omitempty"` + + // Command replaces the template's command when non-empty. + Command []string `json:"command,omitempty"` + + // Args are appended to the template's args. + Args []string `json:"args,omitempty"` + + // Env entries are merged by name, overwriting the template's values. + Env []corev1.EnvVar `json:"env,omitempty"` +} + +// OpenRLWorkerSpec is one worker process's scheduling request. Memory arrives +// already estimated; the controller decides how many devices to spread it +// over and which claim to put it on, and never re-estimates. +type OpenRLWorkerSpec struct { + // Role selects which node pools may host this worker, via the + // openrl.io/trainer and openrl.io/sampler labels. It does not partition + // claims: on a node that accepts both, the two roles share by turns. + Role WorkerRole `json:"role"` + + // ModelID names the model this worker serves: configuration for the + // worker process and the humans reading kubectl get, never identity. + // The worker's identity is metadata.name. + // +kubebuilder:validation:MinLength=1 + ModelID string `json:"modelId"` + + // OwnerID is the unit of fairness: the runtime serves owners round-robin, + // so an owner never gets extra turns for having more processes, requests, + // or adapters. Opaque to the controller, never read by placement -- one + // worker is resident at a time whatever the owners. A worker naming no + // owner is an owner of one. + OwnerID string `json:"ownerId,omitempty"` + + // Memory is the estimator's figure: total peak accelerator memory, an + // aggregate across however many devices it takes. Never re-estimated. + Memory resource.Quantity `json:"memory"` + + // EstimatorVersion records which estimator produced Memory, so a placement + // can be inspected and reproduced later. + EstimatorVersion string `json:"estimatorVersion,omitempty"` + + // PodTemplate names the ConfigMap the worker pod is rendered from. When + // unset the controller falls back to its role-default template. + PodTemplate *PodTemplateRef `json:"podTemplate,omitempty"` + + // Container is what the caller stamps onto the rendered container. + Container *ContainerOverlay `json:"container,omitempty"` +} + +// OpenRLWorkerStatus is what the controller decided and what came of it. +type OpenRLWorkerStatus struct { + // Phase is a coarse summary; Conditions carry the detail. + Phase Phase `json:"phase,omitempty"` + + // DeviceCount is how many accelerators the controller decided this worker + // spans. A decision, not a request: derived from Memory and the capacity + // the pools registered. + DeviceCount int32 `json:"deviceCount,omitempty"` + + // MemoryPerDevice is what each of those devices must provide, i.e. Memory + // divided by DeviceCount and rounded up. + MemoryPerDevice string `json:"memoryPerDevice,omitempty"` + + // HostMemoryWhenParked is the host RAM this worker occupies while + // suspended: its full accelerator footprint, parked in its own host + // address space. This is what actually bounds how many workers a node + // can carry. + HostMemoryWhenParked string `json:"hostMemoryWhenParked,omitempty"` + + // EstimatorVersion echoes which estimator produced the memory figure this + // placement was decided against, so the decision can be reproduced. + EstimatorVersion string `json:"estimatorVersion,omitempty"` + + // ClaimName is the ResourceClaim this worker was assigned to. + ClaimName string `json:"claimName,omitempty"` + + // PodName is the worker pod, once created. + PodName string `json:"podName,omitempty"` + + // NodeName is set only once the claim's allocation is observed. Until then + // the node is genuinely unknown. + NodeName string `json:"nodeName,omitempty"` + + // Reason is the most recent human-readable explanation of Phase. + Reason string `json:"reason,omitempty"` + + // ObservedGeneration is the spec generation this status was computed from. + ObservedGeneration int64 `json:"observedGeneration,omitempty"` + + // Conditions holds the Placed condition and its history. + // +listType=map + // +listMapKey=type + Conditions []metav1.Condition `json:"conditions,omitempty"` +} + +// +kubebuilder:object:root=true +// +kubebuilder:subresource:status +// +kubebuilder:resource:shortName=orw +// +kubebuilder:printcolumn:name="Role",type=string,JSONPath=`.spec.role` +// +kubebuilder:printcolumn:name="Owner",type=string,JSONPath=`.spec.ownerId` +// +kubebuilder:printcolumn:name="GPUs",type=integer,JSONPath=`.status.deviceCount` +// +kubebuilder:printcolumn:name="MemEach",type=string,JSONPath=`.status.memoryPerDevice` +// +kubebuilder:printcolumn:name="Phase",type=string,JSONPath=`.status.phase` +// +kubebuilder:printcolumn:name="Claim",type=string,JSONPath=`.status.claimName` +// +kubebuilder:printcolumn:name="Node",type=string,JSONPath=`.status.nodeName` +// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` + +// OpenRLWorker is the scheduling request for a single worker process. The +// caller creates one per worker it wants; the controller turns it into a +// ResourceClaim (matched or created) and a pod, and records what it picked in +// status -- so `kubectl get openrlworkers` shows what was asked for, where it +// went, and why a pending worker is still pending. +type OpenRLWorker struct { + metav1.TypeMeta `json:",inline"` + metav1.ObjectMeta `json:"metadata,omitempty"` + + // The spec is immutable: every field either places the worker or renders + // its pod, and V1 does not re-place or re-render a live worker. Change by + // deleting and recreating. + // +kubebuilder:validation:XValidation:rule="self == oldSelf",message="OpenRLWorker spec is immutable; delete and recreate the worker" + Spec OpenRLWorkerSpec `json:"spec"` + Status OpenRLWorkerStatus `json:"status,omitempty"` +} + +// +kubebuilder:object:root=true + +// OpenRLWorkerList is a list of OpenRLWorker. +type OpenRLWorkerList struct { + metav1.TypeMeta `json:",inline"` + metav1.ListMeta `json:"metadata,omitempty"` + Items []OpenRLWorker `json:"items"` +} + +func init() { + SchemeBuilder.Register(&OpenRLWorker{}, &OpenRLWorkerList{}) +} diff --git a/controller/api/v1alpha1/zz_generated.deepcopy.go b/controller/api/v1alpha1/zz_generated.deepcopy.go new file mode 100644 index 00000000..6d028042 --- /dev/null +++ b/controller/api/v1alpha1/zz_generated.deepcopy.go @@ -0,0 +1,165 @@ +//go:build !ignore_autogenerated + +// Code generated by controller-gen. DO NOT EDIT. + +package v1alpha1 + +import ( + "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + runtime "k8s.io/apimachinery/pkg/runtime" +) + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *ContainerOverlay) DeepCopyInto(out *ContainerOverlay) { + *out = *in + if in.Command != nil { + in, out := &in.Command, &out.Command + *out = make([]string, len(*in)) + copy(*out, *in) + } + if in.Args != nil { + in, out := &in.Args, &out.Args + *out = make([]string, len(*in)) + copy(*out, *in) + } + if in.Env != nil { + in, out := &in.Env, &out.Env + *out = make([]v1.EnvVar, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ContainerOverlay. +func (in *ContainerOverlay) DeepCopy() *ContainerOverlay { + if in == nil { + return nil + } + out := new(ContainerOverlay) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *OpenRLWorker) DeepCopyInto(out *OpenRLWorker) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ObjectMeta.DeepCopyInto(&out.ObjectMeta) + in.Spec.DeepCopyInto(&out.Spec) + in.Status.DeepCopyInto(&out.Status) +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OpenRLWorker. +func (in *OpenRLWorker) DeepCopy() *OpenRLWorker { + if in == nil { + return nil + } + out := new(OpenRLWorker) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *OpenRLWorker) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *OpenRLWorkerList) DeepCopyInto(out *OpenRLWorkerList) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ListMeta.DeepCopyInto(&out.ListMeta) + if in.Items != nil { + in, out := &in.Items, &out.Items + *out = make([]OpenRLWorker, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OpenRLWorkerList. +func (in *OpenRLWorkerList) DeepCopy() *OpenRLWorkerList { + if in == nil { + return nil + } + out := new(OpenRLWorkerList) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *OpenRLWorkerList) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *OpenRLWorkerSpec) DeepCopyInto(out *OpenRLWorkerSpec) { + *out = *in + out.Memory = in.Memory.DeepCopy() + if in.PodTemplate != nil { + in, out := &in.PodTemplate, &out.PodTemplate + *out = new(PodTemplateRef) + **out = **in + } + if in.Container != nil { + in, out := &in.Container, &out.Container + *out = new(ContainerOverlay) + (*in).DeepCopyInto(*out) + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OpenRLWorkerSpec. +func (in *OpenRLWorkerSpec) DeepCopy() *OpenRLWorkerSpec { + if in == nil { + return nil + } + out := new(OpenRLWorkerSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *OpenRLWorkerStatus) DeepCopyInto(out *OpenRLWorkerStatus) { + *out = *in + if in.Conditions != nil { + in, out := &in.Conditions, &out.Conditions + *out = make([]metav1.Condition, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OpenRLWorkerStatus. +func (in *OpenRLWorkerStatus) DeepCopy() *OpenRLWorkerStatus { + if in == nil { + return nil + } + out := new(OpenRLWorkerStatus) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *PodTemplateRef) DeepCopyInto(out *PodTemplateRef) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PodTemplateRef. +func (in *PodTemplateRef) DeepCopy() *PodTemplateRef { + if in == nil { + return nil + } + out := new(PodTemplateRef) + in.DeepCopyInto(out) + return out +} diff --git a/controller/cmd/manager/main.go b/controller/cmd/manager/main.go new file mode 100644 index 00000000..af018d0f --- /dev/null +++ b/controller/cmd/manager/main.go @@ -0,0 +1,149 @@ +// Command manager runs the Open-RL placement controller. +// +// It watches OpenRLWorker resources and reconciles each into a DRA +// ResourceClaim and a worker pod, letting Kubernetes pick the devices and the +// node. See docs/designs/012-dynamic-placement.md. +package main + +import ( + "flag" + "os" + "time" + + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/runtime" + utilruntime "k8s.io/apimachinery/pkg/util/runtime" + clientgoscheme "k8s.io/client-go/kubernetes/scheme" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/cache" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/healthz" + "sigs.k8s.io/controller-runtime/pkg/log/zap" + metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server" + + openrlv1alpha1 "github.com/gke-labs/open-rl/controller/api/v1alpha1" + "github.com/gke-labs/open-rl/controller/internal/controller" +) + +var scheme = runtime.NewScheme() + +func init() { + utilruntime.Must(clientgoscheme.AddToScheme(scheme)) + utilruntime.Must(openrlv1alpha1.AddToScheme(scheme)) + // +kubebuilder:scaffold:scheme +} + +func main() { + var ( + metricsAddr string + probeAddr string + leaderElection bool + namespace string + deviceClass string + deviceDriver string + trainerTemplate string + samplerTemplate string + retryInterval time.Duration + placementTimeout time.Duration + reclaimInterval time.Duration + ) + + flag.StringVar(&metricsAddr, "metrics-bind-address", "0", "Address the metric endpoint binds to; 0 disables it.") + flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "Address the probe endpoint binds to.") + flag.BoolVar(&leaderElection, "leader-elect", true, + "Hold a lease before placing. Two controllers placing at once would each decide against a fleet missing the other's bookings.") + flag.StringVar(&namespace, "namespace", env("OPEN_RL_WORKER_NAMESPACE", "default"), "Namespace holding workers, claims and pods.") + flag.StringVar(&deviceClass, "device-class", env("OPEN_RL_DEVICE_CLASS", "gpu.nvidia.com"), "DeviceClass generated claims request.") + flag.StringVar(&deviceDriver, "device-driver", env("OPEN_RL_DEVICE_DRIVER", ""), "Driver publishing the ResourceSlices. Defaults to the device class.") + flag.StringVar(&trainerTemplate, "trainer-pod-template", env("OPEN_RL_TRAINER_POD_TEMPLATE_CONFIGMAP", ""), "ConfigMap holding the default trainer pod template.") + flag.StringVar(&samplerTemplate, "sampler-pod-template", env("OPEN_RL_SAMPLER_POD_TEMPLATE_CONFIGMAP", ""), "ConfigMap holding the default sampler pod template.") + flag.DurationVar(&retryInterval, "retry-interval", envDuration("OPEN_RL_RECONCILE_INTERVAL", 10*time.Second), "How often an unplaced worker is retried.") + flag.DurationVar(&placementTimeout, "placement-timeout", envDuration("OPEN_RL_PLACEMENT_TIMEOUT", 15*time.Minute), + "How long a worker may go unplaced before the request is declared unsatisfiable. 0 waits forever.") + flag.DurationVar(&reclaimInterval, "reclaim-interval", envDuration("OPEN_RL_RECLAIM_INTERVAL", time.Minute), "How often idle claims are swept.") + + opts := zap.Options{Development: false} + opts.BindFlags(flag.CommandLine) + flag.Parse() + + ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts))) + setupLog := ctrl.Log.WithName("setup") + + if deviceDriver == "" { + deviceDriver = deviceClass + } + + mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{ + Scheme: scheme, + Metrics: metricsserver.Options{BindAddress: metricsAddr}, + HealthProbeBindAddress: probeAddr, + LeaderElection: leaderElection, + LeaderElectionID: "placement.openrl.io", + // Workers, claims, pods and ConfigMaps are all namespaced; nodes and + // ResourceSlices are cluster-scoped. Nodes are cached only if the + // operator opted them in: placement never reads any other node, and an + // unfiltered informer would deliver every kubelet heartbeat in the + // cluster to this controller's watch. + Cache: cache.Options{ + DefaultNamespaces: map[string]cache.Config{namespace: {}}, + ByObject: map[client.Object]cache.ByObject{ + &corev1.Node{}: {Label: labels.SelectorFromSet(labels.Set{controller.NodeLabelEnabled: "true"})}, + }, + }, + }) + if err != nil { + setupLog.Error(err, "cannot start manager") + os.Exit(1) + } + + reconciler := &controller.OpenRLWorkerReconciler{ + Client: mgr.GetClient(), + Recorder: mgr.GetEventRecorderFor("placement-controller"), + Namespace: namespace, + DeviceClass: deviceClass, + DeviceDriver: deviceDriver, + DefaultPodTemplates: map[openrlv1alpha1.WorkerRole]string{ + openrlv1alpha1.RoleTrainer: trainerTemplate, + openrlv1alpha1.RoleSampler: samplerTemplate, + }, + RetryInterval: retryInterval, + PlacementTimeout: placementTimeout, + ReclaimInterval: reclaimInterval, + } + if err := reconciler.SetupWithManager(mgr); err != nil { + setupLog.Error(err, "cannot set up the OpenRLWorker controller") + os.Exit(1) + } + // +kubebuilder:scaffold:builder + + if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil { + setupLog.Error(err, "cannot add health check") + os.Exit(1) + } + if err := mgr.AddReadyzCheck("readyz", healthz.Ping); err != nil { + setupLog.Error(err, "cannot add ready check") + os.Exit(1) + } + + setupLog.Info("placing workers", "namespace", namespace, "deviceClass", deviceClass, "deviceDriver", deviceDriver) + if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil { + setupLog.Error(err, "manager exited") + os.Exit(1) + } +} + +func env(key, fallback string) string { + if value := os.Getenv(key); value != "" { + return value + } + return fallback +} + +// envDuration reads a Go duration ("30s", "15m") from the environment. +func envDuration(key string, fallback time.Duration) time.Duration { + if parsed, err := time.ParseDuration(os.Getenv(key)); err == nil { + return parsed + } + return fallback +} diff --git a/controller/docs/design.md b/controller/docs/design.md new file mode 100644 index 00000000..4b79460f --- /dev/null +++ b/controller/docs/design.md @@ -0,0 +1,505 @@ +**Status:** Draft +**Author:** Shuby Mishra +**Reviewers:** — +**Last updated:** August 12, 2026 + +## Introduction + +OpenRL is a self-hosted training platform that runs trainer and sampler +workers on accelerator hardware in Kubernetes. Today, its Kubernetes +deployment uses pre-created Dynamic Resource Allocation (DRA) ResourceClaim +objects: trainers share one fixed claim, samplers share another, and OpenRL's +time-slicer coordinates access to the devices behind them. This works for a +fixed deployment, but the accelerator capacity and its assignment to workers +must be decided before workloads arrive. + +This design makes placement dynamic. OpenRL estimates how much accelerator +memory each trainer or sampler process needs. A Go controller places that +process onto a DRA claim and creates its Pod. Kubernetes and the DRA driver +choose the exact node and devices. + +The placement rule is simple: use a free accelerator when possible; otherwise +assign another worker to an existing allocation. Several workers may be +assigned to one claim, but V1 allows exactly one of them to be loaded in +accelerator memory at a time. Workers take turns through suspend and restore. + +## Background + +The merged fft branch creates trainer and sampler Pods directly from the +gateway using role-specific templates. Those templates reference fixed +trainer and sampler claims and select a predetermined accelerator node type. +OpenRL creates and deletes Pods, but it does not estimate their memory or +create, select, and reclaim claims. + +DRA already reports the hardware available to Kubernetes and allocates +devices to claims. OpenRL only needs to add workload policy: determine what +each worker needs, choose an eligible claim, and coordinate workers that +intentionally share it. + +## Design + +### Definitions + +| Term | Meaning | +| --- | --- | +| Worker | One trainer or sampler process. This is the unit placed onto hardware. | +| Owner | The job or shared runtime that receives one turn when an allocation is contested. | +| Claim | A DRA accelerator allocation on one node. Multiple worker Pods may reference it. | +| Resident | The one worker on a claim whose state is currently loaded in accelerator memory. | + +For every claim: + +- Several workers may be assigned to the claim. +- At most one worker is resident and uses the allocation's compute. +- Every other assigned worker is suspended outside accelerator memory. +- A handoff completes suspension of the current resident before restoring the + next worker. + +### V1 limitation: no GPU-memory co-residency + +V1 does not keep two independent worker processes loaded on the same +accelerator allocation, even when their estimated memory would fit together. +If a 45 GiB trainer and a 25 GiB sampler share an 80 GiB GPU, OpenRL still +suspends one before restoring the other. + +This leaves usable accelerator memory idle and may add transfer overhead for +small workers. The tradeoff is a much smaller correctness surface: the +controller checks only whether each worker fits independently, and the +node-local scheduler tracks only one resident instead of maintaining and +recovering a packed resident set. A later version may retain multiple workers +as a node-local memory cache without changing worker identity or claim +membership. + +The end-to-end flow is: + +``` +worker configuration ──> memory estimate ──> OpenRLWorker + │ +ResourceSlices + operator labels ──> Go controller + │ + └──> ResourceClaim + Pod + │ + └──> node-local scheduling +``` + +The gateway decides worker identity and produces the memory estimate. The Go +controller selects capacity and reconciles claims and Pods. Kubernetes and +DRA allocate the hardware. The node-local scheduler coordinates execution on +a shared claim. + +### Hardware visible to OpenRL + +ResourceSlice objects tell OpenRL what hardware exists. Node labels tell +OpenRL which of that hardware it is allowed to use. + +The DRA driver publishes ResourceSlice objects describing devices, their +capacity and attributes, and the nodes that can access them. Platform +engineers opt nodes into OpenRL with labels: + +``` +openrl.io/enabled: "true" +openrl.io/trainer: "true" +openrl.io/sampler: "true" +openrl.io/max-workers-per-claim: "2" +``` + +The controller considers the intersection: + +``` +devices reported through ResourceSlices +∩ devices accessible from openrl.io/enabled nodes += hardware visible to OpenRL +``` + +The role labels may restrict a node to trainers, samplers, or both. If +neither role label is present, both are allowed. `max-workers-per-claim` +limits how many workers may be assigned to one claim and defaults to 1. It +limits queueing and switch overhead; it does not permit multiple residents. + +Labels express operator policy, not hardware facts. For example, +`openrl.io/trainer=true` allows trainers on a node; it does not assert that +the node contains an H100. The DRA inventory remains the source of truth for +the devices. + +OpenRL can place different workers across any number of labeled nodes. A +single worker's claim must still be satisfiable on one node; V1 does not +combine GPUs from different nodes into one distributed worker. + +### Worker identity + +The gateway decides whether an API request reuses an existing worker or +creates a new one: + +| Training path | Worker ownership | Why | +| --- | --- | --- | +| FFT | Each job gets its own owner and its own trainer and sampler workers. | The trainer changes the job's complete model weights and optimizer history. Sharing a trainer process would mix the jobs' parameters. | +| LoRA | Requests for the same base model reuse its workers until adapter capacity is full. | The base model stays fixed; each request changes only its adapter. | + +An FFT job's trainer and sampler are separate workers with the same owner ID. +The trainer updates that job's weights, while the sampler serves snapshots of +those evolving weights. + +Placement operates on workers because each process consumes memory. +Scheduling accounts by owner so that an owner does not receive extra turns +merely because it has more processes, requests, or adapters. The placement +controller treats the owner ID as opaque. + +### Memory estimation + +The estimator calculates the peak accelerator memory required by one worker. +Its inputs include the model, role, training kind, optimizer and offload +settings, context length, peak packed tokens in one forward pass, and sampler +KV-cache or adapter-slot configuration. + +Peak packed tokens is not the same as model context length. If two +131k-token datums are packed into one padded forward, the activation estimate +must use the resulting 262k-token shape. + +The estimator returns one placement quantity: + +```yaml +memory: 50Gi +``` + +It does not choose a GPU model, node, device count, or sharding strategy. The +controller compares the estimate with the eligible DRA inventory and derives +an allocation. Because V1 has one resident per claim, estimates for workers +assigned to the same claim are never added together. OpenRL records the +estimator inputs, result, and version so the decision can be inspected later. + +V1 may begin with a conservative table of measured configurations. A formula +or profiler can replace it without changing placement. + +### Placement decision + +The controller applies the same decision tree to FFT and LoRA workers: + +| Cluster state | Placement | Runtime behavior | +| --- | --- | --- | +| A suitable free allocation exists. | Create a new claim. | The worker runs independently. | +| No free allocation exists, and the worker fits an existing claim independently. | Join that claim if its worker and host-memory limits allow. | Assigned workers suspend and restore as turns change. | +| The worker cannot fit on any eligible one-node allocation. | Leave it pending. | No hardware substitution or partial placement occurs. | + +This policy spreads work while free accelerators exist and shares only under +contention. It does not bin-pack accelerator memory. When several existing +claims are eligible, V1 prefers the claim with the fewest assigned workers, +followed by claim name for deterministic placement. + +The controller owns the claims it creates. It removes a worker's Pod and +claim membership when the worker leaves and deletes a claim after its final +member is gone. V1 does not migrate a running worker after placement. + +### End-to-end scenarios + +#### 1. Free GPUs on different nodes + +Two labeled nodes each expose one free 80 GiB GPU. Two 50 GiB workers arrive. +The controller creates a separate claim for each worker, and DRA allocates +them across the two nodes. Both workers run independently. + +This is normal multi-node cluster support. The limitation is only that one +worker cannot combine the two GPUs across the two nodes. + +#### 2. Two workers would fit together + +Only one 80 GiB GPU remains. An FFT trainer needs 45 GiB and its sampler +needs 25 GiB. Each worker fits independently, so the controller may assign +both Pods to the same claim when no free GPU exists. + +Their combined 70 GiB would fit, but V1 deliberately does not use that fact. +Sampling and training alternate through an exclusive handoff: + +``` +trainer turn: GPU [ trainer ] host [ sampler ] +handoff: suspend trainer -> restore sampler +sampler turn: GPU [ sampler ] host [ trainer ] +``` + +If no other worker is waiting, the current worker may remain resident. +Suspension is required when the scheduler hands the claim to a different +worker. + +#### 3. Two workers do not fit together + +The trainer instead needs 55 GiB and the sampler needs 35 GiB. Each fits the +80 GiB GPU, but their combined 90 GiB does not. They may still be assigned to +the same claim if its worker limit and host-memory budget allow it. + +Runtime behavior is identical to scenario 2: at a safe boundary, the trainer +moves its model, gradients, and optimizer state to host memory before the +sampler wakes. When sampling finishes, vLLM sleep level 1 offloads the +sampler's weights and discards its KV cache so the trainer can restore. V1 +never needs to distinguish whether two workers would fit together. + +#### 4. Several jobs share one GPU + +Two FFT jobs each have a 50 GiB worker on one 80 GiB claim. Their owners take +turns round-robin at safe batch or optimizer-step boundaries. If only one +owner has work, it keeps running. + +A third worker first looks for free hardware elsewhere in the visible fleet. +If none exists, it may join a claim with an available worker slot. If every +eligible claim has reached `max-workers-per-claim`, it remains pending. + +#### 5. Samplers use different model families + +A Qwen sampler and a Llama sampler each require 50 GiB and are assigned to +one 80 GiB claim. They take turns because each runs in its own vLLM process +and sleeps its own engine before the other becomes resident. + +Placement does not require a same-family or same-base-model rule. The workers +only need to fit independently and support suspension. + +#### 6. LoRA and FFT workers share an allocation + +An FFT worker needs 50 GiB and a LoRA worker needs 20 GiB. With no free GPUs +elsewhere, both may be assigned to one 80 GiB claim because each fits +independently. + +They do not remain resident together even though their combined 70 GiB would +fit. They take turns through suspension under the same rule as any other +pair. Placement does not create a special boundary between LoRA and FFT. + +#### 7. LoRA requests reuse workers before placement + +Two LoRA jobs target the same Qwen base model. The gateway assigns both +adapters to the existing Qwen trainer and sampler processes, so the second +request does not create another OpenRLWorker or consume another worker slot +on a claim. + +If the runtime reaches its adapter capacity, or a request uses a different +base model, the gateway creates another set of workers. Those new processes +then enter the normal placement flow. + +#### 8. Node labels exclude otherwise free hardware + +The DRA inventory reports a free H100 on a trainer-only node and a busy L4 on +a sampler node. A new sampler cannot use the H100 even though it is +technically capable of running there: the operator has not allowed samplers +on that node. + +The sampler may join an eligible claim on the L4 if it fits independently and +a worker slot is available. Otherwise it remains pending. ResourceSlice +objects describe what exists; labels decide what OpenRL may use. + +#### 9. max-workers-per-claim limits queue depth + +Three 20 GiB workers arrive for one 80 GiB GPU configured with +`max-workers-per-claim: 2`. The first two may be assigned to the claim and +take exclusive turns. The third cannot join, even though it also fits the GPU +independently. + +The controller looks for another free or shared allocation. If none exists, +the third worker remains pending. `max-workers-per-claim` is an operator +safety limit on wait time and switching overhead, not a memory calculation. + +#### 10. GPU memory fits, but host memory does not + +Two 50 GiB workers can each fit on an 80 GiB GPU. Time-slicing requires +parking one worker's state in host memory while the other is resident. If the +node lacks enough host-memory headroom, the controller does not place the +second worker on that claim. + +The worker tries another eligible allocation or remains pending. This +prevents accelerator oversubscription from causing a node-level out-of-memory +failure. + +#### 11. Capacity exists only across nodes + +Two nodes each expose one 80 GiB GPU, but one worker requires 140 GiB. The +cluster has 160 GiB in aggregate, yet no single node can satisfy the worker, +so it remains pending. + +If one node instead exposes two 80 GiB GPUs, the controller may create a +two-device claim on that node. Placement establishes that the memory exists; +the training runtime must still support that device shape. + +#### 12. Capacity becomes free after placement + +Two workers share one GPU because every other eligible accelerator was busy +when they arrived. Later, another GPU becomes free. V1 leaves the existing +workers on their original claim, so they continue taking turns while the +newly free GPU is available to new workers. + +Moving one of the existing workers would improve throughput, but requires +rebalancing and is a future optimization. + +#### 13. Workers leave a shared claim + +Two workers share one claim. When the first worker exits, the controller +deletes its Pod and removes it from the claim, but the claim remains +allocated for the second worker. When the final worker exits, the controller +deletes the claim and returns the accelerator to the visible fleet. + +#### 14. The memory estimate is wrong + +A worker is estimated at 40 GiB and placed on an 80 GiB GPU. If it actually +peaks above 80 GiB, it may OOM during its turn. If the estimate is too high +instead, OpenRL may reject a placement that would have worked or request a +larger allocation than necessary. + +V1 records the estimate and estimator version for diagnosis but does not +automatically correct future estimates. Feeding observed usage and OOMs back +into the estimator is a future optimization. + +## Future optimizations + +V1 chooses a predictable policy rather than a global optimum. Later versions +may add: + +- node-local co-residency that keeps multiple inactive workers loaded when + their memory fits; +- rebalancing when an accelerator becomes free; +- placement that accounts for suspend and restore cost; +- memory estimates informed by observed usage and OOMs; +- priorities, owner weights, or minimum turn durations; +- explicit device count and topology for TP or FSDP; +- distributed workers spanning multiple nodes; and +- concurrent execution through an isolation mechanism such as MIG or MPS. + +These changes do not alter the V1 boundaries: DRA reports hardware, operators +choose the nodes OpenRL may use, the estimator describes a worker, the Go +controller places it, and the node-local scheduler coordinates shared access. + +## Appendix A: Safety and accounting + +Three checks govern whether a worker can join a claim: + +1. **Independent fit:** the worker must fit the allocation by itself. +2. **Worker limit:** the claim must remain below `max-workers-per-claim`. +3. **Suspended fit:** the node must have enough host memory for every worker + that may be parked while another is resident. + +There is no resident-memory sum in V1. Each worker's memory estimate is +checked against the full allocation independently. For host-memory admission, +the conservative case parks every assigned worker except the smallest one, +which could be the current resident. + +Pods sharing a claim intentionally receive the same devices, so OpenRL must +coordinate their execution. A worker may stay resident while no other worker +needs the claim, but a handoff must finish suspending it before restoration +of the next worker begins. A failed handoff does not grant the next worker +access. + +The preferred suspension mechanism depends on the runtime: + +| Runtime | Suspension mechanism | +| --- | --- | +| vLLM sampler | Sleep level 1 and wake | +| FFT trainer | Application-level host offload and restore | +| Other CUDA worker | CUDA process checkpoint and restore | + +## Appendix B: Alternatives considered + +**Pre-created claims.** This removes dynamic claim lifecycle from OpenRL but +preserves the up-front capacity partitioning the design is intended to +remove. + +**Pack before spreading.** This keeps whole GPUs free for future jobs but +reduces current throughput. V1 spreads first and shares under contention. + +**Co-resident GPU-memory packing.** Keeping several workers loaded can +eliminate transfers when their memory fits together. V1 rejects this +optimization because it requires resident-set accounting, eviction policy, +and recovery from partially completed evictions. It can be added later inside +the node-local scheduler without changing claim membership. + +**Always suspend after a turn.** This simplifies handoff logic but pays +transfer cost even when no other worker is waiting. V1 requires suspension +before a different worker becomes resident, not merely because the current +turn ended. + +**Process checkpointing for every worker.** This is a useful fallback, but +application-aware sampler sleep and trainer offload move less unnecessary +state. + +**Model compatibility groups.** These are unnecessary for V1 placement +because workers never remain resident together. Each only needs to fit +independently and support suspension. LoRA process reuse is decided before +placement. + +**Native DRA without OpenRLWorker.** A Pod can request a claim directly, but +OpenRL would have no durable object for worker identity, the memory estimate, +placement status, and cleanup. + +## Appendix C: Inputs, outputs, and the reconcile sequence + +### Inputs, read fresh on every pass + +The worker itself: + +```yaml +apiVersion: openrl.io/v1alpha1 +kind: OpenRLWorker +metadata: + name: trainer-job-123 # identity; pod and claim names derive from it +spec: + role: trainer # which node pools may host it + modelId: job-123 # configuration for the process, never identity + ownerId: job-123 # fairness unit; opaque, passed through raw + memory: 60Gi # the estimator's figure; must be positive + estimatorVersion: v1 +``` + +The spec is immutable (CEL-enforced): change by deleting and recreating. + +The standing inputs: node labels (`openrl.io/enabled`, `/trainer`, +`/sampler`, `/max-workers-per-claim`), ResourceSlices from the configured +driver -- only the latest complete generation of each pool counts -- giving +device count and per-device memory, and each node's allocatable host memory. + +### Outputs, at most three writes per pass + +The ResourceClaim, when one is cut: + +```yaml +metadata: + name: claim-trainer-job-123-6b2f01a # worker name + UID: unique per + # incarnation, stable within one + labels: + openrl.io/managed: "true" + openrl.io/accelerator-count: "1" # the claim's shape contract + openrl.io/device-memory: 80Gi + openrl.io/sized-against: node-a # reserves the pool until DRA decides +spec: # one request, exact count, CEL bounds: + # floor = the worker's per-device share + # ceiling = the device size the claim was priced against, so DRA cannot + # substitute a bigger device placement never chose +``` + +The Pod: rendered from the operator's template, with the controller's stamps +-- the claim reference, two ORed node-affinity terms (explicit role label, or +no role labels at all: the documented default), an owner reference to the +worker, and the time-slice contract (group = claim name, owner, job id) as +env vars carrying exact values and labels carrying sanitized copies. + +The status: `phase`, `deviceCount`, `memoryPerDevice`, `hostMemoryWhenParked`, +`estimatorVersion`, `claimName`, `podName`, `nodeName`, `reason`, and the +`Placed` condition. + +### The sequence + +1. A watch event names one worker; reconciles run one at a time, because + placement is a fleet-wide decision. +2. **Deleting?** Tear down: delete the pod, requeue until it is verifiably + gone, then release the finalizer. The CR -- and its memory booking -- + survives the pod's termination grace, so a seat can never free while the + process still holds the device. +3. Ensure the finalizer; reject non-positive memory as Failed. +4. **Read the fleet:** pools from slices x labels, managed claims, and + occupancy rebuilt from every worker's status. +5. **Adopt reality:** the pod's claim label outranks a lost status; a pod + owned by a previous incarnation (different ownerRef UID) is replaced, + never adopted; a vanished claim means re-place. +6. **No claim yet?** Decide: a free pool cuts a new claim (unallocated, its + pool reserved) -> Placing; otherwise an allocated claim that passes the + three admission checks is joined -> Placing; otherwise Pending with a + reason that distinguishes busy from impossible, retried each interval + until the placement timeout turns "not yet" into Failed. +7. **Pod:** create it if missing; a pod bound to a stale claim is deleted + and rebuilt, because spec.resourceClaims is immutable. +8. **Report:** Running once the pod runs, the node read from the claim's + allocation; failures carry the scheduler's own words. +9. In the background, the reclaim sweep deletes managed claims that no + worker names, no live pod uses, and DRA no longer reserves, after a + grace period. diff --git a/controller/docs/layout.md b/controller/docs/layout.md new file mode 100644 index 00000000..2c637744 --- /dev/null +++ b/controller/docs/layout.md @@ -0,0 +1,73 @@ +# Where everything is + +The short version: two structs define the API, three functions make the +decision, one function drives Kubernetes, and everything else is glue you can +read once and trust. + +``` +controller/ +├── api/v1alpha1/ +│ ├── openrlworker_types.go the contract: Spec (what the caller asks) and +│ │ Status (what was decided). The +kubebuilder +│ │ comments generate the CRD and its validation; +│ │ doc comments here become `kubectl explain` text. +│ ├── groupversion_info.go scheme registration boilerplate +│ └── zz_generated.deepcopy.go generated; never edit (make generate) +│ +├── internal/placement/ THE DECISION. Pure functions, no Kubernetes. +│ ├── placement.go Decide = ChoosePool (spread onto free devices) +│ │ then SelectClaim (share under contention). +│ │ Claim/Node/Fleet are plain structs; the three +│ │ admission checks and the parked-memory formula +│ │ live here and nowhere else. +│ ├── behavior_test.go the end-to-end behaviors, written as arrivals +│ │ and departures with the estimator's real tier +│ │ figures. Start reading tests here. +│ └── placement_test.go unit tests for the arithmetic, tie-breaks, +│ and the pending-claim reservation rules +│ +├── internal/controller/ THE KUBERNETES GLUE. +│ ├── openrlworker_controller.go Reconcile -> place(): one decision tree per +│ │ pass (has claim? ensure pod : decide, create, +│ │ or mark Pending). Also the watch wiring +│ │ (SetupWithManager) and status writing. +│ ├── fleet.go reads the world into placement's Fleet: +│ │ ResourceSlices x node labels = pools, managed +│ │ claims + worker statuses = occupancy +│ ├── pod.go renders the worker pod: operator template + +│ │ the controller's stamps (claim, affinity, +│ │ time-slice env) and builds the ResourceClaim +│ ├── reclaim.go the sweep that deletes claims nobody uses +│ └── openrlworker_controller_test.go fake-client tests for the glue +│ +├── cmd/manager/main.go flags/env -> Manager -> run. Boilerplate. +│ +├── hack/ +│ ├── kind-smoke.sh the pipeline on kind: fake GPUs by default, +│ │ real DRA via env (see header). `make smoke`. +│ ├── stress.sh churn many workers against a live cluster and +│ │ assert the seat arithmetic every round +│ └── retest-on-box.sh dev helper: sync + test on the GPU box +│ +└── docs/ + ├── design.md the spec. If code and spec disagree, one of + │ them is a bug. + └── layout.md this file + +k8s/deploy/ +├── scheduler/ CRD + RBAC + Deployment. Inert until someone +│ creates OpenRLWorker objects. +└── scheduler-smoke/ kustomize overlay the smoke test applies: + local image, fake-driver env, sleep templates +``` + +Reading order for a first pass: `api/v1alpha1/openrlworker_types.go` (the +contract), then `internal/placement/behavior_test.go` (what it promises), then +`placement.go`'s `Decide`/`SelectClaim`/`ChoosePool` (how), then +`openrlworker_controller.go`'s `place()` (how it touches Kubernetes). That is +~500 lines and everything else is in service of it. + +Not in this module: the node-local time-slicer (who is *resident* right now) +is Python, in `src/accel_timeslicer/`, and ships with the FFT line along with +the gateway's `src/server/scheduler_worker_manager.py`, which turns API +requests into OpenRLWorker objects using the estimator's memory figure. diff --git a/controller/go.mod b/controller/go.mod new file mode 100644 index 00000000..64a4770f --- /dev/null +++ b/controller/go.mod @@ -0,0 +1,64 @@ +module github.com/gke-labs/open-rl/controller + +go 1.26.0 + +require ( + k8s.io/api v0.36.3 + k8s.io/apimachinery v0.36.3 + k8s.io/client-go v0.36.3 + sigs.k8s.io/controller-runtime v0.24.1 + sigs.k8s.io/yaml v1.6.0 +) + +require ( + github.com/beorn7/perks v1.0.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/emicklei/go-restful/v3 v3.13.0 // indirect + github.com/evanphx/json-patch/v5 v5.9.11 // indirect + github.com/fsnotify/fsnotify v1.9.0 // indirect + github.com/fxamacker/cbor/v2 v2.9.0 // indirect + github.com/go-logr/logr v1.4.3 // indirect + github.com/go-logr/zapr v1.3.0 // indirect + github.com/go-openapi/jsonpointer v0.21.0 // indirect + github.com/go-openapi/jsonreference v0.20.2 // indirect + github.com/go-openapi/swag v0.23.0 // indirect + github.com/google/gnostic-models v0.7.0 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/josharian/intern v1.0.0 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect + github.com/prometheus/client_golang v1.23.2 // indirect + github.com/prometheus/client_model v0.6.2 // indirect + github.com/prometheus/common v0.67.5 // indirect + github.com/prometheus/procfs v0.19.2 // indirect + github.com/spf13/pflag v1.0.9 // indirect + github.com/x448/float16 v0.8.4 // indirect + go.uber.org/multierr v1.11.0 // indirect + go.uber.org/zap v1.27.1 // indirect + go.yaml.in/yaml/v2 v2.4.3 // indirect + go.yaml.in/yaml/v3 v3.0.4 // indirect + golang.org/x/net v0.49.0 // indirect + golang.org/x/oauth2 v0.34.0 // indirect + golang.org/x/sync v0.19.0 // indirect + golang.org/x/sys v0.40.0 // indirect + golang.org/x/term v0.39.0 // indirect + golang.org/x/text v0.33.0 // indirect + golang.org/x/time v0.14.0 // indirect + gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect + gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect + gopkg.in/inf.v0 v0.9.1 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect + k8s.io/apiextensions-apiserver v0.36.0 // indirect + k8s.io/klog/v2 v2.140.0 // indirect + k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a // indirect + k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 // indirect + sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect + sigs.k8s.io/randfill v1.0.0 // indirect + sigs.k8s.io/structured-merge-diff/v6 v6.3.3 // indirect +) diff --git a/controller/go.sum b/controller/go.sum new file mode 100644 index 00000000..033b9dc1 --- /dev/null +++ b/controller/go.sum @@ -0,0 +1,171 @@ +github.com/Masterminds/semver/v3 v3.4.0 h1:Zog+i5UMtVoCU8oKka5P7i9q9HgrJeGzI9SA1Xbatp0= +github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/emicklei/go-restful/v3 v3.13.0 h1:C4Bl2xDndpU6nJ4bc1jXd+uTmYPVUwkD6bFY/oTyCes= +github.com/emicklei/go-restful/v3 v3.13.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= +github.com/evanphx/json-patch v0.5.2 h1:xVCHIVMUu1wtM/VkR9jVZ45N3FhZfYMMYGorLCR8P3k= +github.com/evanphx/json-patch v0.5.2/go.mod h1:ZWS5hhDbVDyob71nXKNL0+PWn6ToqBHMikGIFbs31qQ= +github.com/evanphx/json-patch/v5 v5.9.11 h1:/8HVnzMq13/3x9TPvjG08wUGqBTmZBsCWzjTM0wiaDU= +github.com/evanphx/json-patch/v5 v5.9.11/go.mod h1:3j+LviiESTElxA4p3EMKAB9HXj3/XEtnUf6OZxqIQTM= +github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k= +github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0= +github.com/fxamacker/cbor/v2 v2.9.0 h1:NpKPmjDBgUfBms6tr6JZkTHtfFGcMKsw3eGcmD/sapM= +github.com/fxamacker/cbor/v2 v2.9.0/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/zapr v1.3.0 h1:XGdV8XW8zdwFiwOA2Dryh1gj2KRQyOOoNmBy4EplIcQ= +github.com/go-logr/zapr v1.3.0/go.mod h1:YKepepNBd1u/oyhd/yQmtjVXmm9uML4IXUgMOwR8/Gg= +github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= +github.com/go-openapi/jsonpointer v0.21.0 h1:YgdVicSA9vH5RiHs9TZW5oyafXZFc6+2Vc1rr/O9oNQ= +github.com/go-openapi/jsonpointer v0.21.0/go.mod h1:IUyH9l/+uyhIYQ/PXVA41Rexl+kOkAPDdXEYns6fzUY= +github.com/go-openapi/jsonreference v0.20.2 h1:3sVjiK66+uXK/6oQ8xgcRKcFgQ5KXa2KvnJRumpMGbE= +github.com/go-openapi/jsonreference v0.20.2/go.mod h1:Bl1zwGIM8/wsvqjsOQLJ/SH+En5Ap4rVB5KVcIDZG2k= +github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= +github.com/go-openapi/swag v0.23.0 h1:vsEVJDUo2hPJ2tu0/Xc+4noaxyEffXNIs3cOULZ+GrE= +github.com/go-openapi/swag v0.23.0/go.mod h1:esZ8ITTYEsH1V2trKHjAN8Ai7xHb8RV+YSZ577vPjgQ= +github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI= +github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8= +github.com/google/gnostic-models v0.7.0 h1:qwTtogB15McXDaNqTZdzPJRHvaVJlAl+HVQnLmJEJxo= +github.com/google/gnostic-models v0.7.0/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0= +github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/pprof v0.0.0-20250403155104-27863c87afa6 h1:BHT72Gu3keYf3ZEu2J0b1vyeLSOYI8bm5wbJM/8yDe8= +github.com/google/pprof v0.0.0-20250403155104-27863c87afa6/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= +github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee h1:W5t00kpgFdJifH4BDsTlE89Zl93FEloxaWZfGcifgq8= +github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= +github.com/onsi/ginkgo/v2 v2.27.4 h1:fcEcQW/A++6aZAZQNUmNjvA9PSOzefMJBerHJ4t8v8Y= +github.com/onsi/ginkgo/v2 v2.27.4/go.mod h1:ArE1D/XhNXBXCBkKOLkbsb2c81dQHCRcF5zwn/ykDRo= +github.com/onsi/gomega v1.39.0 h1:y2ROC3hKFmQZJNFeGAMeHZKkjBL65mIZcvrLQBF9k6Q= +github.com/onsi/gomega v1.39.0/go.mod h1:ZCU1pkQcXDO5Sl9/VVEGlDyp+zm0m1cmeG5TOzLgdh4= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= +github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg= +github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= +github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= +github.com/prometheus/common v0.67.5 h1:pIgK94WWlQt1WLwAC5j2ynLaBRDiinoAb86HZHTUGI4= +github.com/prometheus/common v0.67.5/go.mod h1:SjE/0MzDEEAyrdr5Gqc6G+sXI67maCxzaT3A2+HqjUw= +github.com/prometheus/procfs v0.19.2 h1:zUMhqEW66Ex7OXIiDkll3tl9a1ZdilUOd/F6ZXw4Vws= +github.com/prometheus/procfs v0.19.2/go.mod h1:M0aotyiemPhBCM0z5w87kL22CxfcH05ZpYlu+b4J7mw= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= +github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= +go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= +go.uber.org/zap v1.27.1 h1:08RqriUEv8+ArZRYSTXy1LeBScaMpVSTBhCeaZYfMYc= +go.uber.org/zap v1.27.1/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E= +go.yaml.in/yaml/v2 v2.4.3 h1:6gvOSjQoTB3vt1l+CU+tSyi/HOjfOjRLJ4YwYZGwRO0= +go.yaml.in/yaml/v2 v2.4.3/go.mod h1:zSxWcmIDjOzPXpjlTTbAsKokqkDNAVtZO0WOMiT90s8= +go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/mod v0.32.0 h1:9F4d3PHLljb6x//jOyokMv3eX+YDeepZSEo3mFJy93c= +golang.org/x/mod v0.32.0/go.mod h1:SgipZ/3h2Ci89DlEtEXWUk/HteuRin+HHhN+WbNhguU= +golang.org/x/net v0.49.0 h1:eeHFmOGUTtaaPSGNmjBKpbng9MulQsJURQUAfUwY++o= +golang.org/x/net v0.49.0/go.mod h1:/ysNB2EvaqvesRkuLAyjI1ycPZlQHM3q01F02UY/MV8= +golang.org/x/oauth2 v0.34.0 h1:hqK/t4AKgbqWkdkcAeI8XLmbK+4m4G5YeQRrmiotGlw= +golang.org/x/oauth2 v0.34.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= +golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= +golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/sys v0.40.0 h1:DBZZqJ2Rkml6QMQsZywtnjnnGvHza6BTfYFWY9kjEWQ= +golang.org/x/sys v0.40.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/term v0.39.0 h1:RclSuaJf32jOqZz74CkPA9qFuVTX7vhLlpfj/IGWlqY= +golang.org/x/term v0.39.0/go.mod h1:yxzUCTP/U+FzoxfdKmLaA0RV1WgE0VY7hXBwKtY/4ww= +golang.org/x/text v0.33.0 h1:B3njUFyqtHDUI5jMn1YIr5B0IE2U0qck04r6d4KPAxE= +golang.org/x/text v0.33.0/go.mod h1:LuMebE6+rBincTi9+xWTY8TztLzKHc/9C1uBCG27+q8= +golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI= +golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4= +golang.org/x/tools v0.41.0 h1:a9b8iMweWG+S0OBnlU36rzLp20z1Rp10w+IY2czHTQc= +golang.org/x/tools v0.41.0/go.mod h1:XSY6eDqxVNiYgezAVqqCeihT4j1U2CCsqvH3WhQpnlg= +gomodules.xyz/jsonpatch/v2 v2.4.0 h1:Ci3iUJyx9UeRx7CeFN8ARgGbkESwJK+KB9lLcWxY/Zw= +gomodules.xyz/jsonpatch/v2 v2.4.0/go.mod h1:AH3dM2RI6uoBZxn3LVrfvJ3E0/9dG4cSrbuBJT4moAY= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/evanphx/json-patch.v4 v4.13.0 h1:czT3CmqEaQ1aanPc5SdlgQrrEIb8w/wwCvWWnfEbYzo= +gopkg.in/evanphx/json-patch.v4 v4.13.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M= +gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc= +gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +k8s.io/api v0.36.3 h1:NxB+05W2UGqXWFXcLO0RB5cnqnUPP5v5sVlaOH0Iz4w= +k8s.io/api v0.36.3/go.mod h1:JzLQKqRHC5+I8RVj/lS3lCg0mg6nWI9Fo/Sk3ElxHzg= +k8s.io/apiextensions-apiserver v0.36.0 h1:Wt7E8J+VBCbj4FjiBfDTK/neXDDjyJVJc7xfuOHImZ0= +k8s.io/apiextensions-apiserver v0.36.0/go.mod h1:kGDjH0msuiIB3tgsYRV0kS9GqpMYMUsQ3GHv7TApyug= +k8s.io/apimachinery v0.36.3 h1:PkzMRBRG8joFD8EhCuQAtNPvJlxb82FwplP26HIzvAM= +k8s.io/apimachinery v0.36.3/go.mod h1:cTSjBWgPe/6CQyBKzY/hDIRWCQQQeK0mfLbml0UYFHE= +k8s.io/client-go v0.36.3 h1:M4JdVzXxYcZk4fGpfDdYnxSwhLKWCFoQsHW6t+z8Hfg= +k8s.io/client-go v0.36.3/go.mod h1:gcPwr0c87vjjG6HB6pWEqOeuYVoXSsREjzux2j6GF30= +k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= +k8s.io/klog/v2 v2.140.0/go.mod h1:o+/RWfJ6PwpnFn7OyAG3QnO47BFsymfEfrz6XyYSSp0= +k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a h1:xCeOEAOoGYl2jnJoHkC3hkbPJgdATINPMAxaynU2Ovg= +k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a/go.mod h1:uGBT7iTA6c6MvqUvSXIaYZo9ukscABYi2btjhvgKGZ0= +k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 h1:AZYQSJemyQB5eRxqcPky+/7EdBj0xi3g0ZcxxJ7vbWU= +k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2/go.mod h1:xDxuJ0whA3d0I4mf/C4ppKHxXynQ+fxnkmQH0vTHnuk= +sigs.k8s.io/controller-runtime v0.24.1 h1:miPEwrmirImAvgME1L9qebGHrOnGJoVmVdtOU9fRfo4= +sigs.k8s.io/controller-runtime v0.24.1/go.mod h1:vFkfY5fGt5xAC/sKb8IBFKgWPNKG9OUG29dR8Y2wImw= +sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 h1:IpInykpT6ceI+QxKBbEflcR5EXP7sU1kvOlxwZh5txg= +sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730/go.mod h1:mdzfpAEoE6DHQEN0uh9ZbOCuHbLK5wOm7dK4ctXE9Tg= +sigs.k8s.io/randfill v1.0.0 h1:JfjMILfT8A6RbawdsK2JXGBR5AQVfd+9TbzrlneTyrU= +sigs.k8s.io/randfill v1.0.0/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY= +sigs.k8s.io/structured-merge-diff/v6 v6.3.3 h1:u08YRbVUi59ri4YD6cg0UqNM4Dimn0sIl+wldcx5PYw= +sigs.k8s.io/structured-merge-diff/v6 v6.3.3/go.mod h1:M3W8sfWvn2HhQDIbGWj3S099YozAsymCo/wrT5ohRUE= +sigs.k8s.io/yaml v1.6.0 h1:G8fkbMSAFqgEFgh4b1wmtzDnioxFCUgTZhlbj5P9QYs= +sigs.k8s.io/yaml v1.6.0/go.mod h1:796bPqUfzR/0jLAl6XjHl3Ck7MiyVv8dbTdyT3/pMf4= diff --git a/controller/hack/kind-smoke.sh b/controller/hack/kind-smoke.sh new file mode 100755 index 00000000..468fa9e5 --- /dev/null +++ b/controller/hack/kind-smoke.sh @@ -0,0 +1,151 @@ +#!/usr/bin/env bash +# The pipeline smoke test: OpenRLWorker -> ResourceClaim -> allocation -> Pod +# Running, against a real API server, real kube-scheduler, and real DRA. +# +# Fake GPUs by default: a kind cluster plus the upstream DRA example driver, +# which publishes synthetic ResourceSlices (8 GPUs x 80Gi per node). The +# controller cannot tell the difference -- everything it reads comes from +# ResourceSlices and node labels -- so the same script against a cluster with +# real hardware exercises the identical path: +# +# ./hack/kind-smoke.sh # kind + fake GPUs +# KEEP=1 ./hack/kind-smoke.sh # leave the cluster up afterwards +# USE_EXISTING_CLUSTER=1 LOAD_INTO= \ +# DEVICE_CLASS=gpu.nvidia.com GPUS=2 MEMORY=10Gi ./hack/kind-smoke.sh +# # current kubectl context, real DRA +# +# GPUS is how many devices the target node exposes (the fake driver's 8 by +# default); MEMORY must fit one of its devices. The run asks for GPUS+1 +# workers and asserts spread-then-share. LOAD_INTO names a kind cluster to +# load the locally built image into; a non-kind cluster needs the image +# pushed somewhere it can pull. +# +# What the behavior tests cannot verify and this does: claim allocation and +# immutability, pod binding, and watch-driven status flow. +set -euo pipefail + +CLUSTER=${CLUSTER:-openrl-smoke} +DEVICE_CLASS=${DEVICE_CLASS:-gpu.example.com} +DEVICE_DRIVER=${DEVICE_DRIVER:-$DEVICE_CLASS} +USE_EXISTING_CLUSTER=${USE_EXISTING_CLUSTER:-0} +LOAD_INTO=${LOAD_INTO:-} +KEEP=${KEEP:-0} +GPUS=${GPUS:-8} +MEMORY=${MEMORY:-50Gi} +# The smoke overlay pins this tag; it is not overridable for that reason. +IMG=open-rl/placement-controller:smoke + +controller=$(cd "$(dirname "$0")/.." && pwd) +repo=$(cd "$controller/.." && pwd) + +say() { printf '\n== %s\n' "$*"; } + +if [ "$USE_EXISTING_CLUSTER" != 1 ]; then + say "creating kind cluster $CLUSTER" + kind create cluster --name "$CLUSTER" --wait 120s + if [ "$KEEP" != 1 ]; then + trap 'kind delete cluster --name "$CLUSTER"' EXIT + fi +fi + +# Fake GPUs. The example driver is the reference DRA implementation: it +# publishes ResourceSlices whose devices carry a memory capacity, exactly like +# the NVIDIA driver, without needing hardware. +if [ "$DEVICE_CLASS" = gpu.example.com ]; then + say "installing the DRA example driver (fake GPUs)" + tmp=$(mktemp -d) + git clone --quiet --depth 1 https://github.com/kubernetes-sigs/dra-example-driver "$tmp/driver" + helm upgrade --install dra-example-driver "$tmp/driver/deployments/helm/dra-example-driver" \ + --create-namespace --namespace dra-example-driver --wait --timeout 180s +fi + +say "waiting for ResourceSlices from $DEVICE_DRIVER" +for _ in $(seq 60); do + kubectl get resourceslices -o jsonpath='{.items[*].spec.driver}' 2>/dev/null | grep -q "$DEVICE_DRIVER" && break + sleep 2 +done +kubectl get resourceslices + +say "building and loading the controller image" +docker build -q -t "$IMG" "$controller" +if [ "$USE_EXISTING_CLUSTER" != 1 ]; then + kind load docker-image "$IMG" --name "$CLUSTER" +elif [ -n "$LOAD_INTO" ]; then + kind load docker-image "$IMG" --name "$LOAD_INTO" +fi + +say "deploying the scheduler (smoke overlay)" +kubectl apply -k "$repo/k8s/deploy/scheduler-smoke" +if [ "$DEVICE_CLASS" != gpu.example.com ]; then + # The overlay defaults to the fake driver; real-hardware runs override it. + kubectl -n default set env deployment/open-rl-scheduler \ + OPEN_RL_DEVICE_CLASS="$DEVICE_CLASS" OPEN_RL_DEVICE_DRIVER="$DEVICE_DRIVER" +fi + +say "opting nodes in" +for node in $(kubectl get nodes -o name); do + kubectl label --overwrite "$node" \ + openrl.io/enabled=true openrl.io/trainer=true openrl.io/sampler=true \ + openrl.io/max-workers-per-claim=2 +done + +if ! kubectl -n default rollout status deployment/open-rl-scheduler --timeout=240s; then + echo "FAIL: the scheduler never became ready" + kubectl -n default describe pods -l app=open-rl-scheduler | tail -30 + kubectl -n default logs deployment/open-rl-scheduler --tail=40 || true + exit 1 +fi + +# One more worker than the node has GPUs, so the pipeline shows both halves +# of the policy: spread while devices are free, then share under contention. +WORKERS=$((GPUS + 1)) + +say "requesting $WORKERS workers of $MEMORY against $GPUS GPUs" +for i in $(seq "$WORKERS"); do + kubectl apply -f - </dev/null || true) + [ "$phase" = Running ] && ok=1 && break + sleep 2 + done + if [ "$ok" != 1 ]; then + echo "FAIL: smoke-$i never reached Running (phase: ${phase:-none})" + kubectl -n default get openrlworkers + kubectl -n default get pods,resourceclaims + exit 1 + fi +done + +say "asserting spread, sharing, and real allocations" +claims=$(kubectl -n default get openrlworkers -o jsonpath='{range .items[*]}{.status.claimName}{"\n"}{end}') +distinct=$(echo "$claims" | sort -u | grep -c .) +if [ "$distinct" != "$GPUS" ]; then + echo "FAIL: $WORKERS workers hold $distinct claims, want $GPUS: one per GPU, with the extra worker sharing" + kubectl -n default get openrlworkers + exit 1 +fi +for claim in $(echo "$claims" | sort -u); do + driver=$(kubectl -n default get resourceclaim "$claim" -o jsonpath='{.status.allocation.devices.results[0].driver}') + if [ "$driver" != "$DEVICE_DRIVER" ]; then + echo "FAIL: claim $claim was not allocated by $DEVICE_DRIVER (got: ${driver:-nothing})" + exit 1 + fi +done +echo "$WORKERS workers on $distinct claims, every claim allocated by $DEVICE_DRIVER" + +say "the run, as an operator would see it" +kubectl -n default get openrlworkers +kubectl -n default get resourceclaims +kubectl -n default get pods -o wide + +say "PASS" diff --git a/controller/hack/retest-on-box.sh b/controller/hack/retest-on-box.sh new file mode 100755 index 00000000..9ac0a092 --- /dev/null +++ b/controller/hack/retest-on-box.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# Dev-loop helper: wait for the dev box to be reachable, sync the controller +# over, and run the Go suite there (the Mac kills freshly built binaries). +set -u +controller=$(cd "$(dirname "$0")/.." && pwd) + +for _ in $(seq "${TRIES:-40}"); do + if ssh -o ConnectTimeout=15 box 'echo alive' >/dev/null 2>&1; then + rsync -a --delete --exclude bin "$controller/" box:~/sched/ || exit 1 + exec ssh box 'export PATH=$PATH:/usr/local/go/bin && cd ~/sched && go vet ./... && go test -count=1 ./... 2>&1 | tail -6' + fi + sleep 30 +done +echo "box never came back" +exit 1 diff --git a/controller/hack/stress.sh b/controller/hack/stress.sh new file mode 100755 index 00000000..00e59fae --- /dev/null +++ b/controller/hack/stress.sh @@ -0,0 +1,107 @@ +#!/usr/bin/env bash +# Churn workers against a live scheduler and check the books every round. +# +# Run it after kind-smoke.sh with KEEP=1 (or against any cluster where the +# scheduler is deployed and nodes are labeled). Each round creates workers up +# to WORKERS, waits for the fleet to settle, asserts the seat arithmetic -- +# running == min(live, GPUS x MAXW), everyone else Pending with a reason, +# never more claims than GPUs -- then deletes a random third and goes again. +# The last round deletes everything and waits for the reclaim sweep to return +# the namespace to empty: no workers, no managed claims, no worker pods. +# +# WORKERS=24 ROUNDS=5 GPUS=8 MAXW=2 MEMORY=50Gi ./hack/stress.sh +set -euo pipefail + +WORKERS=${WORKERS:-24} +ROUNDS=${ROUNDS:-5} +GPUS=${GPUS:-8} +MAXW=${MAXW:-2} +MEMORY=${MEMORY:-50Gi} +SEED=${SEED:-42} +RANDOM=$SEED + +say() { printf '\n== %s\n' "$*"; } +count() { kubectl -n default get "$1" --no-headers 2>/dev/null | grep -c . || true; } + +phase_counts() { + kubectl -n default get openrlworkers -o jsonpath='{range .items[*]}{.status.phase}{"\n"}{end}' +} + +settle() { + local live=$1 seats=$((GPUS * MAXW)) want_running want_pending + want_running=$((live < seats ? live : seats)) + want_pending=$((live - want_running)) + for _ in $(seq 120); do + local running pending failed + running=$(phase_counts | grep -c '^Running$' || true) + pending=$(phase_counts | grep -c '^Pending$' || true) + failed=$(phase_counts | grep -c '^Failed$' || true) + if [ "$failed" != 0 ]; then + echo "FAIL: $failed workers Failed" + kubectl -n default get openrlworkers + exit 1 + fi + if [ "$running" = "$want_running" ] && [ "$pending" = "$want_pending" ]; then + local claims + claims=$(count resourceclaims) + if [ "$claims" -gt "$GPUS" ]; then + echo "FAIL: $claims claims for $GPUS GPUs" + kubectl -n default get resourceclaims + exit 1 + fi + echo "settled: $running running, $pending pending, $claims claims" + # Every pending worker must say why. + kubectl -n default get openrlworkers -o jsonpath='{range .items[?(@.status.phase=="Pending")]}{.metadata.name}: {.status.reason}{"\n"}{end}' + return 0 + fi + sleep 2 + done + echo "FAIL: never settled at $want_running running / $want_pending pending" + kubectl -n default get openrlworkers + exit 1 +} + +for round in $(seq "$ROUNDS"); do + say "round $round: topping up to $WORKERS workers" + for i in $(seq "$WORKERS"); do + kubectl -n default get openrlworker "stress-$i" >/dev/null 2>&1 && continue + kubectl apply -f - >/dev/null </dev/null 2>&1 || true + deleted=$((deleted + 1)) + fi + done + # Deletion is asynchronous; wait for the census to match before asserting. + for _ in $(seq 60); do + [ "$(count openrlworkers)" = "$((WORKERS - deleted))" ] && break + sleep 2 + done + settle "$((WORKERS - deleted))" +done + +say "final: deleting everything and waiting for the reclaim sweep" +kubectl -n default delete openrlworkers --all --wait=true >/dev/null +# Claims outlive their workers by design: a 2m grace plus the sweep interval. +for _ in $(seq 100); do + claims=$(count resourceclaims) + pods=$(kubectl -n default get pods --no-headers 2>/dev/null | grep -c '^orw-' || true) + if [ "$claims" = 0 ] && [ "$pods" = 0 ]; then + say "PASS: namespace is clean -- 0 workers, 0 claims, 0 worker pods" + exit 0 + fi + sleep 5 +done +echo "FAIL: leftovers after deletion: $claims claims, $pods worker pods" +kubectl -n default get pods,resourceclaims +exit 1 diff --git a/controller/internal/controller/fleet.go b/controller/internal/controller/fleet.go new file mode 100644 index 00000000..22091c3b --- /dev/null +++ b/controller/internal/controller/fleet.go @@ -0,0 +1,278 @@ +package controller + +import ( + "context" + "fmt" + "sort" + "strconv" + + corev1 "k8s.io/api/core/v1" + resourcev1 "k8s.io/api/resource/v1" + "k8s.io/apimachinery/pkg/api/resource" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/log" + + openrlv1alpha1 "github.com/gke-labs/open-rl/controller/api/v1alpha1" + "github.com/gke-labs/open-rl/controller/internal/placement" +) + +// Labels the controller stamps on the objects it owns, and reads back to +// rebuild its own state after a restart. +const ( + LabelManaged = "openrl.io/managed" + LabelRole = "openrl.io/role" + LabelAccelCount = "openrl.io/accelerator-count" + LabelDeviceMemory = "openrl.io/device-memory" + LabelClaim = "openrl.io/claim" + LabelWorker = "openrl.io/worker" + LabelSizedAgainst = "openrl.io/sized-against" +) + +// Node labels the operator sets to opt a pool in. Policy, not hardware: the +// DRA driver's ResourceSlices report what the devices actually are. +// +// Opting a node in means opting it in exclusively: the controller counts a +// device as free unless one of its own claims holds it, so GPU workloads it +// does not manage on an enabled node would be invisible to placement. +const ( + NodeLabelEnabled = "openrl.io/enabled" + NodeLabelMaxWorkersPerClaim = "openrl.io/max-workers-per-claim" + NodeLabelTrainer = "openrl.io/trainer" + NodeLabelSampler = "openrl.io/sampler" +) + +var nodeRoleLabel = map[openrlv1alpha1.WorkerRole]string{ + openrlv1alpha1.RoleTrainer: NodeLabelTrainer, + openrlv1alpha1.RoleSampler: NodeLabelSampler, +} + +// readFleet folds ResourceSlices, node labels, managed ResourceClaims and the +// workers already assigned to them into one picture to decide against. +func (r *OpenRLWorkerReconciler) readFleet(ctx context.Context, workers []openrlv1alpha1.OpenRLWorker) (*placement.Fleet, error) { + var nodes corev1.NodeList + if err := r.List(ctx, &nodes, client.MatchingLabels{NodeLabelEnabled: "true"}); err != nil { + return nil, fmt.Errorf("list nodes: %w", err) + } + + var slices resourcev1.ResourceSliceList + if err := r.List(ctx, &slices); err != nil { + return nil, fmt.Errorf("list resourceslices: %w", err) + } + + fleet := placement.NewFleet() + fleet.Nodes = r.poolsFrom(ctx, slices.Items, nodes.Items) + + // Consistent reader: a claim created for the previous worker in a burst + // must be joinable by this one, and the cache may not have caught up. + var claims resourcev1.ResourceClaimList + if err := r.fleetReader().List(ctx, &claims, client.InNamespace(r.Namespace), client.MatchingLabels{LabelManaged: "true"}); err != nil { + return nil, fmt.Errorf("list resourceclaims: %w", err) + } + for i := range claims.Items { + if c := r.claimFrom(ctx, &claims.Items[i]); c != nil { + fleet.Claims[c.Name] = c + } + } + + // Occupancy comes from worker statuses, not pods -- and the finalizer is + // what makes that sound: a deleting worker keeps its CR, and with it the + // real memory booking, until its pod is verifiably gone. + for i := range workers { + bookWorker(fleet, &workers[i]) + } + return fleet, nil +} + +// poolsFrom merges what the driver publishes with what the operator allowed. +// Devices accumulate across slices; where memory differs the smallest wins, +// because the fit must hold for whichever devices DRA picks. +func (r *OpenRLWorkerReconciler) poolsFrom(ctx context.Context, slices []resourcev1.ResourceSlice, nodes []corev1.Node) map[string]*placement.Node { + logger := log.FromContext(ctx) + + devices := map[string]*placement.Node{} + for _, i := range latestCompletePools(ctx, slices, r.DeviceDriver) { + spec := slices[i].Spec + name := *spec.NodeName + for j := range spec.Devices { + device := spec.Devices[j] + capacity, ok := device.Capacity["memory"] + if !ok { + continue + } + memory := capacity.Value.Value() + pool, seen := devices[name] + if !seen { + product := "" + if attr, ok := device.Attributes["productName"]; ok && attr.StringValue != nil { + product = *attr.StringValue + } + devices[name] = &placement.Node{Name: name, DeviceCount: 1, DeviceMemoryBytes: memory, Product: product} + continue + } + pool.DeviceCount++ + pool.DeviceMemoryBytes = min(pool.DeviceMemoryBytes, memory) + } + } + + pools := map[string]*placement.Node{} + for i := range nodes { + node := &nodes[i] + pool, ok := devices[node.Name] + if !ok { + logger.Info("node is enabled but no ResourceSlice from this driver describes it; skipping it for placement", + "node", node.Name, "driver", r.DeviceDriver) + continue + } + // A node naming no role labels allows both roles. + pool.Roles = map[string]bool{} + for role, label := range nodeRoleLabel { + if node.Labels[label] == "true" { + pool.Roles[string(role)] = true + } + } + if len(pool.Roles) == 0 { + for role := range nodeRoleLabel { + pool.Roles[string(role)] = true + } + } + pool.MaxWorkersPerClaim = max(1, labelInt(ctx, node.Labels, NodeLabelMaxWorkersPerClaim, 1, "node "+node.Name)) + pool.HostMemoryBytes = node.Status.Allocatable.Memory().Value() + if pool.HostMemoryBytes == 0 { + logger.Info("node reports no allocatable memory; parked-worker capacity cannot be checked here", + "node", node.Name) + } + pools[node.Name] = pool + } + return pools +} + +// latestCompletePools filters slices to the latest complete generation of +// each (node, pool) -- what the ResourceSlice contract requires of consumers. +// During a driver update, mixing generations double-counts devices and a +// partial generation under-counts them; an incomplete pool is skipped and +// picked up on a later reconcile. +func latestCompletePools(ctx context.Context, slices []resourcev1.ResourceSlice, driver string) []int { + type poolKey struct{ node, pool string } + byPool := map[poolKey][]int{} + for i := range slices { + spec := slices[i].Spec + if spec.NodeName == nil || *spec.NodeName == "" || spec.Driver != driver { + continue + } + key := poolKey{*spec.NodeName, spec.Pool.Name} + byPool[key] = append(byPool[key], i) + } + + var keep []int + for key, indices := range byPool { + latest := int64(-1) + for _, i := range indices { + if g := slices[i].Spec.Pool.Generation; g > latest { + latest = g + } + } + var current []int + for _, i := range indices { + if slices[i].Spec.Pool.Generation == latest { + current = append(current, i) + } + } + if int64(len(current)) != slices[current[0]].Spec.Pool.ResourceSliceCount { + log.FromContext(ctx).Info("skipping incomplete ResourceSlice pool", + "node", key.node, "pool", key.pool, "have", len(current), "want", slices[current[0]].Spec.Pool.ResourceSliceCount) + continue + } + keep = append(keep, current...) + } + sort.Ints(keep) + return keep +} + +// claimFrom reads back the shape the controller stamped on a claim it created. +func (r *OpenRLWorkerReconciler) claimFrom(ctx context.Context, claim *resourcev1.ResourceClaim) *placement.Claim { + count := labelInt(ctx, claim.Labels, LabelAccelCount, 0, "claim "+claim.Name) + if count < 1 { + log.FromContext(ctx).Info("skipping claim with unusable accelerator-count label", "claim", claim.Name) + return nil + } + return &placement.Claim{ + Name: claim.Name, + DeviceCount: count, + Node: allocatedNode(claim), + // Only read while unallocated: the reservation that keeps a burst from + // cutting more claims than the pool has devices. + SizedAgainst: claim.Labels[LabelSizedAgainst], + } +} + +// bookWorker charges a worker's assignment against the claim it names, +// re-deriving memory from the spec (the status string is Gi-rounded). +func bookWorker(fleet *placement.Fleet, worker *openrlv1alpha1.OpenRLWorker) { + if claim, ok := fleet.Claims[worker.Status.ClaimName]; ok { + request := requestFrom(worker) + claim.Book(request.WorkerID, request.PerDeviceBytes(claim.DeviceCount)) + } +} + +// allocatedNode is the node a claim was allocated to, or "" if DRA has not +// decided. DRA reports placement as a node selector; a GPU allocation pins to +// exactly one hostname. +func allocatedNode(claim *resourcev1.ResourceClaim) string { + if claim.Status.Allocation == nil || claim.Status.Allocation.NodeSelector == nil { + return "" + } + for _, term := range claim.Status.Allocation.NodeSelector.NodeSelectorTerms { + for _, expr := range term.MatchExpressions { + if isHostnameKey(expr.Key) && len(expr.Values) > 0 { + return expr.Values[0] + } + } + for _, field := range term.MatchFields { + if isHostnameKey(field.Key) && len(field.Values) > 0 { + return field.Values[0] + } + } + } + return "" +} + +func isHostnameKey(key string) bool { + return key == corev1.LabelHostname || key == "metadata.name" +} + +// requestFrom is the placement Request an OpenRLWorker spec is asking for. +// Validation is the CRD schema's job. +func requestFrom(worker *openrlv1alpha1.OpenRLWorker) placement.Request { + spec := worker.Spec + return placement.Request{ + Role: string(spec.Role), + Memory: spec.Memory.Value(), + // Raw: the spec calls the owner ID opaque, and sanitizing here would + // merge distinct owners ("A/B" and "a-b") into one fairness slot. + // Labels sanitize at the stamping site; env vars carry this exactly. + Owner: spec.OwnerID, + // The CR name: the one identity Kubernetes already guarantees unique. + // Model id is model configuration, not object identity. + WorkerID: worker.Name, + } +} + +// labelInt reads a non-negative integer label, falling back to a default if it +// is missing or nonsense. +func labelInt(ctx context.Context, labels map[string]string, key string, fallback int, subject string) int { + raw, ok := labels[key] + if !ok { + return fallback + } + value, err := strconv.Atoi(raw) + if err != nil || value < 0 { + log.FromContext(ctx).Info("ignoring unparseable label", "subject", subject, "label", key, "value", raw, "using", fallback) + return fallback + } + return value +} + +// gibQuantity renders a byte count as the Gi string the CRD reports. +func gibQuantity(bytes int64) string { + return resource.NewQuantity(placement.CeilGiB(bytes)*placement.GiB, resource.BinarySI).String() +} diff --git a/controller/internal/controller/openrlworker_controller.go b/controller/internal/controller/openrlworker_controller.go new file mode 100644 index 00000000..85145709 --- /dev/null +++ b/controller/internal/controller/openrlworker_controller.go @@ -0,0 +1,498 @@ +package controller + +import ( + "context" + "fmt" + "time" + + corev1 "k8s.io/api/core/v1" + resourcev1 "k8s.io/api/resource/v1" + apiequality "k8s.io/apimachinery/pkg/api/equality" + apierrors "k8s.io/apimachinery/pkg/api/errors" + apimeta "k8s.io/apimachinery/pkg/api/meta" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/tools/record" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/builder" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/controller-runtime/pkg/event" + "sigs.k8s.io/controller-runtime/pkg/handler" + "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/manager" + "sigs.k8s.io/controller-runtime/pkg/predicate" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + openrlv1alpha1 "github.com/gke-labs/open-rl/controller/api/v1alpha1" + "github.com/gke-labs/open-rl/controller/internal/placement" +) + +// OpenRLWorkerReconciler turns OpenRLWorker requests into ResourceClaims and +// pods. The decision lives in internal/placement; this is the part that reads +// and writes Kubernetes objects. Concurrency is one: two reconciles at once +// would each decide against a fleet missing the other's booking. +type OpenRLWorkerReconciler struct { + client.Client + Recorder record.EventRecorder + + // Namespace is where workers, claims and pods live. + Namespace string + // DeviceClass is the DRA DeviceClass generated claims request. + DeviceClass string + // DeviceDriver is the driver publishing the ResourceSlices, and the CEL + // domain its capacities live under. Distinct from DeviceClass in principle, + // identical for NVIDIA's driver. + DeviceDriver string + // DefaultPodTemplates names the ConfigMap per role used when a worker does + // not name one itself. + DefaultPodTemplates map[openrlv1alpha1.WorkerRole]string + // RetryInterval is how often a worker that could not be placed is retried. + RetryInterval time.Duration + // PlacementTimeout is how long a worker may go unplaced before the request + // is declared unsatisfiable. Without it an impossible request waits + // forever, indistinguishable from one that is merely queued. + PlacementTimeout time.Duration + // ReclaimInterval is how often idle claims are swept. + ReclaimInterval time.Duration + + // reader reads straight from the API server, past the informer cache: a + // placement is three writes the cache reflects only eventually, and a + // burst of workers reconciled back to back must each see the previous + // one's booking. Nil (in tests) falls back to the regular client. + reader client.Reader +} + +// fleetReader is the consistent reader for fleet state. +func (r *OpenRLWorkerReconciler) fleetReader() client.Reader { + if r.reader != nil { + return r.reader + } + return r.Client +} + +// +kubebuilder:rbac:groups=openrl.io,resources=openrlworkers,verbs=get;list;watch;update +// +kubebuilder:rbac:groups=openrl.io,resources=openrlworkers/status,verbs=get;update;patch +// +kubebuilder:rbac:groups=resource.k8s.io,resources=resourceclaims,verbs=get;list;watch;create;delete +// +kubebuilder:rbac:groups=resource.k8s.io,resources=resourceslices,verbs=get;list;watch +// +kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;watch;create;delete +// +kubebuilder:rbac:groups=core,resources=nodes,verbs=get;list;watch +// +kubebuilder:rbac:groups=core,resources=configmaps,verbs=get;list;watch +// +kubebuilder:rbac:groups=core,resources=events,verbs=create;patch +// +kubebuilder:rbac:groups=coordination.k8s.io,resources=leases,verbs=get;list;watch;create;update;patch;delete + +// Reconcile places one worker, deciding against a fresh read of the fleet. +func (r *OpenRLWorkerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + // Read through the consistent reader: deciding against a stale copy of + // your own status is how seats get handed out twice. + var worker openrlv1alpha1.OpenRLWorker + if err := r.fleetReader().Get(ctx, req.NamespacedName, &worker); err != nil { + if apierrors.IsNotFound(err) { + return ctrl.Result{}, nil + } + return ctrl.Result{}, err + } + if !worker.DeletionTimestamp.IsZero() { + return r.teardown(ctx, &worker) + } + // The finalizer is the seat guarantee: occupancy is rebuilt from worker + // statuses, so the CR must outlive its pod or the seat frees while the + // process still holds the device. + if controllerutil.AddFinalizer(&worker, workerFinalizer) { + if err := r.Update(ctx, &worker); err != nil { + return ctrl.Result{}, err + } + } + + request := requestFrom(&worker) + + var workers openrlv1alpha1.OpenRLWorkerList + if err := r.fleetReader().List(ctx, &workers, client.InNamespace(r.Namespace)); err != nil { + return ctrl.Result{}, fmt.Errorf("list workers: %w", err) + } + fleet, err := r.readFleet(ctx, workers.Items) + if err != nil { + return ctrl.Result{}, fmt.Errorf("cannot read the fleet, placing nothing this pass: %w", err) + } + + return r.place(ctx, &worker, request, fleet) +} + +func (r *OpenRLWorkerReconciler) place(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, request placement.Request, fleet *placement.Fleet) (ctrl.Result, error) { + logger := log.FromContext(ctx) + + if request.Memory <= 0 { + return ctrl.Result{}, r.fail(ctx, worker, "InvalidSpec", "spec.memory must be a positive quantity") + } + + podName := workerPodName(worker) + pod, err := r.findPod(ctx, podName) + if err != nil { + return ctrl.Result{}, err + } + if pod != nil && pod.Labels[LabelWorker] != "" && pod.Labels[LabelWorker] != sanitizeLabel(worker.Name) { + // A pod wearing this name but another worker's label is a collision, + // not something to adopt or delete. + return ctrl.Result{}, r.fail(ctx, worker, "PodConflict", + fmt.Sprintf("pod %s belongs to worker %s", podName, pod.Labels[LabelWorker])) + } + if pod != nil { + if owner := metav1.GetControllerOf(pod); owner != nil && owner.UID != worker.UID { + // A previous incarnation's pod: a worker deleted and recreated + // under the same name reaches here before garbage collection has + // caught up. Adopting it would inherit a claim seat this CR never + // booked -- the over-admission a recreate storm produces -- and + // its dying phase would be reported as this worker's. Replace it. + logger.Info("pod belongs to a previous incarnation; replacing it", "pod", podName, "worker", worker.Name) + if err := r.Delete(ctx, pod, client.Preconditions{UID: &pod.UID}); err != nil && !apierrors.IsNotFound(err) { + return ctrl.Result{}, err + } + // No requeue: the deletion event re-enqueues through Owns. + return ctrl.Result{}, r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase = openrlv1alpha1.PhasePlacing + s.Reason = "ReplacingPredecessorPod" + }) + } + } + + claimName := worker.Status.ClaimName + if pod != nil && pod.Labels[LabelClaim] != "" && pod.Labels[LabelClaim] != claimName { + // The running pod is the truth; adopting it recovers from a restart + // that lost an unpatched status. + claimName = pod.Labels[LabelClaim] + } + if claimName != "" { + if _, live := fleet.Claims[claimName]; !live { + logger.Info("assigned claim no longer exists; re-placing", "claim", claimName, "worker", worker.Name) + claimName = "" + } + } + + if claimName == "" { + claim, created, err := r.assign(ctx, worker, request, fleet) + if err != nil { + return ctrl.Result{}, err + } + if claim == nil { + reason := placement.Explain(request, fleet, "") + if r.expired(worker) { + return ctrl.Result{}, r.fail(ctx, worker, "Unsatisfiable", reason) + } + return ctrl.Result{RequeueAfter: r.RetryInterval}, r.markPending(ctx, worker, reason) + } + claimName = claim.Name + perDevice := request.PerDeviceBytes(claim.DeviceCount) + + verb := "SharedExistingClaim" + if created { + verb = fmt.Sprintf("CreatedClaim: %dx%s", claim.DeviceCount, gibQuantity(perDevice)) + } + if err := r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase = openrlv1alpha1.PhasePlacing + s.ClaimName = claimName + s.Reason = verb + recordFootprint(s, worker, request, claim.DeviceCount, perDevice) + }); err != nil { + return ctrl.Result{}, err + } + } + + if pod != nil && pod.Labels[LabelClaim] != "" && pod.Labels[LabelClaim] != claimName { + // A pod's spec.resourceClaims is immutable, so a re-placed worker's + // old pod can never reach the new claim: delete it and rebuild next + // pass. No requeue -- the deletion event re-enqueues via Owns. + logger.Info("pod is bound to a stale claim; recreating it", "pod", podName, "was", pod.Labels[LabelClaim], "now", claimName, "worker", worker.Name) + if err := r.Delete(ctx, pod); err != nil && !apierrors.IsNotFound(err) { + return ctrl.Result{}, err + } + return ctrl.Result{}, r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase = openrlv1alpha1.PhasePlacing + s.ClaimName = claimName + s.Reason = "RecreatingPodOnNewClaim" + }) + } + + if pod == nil { + if err := r.createPod(ctx, worker, podName, claimName); err != nil { + return ctrl.Result{}, err + } + return ctrl.Result{}, r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase = openrlv1alpha1.PhasePlacing + s.ClaimName = claimName + s.PodName = podName + s.Reason = "PodCreated" + }) + } + + if detail := unschedulableMessage(pod); detail != "" { + reason := placement.Explain(request, fleet, detail) + if r.expired(worker) { + return ctrl.Result{}, r.fail(ctx, worker, "Unschedulable", reason) + } + return ctrl.Result{RequeueAfter: r.RetryInterval}, r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase = openrlv1alpha1.PhasePending + s.ClaimName, s.PodName, s.Reason = claimName, podName, reason + setCondition(s, metav1.ConditionFalse, "Unschedulable", reason) + }) + } + + return ctrl.Result{}, r.reportPod(ctx, worker, fleet, pod, claimName, podName) +} + +// workerFinalizer is the deleted-worker seat guarantee: the CR -- and with +// it, the memory booking -- survives until its pod is verifiably gone. +const workerFinalizer = "openrl.io/placement" + +// teardown drives a deleting worker: delete its pod, wait for the process to +// actually exit, then let the CR go. Claims are not touched here -- they are +// shared, and the reclaim sweep owns their end of life. +func (r *OpenRLWorkerReconciler) teardown(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker) (ctrl.Result, error) { + if !controllerutil.ContainsFinalizer(worker, workerFinalizer) { + return ctrl.Result{}, nil + } + pod, err := r.findPod(ctx, workerPodName(worker)) + if err != nil { + return ctrl.Result{}, err + } + if pod != nil && pod.Labels[LabelWorker] == sanitizeLabel(worker.Name) { + if pod.DeletionTimestamp.IsZero() { + if err := r.Delete(ctx, pod); err != nil && !apierrors.IsNotFound(err) { + return ctrl.Result{}, err + } + } + // Still terminating: the seat stays booked until it is gone. + return ctrl.Result{RequeueAfter: r.RetryInterval}, nil + } + controllerutil.RemoveFinalizer(worker, workerFinalizer) + return ctrl.Result{}, r.Update(ctx, worker) +} + +// assign cuts a new claim, or joins an existing one; the returned bool says +// which. The spread-before-share policy itself lives in placement.Decide. +func (r *OpenRLWorkerReconciler) assign(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, request placement.Request, fleet *placement.Fleet) (*placement.Claim, bool, error) { + pool, join := placement.Decide(request, fleet) + if join != nil { + // Book immediately, so a worker reconciled straight after this one + // sees the seat taken. + join.Book(request.WorkerID, request.PerDeviceBytes(join.DeviceCount)) + return join, false, nil + } + if pool == nil { + return nil, false, nil + } + + perDevice := request.PerDeviceBytes(pool.DeviceCount) + claim := &placement.Claim{ + Name: claimNameFor(worker), + DeviceCount: pool.DeviceCount, + // Node stays empty until DRA decides; SizedAgainst reserves the + // pool's devices in the meantime. + SizedAgainst: pool.Node.Name, + } + claim.Book(request.WorkerID, perDevice) + + log.FromContext(ctx).Info("cutting a claim", + "claim", claim.Name, "devices", pool.DeviceCount, "perDevice", gibQuantity(perDevice), "sizedAgainst", pool.Node.Name) + + body := r.buildClaim(worker, claim, perDevice, pool.Node.DeviceMemoryBytes) + if err := r.Create(ctx, body); err != nil { + if !apierrors.IsAlreadyExists(err) { + return nil, false, fmt.Errorf("create claim %s: %w", claim.Name, err) + } + // Claim names are UID-derived, so an existing claim is this same + // incarnation's earlier create. Adopt the cluster's copy -- it may + // already be allocated. + var existing resourcev1.ResourceClaim + if err := r.fleetReader().Get(ctx, types.NamespacedName{Namespace: r.Namespace, Name: claim.Name}, &existing); err != nil { + return nil, false, fmt.Errorf("read existing claim %s: %w", claim.Name, err) + } + adopted := r.claimFrom(ctx, &existing) + if adopted == nil { + return nil, false, fmt.Errorf("claim %s exists but its shape is unreadable", claim.Name) + } + adopted.Book(request.WorkerID, perDevice) + fleet.Claims[adopted.Name] = adopted + return adopted, true, nil + } + fleet.Claims[claim.Name] = claim + return claim, true, nil +} + +func (r *OpenRLWorkerReconciler) createPod(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, podName, claimName string) error { + pod, err := r.renderPod(ctx, worker, podName, claimName) + if err != nil { + // Record the failure, then still return the error: a nil return here + // would read as "pod created" and nothing would retry the render. + if patchErr := r.fail(ctx, worker, "TemplateError", err.Error()); patchErr != nil { + return patchErr + } + return err + } + if err := r.Create(ctx, pod); err != nil && !apierrors.IsAlreadyExists(err) { + return fmt.Errorf("create pod %s: %w", podName, err) + } + return nil +} + +// findPod returns the worker's pod, or nil if it has none. A terminal pod is +// reported, not replaced: whether a finished model still wants a worker is +// the gateway's call. +func (r *OpenRLWorkerReconciler) findPod(ctx context.Context, podName string) (*corev1.Pod, error) { + var pod corev1.Pod + err := r.Get(ctx, types.NamespacedName{Namespace: r.Namespace, Name: podName}, &pod) + if apierrors.IsNotFound(err) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("read pod %s: %w", podName, err) + } + return &pod, nil +} + +func (r *OpenRLWorkerReconciler) reportPod(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, fleet *placement.Fleet, pod *corev1.Pod, claimName, podName string) error { + phase := openrlv1alpha1.PhasePlacing + reason := "" + switch pod.Status.Phase { + case corev1.PodRunning, corev1.PodSucceeded: + phase = openrlv1alpha1.PhaseRunning + case corev1.PodFailed: + phase, reason = openrlv1alpha1.PhaseFailed, "PodFailed" + } + + node := pod.Spec.NodeName + if node == "" { + if claim, ok := fleet.Claims[claimName]; ok { + node = claim.Node + } + } + + return r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase, s.ClaimName, s.PodName, s.NodeName, s.Reason = phase, claimName, podName, node, reason + if phase == openrlv1alpha1.PhaseRunning { + setCondition(s, metav1.ConditionTrue, "Placed", "worker is running on "+claimName) + } + }) +} + +// recordFootprint writes down the memory this placement implies. Parking a +// worker costs its whole accelerator footprint in host RAM, however it is +// spread, so the parked figure is the request's memory itself. +func recordFootprint(status *openrlv1alpha1.OpenRLWorkerStatus, worker *openrlv1alpha1.OpenRLWorker, request placement.Request, deviceCount int, perDevice int64) { + status.DeviceCount = int32(deviceCount) + status.MemoryPerDevice = gibQuantity(perDevice) + status.HostMemoryWhenParked = gibQuantity(request.Memory) + status.EstimatorVersion = worker.Spec.EstimatorVersion +} + +// -- status ------------------------------------------------------------------- + +func (r *OpenRLWorkerReconciler) patchStatus(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, mutate func(*openrlv1alpha1.OpenRLWorkerStatus)) error { + before := worker.Status.DeepCopy() + mutate(&worker.Status) + worker.Status.ObservedGeneration = worker.Generation + // Skip no-op writes: they would re-trigger the watch forever. + if apiequality.Semantic.DeepEqual(before, &worker.Status) { + return nil + } + if err := r.Status().Update(ctx, worker); err != nil { + if apierrors.IsConflict(err) { + // Someone else wrote first; the next reconcile recomputes from scratch. + return nil + } + return fmt.Errorf("patch status of %s: %w", worker.Name, err) + } + return nil +} + +func (r *OpenRLWorkerReconciler) markPending(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, reason string) error { + return r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase, s.Reason = openrlv1alpha1.PhasePending, reason + setCondition(s, metav1.ConditionFalse, "WaitingForCapacity", reason) + }) +} + +func (r *OpenRLWorkerReconciler) fail(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, reason, message string) error { + if r.Recorder != nil { + r.Recorder.Event(worker, corev1.EventTypeWarning, reason, message) + } + return r.patchStatus(ctx, worker, func(s *openrlv1alpha1.OpenRLWorkerStatus) { + s.Phase, s.Reason = openrlv1alpha1.PhaseFailed, message + setCondition(s, metav1.ConditionFalse, reason, message) + }) +} + +// expired reports whether this worker has been waiting past the point where +// "not yet" should be called "no". +func (r *OpenRLWorkerReconciler) expired(worker *openrlv1alpha1.OpenRLWorker) bool { + if r.PlacementTimeout <= 0 { + return false + } + since := worker.CreationTimestamp.Time + if condition := apimeta.FindStatusCondition(worker.Status.Conditions, openrlv1alpha1.ConditionPlaced); condition != nil { + since = condition.LastTransitionTime.Time + } + return time.Since(since) > r.PlacementTimeout +} + +func setCondition(status *openrlv1alpha1.OpenRLWorkerStatus, state metav1.ConditionStatus, reason, message string) { + // Kubernetes rejects condition messages over 32KiB; SetStatusCondition + // does not truncate. + if len(message) > 32768 { + message = message[:32768] + } + apimeta.SetStatusCondition(&status.Conditions, metav1.Condition{ + Type: openrlv1alpha1.ConditionPlaced, + Status: state, + Reason: reason, + Message: message, + }) +} + +// -- wiring -------------------------------------------------------------------- + +// SetupWithManager registers the reconciler and the claim-reclaim sweep. +func (r *OpenRLWorkerReconciler) SetupWithManager(mgr ctrl.Manager) error { + r.reader = mgr.GetAPIReader() + + if err := mgr.Add(manager.RunnableFunc(r.runReclaim)); err != nil { + return err + } + + // Capacity changes are fleet-wide: a freed claim might unblock any + // pending worker, so these events wake every worker rather than one. + wakeAll := handler.EnqueueRequestsFromMapFunc(func(ctx context.Context, _ client.Object) []reconcile.Request { + var workers openrlv1alpha1.OpenRLWorkerList + if err := mgr.GetClient().List(ctx, &workers, client.InNamespace(r.Namespace)); err != nil { + return nil + } + requests := make([]reconcile.Request, 0, len(workers.Items)) + for i := range workers.Items { + requests = append(requests, reconcile.Request{ + NamespacedName: types.NamespacedName{Namespace: workers.Items[i].Namespace, Name: workers.Items[i].Name}, + }) + } + return requests + }) + + // Kubelet heartbeats rewrite node status every few seconds; only label + // and allocatable-memory changes are capacity events worth waking for. + nodeCapacityChanged := predicate.Or( + predicate.LabelChangedPredicate{}, + predicate.Funcs{ + UpdateFunc: func(e event.UpdateEvent) bool { + before, okBefore := e.ObjectOld.(*corev1.Node) + after, okAfter := e.ObjectNew.(*corev1.Node) + return okBefore && okAfter && !before.Status.Allocatable.Memory().Equal(*after.Status.Allocatable.Memory()) + }, + }, + ) + + return ctrl.NewControllerManagedBy(mgr). + For(&openrlv1alpha1.OpenRLWorker{}). + Owns(&corev1.Pod{}). + Watches(&resourcev1.ResourceClaim{}, wakeAll, builder.WithPredicates(managedClaims())). + Watches(&corev1.Node{}, wakeAll, builder.WithPredicates(nodeCapacityChanged)). + WithOptions(controller.Options{MaxConcurrentReconciles: 1}). + Complete(r) +} diff --git a/controller/internal/controller/openrlworker_controller_test.go b/controller/internal/controller/openrlworker_controller_test.go new file mode 100644 index 00000000..0d2984a8 --- /dev/null +++ b/controller/internal/controller/openrlworker_controller_test.go @@ -0,0 +1,728 @@ +package controller + +import ( + "context" + "strconv" + "strings" + "testing" + "time" + + corev1 "k8s.io/api/core/v1" + resourcev1 "k8s.io/api/resource/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + clientgoscheme "k8s.io/client-go/kubernetes/scheme" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + openrlv1alpha1 "github.com/gke-labs/open-rl/controller/api/v1alpha1" +) + +const ( + testNamespace = "open-rl" + testDriver = "gpu.nvidia.com" + testNode = "node-a" + testTemplate = "trainer-pod-template" +) + +// podTemplateYAML is deliberately hostile to the controller: it pins a +// ResourceClaim and a nodeSelector of its own. Both must be overwritten, since +// picking hardware is the whole job of this controller. +const podTemplateYAML = ` +apiVersion: v1 +kind: Pod +spec: + nodeSelector: + cloud.google.com/gke-accelerator: nvidia-l4 + resourceClaims: + - name: gpu + resourceClaimName: someone-elses-claim + containers: + - name: worker + image: template-image + env: + - name: KEEP_ME + value: "1" +` + +func testScheme(t *testing.T) *runtime.Scheme { + t.Helper() + scheme := runtime.NewScheme() + if err := clientgoscheme.AddToScheme(scheme); err != nil { + t.Fatalf("add client-go scheme: %v", err) + } + if err := openrlv1alpha1.AddToScheme(scheme); err != nil { + t.Fatalf("add openrl scheme: %v", err) + } + return scheme +} + +// enabledNode is a pool the operator has opted in for both roles, described by +// a ResourceSlice with two 96Gi devices. maxWorkers is how many time-sliced +// workers the operator will let share one claim here; 1 means no sharing, which +// is also what an unlabelled node gets. +func enabledNode(maxWorkers int) []client.Object { + node := &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: testNode, + Labels: map[string]string{ + NodeLabelEnabled: "true", + NodeLabelTrainer: "true", + NodeLabelSampler: "true", + NodeLabelMaxWorkersPerClaim: strconv.Itoa(maxWorkers), + }, + }, + Status: corev1.NodeStatus{ + Allocatable: corev1.ResourceList{corev1.ResourceMemory: resource.MustParse("340Gi")}, + }, + } + + nodeName := testNode + product := "NVIDIA RTX PRO 6000 Blackwell" + device := func(name string) resourcev1.Device { + return resourcev1.Device{ + Name: name, + Attributes: map[resourcev1.QualifiedName]resourcev1.DeviceAttribute{ + "productName": {StringValue: &product}, + }, + Capacity: map[resourcev1.QualifiedName]resourcev1.DeviceCapacity{ + "memory": {Value: resource.MustParse("96Gi")}, + }, + } + } + slice := &resourcev1.ResourceSlice{ + ObjectMeta: metav1.ObjectMeta{Name: "slice-a"}, + Spec: resourcev1.ResourceSliceSpec{ + Driver: testDriver, + NodeName: &nodeName, + Pool: resourcev1.ResourcePool{Name: testNode, ResourceSliceCount: 1}, + Devices: []resourcev1.Device{device("gpu-0"), device("gpu-1")}, + }, + } + + // The key is deliberately not the controller's "pod.yaml" default: the + // deployed templates are shared with the static worker manager and keep + // its key names, so every test also exercises the lone-key fallback. + return []client.Object{node, slice, &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{Name: testTemplate, Namespace: testNamespace}, + Data: map[string]string{"trainer-worker-pod.yaml": podTemplateYAML}, + }} +} + +// worker is the whole of a request: a role, an id, how much accelerator memory +// it needs, and -- if it shares weights with anyone -- the owner it shares +// them with. Everything else the scheduler derives. +func worker(name, modelID string, role openrlv1alpha1.WorkerRole, memory string) *openrlv1alpha1.OpenRLWorker { + return &openrlv1alpha1.OpenRLWorker{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: testNamespace, + CreationTimestamp: metav1.Now(), + }, + Spec: openrlv1alpha1.OpenRLWorkerSpec{ + Role: role, + ModelID: modelID, + Memory: resource.MustParse(memory), + }, + } +} + +// trainerWorker shares nothing, so it is an owner of one: it competes for +// turns alone against every other worker on its claim. +func trainerWorker(name, modelID string) *openrlv1alpha1.OpenRLWorker { + return worker(name, modelID, openrlv1alpha1.RoleTrainer, "24Gi") +} + +// ownedWorker names the base model it serves. Workers naming the same owner +// share one fairness slot; they still take turns one at a time. +func ownedWorker(name, modelID, owner string) *openrlv1alpha1.OpenRLWorker { + w := worker(name, modelID, openrlv1alpha1.RoleTrainer, "24Gi") + w.Spec.OwnerID = owner + return w +} + +// fillerWorker occupies one whole 96Gi device, forcing later workers to +// contend for the other. +func fillerWorker(name, modelID string) *openrlv1alpha1.OpenRLWorker { + return worker(name, modelID, openrlv1alpha1.RoleTrainer, "90Gi") +} + +func newReconciler(t *testing.T, objects ...client.Object) *OpenRLWorkerReconciler { + t.Helper() + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(objects...). + WithStatusSubresource(&openrlv1alpha1.OpenRLWorker{}). + Build() + return &OpenRLWorkerReconciler{ + Client: c, + Namespace: testNamespace, + DeviceClass: testDriver, + DeviceDriver: testDriver, + DefaultPodTemplates: map[openrlv1alpha1.WorkerRole]string{ + openrlv1alpha1.RoleTrainer: testTemplate, + openrlv1alpha1.RoleSampler: testTemplate, + }, + RetryInterval: time.Second, + PlacementTimeout: time.Hour, + ReclaimInterval: time.Minute, + } +} + +func runReconcile(t *testing.T, r *OpenRLWorkerReconciler, name string) ctrl.Result { + t.Helper() + result, err := r.Reconcile(context.Background(), ctrl.Request{ + NamespacedName: types.NamespacedName{Namespace: testNamespace, Name: name}, + }) + if err != nil { + t.Fatalf("reconcile %s: %v", name, err) + } + return result +} + +// settle places the worker and creates its pod: the claim is cut on one pass +// and the pod built on the next. +func settle(t *testing.T, r *OpenRLWorkerReconciler, names ...string) { + t.Helper() + for _, name := range names { + runReconcile(t, r, name) + runReconcile(t, r, name) + } +} + +func getWorker(t *testing.T, r *OpenRLWorkerReconciler, name string) *openrlv1alpha1.OpenRLWorker { + t.Helper() + var w openrlv1alpha1.OpenRLWorker + if err := r.Get(context.Background(), types.NamespacedName{Namespace: testNamespace, Name: name}, &w); err != nil { + t.Fatalf("get worker %s: %v", name, err) + } + return &w +} + +func claimOf(t *testing.T, r *OpenRLWorkerReconciler, name string) string { + t.Helper() + claim := getWorker(t, r, name).Status.ClaimName + if claim == "" { + t.Fatalf("worker %s was not placed", name) + } + return claim +} + +func getPod(t *testing.T, r *OpenRLWorkerReconciler, name string) *corev1.Pod { + t.Helper() + var pod corev1.Pod + err := r.Get(context.Background(), types.NamespacedName{Namespace: testNamespace, Name: name}, &pod) + if apierrors.IsNotFound(err) { + return nil + } + if err != nil { + t.Fatalf("get pod %s: %v", name, err) + } + return &pod +} + +func envOf(container corev1.Container, name string) string { + for _, env := range container.Env { + if env.Name == name { + return env.Value + } + } + return "" +} + +// allocateClaim plays DRA: pins the claim to the test node, which is what +// makes it joinable. Sharing decisions only ever run against allocated claims. +func allocateClaim(t *testing.T, r *OpenRLWorkerReconciler, name string) { + t.Helper() + var claim resourcev1.ResourceClaim + if err := r.Get(context.Background(), types.NamespacedName{Namespace: testNamespace, Name: name}, &claim); err != nil { + t.Fatalf("get claim %s: %v", name, err) + } + claim.Status.Allocation = &resourcev1.AllocationResult{ + NodeSelector: &corev1.NodeSelector{NodeSelectorTerms: []corev1.NodeSelectorTerm{{ + MatchFields: []corev1.NodeSelectorRequirement{{ + Key: "metadata.name", Operator: corev1.NodeSelectorOpIn, Values: []string{testNode}, + }}, + }}}, + } + if err := r.Update(context.Background(), &claim); err != nil { + t.Fatalf("allocate claim %s: %v", name, err) + } +} + +// requiresNodeLabels reports whether some affinity term demands exactly these +// label keys set to "true". +func requiresNodeLabels(pod *corev1.Pod, keys ...string) bool { + if pod.Spec.Affinity == nil || pod.Spec.Affinity.NodeAffinity == nil || + pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil { + return false + } + for _, term := range pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms { + matched := 0 + for _, want := range keys { + for _, expr := range term.MatchExpressions { + if expr.Key == want && expr.Operator == corev1.NodeSelectorOpIn && len(expr.Values) == 1 && expr.Values[0] == "true" { + matched++ + } + } + } + if matched == len(keys) && len(term.MatchExpressions) == len(keys) { + return true + } + } + return false +} + +// A worker with nothing placed yet gets a claim and a pod, and the pod is +// rendered against the claim rather than against whatever the template said. +func TestReconcilePlacesUnplacedWorker(t *testing.T) { + r := newReconciler(t, append(enabledNode(4), trainerWorker("w-a", "model-a"))...) + + runReconcile(t, r, "w-a") + + placed := getWorker(t, r, "w-a") + if placed.Status.Phase != openrlv1alpha1.PhasePlacing { + t.Fatalf("phase = %q, want Placing", placed.Status.Phase) + } + claimName := placed.Status.ClaimName + if claimName == "" { + t.Fatal("no claim recorded on status") + } + + var claim resourcev1.ResourceClaim + if err := r.Get(context.Background(), types.NamespacedName{Namespace: testNamespace, Name: claimName}, &claim); err != nil { + t.Fatalf("claim %s was not created: %v", claimName, err) + } + if claim.Labels[LabelManaged] != "true" { + t.Errorf("claim is missing the managed label: %v", claim.Labels) + } + if claim.Labels[LabelDeviceMemory] != "96Gi" { + t.Errorf("device-memory label = %q, want the sized device's 96Gi", claim.Labels[LabelDeviceMemory]) + } + // The CEL carries both bounds: the floor is the worker's share, the + // ceiling is the device size the claim was priced against, so DRA cannot + // satisfy this claim with a bigger device placement never chose. + cel := claim.Spec.Devices.Requests[0].Exactly.Selectors[0].CEL.Expression + if !strings.Contains(cel, `quantity("24Gi")) >= 0`) || !strings.Contains(cel, `quantity("96Gi")) <= 0`) { + t.Errorf("claim CEL = %q, want a 24Gi floor and a 96Gi ceiling", cel) + } + + // The pod is created on the pass after the claim, so reconcile again. + runReconcile(t, r, "w-a") + pod := getPod(t, r, "orw-w-a") + if pod == nil { + t.Fatal("pod was not created") + } + if got := pod.Labels[LabelClaim]; got != claimName { + t.Errorf("pod claim label = %q, want %q", got, claimName) + } + if len(pod.Spec.ResourceClaims) != 1 || *pod.Spec.ResourceClaims[0].ResourceClaimName != claimName { + t.Errorf("pod resourceClaims = %+v, want the template's claim replaced by %q", pod.Spec.ResourceClaims, claimName) + } + if len(pod.Spec.NodeSelector) != 0 { + t.Errorf("nodeSelector = %v, want none: the template's pin is dropped and affinity rules instead", pod.Spec.NodeSelector) + } + // One term for explicitly-labeled trainer nodes, one for nodes naming no + // roles at all -- the documented default that both roles are allowed. + if !requiresNodeLabels(pod, NodeLabelEnabled, NodeLabelTrainer) { + t.Errorf("affinity %+v lacks the enabled+trainer term", pod.Spec.Affinity) + } + if terms := pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms; len(terms) != 2 { + t.Errorf("affinity has %d terms, want 2: the second admits role-unlabeled nodes", len(terms)) + } + + // The group is the claim -- not a cluster-wide "trainers" bucket -- and a + // worker that named no owner is an owner of one: it competes for turns + // alone, under its own name. + if pod.Labels[timeSliceEnabledLabel] != "true" || pod.Labels[timeSliceGroupLabel] != claimName { + t.Errorf("time-slice labels = %v, want enabled with group %q", pod.Labels, claimName) + } + if got := pod.Labels[timeSliceOwnerLabel]; got != "w-a" { + t.Errorf("owner label = %q, want the worker's own name", got) + } + container := pod.Spec.Containers[0] + if got := envOf(container, timeSliceGroupEnv); got != claimName { + t.Errorf("%s = %q, want %q", timeSliceGroupEnv, got, claimName) + } + if got := envOf(container, timeSliceOwnerEnv); got != "w-a" { + t.Errorf("%s = %q, want %q", timeSliceOwnerEnv, got, "w-a") + } + if got := envOf(container, "KEEP_ME"); got != "1" { + t.Errorf("the template's own env was dropped: %v", container.Env) + } +} + +// The device count is derived from memory, never asked for. There is no +// sharding: the model is laid out layer by layer over whatever it is given, so +// 120Gi on 96Gi devices needs two of them and half the model sits on each. +func TestReconcileDerivesTheDeviceCountFromMemory(t *testing.T) { + big := worker("w-big", "model-big", openrlv1alpha1.RoleTrainer, "120Gi") + r := newReconciler(t, append(enabledNode(4), big)...) + + runReconcile(t, r, "w-big") + + status := getWorker(t, r, "w-big").Status + if status.DeviceCount != 2 { + t.Errorf("deviceCount = %d, want 2: 120Gi does not fit one 96Gi device", status.DeviceCount) + } + if status.MemoryPerDevice != "60Gi" { + t.Errorf("memoryPerDevice = %q, want 60Gi", status.MemoryPerDevice) + } + // Parking moves the whole footprint to host RAM, however it was spread. + if status.HostMemoryWhenParked != "120Gi" { + t.Errorf("hostMemoryWhenParked = %q, want 120Gi", status.HostMemoryWhenParked) + } +} + +// The regression this test exists for: spec.resourceClaims is immutable, so a +// worker re-placed onto a different claim can never be reached by its existing +// pod. The controller has to delete it, and the next pass has to build one +// against the new claim. +func TestReconcileRecreatesPodBoundToAStaleClaim(t *testing.T) { + w := trainerWorker("w-a", "model-a") + w.Status = openrlv1alpha1.OpenRLWorkerStatus{ + Phase: openrlv1alpha1.PhaseRunning, + ClaimName: "claim-gone", + PodName: "orw-w-a", + } + stale := "claim-gone" + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "orw-w-a", + Namespace: testNamespace, + Labels: map[string]string{LabelClaim: stale, LabelWorker: "w-a"}, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{Name: "worker", Image: "template-image"}}, + ResourceClaims: []corev1.PodResourceClaim{{Name: podClaimName, ResourceClaimName: &stale}}, + }, + } + r := newReconciler(t, append(enabledNode(4), w, pod)...) + + runReconcile(t, r, "w-a") + + if getPod(t, r, "orw-w-a") != nil { + t.Fatal("the pod bound to the vanished claim was left in place") + } + after := getWorker(t, r, "w-a") + if after.Status.Reason != "RecreatingPodOnNewClaim" { + t.Errorf("reason = %q, want RecreatingPodOnNewClaim", after.Status.Reason) + } + if after.Status.ClaimName == stale || after.Status.ClaimName == "" { + t.Fatalf("claim = %q, want a freshly cut one", after.Status.ClaimName) + } + + // Converge: the next pass builds a pod against the claim the worker now holds. + runReconcile(t, r, "w-a") + rebuilt := getPod(t, r, "orw-w-a") + if rebuilt == nil { + t.Fatal("no pod was rebuilt") + } + if got := *rebuilt.Spec.ResourceClaims[0].ResourceClaimName; got != after.Status.ClaimName { + t.Errorf("rebuilt pod names claim %q, want %q", got, after.Status.ClaimName) + } +} + +// A worker deleted and recreated under the same name must not adopt its +// predecessor's still-terminating pod: the controller ownerRef UID tells the +// incarnations apart. Adoption would inherit a claim seat this CR never +// booked -- the over-admission a delete/recreate storm produces. +func TestReconcileReplacesAPredecessorsPod(t *testing.T) { + w := trainerWorker("w-a", "model-a") + w.UID = "new-incarnation" + isController := true + old := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "orw-w-a", + Namespace: testNamespace, + Labels: map[string]string{LabelClaim: "claim-old", LabelWorker: "w-a"}, + OwnerReferences: []metav1.OwnerReference{{ + APIVersion: "openrl.io/v1alpha1", Kind: "OpenRLWorker", + Name: "w-a", UID: "old-incarnation", Controller: &isController, + }}, + }, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "worker", Image: "template-image"}}}, + } + r := newReconciler(t, append(enabledNode(4), w, old)...) + + runReconcile(t, r, "w-a") + + if getPod(t, r, "orw-w-a") != nil { + t.Fatal("the predecessor's pod was adopted or left in place") + } + after := getWorker(t, r, "w-a") + if after.Status.Reason != "ReplacingPredecessorPod" { + t.Errorf("reason = %q, want ReplacingPredecessorPod", after.Status.Reason) + } + if after.Status.ClaimName == "claim-old" { + t.Errorf("inherited the predecessor's claim %q", after.Status.ClaimName) + } +} + +// A pod that is already on the right claim is left alone. Without this, the +// delete branch above would restart every worker on every reconcile. +func TestReconcileLeavesAMatchingPodAlone(t *testing.T) { + r := newReconciler(t, append(enabledNode(4), trainerWorker("w-a", "model-a"))...) + + settle(t, r, "w-a") + created := getPod(t, r, "orw-w-a") + if created == nil { + t.Fatal("pod was not created") + } + + for i := 0; i < 3; i++ { + runReconcile(t, r, "w-a") + } + still := getPod(t, r, "orw-w-a") + if still == nil { + t.Fatal("a settled pod was deleted") + } + if still.UID != created.UID { + t.Errorf("pod was recreated (uid %q -> %q)", created.UID, still.UID) + } +} + +// Spread onto free capacity first: while unclaimed accelerators exist, each +// worker gets its own claim, because sharing under no contention just costs +// throughput. Only when the devices run out does a worker join an existing +// claim. +func TestReconcileSpreadsThenSharesUnderContention(t *testing.T) { + r := newReconciler(t, append(enabledNode(4), + trainerWorker("w-a", "model-a"), trainerWorker("w-b", "model-b"), trainerWorker("w-c", "model-c"))...) + + runReconcile(t, r, "w-a") + allocateClaim(t, r, claimOf(t, r, "w-a")) + runReconcile(t, r, "w-b") + allocateClaim(t, r, claimOf(t, r, "w-b")) + runReconcile(t, r, "w-c") + + a, b, c := claimOf(t, r, "w-a"), claimOf(t, r, "w-b"), claimOf(t, r, "w-c") + if a == b { + t.Errorf("both trainers landed on claim %q while a device was still free", a) + } + if c != a && c != b { + t.Errorf("third worker cut claim %q, but both devices were taken and it should have shared", c) + } + + var claims resourcev1.ResourceClaimList + if err := r.List(context.Background(), &claims, client.InNamespace(testNamespace)); err != nil { + t.Fatalf("list claims: %v", err) + } + if len(claims.Items) != 2 { + t.Errorf("cut %d claims, want 2: one per device, and the third worker sharing", len(claims.Items)) + } +} + +// A trainer and a sampler on one accelerator, once the devices are contended. +// Role selects which node pools may host a worker and nothing else -- in +// particular it does not partition claims -- so on a node labelled for both, +// the two halves of the loop take turns on one GPU instead of one of them +// going unplaced. +// +// The filler holds the other device. Both existing claims hold one worker, so +// the tie breaks by name and the sampler joins the trainer's claim. +func TestReconcileRunsATrainerAndASamplerOnOneClaim(t *testing.T) { + filler := fillerWorker("w-x", "model-x") + trainer := trainerWorker("w-t", "model-t") + sampler := worker("w-s", "model-s", openrlv1alpha1.RoleSampler, "24Gi") + r := newReconciler(t, append(enabledNode(4), filler, trainer, sampler)...) + + settle(t, r, "w-x") + allocateClaim(t, r, claimOf(t, r, "w-x")) + settle(t, r, "w-t") + allocateClaim(t, r, claimOf(t, r, "w-t")) + settle(t, r, "w-s") + + claim := claimOf(t, r, "w-t") + if got := claimOf(t, r, "w-s"); got != claim { + t.Fatalf("sampler landed on %q and trainer on %q, want one claim", got, claim) + } + + // Same group, different owners: they share the bundle, not the weights. + trainerPod, samplerPod := getPod(t, r, "orw-w-t"), getPod(t, r, "orw-w-s") + if trainerPod == nil || samplerPod == nil { + t.Fatalf("pods missing: trainer=%v sampler=%v", trainerPod != nil, samplerPod != nil) + } + if trainerPod.Labels[timeSliceGroupLabel] != samplerPod.Labels[timeSliceGroupLabel] { + t.Errorf("time-slice groups differ: %q and %q", trainerPod.Labels[timeSliceGroupLabel], samplerPod.Labels[timeSliceGroupLabel]) + } + if trainerPod.Labels[timeSliceOwnerLabel] == samplerPod.Labels[timeSliceOwnerLabel] { + t.Errorf("both pods are in owner %q, but they share nothing", trainerPod.Labels[timeSliceOwnerLabel]) + } + // The sampler still may not land on a pool the operator closed to samplers. + if !requiresNodeLabels(samplerPod, NodeLabelEnabled, NodeLabelSampler) { + t.Errorf("sampler affinity %+v lacks the enabled+sampler term", samplerPod.Spec.Affinity) + } +} + +// Workers under contention, two owners between them. The owner is a string +// the caller chooses and the scheduler only ever compares: same string means +// one fairness slot at runtime, different string means separate turns. It +// never shapes placement -- the controller has no table of which workload +// kinds may share, and contention spreads to the claim with the fewest +// assigned workers. +// +// The filler takes the second device, so everything after w-a shares. +func TestReconcileGroupsWorkersByTheirOwnerID(t *testing.T) { + workers := []client.Object{ + fillerWorker("w-x", "model-x"), + ownedWorker("w-a", "model-a", "Qwen/Qwen3-0.6B"), + ownedWorker("w-b", "model-b", "Qwen/Qwen3-0.6B"), + ownedWorker("w-c", "model-c", "meta-llama/Llama-3-8B"), + } + r := newReconciler(t, append(enabledNode(4), workers...)...) + + settle(t, r, "w-x") + allocateClaim(t, r, claimOf(t, r, "w-x")) + settle(t, r, "w-a") + allocateClaim(t, r, claimOf(t, r, "w-a")) + settle(t, r, "w-b", "w-c") + + // w-b joins the least-loaded claim by name; w-c then finds the filler's + // claim emptier. Owner never enters the placement decision. + claim := claimOf(t, r, "w-a") + if b := claimOf(t, r, "w-b"); b != claim { + t.Fatalf("w-b landed on %q and w-a on %q, want the shared claim", b, claim) + } + if c := claimOf(t, r, "w-c"); c != claimOf(t, r, "w-x") { + t.Fatalf("w-c landed on %q, want the filler's claim -- the one with the fewest workers", c) + } + + owner := func(name string) string { + pod := getPod(t, r, "orw-"+name) + if pod == nil { + t.Fatalf("no pod for %s", name) + } + return pod.Labels[timeSliceOwnerLabel] + } + if a, b := owner("w-a"), owner("w-b"); a != b { + t.Errorf("owners %q and %q differ, but both named the same base model", a, b) + } + if a, c := owner("w-a"), owner("w-c"); a == c { + t.Errorf("both pods are in owner %q, but they serve different base models", a) + } +} + +// The same two trainers on a pool that seats one worker per claim get a claim +// each. openrl.io/max-workers-per-claim is what turns sharing on, and a node +// without the label defaults to no sharing rather than to unlimited. +func TestReconcileDoesNotShareWhenTheNodeSeatsOneResident(t *testing.T) { + r := newReconciler(t, append(enabledNode(1), trainerWorker("w-a", "model-a"), trainerWorker("w-b", "model-b"))...) + + runReconcile(t, r, "w-a") + runReconcile(t, r, "w-b") + + if a, b := claimOf(t, r, "w-a"), claimOf(t, r, "w-b"); a == b { + t.Errorf("both workers landed on claim %q, but the pool seats one resident", a) + } +} + +// Deleting a worker frees its seat only when its pod is verifiably gone: +// the finalizer holds the CR -- and with it the real memory booking -- +// through the pod's termination grace, so max-workers-per-claim cannot break +// for the width of the garbage-collection window. +func TestDeletedWorkerHoldsItsSeatUntilThePodIsGone(t *testing.T) { + r := newReconciler(t, append(enabledNode(1), + trainerWorker("w-a", "model-a"), trainerWorker("w-b", "model-b"), trainerWorker("w-c", "model-c"))...) + + settle(t, r, "w-a") + claimA := claimOf(t, r, "w-a") + allocateClaim(t, r, claimA) + settle(t, r, "w-b") + allocateClaim(t, r, claimOf(t, r, "w-b")) + + // Pin w-a's pod the way a kubelet mid-termination would, then delete the + // worker. The finalizer keeps the CR while the pod drains. + pod := getPod(t, r, "orw-w-a") + pod.Finalizers = append(pod.Finalizers, "test.openrl.io/hold") + if err := r.Update(context.Background(), pod); err != nil { + t.Fatal(err) + } + if err := r.Delete(context.Background(), getWorker(t, r, "w-a")); err != nil { + t.Fatal(err) + } + runReconcile(t, r, "w-a") + + if getWorker(t, r, "w-a").DeletionTimestamp.IsZero() { + t.Fatal("the worker should be terminating, held by its finalizer") + } + // Both devices claimed, both single seats booked (w-a's still counts): + // w-c has nowhere to go while the process drains. + runReconcile(t, r, "w-c") + if phase := getWorker(t, r, "w-c").Status.Phase; phase != openrlv1alpha1.PhasePending { + t.Fatalf("w-c is %q, want Pending: the terminating worker still holds its seat", phase) + } + + // The process exits: the pod goes, then the worker, then the seat. + pod = getPod(t, r, "orw-w-a") + pod.Finalizers = nil + if err := r.Update(context.Background(), pod); err != nil { + t.Fatal(err) + } + runReconcile(t, r, "w-a") + var gone openrlv1alpha1.OpenRLWorker + if err := r.Get(context.Background(), types.NamespacedName{Namespace: testNamespace, Name: "w-a"}, &gone); !apierrors.IsNotFound(err) { + t.Fatalf("worker still present after its pod died: %v", err) + } + runReconcile(t, r, "w-c") + if got := claimOf(t, r, "w-c"); got != claimA { + t.Fatalf("w-c landed on %q, want the freed seat on %q", got, claimA) + } +} + +// Zero or negative memory is a broken request, not a free placement. +func TestReconcileFailsNonPositiveMemory(t *testing.T) { + r := newReconciler(t, append(enabledNode(4), worker("w-a", "model-a", openrlv1alpha1.RoleTrainer, "0"))...) + + runReconcile(t, r, "w-a") + + after := getWorker(t, r, "w-a") + if after.Status.Phase != openrlv1alpha1.PhaseFailed { + t.Fatalf("phase = %q, want Failed for memory: 0", after.Status.Phase) + } +} + +// Claim names derive from the worker's UID -- unique per incarnation, stable +// within one -- and stay label-legal however long the worker's name is. +func TestClaimNamesAreUIDUniqueAndLabelSafe(t *testing.T) { + first := trainerWorker(strings.Repeat("w", 100), "model-a") + first.UID = "11111111-aaaa-bbbb-cccc-dddddddddddd" + name := claimNameFor(first) + if len(name) > 63 { + t.Fatalf("claim name %q is %d chars; labels cap at 63", name, len(name)) + } + if claimNameFor(first) != name { + t.Error("one incarnation must converge on one claim name") + } + + reborn := trainerWorker(strings.Repeat("w", 100), "model-a") + reborn.UID = "22222222-aaaa-bbbb-cccc-dddddddddddd" + if claimNameFor(reborn) == name { + t.Error("a recreated worker collided with its predecessor's claim name") + } +} + +// A worker asking for more than any registered pool can provide is reported as +// pending with an explanation, not silently dropped or placed anyway. +func TestReconcileReportsAnUnplaceableWorkerAsPending(t *testing.T) { + r := newReconciler(t, append(enabledNode(4), worker("w-a", "model-a", openrlv1alpha1.RoleTrainer, "4000Gi"))...) + + result := runReconcile(t, r, "w-a") + if result.RequeueAfter != r.RetryInterval { + t.Errorf("requeueAfter = %v, want %v", result.RequeueAfter, r.RetryInterval) + } + after := getWorker(t, r, "w-a") + if after.Status.Phase != openrlv1alpha1.PhasePending { + t.Fatalf("phase = %q, want Pending", after.Status.Phase) + } + if after.Status.Reason == "" { + t.Error("a pending worker should carry an explanation") + } + if getPod(t, r, "orw-w-a") != nil { + t.Error("a pod was created for a worker that was never placed") + } +} diff --git a/controller/internal/controller/pod.go b/controller/internal/controller/pod.go new file mode 100644 index 00000000..aa65aba5 --- /dev/null +++ b/controller/internal/controller/pod.go @@ -0,0 +1,309 @@ +package controller + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "regexp" + "strings" + + corev1 "k8s.io/api/core/v1" + resourcev1 "k8s.io/api/resource/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/yaml" + + openrlv1alpha1 "github.com/gke-labs/open-rl/controller/api/v1alpha1" + "github.com/gke-labs/open-rl/controller/internal/placement" +) + +// Time-slicer contract, mirrored from src/accel_timeslicer/workload.py. A pod +// label so other pods can discover the workload, an env var so the process +// itself knows which group and owner it belongs to. All of them must agree. +const ( + timeSliceEnabledLabel = "accel-timeslicer" + timeSliceGroupLabel = "timeslice.io/group" + timeSliceOwnerLabel = "timeslice.io/owner" + timeSliceJobIDLabel = "timeslice.io/job-id" + timeSliceJobIDEnv = "OPEN_RL_TIME_SLICE_JOB_ID" + timeSliceGroupEnv = "OPEN_RL_TIME_SLICE_GROUP" + timeSliceOwnerEnv = "OPEN_RL_TIME_SLICE_OWNER" +) + +// The name the pod and the claim agree to call the allocation. +const podClaimName = "gpu" + +// labelUnsafe matches everything a DNS-1123 label value may not contain. +var labelUnsafe = regexp.MustCompile(`[^a-z0-9-]+`) + +// sanitizeLabel reduces an arbitrary id to something usable as a label value. +// Empty in, empty out. +func sanitizeLabel(value string) string { + cleaned := strings.Trim(labelUnsafe.ReplaceAllString(strings.ToLower(value), "-"), "-") + if len(cleaned) > 63 { + cleaned = strings.TrimRight(cleaned[:63], "-") + } + return cleaned +} + +// workerPodName derives the pod from the CR name -- the one identity +// Kubernetes already guarantees unique. A truncated name keeps a hash of the +// full one so two long names cannot collide. +func workerPodName(worker *openrlv1alpha1.OpenRLWorker) string { + name := "orw-" + worker.Name + if len(name) <= 253 { + return name + } + sum := sha256.Sum256([]byte(name)) + return name[:245] + "-" + hex.EncodeToString(sum[:])[:7] +} + +// claimNameFor derives the claim from the worker's UID: unique per +// incarnation, so a recreated worker can never collide with -- or blindly +// adopt -- its predecessor's claim, while repeated reconciles of one +// incarnation converge on one name. (The predecessor's claim stays reusable +// the honest way: once allocated, SelectClaim can join it with real checks.) +// The worker's name rides along for operators, truncated because the claim +// name travels as a label value (the pod's time-slice group, max 63). +func claimNameFor(worker *openrlv1alpha1.OpenRLWorker) string { + name := worker.Name + if len(name) > 48 { + name = strings.TrimRight(name[:48], "-.") + } + uid := string(worker.UID) + if uid == "" { + // Objects always carry a UID in a real cluster; bare fixtures don't. + return "claim-" + name + } + if len(uid) > 8 { + uid = uid[:8] + } + return "claim-" + name + "-" + uid +} + +// buildClaim builds a ResourceClaim: the device count we decided, plus CEL +// bounds on per-device memory. The floor is the worker's share; the ceiling +// is the device size the claim was priced against -- without it, DRA could +// satisfy an L4-sized claim with an H100 and strand a later big worker. +// Deliberately no node selector: which node satisfies this is +// kube-scheduler's decision, steered by the pod. +func (r *OpenRLWorkerReconciler) buildClaim(worker *openrlv1alpha1.OpenRLWorker, claim *placement.Claim, perDeviceBytes, deviceMemoryBytes int64) *resourcev1.ResourceClaim { + // No role or owner label: which workers sit on a claim is rebuilt from + // the workers that reference it. Count and device memory are the claim's + // shape contract, checked when a create collides with an existing claim. + labels := map[string]string{ + LabelManaged: "true", + LabelAccelCount: fmt.Sprint(claim.DeviceCount), + LabelDeviceMemory: gibQuantity(deviceMemoryBytes), + LabelSizedAgainst: claim.SizedAgainst, + } + + floor := fmt.Sprintf(`device.capacity["%s"].memory.compareTo(quantity("%dGi")) >= 0 && device.capacity["%s"].memory.compareTo(quantity("%dGi")) <= 0`, + r.DeviceDriver, placement.CeilGiB(perDeviceBytes), r.DeviceDriver, placement.CeilGiB(deviceMemoryBytes)) + + return &resourcev1.ResourceClaim{ + ObjectMeta: metav1.ObjectMeta{ + Name: claim.Name, + Namespace: r.Namespace, + Labels: labels, + }, + Spec: resourcev1.ResourceClaimSpec{ + Devices: resourcev1.DeviceClaim{ + Requests: []resourcev1.DeviceRequest{{ + Name: podClaimName, + Exactly: &resourcev1.ExactDeviceRequest{ + DeviceClassName: r.DeviceClass, + Count: int64(claim.DeviceCount), + AllocationMode: resourcev1.DeviceAllocationModeExactCount, + Selectors: []resourcev1.DeviceSelector{{ + CEL: &resourcev1.CELDeviceSelector{Expression: floor}, + }}, + }, + }}, + }, + }, + } +} + +// renderPod builds the worker pod: the operator's template for everything +// placement has no opinion about, the controller's decision for everything it +// does. +func (r *OpenRLWorkerReconciler) renderPod(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker, podName, claimName string) (*corev1.Pod, error) { + pod, err := r.loadTemplate(ctx, worker) + if err != nil { + return nil, err + } + if len(pod.Spec.Containers) == 0 { + return nil, fmt.Errorf("pod template for role %s declares no containers", worker.Spec.Role) + } + + pod.Name = podName + pod.Namespace = r.Namespace + applyOverlay(&pod.Spec.Containers[0], worker.Spec.Container) + attachClaim(pod, claimName) + attachTimeSliceGroup(pod, worker, claimName) + + if pod.Labels == nil { + pod.Labels = map[string]string{} + } + pod.Labels["app"] = "open-rl-" + string(worker.Spec.Role) + "-worker" + pod.Labels[LabelClaim] = claimName + // Label values cap at 63 characters and forbid dots; names do neither. + // Labels carry sanitized identities, env vars carry the full ones. + pod.Labels[LabelWorker] = sanitizeLabel(worker.Name) + pod.Labels[LabelRole] = string(worker.Spec.Role) + + // Constraining the pod is what constrains where its claim can land; + // whatever SKU or affinity the template pinned is dropped on purpose. + // Two ORed terms, because a node selector cannot say "or unlabeled": + // the documented default is that a node naming no role labels takes both. + pod.Spec.NodeSelector = nil + pod.Spec.Affinity = &corev1.Affinity{NodeAffinity: &corev1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ + NodeSelectorTerms: []corev1.NodeSelectorTerm{ + {MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: NodeLabelEnabled, Operator: corev1.NodeSelectorOpIn, Values: []string{"true"}}, + {Key: nodeRoleLabel[worker.Spec.Role], Operator: corev1.NodeSelectorOpIn, Values: []string{"true"}}, + }}, + {MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: NodeLabelEnabled, Operator: corev1.NodeSelectorOpIn, Values: []string{"true"}}, + {Key: NodeLabelTrainer, Operator: corev1.NodeSelectorOpDoesNotExist}, + {Key: NodeLabelSampler, Operator: corev1.NodeSelectorOpDoesNotExist}, + }}, + }, + }, + }} + + if err := controllerutil.SetControllerReference(worker, pod, r.Scheme()); err != nil { + return nil, fmt.Errorf("set owner of pod %s: %w", podName, err) + } + return pod, nil +} + +// loadTemplate reads the pod YAML the worker asked for, or the controller's +// role default. +func (r *OpenRLWorkerReconciler) loadTemplate(ctx context.Context, worker *openrlv1alpha1.OpenRLWorker) (*corev1.Pod, error) { + name, key := r.DefaultPodTemplates[worker.Spec.Role], "pod.yaml" + if ref := worker.Spec.PodTemplate; ref != nil { + name = ref.Name + if ref.Key != "" { + key = ref.Key + } + } + if name == "" { + return nil, fmt.Errorf("no pod template configured for role %s and none given in spec.podTemplate", worker.Spec.Role) + } + + var cm corev1.ConfigMap + if err := r.Get(ctx, types.NamespacedName{Namespace: r.Namespace, Name: name}, &cm); err != nil { + return nil, fmt.Errorf("read pod template ConfigMap %s: %w", name, err) + } + raw, ok := cm.Data[key] + if !ok && len(cm.Data) == 1 { + // A single-key ConfigMap is unambiguous; the deployed templates keep + // the static worker manager's key names. + for _, raw = range cm.Data { + ok = true + } + } + if !ok { + return nil, fmt.Errorf("pod template ConfigMap %s has no key %q and does not have exactly one key", name, key) + } + + var pod corev1.Pod + if err := yaml.Unmarshal([]byte(raw), &pod); err != nil { + return nil, fmt.Errorf("parse pod template %s/%s: %w", name, key, err) + } + return &pod, nil +} + +// applyOverlay stamps the gateway's per-model decisions onto the container, +// so the controller never reads the gateway's metadata store. +func applyOverlay(container *corev1.Container, overlay *openrlv1alpha1.ContainerOverlay) { + if overlay == nil { + return + } + if overlay.Image != "" { + container.Image = overlay.Image + } + if len(overlay.Command) > 0 { + container.Command = overlay.Command + } + container.Args = append(container.Args, overlay.Args...) + for _, env := range overlay.Env { + setEnv(container, env) + } +} + +// setEnv merges one variable by name, overwriting whatever the template had. +func setEnv(container *corev1.Container, want corev1.EnvVar) { + for i := range container.Env { + if container.Env[i].Name == want.Name { + container.Env[i] = want + return + } + } + container.Env = append(container.Env, want) +} + +// attachClaim points the pod at a specific ResourceClaim, replacing whatever +// the template carried -- two GPU claims would pin the pod to the +// intersection of two allocations. +func attachClaim(pod *corev1.Pod, claimName string) { + pod.Spec.ResourceClaims = []corev1.PodResourceClaim{{ + Name: podClaimName, + ResourceClaimName: &claimName, + }} + for i := range pod.Spec.Containers { + pod.Spec.Containers[i].Resources.Claims = []corev1.ResourceClaim{{Name: podClaimName}} + } +} + +// attachTimeSliceGroup tells the node-local time-slicer which accelerator +// bundle this worker shares (group = the claim, so unrelated allocations +// never wait on each other) and which owner it is served under (turns rotate +// between owners; one worker resident at a time regardless). +func attachTimeSliceGroup(pod *corev1.Pod, worker *openrlv1alpha1.OpenRLWorker, claimName string) { + if pod.Labels == nil { + pod.Labels = map[string]string{} + } + // The time-slice job id is the placement WorkerID: one identity formula, + // so the booking key and the runtime key can never drift apart. Env vars + // carry the exact values; the labels are sanitized copies for discovery. + request := requestFrom(worker) + jobID := request.WorkerID + owner := request.OwnerKey() + + pod.Labels[timeSliceEnabledLabel] = "true" + pod.Labels[timeSliceGroupLabel] = claimName + pod.Labels[timeSliceOwnerLabel] = sanitizeLabel(owner) + pod.Labels[timeSliceJobIDLabel] = sanitizeLabel(jobID) + for i := range pod.Spec.Containers { + setEnv(&pod.Spec.Containers[i], corev1.EnvVar{Name: timeSliceGroupEnv, Value: claimName}) + setEnv(&pod.Spec.Containers[i], corev1.EnvVar{Name: timeSliceOwnerEnv, Value: owner}) + setEnv(&pod.Spec.Containers[i], corev1.EnvVar{Name: timeSliceJobIDEnv, Value: jobID}) + } +} + +// unschedulableMessage is the scheduler's reason a pod cannot be placed, if it +// gave one. +func unschedulableMessage(pod *corev1.Pod) string { + if pod.Status.Phase != corev1.PodPending && pod.Status.Phase != "" { + return "" + } + for _, condition := range pod.Status.Conditions { + if condition.Type != corev1.PodScheduled || condition.Status != corev1.ConditionFalse { + continue + } + detail := condition.Message + if detail == "" { + detail = condition.Reason + } + if detail != "" { + return "Unschedulable: " + detail + } + } + return "" +} diff --git a/controller/internal/controller/reclaim.go b/controller/internal/controller/reclaim.go new file mode 100644 index 00000000..dbf83ad8 --- /dev/null +++ b/controller/internal/controller/reclaim.go @@ -0,0 +1,104 @@ +package controller + +import ( + "context" + "time" + + corev1 "k8s.io/api/core/v1" + resourcev1 "k8s.io/api/resource/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/predicate" + + openrlv1alpha1 "github.com/gke-labs/open-rl/controller/api/v1alpha1" +) + +// claimGracePeriod is how long a newly cut claim is safe from the reclaim +// sweep. A claim is created before the worker status that names it, so +// without a grace period the sweep could delete a claim out from under a +// worker that is still being placed. +const claimGracePeriod = 2 * time.Minute + +// managedClaims restricts the ResourceClaim watch to the ones this controller +// created, so unrelated DRA traffic does not wake every worker. +func managedClaims() predicate.Predicate { + return predicate.NewPredicateFuncs(func(obj client.Object) bool { + return obj.GetLabels()[LabelManaged] == "true" + }) +} + +// runReclaim sweeps idle claims until the manager stops. +func (r *OpenRLWorkerReconciler) runReclaim(ctx context.Context) error { + ticker := time.NewTicker(r.ReclaimInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return nil + case <-ticker.C: + if err := r.reclaimIdleClaims(ctx); err != nil { + log.FromContext(ctx).Error(err, "reclaim sweep failed") + } + } + } +} + +// reclaimIdleClaims deletes managed claims that no longer back any worker. +// Claims carry no owner reference -- a shared claim belongs to no single +// worker, so GC'ing it with its creator would pull the allocation out from +// under everyone else -- which is why this sweep must exist. Four stays of +// execution: an OpenRLWorker names it, a live pod sits on it, DRA still +// reserves it, or it is too young to judge. +func (r *OpenRLWorkerReconciler) reclaimIdleClaims(ctx context.Context) error { + logger := log.FromContext(ctx) + + var claims resourcev1.ResourceClaimList + if err := r.List(ctx, &claims, client.InNamespace(r.Namespace), client.MatchingLabels{LabelManaged: "true"}); err != nil { + return err + } + if len(claims.Items) == 0 { + return nil + } + + spokenFor := map[string]bool{} + + var workers openrlv1alpha1.OpenRLWorkerList + if err := r.fleetReader().List(ctx, &workers, client.InNamespace(r.Namespace)); err != nil { + return err + } + for i := range workers.Items { + if name := workers.Items[i].Status.ClaimName; name != "" { + spokenFor[name] = true + } + } + + var pods corev1.PodList + if err := r.List(ctx, &pods, client.InNamespace(r.Namespace), client.HasLabels{LabelClaim}); err != nil { + return err + } + for i := range pods.Items { + pod := &pods.Items[i] + if pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { + continue + } + spokenFor[pod.Labels[LabelClaim]] = true + } + + for i := range claims.Items { + claim := &claims.Items[i] + switch { + case spokenFor[claim.Name]: + continue + case len(claim.Status.ReservedFor) > 0: + continue + case time.Since(claim.CreationTimestamp.Time) < claimGracePeriod: + continue + } + logger.Info("reclaiming idle claim", "claim", claim.Name) + if err := r.Delete(ctx, claim); err != nil && !apierrors.IsNotFound(err) { + logger.Error(err, "failed to delete idle claim", "claim", claim.Name) + } + } + return nil +} diff --git a/controller/internal/placement/behavior_test.go b/controller/internal/placement/behavior_test.go new file mode 100644 index 00000000..1a25c5c0 --- /dev/null +++ b/controller/internal/placement/behavior_test.go @@ -0,0 +1,246 @@ +// What we expect of placement, end to end, played through Decide the way the +// controller plays it -- with DRA's answer arriving instantly. The memory +// figures are the estimator's real outputs: the gateway's tier table says +// 10Gi for the L4 tier and 60Gi for the 80Gi tier, and the hardware shapes +// are the ones we run (2x L4 24Gi dev box, single 80Gi devices). +// +// - free GPUs are used before anyone shares +// - workers share a claim whether or not their sums fit; only independent +// fit matters +// - a full claim makes the next worker wait, and a departure frees the seat +// - role labels are policy: they hide hardware from workers +// - host memory bounds how many suspended workers a node carries +// - a worker no single node can hold stays waiting, with a reason +// - placed workers stay put when capacity frees; new capacity serves new +// work +// +// Who is resident at any moment is the timeslicer's job, tested in +// tests/test_accel_timeslicer.py. The pipeline against a real API server is +// hack/kind-smoke.sh. +package placement + +import ( + "strings" + "testing" +) + +// cluster drives Decide exactly as the controller's assign step does, with +// the one simulation liberty that a new claim binds to its pool immediately +// (the real one waits for DRA; pending claims are covered by the unit tests). +type cluster struct { + t *testing.T + fleet *Fleet + placed map[string]*Claim +} + +func newCluster(t *testing.T, nodes ...*Node) *cluster { + t.Helper() + fleet := NewFleet() + for _, n := range nodes { + fleet.Nodes[n.Name] = n + } + return &cluster{t: t, fleet: fleet, placed: map[string]*Claim{}} +} + +// arrive places one worker and reports where it landed: a claim name, or "" +// while it waits. +func (c *cluster) arrive(req Request) string { + c.t.Helper() + pool, join := Decide(req, c.fleet) + switch { + case pool != nil: + claim := &Claim{Name: "claim-" + req.WorkerID, DeviceCount: pool.DeviceCount, Node: pool.Node.Name} + claim.Book(req.WorkerID, req.PerDeviceBytes(pool.DeviceCount)) + c.fleet.Claims[claim.Name] = claim + c.placed[req.WorkerID] = claim + return claim.Name + case join != nil: + join.Book(req.WorkerID, req.PerDeviceBytes(join.DeviceCount)) + c.placed[req.WorkerID] = join + return join.Name + } + return "" +} + +// leave releases a worker's seat; an emptied claim returns its accelerators. +func (c *cluster) leave(workerID string) { + c.t.Helper() + claim := c.placed[workerID] + if claim == nil { + c.t.Fatalf("%s was never placed", workerID) + } + claim.Release(workerID) + delete(c.placed, workerID) + if claim.Workers() == 0 { + delete(c.fleet.Claims, claim.Name) + } +} + +// waitingReason is why an unplaced worker is waiting. +func (c *cluster) waitingReason(req Request) string { + return Explain(req, c.fleet, "") +} + +func l4Box(name string) *Node { + return &Node{ + Name: name, DeviceCount: 2, DeviceMemoryBytes: gib(24), HostMemoryBytes: gib(94), + Roles: map[string]bool{"trainer": true, "sampler": true}, MaxWorkersPerClaim: 4, Product: "NVIDIA L4", + } +} + +func gpu80(name string, maxWorkers int, roles ...string) *Node { + allowed := map[string]bool{} + for _, role := range roles { + allowed[role] = true + } + return &Node{ + Name: name, DeviceCount: 1, DeviceMemoryBytes: gib(80), HostMemoryBytes: gib(180), + Roles: allowed, MaxWorkersPerClaim: maxWorkers, + } +} + +func TestWorkersSpreadAcrossFreeGPUs(t *testing.T) { + c := newCluster(t, gpu80("node-a", 2, "trainer"), gpu80("node-b", 2, "trainer")) + + // Two 80Gi-tier estimates (the gateway's table says 60Gi) while both + // GPUs are free: a claim each, on different nodes. + a := c.arrive(trainer("w1", 60)) + b := c.arrive(trainer("w2", 60)) + if a == "" || b == "" || a == b { + t.Fatalf("placed on %q and %q, want separate claims while devices are free", a, b) + } + if c.placed["w1"].Node == c.placed["w2"].Node { + t.Errorf("both claims landed on %s, want separate nodes", c.placed["w1"].Node) + } +} + +// An FFT trainer and its sampler under contention: whether their sum would +// fit (45+25 on 80Gi) or not (55+35), V1 places them identically, because +// each only has to fit alone -- they take exclusive turns at runtime. +func TestWorkersShareAClaimWhetherOrNotTheirSumFits(t *testing.T) { + for name, pair := range map[string][2]int64{ + "sum fits": {45, 25}, + "sum does not": {55, 35}, + } { + t.Run(name, func(t *testing.T) { + c := newCluster(t, gpu80("gpu", 2, "trainer", "sampler")) + first := c.arrive(trainer("trainer", pair[0])) + second := c.arrive(Request{Role: "sampler", WorkerID: "sampler", Memory: gib(pair[1])}) + if first == "" || first != second { + t.Fatalf("placed on %q and %q, want one shared claim", first, second) + } + }) + } +} + +// max-workers-per-claim is a seat count, not a memory rule: the third worker +// fits the GPU but waits, and takes the seat the moment one frees. +func TestAFullClaimMakesTheNextWorkerWaitForASeat(t *testing.T) { + c := newCluster(t, gpu80("gpu", 2, "trainer")) + c.arrive(trainer("j1", 50)) + shared := c.arrive(trainer("j2", 50)) + + if got := c.arrive(trainer("j3", 20)); got != "" { + t.Fatalf("j3 was seated on %q, but the claim is at max-workers-per-claim", got) + } + c.leave("j1") + if got := c.arrive(trainer("j3", 20)); got != shared { + t.Fatalf("j3 landed on %q, want the seat freed on %q -- the claim outlives its first worker", got, shared) + } +} + +// Labels decide what OpenRL may use: a free GPU on a trainer-only node is +// invisible to a sampler, and the reason names the policy. +func TestRoleLabelsHideFreeHardware(t *testing.T) { + c := newCluster(t, gpu80("trainer-only", 2, "trainer")) + + sampler := Request{Role: "sampler", WorkerID: "sampler", Memory: gib(10)} + if got := c.arrive(sampler); got != "" { + t.Fatalf("sampler landed on %q, but the operator allowed no sampler nodes", got) + } + if reason := c.waitingReason(sampler); !strings.Contains(reason, "no enabled node accepts sampler") { + t.Errorf("reason = %q, want the policy named, not capacity", reason) + } +} + +// GPU memory fits but host memory does not: suspended workers park in host +// RAM, so a thin-host node refuses the worker that would overflow it -- until +// a departure frees the memory. +func TestHostMemoryBoundsAdmission(t *testing.T) { + thin := gpu80("thin-host", 8, "trainer") + thin.HostMemoryBytes = gib(40) // budget 34Gi after headroom + c := newCluster(t, thin) + + c.arrive(trainer("w1", 20)) + if got := c.arrive(trainer("w2", 20)); got == "" { + t.Fatal("w2 parks only 20Gi, inside the budget, and should have joined") + } + if got := c.arrive(trainer("w3", 20)); got != "" { + t.Fatalf("w3 was seated on %q, but parking two 20Gi workers exceeds the 34Gi budget", got) + } + c.leave("w1") + if got := c.arrive(trainer("w3", 20)); got == "" { + t.Fatal("w3 still refused after the departure freed host memory") + } +} + +// Capacity that exists only across nodes places nothing; the same memory on +// one node places as a multi-device claim. +func TestAWorkerNoSingleNodeCanHoldStaysWaiting(t *testing.T) { + across := newCluster(t, gpu80("node-a", 2, "trainer"), gpu80("node-b", 2, "trainer")) + big := trainer("big", 140) + if got := across.arrive(big); got != "" { + t.Fatalf("140Gi landed on %q, but no single node can hold it", got) + } + if reason := across.waitingReason(big); !strings.Contains(reason, "NoCapacity") { + t.Errorf("reason = %q, want NoCapacity: waiting would never help", reason) + } + + oneNode := newCluster(t, &Node{ + Name: "wide", DeviceCount: 2, DeviceMemoryBytes: gib(80), HostMemoryBytes: gib(360), + Roles: map[string]bool{"trainer": true}, MaxWorkersPerClaim: 2, + }) + if oneNode.arrive(big) == "" || oneNode.placed["big"].DeviceCount != 2 { + t.Fatalf("placed = %+v, want a two-device claim on the one node that fits it", oneNode.placed["big"]) + } +} + +// V1 does not rebalance: the pair placed under contention keeps sharing after +// a GPU frees, and the freed GPU serves new work instead. +func TestPlacedWorkersStayPutWhenAGPUFrees(t *testing.T) { + c := newCluster(t, gpu80("gpu-a", 2, "trainer"), gpu80("gpu-b", 2, "trainer")) + c.arrive(trainer("z-blocker", 60)) + shared := c.arrive(trainer("w1", 50)) + if got := c.arrive(trainer("w2", 50)); got != shared { + t.Fatalf("w2 on %q, want it sharing %q while both GPUs were taken", got, shared) + } + + c.leave("z-blocker") + if got := c.placed["w2"]; got.Name != shared { + t.Fatalf("w2 moved to %q; V1 never migrates a placed worker", got.Name) + } + fresh := c.arrive(trainer("fresh", 50)) + if fresh == "" || c.placed["fresh"].Node != "gpu-a" { + t.Fatalf("fresh = %+v, want it on the GPU the blocker freed", c.placed["fresh"]) + } +} + +// The dev box, using the estimator's own tier outputs: 10Gi lora-tier +// workers on 2x L4. Spread first, then share, then wait for a seat. +func TestTierTableWorkloadsOnTheDevBox(t *testing.T) { + c := newCluster(t, l4Box("box")) + + a := c.arrive(trainer("lora-1", 10)) + b := c.arrive(trainer("lora-2", 10)) + if a == b { + t.Fatalf("both loras landed on %q while an L4 was free", a) + } + // The 80Gi tier's 60Gi estimate does not fit any L4, whole box or not. + if got := c.arrive(trainer("fft-8b", 60)); got != "" { + t.Fatalf("a 60Gi estimate landed on %q, but the box's devices are 24Gi", got) + } + // More loras double up on the existing claims instead. + if got := c.arrive(trainer("lora-3", 10)); got != a { + t.Fatalf("lora-3 landed on %q, want the fewest-workers claim %q", got, a) + } +} diff --git a/controller/internal/placement/placement.go b/controller/internal/placement/placement.go new file mode 100644 index 00000000..c8a0e867 --- /dev/null +++ b/controller/internal/placement/placement.go @@ -0,0 +1,332 @@ +// Package placement is the scheduling decision: pure functions over a Request +// and a Fleet, no Kubernetes imports. A claim is a bundle of accelerators; +// several workers may be assigned to it; exactly one is resident at a time. +package placement + +import ( + "fmt" + "sort" +) + +// GiB is the unit every memory figure here is reported in. +const GiB int64 = 1 << 30 + +// CeilGiB rounds a byte count up to whole GiB. The one rounding rule for +// every figure the scheduler reports or writes into a CEL selector. +func CeilGiB(bytes int64) int64 { + return (bytes + GiB - 1) / GiB +} + +// HostMemoryHeadroom is the share of a node's allocatable memory left for the +// kubelet, the DRA driver, page cache, and the resident worker's own host-side +// allocations. Only the remainder may hold suspended workers. +const HostMemoryHeadroom = 0.15 + +// Node is one accelerator pool: hardware from the driver's ResourceSlice, +// policy from the operator's node labels. +type Node struct { + Name string + // DeviceCount and DeviceMemoryBytes come from the DRA driver. + DeviceCount int + DeviceMemoryBytes int64 + // HostMemoryBytes is the node's allocatable memory, which bounds how many + // workers can be parked here at once. + HostMemoryBytes int64 + // Roles is the set of worker roles the operator allowed on this pool. + Roles map[string]bool + // MaxWorkersPerClaim is the openrl.io/max-workers-per-claim ceiling. It is + // a policy cap on queueing and switch overhead, not a capacity check, and + // it never permits multiple residents. + MaxWorkersPerClaim int + Product string +} + +// Accepts reports whether the operator allowed this role on this pool. +func (n *Node) Accepts(role string) bool { return n.Roles[role] } + +// Describe renders the pool's hardware for an error message. +func (n *Node) Describe() string { + hardware := fmt.Sprintf("%dGi x %d", n.DeviceMemoryBytes/GiB, n.DeviceCount) + if n.Product == "" { + return hardware + } + return hardware + " " + n.Product +} + +// HostBudget is how much of this node's memory may hold suspended workers. +// Suspension does not spill to disk -- it parks device memory in the process's +// own host address space -- so exceeding this OOM-kills the node. Zero means +// the node did not report allocatable memory; the check is skipped rather +// than guessed at. +func (n *Node) HostBudget() int64 { + return int64(float64(n.HostMemoryBytes) * (1 - HostMemoryHeadroom)) +} + +// Claim is a ResourceClaim, plus what is already sitting on it. Claims are +// not partitioned by role or workload type: anything the node accepts may +// join and take turns. +type Claim struct { + Name string + DeviceCount int + // Node is where the claim was allocated, empty until DRA has decided. + Node string + // SizedAgainst is the pool a not-yet-allocated claim was cut for. Until + // DRA decides, the claim reserves devices there, so a burst stops cutting + // claims once a pool's devices are spoken for and starts joining instead. + SizedAgainst string + // booked is per-device bytes per assigned worker -- deliberately not one + // total, because nothing is ever summed against the device. + booked map[string]int64 +} + +// Allocated reports whether the scheduler has said where this claim landed. +func (c *Claim) Allocated() bool { return c.Node != "" } + +// Workers is how many workers are assigned to this claim. +func (c *Claim) Workers() int { return len(c.booked) } + +// Book accepts a placement: one more assigned worker and its per-device bytes. +func (c *Claim) Book(workerID string, perDeviceBytes int64) { + if c.booked == nil { + c.booked = map[string]int64{} + } + c.booked[workerID] = perDeviceBytes +} + +// Release gives back a worker's seat and the memory that came with it. +func (c *Claim) Release(workerID string) { + delete(c.booked, workerID) +} + +// ParkedBytesWith is the host memory this claim's suspended workers would +// hold if one more worker of perDeviceBytes joined: the conservative case is +// the smallest worker resident and every other one parked, each holding what +// it had on every device. +func (c *Claim) ParkedBytesWith(perDeviceBytes int64) int64 { + total, smallest := perDeviceBytes, perDeviceBytes + for _, bytes := range c.booked { + total += bytes + if bytes < smallest { + smallest = bytes + } + } + return (total - smallest) * int64(c.DeviceCount) +} + +// Fleet is everything placement decides against. +type Fleet struct { + Nodes map[string]*Node + Claims map[string]*Claim +} + +// NewFleet returns an empty Fleet. +func NewFleet() *Fleet { + return &Fleet{Nodes: map[string]*Node{}, Claims: map[string]*Claim{}} +} + +// FreeDevices is how many of a node's accelerators no claim has taken yet. +// An unallocated claim counts against the pool it was sized for: without that +// reservation, a burst reconciling before DRA decides would cut one claim per +// worker and leave most of them unsatisfiable. +func (f *Fleet) FreeDevices(node *Node) int { + free := node.DeviceCount + for _, claim := range f.Claims { + if claim.Node == node.Name || (!claim.Allocated() && claim.SizedAgainst == node.Name) { + free -= claim.DeviceCount + } + } + return free +} + +// Request is one worker's needs, parsed out of its spec once. +type Request struct { + // Role selects node pools and nothing else; it does not partition claims. + Role string + // Memory is the total accelerator memory the worker needs, across however + // many devices it ends up on. + Memory int64 + // Owner is the runtime fairness unit. Placement never reads it; the + // timeslicer does. + Owner string + // WorkerID identifies the worker. Required, and required to be unique. + WorkerID string +} + +// OwnerKey is the fairness unit the timeslicer serves this worker under; a +// worker naming no owner is an owner of one. +func (r Request) OwnerKey() string { + if r.Owner != "" { + return r.Owner + } + return r.WorkerID +} + +// DevicesOn is how many of a node's devices this workload needs, or 0 if the +// pool cannot hold it. Plain ceiling division: there is no sharding, so any +// count will do and aggregate memory is the only thing that has to add up. +func (r Request) DevicesOn(n *Node) int { + if n.DeviceMemoryBytes <= 0 { + return 0 + } + count := int((r.Memory + n.DeviceMemoryBytes - 1) / n.DeviceMemoryBytes) + if count < 1 { + count = 1 + } + if count > n.DeviceCount { + return 0 + } + return count +} + +// PerDeviceBytes is the workload's share of each device when spread over +// deviceCount of them. An even split that a layer-wise layout only +// approximates; erring high is the safe direction. +func (r Request) PerDeviceBytes(deviceCount int) int64 { + if deviceCount < 1 { + panic(fmt.Sprintf("deviceCount must be >= 1, got %d", deviceCount)) + } + return (r.Memory + int64(deviceCount) - 1) / int64(deviceCount) +} + +// candidateNodes is every pool that accepts the role and fits the workload, +// with the device count it would take there. +func candidateNodes(req Request, fleet *Fleet) map[string]int { + fits := map[string]int{} + for name, node := range fleet.Nodes { + if !node.Accepts(req.Role) { + continue + } + if count := req.DevicesOn(node); count > 0 { + fits[name] = count + } + } + return fits +} + +// Decide is the placement policy in one sentence: spread onto a free pool +// while one exists, share an existing claim only under contention. Exactly +// one result is non-nil, or both are nil when nothing will have the worker. +func Decide(req Request, fleet *Fleet) (*Pool, *Claim) { + if pool := ChoosePool(req, fleet); pool != nil { + return pool, nil + } + return nil, SelectClaim(req, fleet) +} + +// SelectClaim is the claim this worker should join, or nil if none will have +// it -- the sharing half of Decide. +// +// Only allocated claims are joinable: an unallocated claim's node is unknown, +// so nothing about it -- device memory, host memory, role policy, the worker +// ceiling -- can be checked against anything real. A burst does not scatter +// while it waits, because pending claims reserve their pool (FreeDevices); +// the joiners simply retry once DRA has decided. +// +// Three checks govern joining: the worker fits a device by itself (nothing +// is summed -- only one worker is ever resident), the claim is below +// max-workers-per-claim, and the node has host memory for the workers that +// may be parked. Preference: fewest assigned workers, then name, so the +// choice is deterministic. +func SelectClaim(req Request, fleet *Fleet) *Claim { + var best *Claim + for _, claim := range fleet.Claims { + node := fleet.Nodes[claim.Node] + if !claim.Allocated() || node == nil || !node.Accepts(req.Role) || req.DevicesOn(node) != claim.DeviceCount { + continue + } + if claim.Workers() >= node.MaxWorkersPerClaim { + continue + } + perDevice := req.PerDeviceBytes(claim.DeviceCount) + if perDevice > node.DeviceMemoryBytes { + continue + } + // A zero budget means the node reported no allocatable memory; skip + // the check rather than refuse every claim on it. + if budget := node.HostBudget(); budget > 0 && claim.ParkedBytesWith(perDevice) > budget { + continue + } + if best == nil || claim.Workers() < best.Workers() || + (claim.Workers() == best.Workers() && claim.Name < best.Name) { + best = claim + } + } + return best +} + +// Pool is the node a new claim is sized against and how wide it would be. +type Pool struct { + Node *Node + DeviceCount int +} + +// ChoosePool picks the pool a new claim is sized against -- best fit by +// wasted memory among pools with unclaimed accelerators -- or nil if none has +// room. Sizing only: DRA picks where the claim actually lands. +func ChoosePool(req Request, fleet *Fleet) *Pool { + var best *Pool + var bestWaste int64 + names := make([]string, 0, len(fleet.Nodes)) + for name := range fleet.Nodes { + names = append(names, name) + } + sort.Strings(names) + + for _, name := range names { + node := fleet.Nodes[name] + if !node.Accepts(req.Role) { + continue + } + count := req.DevicesOn(node) + if count == 0 || count > fleet.FreeDevices(node) { + continue + } + waste := int64(count)*node.DeviceMemoryBytes - req.Memory + if best == nil || waste < bestWaste || (waste == bestWaste && count < best.DeviceCount) { + best, bestWaste = &Pool{Node: node, DeviceCount: count}, waste + } + } + return best +} + +// Explain says why this worker is not running. "Busy, retry" and "too small, +// never" are deliberately different answers; detail carries the caller's own +// words about what failed. +func Explain(req Request, fleet *Fleet, detail string) string { + var pools []*Node + for _, node := range fleet.Nodes { + if node.Accepts(req.Role) { + pools = append(pools, node) + } + } + + var reason string + switch { + case len(pools) == 0: + reason = fmt.Sprintf("NoCapacity: no enabled node accepts %s workers", req.Role) + case len(candidateNodes(req, fleet)) > 0: + // The hardware exists; it is busy. Retrying is the right move. + reason = "WaitingForCapacity: a pool fits this workload but none has a free seat or a free accelerator" + default: + biggest := pools[0] + for _, node := range pools[1:] { + if int64(node.DeviceCount)*node.DeviceMemoryBytes > int64(biggest.DeviceCount)*biggest.DeviceMemoryBytes { + biggest = node + } + } + reason = fmt.Sprintf("NoCapacity: needs %dGi in total; largest pool offers %s", + CeilGiB(req.Memory), biggest.Describe()) + } + + if detail != "" { + reason += ". " + detail + } + return truncate(reason, 1024) +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] +} diff --git a/controller/internal/placement/placement_test.go b/controller/internal/placement/placement_test.go new file mode 100644 index 00000000..e7352c79 --- /dev/null +++ b/controller/internal/placement/placement_test.go @@ -0,0 +1,255 @@ +// The spec's end-to-end scenarios are tested in internal/sim (sim_test.go), +// which drives these same functions. What lives here is only what the sim +// cannot reach: the burst window where claims exist but DRA has not +// allocated them (the sim binds immediately), the deterministic orderings, +// and the arithmetic the spec pins in its appendix. +package placement + +import ( + "strings" + "testing" +) + +func gib(n int64) int64 { return n * GiB } + +// l4Node is the dev box: g2-standard-24 with 2x L4 24Gi and 94Gi allocatable. +func l4Node(name string, maxWorkers int, roles ...string) *Node { + allowed := map[string]bool{} + for _, role := range roles { + allowed[role] = true + } + return &Node{ + Name: name, + DeviceCount: 2, + DeviceMemoryBytes: gib(24), + HostMemoryBytes: gib(94), + Roles: allowed, + MaxWorkersPerClaim: maxWorkers, + Product: "NVIDIA L4", + } +} + +// bigNode is a pool of any shape, for the cases where a workload has to choose. +func bigNode(name string, devices int, deviceGiB int64, maxWorkers int, roles ...string) *Node { + allowed := map[string]bool{} + for _, role := range roles { + allowed[role] = true + } + return &Node{ + Name: name, + DeviceCount: devices, + DeviceMemoryBytes: gib(deviceGiB), + HostMemoryBytes: gib(340), + Roles: allowed, + MaxWorkersPerClaim: maxWorkers, + } +} + +func trainer(id string, memoryGiB int64) Request { + return Request{Role: "trainer", WorkerID: id, Memory: gib(memoryGiB)} +} + +// booked is a claim with workers already assigned to it, charged worker by +// worker the way bookWorker rebuilds one from the workers that reference it. +func booked(c *Claim, perDevice int64, workers ...string) *Claim { + for _, worker := range workers { + c.Book(worker, perDevice) + } + return c +} + +func name(c *Claim) string { + if c == nil { + return "" + } + return c.Name +} + +// There is no sharding: the model is laid out layer by layer over whatever +// devices it gets, so aggregate memory is the only thing that has to add up and +// any device count will do -- including three. +func TestDevicesOnIsPlainCeilingDivision(t *testing.T) { + node := bigNode("n", 8, 24, 1, "trainer") + for _, tc := range []struct { + memoryGiB int64 + want int + }{ + {10, 1}, + {24, 1}, + {25, 2}, + {60, 3}, // not a power of two, and that is fine + {192, 8}, + {193, 0}, // more than the pool holds + } { + if got := trainer("w", tc.memoryGiB).DevicesOn(node); got != tc.want { + t.Errorf("DevicesOn(%dGi) = %d, want %d", tc.memoryGiB, got, tc.want) + } + } +} + +func TestPerDeviceBytesRoundsUp(t *testing.T) { + req := trainer("w", 30) + if got, want := req.PerDeviceBytes(2), gib(15); got != want { + t.Errorf("PerDeviceBytes(2) = %d, want %d", got, want) + } + // 30Gi over 4 devices is 7.5Gi each; erring high is the safe direction. + if got, want := req.PerDeviceBytes(4), (gib(30)+3)/4; got != want { + t.Errorf("PerDeviceBytes(4) = %d, want %d", got, want) + } +} + +func TestOwnerKey(t *testing.T) { + for _, tc := range []struct { + name string + req Request + want string + }{ + {"a named owner is the owner", Request{Owner: "qwen3-0-6b", WorkerID: "job-a"}, "qwen3-0-6b"}, + {"no owner means an owner of one", Request{WorkerID: "job-a"}, "job-a"}, + } { + if got := tc.req.OwnerKey(); got != tc.want { + t.Errorf("%s: OwnerKey() = %q, want %q", tc.name, got, tc.want) + } + } +} + +// A pending claim's node is unknown, so nothing about it can be checked +// against anything real: it is reserved (FreeDevices), never joined. The +// burst waits a retry instead of scattering -- and a 70Gi worker can no +// longer join a claim whose stored selector only guarantees 10Gi devices. +func TestPendingClaimsAreReservedNotJoined(t *testing.T) { + fleet := NewFleet() + fleet.Nodes["n"] = &Node{ + Name: "n", DeviceCount: 1, DeviceMemoryBytes: gib(24), HostMemoryBytes: gib(94), + Roles: map[string]bool{"trainer": true}, MaxWorkersPerClaim: 4, + } + fleet.Claims["pending"] = booked(&Claim{Name: "pending", DeviceCount: 1, SizedAgainst: "n"}, gib(6), "job-a") + + if got := SelectClaim(trainer("job-b", 6), fleet); got != nil { + t.Errorf("joined %q, but an unallocated claim must not be joined", name(got)) + } + if pool, join := Decide(trainer("job-b", 6), fleet); pool != nil || join != nil { + t.Errorf("Decide = (%+v, %v), want the worker to wait for the allocation", pool, name(join)) + } +} + +// max-workers-per-claim is enforced at the join: a full allocated claim +// refuses however much device memory remains. +func TestSelectClaimRejectsFullClaims(t *testing.T) { + fleet := NewFleet() + fleet.Nodes["n"] = l4Node("n", 2, "trainer") + fleet.Claims["full"] = booked(&Claim{Name: "full", DeviceCount: 1, Node: "n"}, gib(2), "job-a", "job-b") + + if got := SelectClaim(trainer("job-c", 6), fleet); got != nil { + t.Errorf("joined %q, but max-workers-per-claim is 2 and both seats are taken", name(got)) + } +} + +// An unallocated claim reserves the pool it was sized against, so a burst +// stops cutting claims once a pool's devices are spoken for. +func TestUnallocatedClaimsReserveTheirPool(t *testing.T) { + fleet := NewFleet() + node := bigNode("n", 2, 80, 2, "trainer") + fleet.Nodes["n"] = node + fleet.Claims["c1"] = &Claim{Name: "c1", DeviceCount: 1, SizedAgainst: "n"} + fleet.Claims["c2"] = &Claim{Name: "c2", DeviceCount: 1, SizedAgainst: "n"} + + if free := fleet.FreeDevices(node); free != 0 { + t.Errorf("FreeDevices = %d, want 0: both devices are reserved by pending claims", free) + } + if pool := ChoosePool(trainer("w", 10), fleet); pool != nil { + t.Errorf("ChoosePool = %+v, want nil so the worker joins instead of cutting a third claim", pool) + } +} + +// Decide is the one place the spread-before-share order lives: a free pool +// wins over a joinable claim; under contention the claim wins over waiting. +func TestDecideSpreadsBeforeSharing(t *testing.T) { + fleet := NewFleet() + fleet.Nodes["n"] = bigNode("n", 2, 24, 4, "trainer") + fleet.Claims["c"] = booked(&Claim{Name: "c", DeviceCount: 1, Node: "n"}, gib(6), "job-a") + + if pool, join := Decide(trainer("job-b", 6), fleet); pool == nil || join != nil { + t.Errorf("Decide = (%+v, %v), want the free device, not the shared claim", pool, name(join)) + } + + fleet.Claims["c2"] = booked(&Claim{Name: "c2", DeviceCount: 1, Node: "n"}, gib(6), "job-b") + if pool, join := Decide(trainer("job-c", 6), fleet); pool != nil || join == nil { + t.Errorf("Decide = (%+v, %v), want a shared claim once no device is free", pool, name(join)) + } +} + +// The spec's preference order: fewest assigned workers, then claim name -- +// deterministic across reconciles. +func TestSelectClaimPrefersFewestWorkersThenName(t *testing.T) { + fleet := NewFleet() + fleet.Nodes["n"] = l4Node("n", 4, "trainer") + fleet.Claims["quiet"] = booked(&Claim{Name: "quiet", DeviceCount: 1, Node: "n"}, gib(2), "job-b") + fleet.Claims["crowded"] = booked(&Claim{Name: "crowded", DeviceCount: 1, Node: "n"}, gib(2), "job-c", "job-d", "job-e") + + if got := SelectClaim(trainer("job-new", 6), fleet); name(got) != "quiet" { + t.Errorf("joined %q, want the claim with the fewest workers", name(got)) + } + + fleet.Claims["a-quiet"] = booked(&Claim{Name: "a-quiet", DeviceCount: 1, Node: "n"}, gib(2), "job-f") + if got := SelectClaim(trainer("job-new", 6), fleet); name(got) != "a-quiet" { + t.Errorf("joined %q, want the tie broken by name", name(got)) + } +} + +// New claims are sized best-fit by wasted memory, so small work stays off the +// big pool while an L4 can hold it. +func TestChoosePoolPrefersTheTightestFit(t *testing.T) { + fleet := NewFleet() + fleet.Nodes["l4"] = bigNode("l4", 4, 24, 4, "trainer") + fleet.Nodes["big"] = bigNode("big", 4, 96, 4, "trainer") + + // 20Gi wastes 4Gi on an L4 and 76Gi on the big pool. + pool := ChoosePool(trainer("w", 20), fleet) + if pool == nil || pool.Node.Name != "l4" || pool.DeviceCount != 1 { + t.Fatalf("ChoosePool picked %+v, want 1 device on the L4 pool", pool) + } + + // 200Gi does not fit four L4s at all, so the big pool is the only answer. + pool = ChoosePool(trainer("w", 200), fleet) + if pool == nil || pool.Node.Name != "big" || pool.DeviceCount != 3 { + t.Fatalf("ChoosePool picked %+v, want 3 devices on the big pool", pool) + } +} + +// Appendix A's host-memory rule, pinned as arithmetic: the conservative case +// parks every assigned worker except the smallest one. +func TestParkedBytesLeaveTheSmallestWorkerResident(t *testing.T) { + claim := booked(&Claim{Name: "c", DeviceCount: 2, Node: "n"}, gib(10), "job-a") + claim.Book("job-b", gib(4)) + + // Joining with 6Gi: workers are 10, 4, 6; the 4Gi one stays resident, so + // 16Gi parks per device, over 2 devices. + if got, want := claim.ParkedBytesWith(gib(6)), gib(32); got != want { + t.Errorf("ParkedBytesWith = %d, want %d", got, want) + } +} + +// A pending worker's reason distinguishes "busy, retry" from "too small, +// never" -- the spec's actionable-reason requirement. +func TestExplain(t *testing.T) { + empty := NewFleet() + if got := Explain(trainer("w", 6), empty, ""); !strings.HasPrefix(got, "NoCapacity") { + t.Errorf("Explain = %q, want NoCapacity for a fleet with no pools", got) + } + + tooSmall := NewFleet() + tooSmall.Nodes["n"] = l4Node("n", 4, "trainer") // 2x24Gi + got := Explain(trainer("w", 200), tooSmall, "") + if !strings.HasPrefix(got, "NoCapacity") || !strings.Contains(got, "200Gi") { + t.Errorf("Explain = %q, want NoCapacity naming the 200Gi it could not fit", got) + } + + full := NewFleet() + full.Nodes["n"] = l4Node("n", 1, "trainer") + full.Claims["c"] = booked(&Claim{Name: "c", DeviceCount: 1, Node: "n"}, gib(20), "job-a") + got = Explain(trainer("w", 6), full, "pod is unschedulable") + if !strings.HasPrefix(got, "WaitingForCapacity") || !strings.Contains(got, "pod is unschedulable") { + t.Errorf("Explain = %q, want WaitingForCapacity carrying the caller's detail", got) + } +} diff --git a/k8s/deploy/scheduler-smoke/kustomization.yaml b/k8s/deploy/scheduler-smoke/kustomization.yaml new file mode 100644 index 00000000..cbe02e59 --- /dev/null +++ b/k8s/deploy/scheduler-smoke/kustomization.yaml @@ -0,0 +1,41 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +# The smoke overlay: the shipped manifests plus exactly what a smoke run +# needs -- the locally built image, the fake DRA driver's device class, and +# sleep-only worker pod templates. Applied by controller/hack/kind-smoke.sh, +# so the smoke test validates the real manifests instead of shell fixups. +resources: + - ../scheduler + +images: + - name: ghcr.io/gke-labs/open-rl/placement-controller + newName: open-rl/placement-controller + newTag: smoke + +patches: + - patch: |- + apiVersion: apps/v1 + kind: Deployment + metadata: + name: open-rl-scheduler + spec: + template: + spec: + containers: + - name: manager + # The kind-loaded image must not be re-pulled from a registry. + imagePullPolicy: IfNotPresent + env: + - name: OPEN_RL_DEVICE_CLASS + value: gpu.example.com + - name: OPEN_RL_DEVICE_DRIVER + value: gpu.example.com + +configMapGenerator: + - name: open-rl-trainer-worker-pod-template + files: [worker-pod.yaml] + options: {disableNameSuffixHash: true} + - name: open-rl-sampler-worker-pod-template + files: [worker-pod.yaml] + options: {disableNameSuffixHash: true} diff --git a/k8s/deploy/scheduler-smoke/worker-pod.yaml b/k8s/deploy/scheduler-smoke/worker-pod.yaml new file mode 100644 index 00000000..f62bf715 --- /dev/null +++ b/k8s/deploy/scheduler-smoke/worker-pod.yaml @@ -0,0 +1,9 @@ +# The smoke test has no model to run; placement neither knows nor cares. +apiVersion: v1 +kind: Pod +spec: + restartPolicy: Never + containers: + - name: worker + image: busybox:1.36 + command: ["sleep", "infinity"] diff --git a/k8s/deploy/scheduler/00-openrlworker-crd.yaml b/k8s/deploy/scheduler/00-openrlworker-crd.yaml new file mode 100644 index 00000000..3c52f6ae --- /dev/null +++ b/k8s/deploy/scheduler/00-openrlworker-crd.yaml @@ -0,0 +1,440 @@ +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + controller-gen.kubebuilder.io/version: v0.19.0 + name: openrlworkers.openrl.io +spec: + group: openrl.io + names: + kind: OpenRLWorker + listKind: OpenRLWorkerList + plural: openrlworkers + shortNames: + - orw + singular: openrlworker + scope: Namespaced + versions: + - additionalPrinterColumns: + - jsonPath: .spec.role + name: Role + type: string + - jsonPath: .spec.ownerId + name: Owner + type: string + - jsonPath: .status.deviceCount + name: GPUs + type: integer + - jsonPath: .status.memoryPerDevice + name: MemEach + type: string + - jsonPath: .status.phase + name: Phase + type: string + - jsonPath: .status.claimName + name: Claim + type: string + - jsonPath: .status.nodeName + name: Node + type: string + - jsonPath: .metadata.creationTimestamp + name: Age + type: date + name: v1alpha1 + schema: + openAPIV3Schema: + description: |- + OpenRLWorker is the scheduling request for a single worker process. The + caller creates one per worker it wants; the controller turns it into a + ResourceClaim (matched or created) and a pod, and records what it picked in + status -- so `kubectl get openrlworkers` shows what was asked for, where it + went, and why a pending worker is still pending. + properties: + apiVersion: + description: |- + APIVersion defines the versioned schema of this representation of an object. + Servers should convert recognized schemas to the latest internal value, and + may reject unrecognized values. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources + type: string + kind: + description: |- + Kind is a string value representing the REST resource this object represents. + Servers may infer this from the endpoint the client submits requests to. + Cannot be updated. + In CamelCase. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds + type: string + metadata: + type: object + spec: + description: |- + The spec is immutable: every field either places the worker or renders + its pod, and V1 does not re-place or re-render a live worker. Change by + deleting and recreating. + properties: + container: + description: Container is what the caller stamps onto the rendered + container. + properties: + args: + description: Args are appended to the template's args. + items: + type: string + type: array + command: + description: Command replaces the template's command when non-empty. + items: + type: string + type: array + env: + description: Env entries are merged by name, overwriting the template's + values. + items: + description: EnvVar represents an environment variable present + in a Container. + properties: + name: + description: |- + Name of the environment variable. + May consist of any printable ASCII characters except '='. + type: string + value: + description: |- + Variable references $(VAR_NAME) are expanded + using the previously defined environment variables in the container and + any service environment variables. If a variable cannot be resolved, + the reference in the input string will be unchanged. Double $$ are reduced + to a single $, which allows for escaping the $(VAR_NAME) syntax: i.e. + "$$(VAR_NAME)" will produce the string literal "$(VAR_NAME)". + Escaped references will never be expanded, regardless of whether the variable + exists or not. + Defaults to "". + type: string + valueFrom: + description: Source for the environment variable's value. + Cannot be used if value is not empty. + properties: + configMapKeyRef: + description: Selects a key of a ConfigMap. + properties: + key: + description: The key to select. + type: string + name: + default: "" + description: |- + Name of the referent. + This field is effectively required, but due to backwards compatibility is + allowed to be empty. Instances of this type with an empty value here are + almost certainly wrong. + More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names + type: string + optional: + description: Specify whether the ConfigMap or its + key must be defined + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + fieldRef: + description: |- + Selects a field of the pod: supports metadata.name, metadata.namespace, `metadata.labels['']`, `metadata.annotations['']`, + spec.nodeName, spec.serviceAccountName, status.hostIP, status.podIP, status.podIPs. + properties: + apiVersion: + description: Version of the schema the FieldPath + is written in terms of, defaults to "v1". + type: string + fieldPath: + description: Path of the field to select in the + specified API version. + type: string + required: + - fieldPath + type: object + x-kubernetes-map-type: atomic + fileKeyRef: + description: |- + FileKeyRef selects a key of the env file. + Requires the EnvFiles feature gate to be enabled. + properties: + key: + description: |- + The key within the env file. An invalid key will prevent the pod from starting. + The keys defined within a source may consist of any printable ASCII characters except '='. + During Alpha stage of the EnvFiles feature gate, the key size is limited to 128 characters. + type: string + optional: + default: false + description: |- + Specify whether the file or its key must be defined. If the file or key + does not exist, then the env var is not published. + If optional is set to true and the specified key does not exist, + the environment variable will not be set in the Pod's containers. + + If optional is set to false and the specified key does not exist, + an error will be returned during Pod creation. + type: boolean + path: + description: |- + The path within the volume from which to select the file. + Must be relative and may not contain the '..' path or start with '..'. + type: string + volumeName: + description: The name of the volume mount containing + the env file. + type: string + required: + - key + - path + - volumeName + type: object + x-kubernetes-map-type: atomic + resourceFieldRef: + description: |- + Selects a resource of the container: only resources limits and requests + (limits.cpu, limits.memory, limits.ephemeral-storage, requests.cpu, requests.memory and requests.ephemeral-storage) are currently supported. + properties: + containerName: + description: 'Container name: required for volumes, + optional for env vars' + type: string + divisor: + anyOf: + - type: integer + - type: string + description: Specifies the output format of the + exposed resources, defaults to "1" + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + resource: + description: 'Required: resource to select' + type: string + required: + - resource + type: object + x-kubernetes-map-type: atomic + secretKeyRef: + description: Selects a key of a secret in the pod's + namespace + properties: + key: + description: The key of the secret to select from. Must + be a valid secret key. + type: string + name: + default: "" + description: |- + Name of the referent. + This field is effectively required, but due to backwards compatibility is + allowed to be empty. Instances of this type with an empty value here are + almost certainly wrong. + More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names + type: string + optional: + description: Specify whether the Secret or its key + must be defined + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + type: object + required: + - name + type: object + type: array + image: + description: Image replaces the template's image when set. + type: string + type: object + estimatorVersion: + description: |- + EstimatorVersion records which estimator produced Memory, so a placement + can be inspected and reproduced later. + type: string + memory: + anyOf: + - type: integer + - type: string + description: |- + Memory is the estimator's figure: total peak accelerator memory, an + aggregate across however many devices it takes. Never re-estimated. + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + modelId: + description: |- + ModelID names the model this worker serves: configuration for the + worker process and the humans reading kubectl get, never identity. + The worker's identity is metadata.name. + minLength: 1 + type: string + ownerId: + description: |- + OwnerID is the unit of fairness: the runtime serves owners round-robin, + so an owner never gets extra turns for having more processes, requests, + or adapters. Opaque to the controller, never read by placement -- one + worker is resident at a time whatever the owners. A worker naming no + owner is an owner of one. + type: string + podTemplate: + description: |- + PodTemplate names the ConfigMap the worker pod is rendered from. When + unset the controller falls back to its role-default template. + properties: + key: + default: pod.yaml + description: Key within the ConfigMap holding the pod YAML. + type: string + name: + description: Name of the ConfigMap in the controller's namespace. + minLength: 1 + type: string + required: + - name + type: object + role: + description: |- + Role selects which node pools may host this worker, via the + openrl.io/trainer and openrl.io/sampler labels. It does not partition + claims: on a node that accepts both, the two roles share by turns. + enum: + - trainer + - sampler + type: string + required: + - memory + - modelId + - role + type: object + x-kubernetes-validations: + - message: OpenRLWorker spec is immutable; delete and recreate the worker + rule: self == oldSelf + status: + description: OpenRLWorkerStatus is what the controller decided and what + came of it. + properties: + claimName: + description: ClaimName is the ResourceClaim this worker was assigned + to. + type: string + conditions: + description: Conditions holds the Placed condition and its history. + items: + description: Condition contains details for one aspect of the current + state of this API Resource. + properties: + lastTransitionTime: + description: |- + lastTransitionTime is the last time the condition transitioned from one status to another. + This should be when the underlying condition changed. If that is not known, then using the time when the API field changed is acceptable. + format: date-time + type: string + message: + description: |- + message is a human readable message indicating details about the transition. + This may be an empty string. + maxLength: 32768 + type: string + observedGeneration: + description: |- + observedGeneration represents the .metadata.generation that the condition was set based upon. + For instance, if .metadata.generation is currently 12, but the .status.conditions[x].observedGeneration is 9, the condition is out of date + with respect to the current state of the instance. + format: int64 + minimum: 0 + type: integer + reason: + description: |- + reason contains a programmatic identifier indicating the reason for the condition's last transition. + Producers of specific condition types may define expected values and meanings for this field, + and whether the values are considered a guaranteed API. + The value should be a CamelCase string. + This field may not be empty. + maxLength: 1024 + minLength: 1 + pattern: ^[A-Za-z]([A-Za-z0-9_,:]*[A-Za-z0-9_])?$ + type: string + status: + description: status of the condition, one of True, False, Unknown. + enum: + - "True" + - "False" + - Unknown + type: string + type: + description: type of condition in CamelCase or in foo.example.com/CamelCase. + maxLength: 316 + pattern: ^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*/)?(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])$ + type: string + required: + - lastTransitionTime + - message + - reason + - status + - type + type: object + type: array + x-kubernetes-list-map-keys: + - type + x-kubernetes-list-type: map + deviceCount: + description: |- + DeviceCount is how many accelerators the controller decided this worker + spans. A decision, not a request: derived from Memory and the capacity + the pools registered. + format: int32 + type: integer + estimatorVersion: + description: |- + EstimatorVersion echoes which estimator produced the memory figure this + placement was decided against, so the decision can be reproduced. + type: string + hostMemoryWhenParked: + description: |- + HostMemoryWhenParked is the host RAM this worker occupies while + suspended: its full accelerator footprint, parked in its own host + address space. This is what actually bounds how many workers a node + can carry. + type: string + memoryPerDevice: + description: |- + MemoryPerDevice is what each of those devices must provide, i.e. Memory + divided by DeviceCount and rounded up. + type: string + nodeName: + description: |- + NodeName is set only once the claim's allocation is observed. Until then + the node is genuinely unknown. + type: string + observedGeneration: + description: ObservedGeneration is the spec generation this status + was computed from. + format: int64 + type: integer + phase: + description: Phase is a coarse summary; Conditions carry the detail. + enum: + - Pending + - Placing + - Running + - Failed + type: string + podName: + description: PodName is the worker pod, once created. + type: string + reason: + description: Reason is the most recent human-readable explanation + of Phase. + type: string + type: object + required: + - spec + type: object + served: true + storage: true + subresources: + status: {} diff --git a/k8s/deploy/scheduler/01-scheduler.yaml b/k8s/deploy/scheduler/01-scheduler.yaml new file mode 100644 index 00000000..037cbb44 --- /dev/null +++ b/k8s/deploy/scheduler/01-scheduler.yaml @@ -0,0 +1,175 @@ +# The GPU scheduler. +# +# It runs beside kube-scheduler rather than replacing it: it decides which +# ResourceClaim a worker joins -- and cuts a new one when none will have it -- +# and the real scheduler decides where that claim lands. Nothing else in the +# cluster changes when this is applied; without any OpenRLWorker objects it +# watches an empty namespace and does nothing. +# +# Placement is a global bin-packing decision: a reconcile reads a snapshot of +# claims, nodes and workers and then writes, so two reconciles running at once +# would each decide against a fleet missing the other's booking. The controller +# runs one reconcile at a time internally, and holds a leader lease so a rolling +# update never has two processes deciding at once. +apiVersion: v1 +kind: ServiceAccount +metadata: + name: open-rl-scheduler +--- +# Cluster-scoped only for reading nodes: the scheduler must see the capacity +# labels (openrl.io/enabled, /trainer, /sampler, /max-workers-per-claim) to know which +# nodes accept which role and how many workers may share a claim there, and the +# node's allocatable memory, which is what actually bounds how many workers can +# be parked on it at once. +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: open-rl-scheduler-nodes +rules: +- apiGroups: [""] + resources: ["nodes"] + verbs: ["get", "list", "watch"] +# ResourceSlices are cluster-scoped and are how the scheduler learns what +# accelerators each node actually has. +- apiGroups: ["resource.k8s.io"] + resources: ["resourceslices"] + verbs: ["get", "list", "watch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: open-rl-scheduler-nodes +subjects: +- kind: ServiceAccount + name: open-rl-scheduler + namespace: default +roleRef: + kind: ClusterRole + name: open-rl-scheduler-nodes + apiGroup: rbac.authorization.k8s.io +--- +# Everything else is namespaced, so the scheduler cannot reach pods or claims +# outside the namespace it was pointed at. +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: open-rl-scheduler +rules: +- apiGroups: ["openrl.io"] + resources: ["openrlworkers"] + # update is for the finalizer only: a deleted worker holds its seat until + # its pod is verifiably gone. The spec itself is CEL-immutable. + verbs: ["get", "list", "watch", "update"] +# The scheduler never edits a request, only reports on it. +- apiGroups: ["openrl.io"] + resources: ["openrlworkers/status"] + verbs: ["get", "update", "patch"] +- apiGroups: ["resource.k8s.io"] + resources: ["resourceclaims"] + verbs: ["get", "list", "watch", "create", "delete"] +- apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "watch", "create", "delete"] +# Worker pods are rendered from templates held in ConfigMaps. +- apiGroups: [""] + resources: ["configmaps"] + verbs: ["get", "list", "watch"] +- apiGroups: [""] + resources: ["events"] + verbs: ["create", "patch"] +# Leader election. +- apiGroups: ["coordination.k8s.io"] + resources: ["leases"] + verbs: ["get", "list", "watch", "create", "update", "patch", "delete"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: open-rl-scheduler +subjects: +- kind: ServiceAccount + name: open-rl-scheduler +roleRef: + kind: Role + name: open-rl-scheduler + apiGroup: rbac.authorization.k8s.io +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: open-rl-scheduler + labels: + app: open-rl-scheduler +spec: + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + app: open-rl-scheduler + template: + metadata: + labels: + app: open-rl-scheduler + spec: + serviceAccountName: open-rl-scheduler + securityContext: + runAsNonRoot: true + containers: + - name: manager + # Built by `make docker-build docker-push` from controller/. + image: ghcr.io/gke-labs/open-rl/placement-controller:latest + args: + - --leader-elect + - --health-probe-bind-address=:8081 + - --metrics-bind-address=:8080 + env: + - name: OPEN_RL_WORKER_NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + - name: OPEN_RL_RECONCILE_INTERVAL + value: "10s" + # The DeviceClass generated claims request, and the driver whose + # ResourceSlices describe the hardware. The driver name is also the CEL + # domain the per-device memory floor is written against. + - name: OPEN_RL_DEVICE_CLASS + value: "gpu.nvidia.com" + - name: OPEN_RL_DEVICE_DRIVER + value: "gpu.nvidia.com" + # Worker pods are rendered from these ConfigMaps, read through the API + # rather than mounted: a worker may name its own template, and the set + # of templates is not known when this pod starts. + - name: OPEN_RL_TRAINER_POD_TEMPLATE_CONFIGMAP + value: "open-rl-trainer-worker-pod-template" + - name: OPEN_RL_SAMPLER_POD_TEMPLATE_CONFIGMAP + value: "open-rl-sampler-worker-pod-template" + # How long a worker may go unplaced before the request is called + # unsatisfiable. Without a deadline an impossible request waits forever, + # indistinguishable from one that is merely queued behind a busy fleet. + - name: OPEN_RL_PLACEMENT_TIMEOUT + value: "15m" + - name: OPEN_RL_RECLAIM_INTERVAL + value: "1m" + ports: + - name: metrics + containerPort: 8080 + livenessProbe: + httpGet: {path: /healthz, port: 8081} + initialDelaySeconds: 15 + periodSeconds: 20 + readinessProbe: + httpGet: {path: /readyz, port: 8081} + initialDelaySeconds: 5 + periodSeconds: 10 + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: ["ALL"] + resources: + requests: + cpu: 100m + memory: 128Mi + limits: + memory: 512Mi diff --git a/k8s/deploy/scheduler/kustomization.yaml b/k8s/deploy/scheduler/kustomization.yaml new file mode 100644 index 00000000..cf5b4894 --- /dev/null +++ b/k8s/deploy/scheduler/kustomization.yaml @@ -0,0 +1,16 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +# The scheduler on its own: a CRD and one deployment, no base overlay. +# +# Applying this changes nothing about a running cluster. The scheduler only acts +# on OpenRLWorker objects, and nothing creates those yet -- the gateway still +# manages its workers itself. That is the point: this can land, be applied, and +# be exercised by hand before anything depends on it. +# +# kubectl apply -k k8s/deploy/scheduler +# kubectl label node openrl.io/enabled=true openrl.io/trainer=true openrl.io/max-workers-per-claim=4 +# kubectl get openrlworkers -w +resources: + - 00-openrlworker-crd.yaml + - 01-scheduler.yaml