diff --git a/scheduler/controller/api/v1alpha1/workload_types.go b/scheduler/controller/api/v1alpha1/workload_types.go index 57c3d8cbf..d06fe427d 100644 --- a/scheduler/controller/api/v1alpha1/workload_types.go +++ b/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 diff --git a/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go b/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go index ef6ccd2c7..51b210a41 100644 --- a/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go +++ b/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go @@ -13,6 +13,11 @@ import ( func (in *AcceleratorSpec) DeepCopyInto(out *AcceleratorSpec) { *out = *in out.Memory = in.Memory.DeepCopy() + if in.PreferredTypes != nil { + in, out := &in.PreferredTypes, &out.PreferredTypes + *out = make([]string, len(*in)) + copy(*out, *in) + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AcceleratorSpec. diff --git a/scheduler/controller/cmd/manager/main.go b/scheduler/controller/cmd/manager/main.go index 6c734ab74..6bce4251c 100644 --- a/scheduler/controller/cmd/manager/main.go +++ b/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 { diff --git a/scheduler/controller/internal/controller/fleet.go b/scheduler/controller/internal/controller/fleet.go index f805aa46b..e6b1da8a8 100644 --- a/scheduler/controller/internal/controller/fleet.go +++ b/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, diff --git a/scheduler/controller/internal/controller/pod.go b/scheduler/controller/internal/controller/pod.go index 868010ad8..4450d9959 100644 --- a/scheduler/controller/internal/controller/pod.go +++ b/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 { diff --git a/scheduler/controller/internal/controller/workload_controller.go b/scheduler/controller/internal/controller/workload_controller.go index d8dfe8856..325ad10a3 100644 --- a/scheduler/controller/internal/controller/workload_controller.go +++ b/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 diff --git a/scheduler/controller/internal/controller/workload_controller_test.go b/scheduler/controller/internal/controller/workload_controller_test.go index a9658d6db..c8b6e9f0e 100644 --- a/scheduler/controller/internal/controller/workload_controller_test.go +++ b/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) { diff --git a/scheduler/deploy/base/00-workload-crd.yaml b/scheduler/deploy/base/00-workload-crd.yaml index 2042ce6bb..38c813e3d 100644 --- a/scheduler/deploy/base/00-workload-crd.yaml +++ b/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 diff --git a/scheduler/docs/design.md b/scheduler/docs/design.md index b86bbbb23..d561b3099 100644 --- a/scheduler/docs/design.md +++ b/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: diff --git a/skills/openrl/references/lifecycle.md b/skills/openrl/references/lifecycle.md index d71eaf4f7..a170887ac 100644 --- a/skills/openrl/references/lifecycle.md +++ b/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 diff --git a/skills/openrl/references/runbook.md b/skills/openrl/references/runbook.md index 5cec0a708..851c69bd9 100644 --- a/skills/openrl/references/runbook.md +++ b/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.: 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 | diff --git a/src/server/api_server.py b/src/server/api_server.py index 095e19ced..92fec4689 100644 --- a/src/server/api_server.py +++ b/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) diff --git a/src/server/model_metadata.py b/src/server/model_metadata.py index eb345a6e5..0a6fbeb92 100644 --- a/src/server/model_metadata.py +++ b/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 diff --git a/src/server/openrl_status.py b/src/server/openrl_status.py index 16acf057a..e5e0ef158 100644 --- a/src/server/openrl_status.py +++ b/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), } diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index 29743bed8..ffad3b3c2 100644 --- a/src/server/scheduler_worker_manager.py +++ b/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]: diff --git a/tests/test_api_server_paths.py b/tests/test_api_server_paths.py index 51a7d9b54..615fdad81 100644 --- a/tests/test_api_server_paths.py +++ b/tests/test_api_server_paths.py @@ -570,6 +570,29 @@ def test_sampler_replicas_need_a_server_that_launches_pods(self) -> None: self.assertEqual(response.status_code, 400) self.assertIn("openrl.sampler_replicas above 1", response.json()["error"]) + def test_gpu_types_are_an_ordered_list_on_the_model(self) -> None: + gpu_types = {"openrl.gpu_types": "nvidia-a100-80gb, nvidia-l4"} + model_id = self.post("create_model", {"base_model": "m", "user_metadata": gpu_types}).json()["request_id"] + self.assertEqual(self.metadata(model_id)["gpu_types"], ["nvidia-a100-80gb", "nvidia-l4"]) + # TINKER_TAGS splits on commas, so a tag separates types with semicolons. + tagged = {"openrl.gpu_types": "nvidia-a100-80gb;nvidia-l4"} + model_id = self.post("create_model", {"base_model": "m", "user_metadata": tagged}).json()["request_id"] + self.assertEqual(self.metadata(model_id)["gpu_types"], ["nvidia-a100-80gb", "nvidia-l4"]) + model_id = self.post("create_model", {"base_model": "m"}).json()["request_id"] + self.assertEqual(self.metadata(model_id)["gpu_types"], []) + + response = self.post("create_model", {"base_model": "m", "user_metadata": {"openrl.gpu_types": "nvidia-l4,nvidia-l4"}}) + self.assertEqual(response.status_code, 400) + self.assertIn("openrl.gpu_types", response.json()["error"]) + response = self.post("create_model", {"base_model": "m", "user_metadata": {"openrl.gpu_types": "nvidia l4"}}) + self.assertEqual(response.status_code, 400) + self.assertIn("not an accelerator type name", response.json()["error"]) + + with patch.object(api_server, "worker_manager", object.__new__(LocalWorkerManager)): + response = self.post("create_model", {"base_model": "m", "user_metadata": gpu_types}) + self.assertEqual(response.status_code, 400) + self.assertIn("openrl.gpu_types needs a server that launches workers as pods", response.json()["error"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index 40becf0e9..984b56374 100644 --- a/tests/test_scheduler_worker_manager.py +++ b/tests/test_scheduler_worker_manager.py @@ -354,6 +354,25 @@ def test_sampler_replicas_are_single_gpu_workloads_released_together(self) -> No self.assertEqual(sampler["spec"]["ownerID"], "job-sr") self.assertEqual(sorted(self.api.deleted), sorted(names)) + def test_preferred_gpu_types_ride_both_workloads_in_order(self) -> None: + meta = {"base_model": "google/gemma-4-E2B", "fine_tuning_type": "lora", "gpu_types": ["nvidia-a100-80gb", "nvidia-l4"]} + s = self.store_with("job-gt", meta) + with patch("server.worker_manager.get_state_store", return_value=s): + self.manager.ensure("job-gt", "trainer") + self.manager.ensure("job-gt", "sampler") + + trainer, sampler = self.api.created + for worker in (trainer, sampler): + self.assertEqual(worker["spec"]["accelerator"]["preferredTypes"], ["nvidia-a100-80gb", "nvidia-l4"]) + self.assertEqual(trainer["spec"]["accelerator"]["mode"], "SingleGPU") + + def test_no_preferred_gpu_types_leaves_the_accelerator_spec_alone(self) -> None: + s = self.store_with("job-fft", {"base_model": "Qwen/Qwen3-8B", "fine_tuning_type": "full"}) + with patch("server.worker_manager.get_state_store", return_value=s): + self.manager.ensure("job-fft", "trainer") + + self.assertNotIn("preferredTypes", self.api.created[0]["spec"]["accelerator"]) + class MixedSamplingSessionTest(unittest.IsolatedAsyncioTestCase): async def test_lora_and_fft_sessions_launch_their_own_sampler_types(self) -> None: