perf: prefer GPU types per job with openrl.gpu_types

scheduler/controller/api/v1alpha1/workload_types.go

@@ -81,6 +81,12 @@ type AcceleratorSpec struct {
 	// +kubebuilder:validation:Minimum=2
 	// +optional
 	Devices int32 `json:"devices,omitempty"`
+
+	// PreferredTypes lists accelerator types, most preferred first, as values
+	// of the controller's accelerator type node label. When set they replace
+	// the tight-fit node preference; a soft preference, never a constraint.
+	// +optional
+	PreferredTypes []string `json:"preferredTypes,omitempty"`
 }
 
 // WorkloadSpec is one worker process's scheduling request. Memory arrives

scheduler/controller/cmd/manager/main.go

@@ -47,6 +47,7 @@ func main() {
 		retryInterval    time.Duration
 		placementTimeout time.Duration
 		strategy         string
+		typeLabel        string
 		maxConcurrent    int
 	)
 
@@ -62,6 +63,8 @@ func main() {
 		"How long a worker may go unplaced before the request is declared unsatisfiable. 0 waits forever.")
 	flag.StringVar(&strategy, "placement-strategy", env("OPEN_RL_PLACEMENT_STRATEGY", string(placement.StrategyBinPack)),
 		"binpack shares eligible FFT claims first; spread requests a GPU first, then falls back to FFT sharing. LoRA never shares.")
+	flag.StringVar(&typeLabel, "accelerator-type-label", env("OPEN_RL_ACCELERATOR_TYPE_LABEL", controller.DefaultAcceleratorTypeLabel),
+		"Node label naming a node's accelerator type. A Workload's preferredTypes are matched against its values.")
 	flag.IntVar(&maxConcurrent, "max-concurrent-reconciles", envInt("OPEN_RL_MAX_CONCURRENT_RECONCILES", 4),
 		"How many workers place at once. Seat booking is CAS-arbitrated, so concurrency risks only transient over-cut claims, which the sharing fallback drains.")
 
@@ -116,6 +119,7 @@ func main() {
 		RetryInterval:           retryInterval,
 		PlacementTimeout:        placementTimeout,
 		PlacementStrategy:       parsedStrategy,
+		AcceleratorTypeLabel:    typeLabel,
 		MaxConcurrentReconciles: maxConcurrent,
 	}
 	if err := reconciler.SetupWithManager(mgr); err != nil {

scheduler/controller/internal/controller/fleet.go

@@ -36,6 +36,10 @@ const (
 	NodeLabelSampler = "openrl.io/sampler"
 )
 
+// DefaultAcceleratorTypeLabel is the node label naming a node's accelerator
+// type, the one a Workload's PreferredTypes match unless overridden.
+const DefaultAcceleratorTypeLabel = "cloud.google.com/gke-accelerator"
+
 var nodeRoleLabel = map[openrlv1alpha1.WorkerRole]string{
 	openrlv1alpha1.RoleTrainer: NodeLabelTrainer,
 	openrlv1alpha1.RoleSampler: NodeLabelSampler,

scheduler/controller/internal/controller/pod.go

@@ -231,6 +231,10 @@ func (r *WorkloadReconciler) renderPod(worker *openrlv1alpha1.Workload, podName,
 		// requires exactly one value, hence one term per node.
 		PreferredDuringSchedulingIgnoredDuringExecution: preferredTerms(prefs),
 	}}
+	// A job naming accelerator types ranks nodes by type instead of by size.
+	if types := worker.Spec.Accelerator.PreferredTypes; len(types) > 0 {
+		pod.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution = typeTerms(r.acceleratorTypeLabel(), types)
+	}
 
 	if err := controllerutil.SetControllerReference(worker, pod, r.Scheme()); err != nil {
 		return nil, fmt.Errorf("set owner of pod %s: %w", podName, err)
@@ -256,6 +260,31 @@ func preferredTerms(prefs []placement.NodePreference) []corev1.PreferredScheduli
 	return terms
 }
 
+// typeTerms renders preferred accelerator types as weighted scheduling
+// terms, 100 for the first and 10 less per position, never below 1.
+func typeTerms(label string, types []string) []corev1.PreferredSchedulingTerm {
+	terms := make([]corev1.PreferredSchedulingTerm, 0, len(types))
+	for i, name := range types {
+		terms = append(terms, corev1.PreferredSchedulingTerm{
+			Weight: int32(max(100-10*i, 1)),
+			Preference: corev1.NodeSelectorTerm{
+				MatchExpressions: []corev1.NodeSelectorRequirement{{
+					Key: label, Operator: corev1.NodeSelectorOpIn, Values: []string{name},
+				}},
+			},
+		})
+	}
+	return terms
+}
+
+// acceleratorTypeLabel is the node label PreferredTypes match against.
+func (r *WorkloadReconciler) acceleratorTypeLabel() string {
+	if r.AcceleratorTypeLabel != "" {
+		return r.AcceleratorTypeLabel
+	}
+	return DefaultAcceleratorTypeLabel
+}
+
 // setEnv merges one variable by name, overwriting whatever the template had.
 func setEnv(container *corev1.Container, want corev1.EnvVar) {
 	for i := range container.Env {

scheduler/controller/internal/controller/workload_controller.go

@@ -54,6 +54,9 @@ type WorkloadReconciler struct {
 	// PlacementStrategy orders the two placement moves; the zero value reads
 	// as binpack.
 	PlacementStrategy placement.Strategy
+	// AcceleratorTypeLabel is the node label a Workload's PreferredTypes
+	// match; empty reads as DefaultAcceleratorTypeLabel.
+	AcceleratorTypeLabel string
 	// MaxConcurrentReconciles is how many workers place at once.
 	MaxConcurrentReconciles int
 

scheduler/controller/internal/controller/workload_controller_test.go

@@ -3,6 +3,7 @@ package controller
 import (
 	"context"
 	"fmt"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -1116,6 +1117,73 @@ func TestPodPrefersTheSmallestAdequateNode(t *testing.T) {
 	}
 }
 
+// A job naming accelerator types swaps the tight-fit terms for one term per
+// type, weighted by list position; naming none leaves the tight-fit terms.
+func TestPodPrefersTheNamedAcceleratorTypes(t *testing.T) {
+	r := newReconciler(t)
+	prefs := []placement.NodePreference{{Weight: 100, Nodes: []string{"node-l4"}}, {Weight: 50, Nodes: []string{"node-a100"}}}
+
+	plain := trainerWorker("w-a", "model-a")
+	pod, err := r.renderPod(plain, "orw-w-a", "claim-a", prefs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := pod.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution, preferredTerms(prefs); !reflect.DeepEqual(got, want) {
+		t.Fatalf("without preferred types terms = %+v, want the tight-fit terms %+v", got, want)
+	}
+	required := pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution
+
+	typed := trainerWorker("w-b", "model-b")
+	typed.Spec.Accelerator.PreferredTypes = []string{"nvidia-a100-80gb", "nvidia-l4"}
+	pod, err = r.renderPod(typed, "orw-w-b", "claim-b", prefs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution, required) {
+		t.Errorf("required affinity changed with preferred types: %+v", pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution)
+	}
+	terms := pod.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution
+	if len(terms) != 2 {
+		t.Fatalf("preferred terms = %+v, want one per type", terms)
+	}
+	for i, want := range []struct {
+		weight int32
+		value  string
+	}{{100, "nvidia-a100-80gb"}, {90, "nvidia-l4"}} {
+		term := terms[i]
+		if len(term.Preference.MatchFields) != 0 || len(term.Preference.MatchExpressions) != 1 {
+			t.Fatalf("term %d = %+v, want one type expression and no node names", i, term)
+		}
+		expr := term.Preference.MatchExpressions[0]
+		if term.Weight != want.weight || expr.Key != DefaultAcceleratorTypeLabel || expr.Operator != corev1.NodeSelectorOpIn ||
+			!reflect.DeepEqual(expr.Values, []string{want.value}) {
+			t.Errorf("term %d = %+v, want %s=%s at weight %d", i, term, DefaultAcceleratorTypeLabel, want.value, want.weight)
+		}
+	}
+
+	r.AcceleratorTypeLabel = "example.com/accelerator"
+	pod, err = r.renderPod(typed, "orw-w-b", "claim-b", prefs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, term := range pod.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution {
+		if key := term.Preference.MatchExpressions[0].Key; key != "example.com/accelerator" {
+			t.Errorf("term keyed on %q, want the overridden label", key)
+		}
+	}
+}
+
+func TestTypeTermWeightsFloorAtOne(t *testing.T) {
+	types := make([]string, 12)
+	for i := range types {
+		types[i] = fmt.Sprintf("type-%d", i)
+	}
+	terms := typeTerms(DefaultAcceleratorTypeLabel, types)
+	if terms[9].Weight != 10 || terms[10].Weight != 1 || terms[11].Weight != 1 {
+		t.Errorf("weights %d, %d, %d at positions 9-11, want 10, 1, 1", terms[9].Weight, terms[10].Weight, terms[11].Weight)
+	}
+}
+
 // 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) {

scheduler/deploy/base/00-workload-crd.yaml

@@ -100,6 +100,14 @@ spec:
                     - SingleGPU
                     - MultiGPU
                     type: string
+                  preferredTypes:
+                    description: |-
+                      PreferredTypes lists accelerator types, most preferred first, as values
+                      of the controller's accelerator type node label. When set they replace
+                      the tight-fit node preference; a soft preference, never a constraint.
+                    items:
+                      type: string
+                    type: array
                 required:
                 - memory
                 type: object

scheduler/docs/design.md

@@ -700,6 +700,17 @@ among nodes that can already host the pod. This stays survey-less: the
 preferences are compiled from which shapes live where, a static hardware
 fact, never from what is free.
 
+A workload can override that order with `accelerator.preferredTypes`, an
+ordered list of accelerator type names (the API server sets it from the
+job's `openrl.gpu_types`). Smallest-fit is the wrong choice for a job that
+runs several times faster on a bigger device. When the list is set, the
+tight-fit terms are replaced by one term per type, matching the node label
+`--accelerator-type-label` (`OPEN_RL_ACCELERATOR_TYPE_LABEL`, default
+`cloud.google.com/gke-accelerator`), weighted 100 for the first type and 10
+less per position, floored at 1. These terms are preferences as well: with
+the named types busy or absent, the pod lands on any device that fits. The
+claim's tiers are unchanged, and joining a shared ledger ignores the list.
+
 The tiers are rendered as DRA's prioritized list (`firstAvailable` in the
 claim's device request); the allocator tries them in list order and takes the
 first that fits on the node under consideration. Concretely, the claim the controller cuts looks like:

skills/openrl/references/lifecycle.md

@@ -65,7 +65,9 @@ full fine-tuning cannot span devices.
 1. **Prefer a bigger device.** `openrl.exclusive=true` gives the job its own trainer and
    sampler. If the cluster has a larger pool, the scheduler places them where they fit;
    nothing else changes. Ask the platform owner which pools exist; the endpoint does not
-   list them yet.
+   list them yet. The scheduler picks the smallest GPU that fits, which is not always the
+   fastest: `openrl.gpu_types=nvidia-a100-80gb;nvidia-l4` makes the job's workers prefer
+   those types in that order, and still fall back to any GPU that fits when they are busy.
 2. **Then spread the trainer over several GPUs.** When no single device holds the
    trainer, set `openrl.trainer_backend=automodel` and `openrl.trainer_gpus=N` (2 to 8).
    The Automodel trainer shards the model with FSDP2 across `N` GPUs on one node, so each

skills/openrl/references/runbook.md

@@ -20,7 +20,7 @@ wait out the normal duration, then ask the platform owner.
 | `Session heartbeat failed for N seconds` in the client log | `curl .../healthz` through the same URL the client uses | if `healthz` fails: the connection, restart the port-forward loop. Under 120 s the run continues; over 120 s see the row above. If `healthz` is fine, the SDK's own thread is blocked: the cookbook `debug` skill, GIL contention | nothing, unless `healthz` itself fails |
 | HTTP 400 `The x-open-rl-fine-tuning-type header was removed; set openrl.fine_tuning_type ...` | the client sets a removed header or `full_config` | pass the setting named in the message as `user_metadata` or `TINKER_TAGS` | nothing |
 | HTTP 400 `openrl.<name>: unknown setting` or `must be ...` | spelling and value against the schema: `curl -s $OPENRL/openapi.json \| jq .components.schemas.OpenRLSettings` | fix the value; values are strings | nothing |
-| HTTP 400 `openrl.exclusive needs a server that launches workers per model` (or `sampler_replicas`, `trainer_backend`, `trainer_gpus`) | `openrl.worker_launcher` in capabilities is not `scheduler` | run against the Kubernetes deployment; these settings need it | nothing |
+| HTTP 400 `openrl.exclusive needs a server that launches workers per model` (or `sampler_replicas`, `trainer_backend`, `trainer_gpus`, `gpu_types`) | `openrl.worker_launcher` in capabilities is not `scheduler` | run against the Kubernetes deployment; these settings need it | nothing |
 | HTTP 400 on `create_model` with `openrl.fine_tuning_type=full` saying FFT is disabled | `openrl.fft_enabled` in capabilities | use LoRA, or ask for the FFT release | "please deploy `openrl-fft.yaml`" |
 | Samples are token soup or a different model's style | `get_info` for your model returns your base model as `tokenizer_id`; if it does, the renderer is the suspect: the cookbook `debug` skill | if `tokenizer_id` is wrong, report; this was a server bug fixed in v0.0.5 | OpenRL version and the model id |
 | Steps are two to three times slower than yesterday | `shares_workers: true` in the model view means other LoRA jobs on the same base model share your trainer; several FFT jobs time-slice a GPU | expected sharing; stagger runs if wall-clock matters. If nothing shares your workers, profile the client: the cookbook `debug` skill | nothing |

src/server/api_server.py

@@ -22,7 +22,18 @@
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from pydantic import AfterValidator, AliasChoices, BaseModel, BeforeValidator, ConfigDict, Field, ValidationError, ValidationInfo, field_validator
+from pydantic import (
+  AfterValidator,
+  AliasChoices,
+  BaseModel,
+  BeforeValidator,
+  ConfigDict,
+  Field,
+  ValidationError,
+  ValidationInfo,
+  WithJsonSchema,
+  field_validator,
+)
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
 from server import openrl_status, proto_codec
@@ -181,6 +192,23 @@ def parse_trainer_backend(value: str) -> str:
   return value
 
 
+# A Kubernetes label value, what a node's accelerator type label holds.
+LABEL_VALUE = re.compile(r"[A-Za-z0-9]([-A-Za-z0-9_.]{0,61}[A-Za-z0-9])?")
+
+
+def parse_gpu_types(value: Any) -> list[str]:
+  # TINKER_TAGS splits on commas, so a tag lists types with semicolons.
+  names = [name.strip() for name in re.split(r"[,;]", value)] if isinstance(value, str) else value
+  if not isinstance(names, list) or not all(isinstance(name, str) for name in names):
+    raise ValueError(f"must be a list of accelerator types separated by ; or ,, got {value!r}")
+  names = [name for name in names if name]
+  if bad := [name for name in names if not LABEL_VALUE.fullmatch(name)]:
+    raise ValueError(f"{bad[0]!r} is not an accelerator type name like nvidia-a100-80gb")
+  if len(set(names)) != len(names):
+    raise ValueError(f"lists an accelerator type twice: {value!r}")
+  return names
+
+
 class Settings(BaseModel):
   """The descriptions are what get_server_capabilities and the skill show, so
   they are written for the client choosing a value. Each field's wire name is
@@ -211,6 +239,14 @@ class Settings(BaseModel):
     le=8,
     description="Trainer GPUs that share one sequence, as context parallelism. The rest of trainer_gpus is data parallel.",
   )
+  gpu_types: Annotated[list[str], BeforeValidator(parse_gpu_types), WithJsonSchema({"type": "string"})] = Field(
+    default_factory=list,
+    description=(
+      "Accelerator types the model's workers prefer, most preferred first, separated by semicolons: nvidia-a100-80gb;nvidia-l4. "
+      "A preference, not a requirement: a worker still lands on any GPU that fits when those are busy. "
+      "Unset, workers prefer the smallest GPU that fits."
+    ),
+  )
   fine_tuning_type: FineTuningType | None = Field(
     default=None,
     description="lora unless set. A model restored from a checkpoint takes the checkpoint's type.",
@@ -498,6 +534,8 @@ async def _extract_and_persist_model_metadata(
     raise ValueError("openrl.exclusive needs a server that launches workers per model")
   if settings.sampler_replicas > 1 and isinstance(worker_manager, LocalWorkerManager):
     raise ValueError("openrl.sampler_replicas above 1 needs a server that launches workers as pods")
+  if settings.gpu_types and (worker_manager is None or isinstance(worker_manager, LocalWorkerManager)):
+    raise ValueError("openrl.gpu_types needs a server that launches workers as pods")
   if settings.trainer_backend != "pytorch" and worker_manager is None:
     raise ValueError(f"openrl.trainer_backend={settings.trainer_backend} needs a server that launches workers per model")
   if settings.trainer_backend not in TRAINER_BACKENDS and isinstance(worker_manager, LocalWorkerManager):
@@ -523,6 +561,7 @@ async def _extract_and_persist_model_metadata(
     trainer_backend=settings.trainer_backend,
     trainer_gpus=settings.trainer_gpus,
     trainer_cp=settings.trainer_cp,
+    gpu_types=settings.gpu_types,
   )
   await persist_model_metadata(state, model_id, meta_obj)
 

src/server/model_metadata.py

@@ -48,6 +48,8 @@ class TrainingModelMetadata(BaseModel):
   trainer_backend: str = "pytorch"
   trainer_gpus: int = 1
   trainer_cp: int = 1
+  # Accelerator types the workers prefer, most preferred first.
+  gpu_types: list[str] = Field(default_factory=list)
   status: str = "active"
   updated_at: float = 0.0
   completed_at: float | None = None

src/server/openrl_status.py

@@ -89,6 +89,7 @@ def settings_of(meta: TrainingModelMetadata) -> dict[str, Any]:
     "openrl.trainer_backend": meta.trainer_backend,
     "openrl.trainer_gpus": meta.trainer_gpus,
     "openrl.trainer_cp": meta.trainer_cp,
+    "openrl.gpu_types": ",".join(meta.gpu_types),
   }
 
 

src/server/scheduler_worker_manager.py

@@ -185,8 +185,14 @@ def pod_template(worker: Worker) -> dict[str, Any]:
 def accelerator_spec(worker: Worker) -> dict[str, Any]:
   """A torchrun group asks for its devices, each sized for a whole replica."""
   if worker.devices > 1:
-    return {"mode": "MultiGPU", "devices": worker.devices, "memory": worker.footprint.accelerator}
-  return {"mode": "SingleGPU", "memory": worker.footprint.accelerator}
+    spec = {"mode": "MultiGPU", "devices": worker.devices, "memory": worker.footprint.accelerator}
+  else:
+    spec = {"mode": "SingleGPU", "memory": worker.footprint.accelerator}
+  # A shared LoRA runtime's Workloads carry the preference of the job whose
+  # request created them; later jobs reuse them as they are.
+  if worker.meta.gpu_types:
+    spec["preferredTypes"] = list(worker.meta.gpu_types)
+  return spec
 
 
 def workload_body(worker: Worker) -> dict[str, Any]: