Drop piecewise_constant from the trainable lr_scheduler allow-list

piecewise_constant is the only diffusers scheduler that needs a step_rules string, and
neither diffusion trainer passes one (get_scheduler is called with only warmup/training
steps, and there is no config field for it). Accepting it let /diffusion/start pass
normalized(), free the resident GPU workloads, spawn the trainer, and only then crash in
the subprocess (get_piecewise_constant_schedule does step_rules.split(",") on None) -- the
exact evict-then-fail the up-front validation exists to prevent. Reject it now with a clear
400. The remaining six schedulers all run with only warmup/training steps.

studio/backend/core/training/diffusion_train_common.py

@@ -35,7 +35,13 @@ from core.inference.diffusion_families import (
 # set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback.
 DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
 
-# diffusers' SchedulerType names (diffusers.optimization.get_scheduler).
+# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). piecewise_constant is
+# intentionally excluded: it is the only scheduler that needs a `step_rules` string, which the
+# trainers never pass (get_scheduler is called with only warmup/training steps, and there is no
+# config field for it). Accepting it would pass normalized(), free the resident GPU workloads,
+# then crash in the trainer subprocess (get_piecewise_constant_schedule does step_rules.split(",")
+# on None) -- the exact evict-then-fail the up-front validation exists to prevent. The remaining
+# six all run with only warmup/training steps.
 _LR_SCHEDULERS: frozenset[str] = frozenset(
     {
         "linear",
@@ -44,7 +50,6 @@ _LR_SCHEDULERS: frozenset[str] = frozenset(
         "polynomial",
         "constant",
         "constant_with_warmup",
-        "piecewise_constant",
     }
 )
 

studio/backend/tests/test_diffusion_lora_trainer.py

@@ -114,6 +114,32 @@ def test_config_normalized_validation(kw):
         DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized()
 
 
+def test_normalized_rejects_piecewise_constant():
+    # piecewise_constant needs a step_rules string the trainers never supply, so get_scheduler()
+    # would crash in the trainer subprocess AFTER the resident GPU workloads are freed. It must be
+    # rejected up front (a clean ValueError -> 400), not accepted like the other schedulers.
+    with pytest.raises(ValueError, match = "lr_scheduler"):
+        DiffusionLoraConfig(
+            base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "piecewise_constant"
+        ).normalized()
+
+
+def test_normalized_accepts_supported_schedulers():
+    # Every scheduler in the allow-list runs with only warmup/training steps (no extra required arg).
+    for sched in (
+        "linear",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "constant",
+        "constant_with_warmup",
+    ):
+        cfg = DiffusionLoraConfig(
+            base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = sched
+        ).normalized()
+        assert cfg.lr_scheduler == sched
+
+
 def test_compute_sdxl_add_time_ids():
     assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024)