fix: training never replaces cold-start model _drop_timestamp called as instance method

tests/test_dual_server_client.py

@@ -1307,6 +1307,10 @@ def test_dual_server_quantile_regression_learns_distribution_stress():
 
     # 4) Wait for training to complete
     time.sleep(30)
+    training_status = requests.get(f"{TRAINING_URL}/data/status", timeout=10)
+    assert training_status.status_code == 200, "training status endpoint failed"
+    assert training_status.json().get("last_retrain") is not None, "training did not complete"
+
     # 5) Sync models to prediction server
     synced = False
     for _ in range(10):

tests/test_training_server.py

@@ -0,0 +1,50 @@
+def test_training_replaces_cold_start_models(tmp_path, monkeypatch):
+    """A retrain must replace the one-point cold-start models."""
+    monkeypatch.setenv("LATENCY_MODEL_TYPE", "xgboost")
+    monkeypatch.setenv("LATENCY_ENSEMBLE_MODE", "false")
+
+    from training import training_server as server
+
+    monkeypatch.setattr(server.settings, "TTFT_MODEL_PATH", str(tmp_path / "ttft.joblib"))
+    monkeypatch.setattr(server.settings, "TPOT_MODEL_PATH", str(tmp_path / "tpot.joblib"))
+    monkeypatch.setattr(server.settings, "TTFT_SCALER_PATH", str(tmp_path / "ttft_scaler.joblib"))
+    monkeypatch.setattr(server.settings, "TPOT_SCALER_PATH", str(tmp_path / "tpot_scaler.joblib"))
+    monkeypatch.setattr(server.settings, "MIN_SAMPLES_FOR_RETRAIN", 10)
+    monkeypatch.setattr(server.settings, "MIN_SAMPLES_FOR_RETRAIN_FRESH", 10)
+    monkeypatch.setattr(server.settings, "ENSEMBLE_MODE", False)
+
+    predictor = server.LatencyPredictor()
+    predictor.load_models()
+
+    features = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 400,
+        "num_request_waiting": 3,
+        "num_request_running": 2,
+        "num_tokens_generated": 10,
+        "prefix_cache_score": 0.5,
+    }
+    initial = predictor.predict(features)[:2]
+
+    for i in range(100):
+        input_tokens = 50 + i
+        predictor.add_training_sample(
+            {
+                "kv_cache_percentage": 0.2 + (i % 5) * 0.1,
+                "input_token_length": input_tokens,
+                "num_request_waiting": i % 4,
+                "num_request_running": 1 + i % 3,
+                "actual_ttft_ms": 100.0 + input_tokens * 2.0,
+                "actual_tpot_ms": 20.0 + input_tokens * 0.5,
+                "num_tokens_generated": 1 + i % 10,
+                "prefix_cache_score": (i % 10) / 10.0,
+            }
+        )
+
+    predictor.train()
+    updated = predictor.predict(features)[:2]
+
+    assert predictor.last_retrain_time is not None
+    assert updated != initial
+    assert updated != (10.0, 10.0)
+    assert predictor.is_ready

training/training_server.py

@@ -56,10 +56,10 @@
 from common.types import ModelType, ObjectiveType, QueueGatedModel, RandomDropDeque
 
 
-@staticmethod
 def _drop_timestamp(rows: list[dict]) -> list[dict]:
     return [{k: v for k, v in row.items() if k != "timestamp"} for row in rows]
 
+
 # --- Configuration ---
 class Settings:
     """
@@ -711,7 +711,7 @@ def _calculate_metrics_on_test(self, model, scaler, test_data, model_name, targe
         Returns (None, None, None) on failure.
         """
         try:
-            clean_metrics = self._drop_timestamp(test_data)
+            clean_metrics = _drop_timestamp(test_data)
             df_raw = pd.DataFrame(clean_metrics).dropna()
             df_raw = df_raw[df_raw[target_col] > 0]
 
@@ -856,7 +856,7 @@ def train(self):
 
             # Train TTFT
             if ttft_snap:
-                clean_ttft = self._drop_timestamp(ttft_snap)
+                clean_ttft = _drop_timestamp(ttft_snap)
                 raw_ttft = pd.DataFrame(clean_ttft).dropna()
                 raw_ttft = raw_ttft[raw_ttft["actual_ttft_ms"] > 0]
                 df_ttft = self._prepare_features_with_interaction(raw_ttft.copy(), model_type="ttft")
@@ -965,7 +965,7 @@ def train(self):
 
             # Train TPOT
             if tpot_snap:
-                clean_tpot = self._drop_timestamp(tpot_snap)
+                clean_tpot = _drop_timestamp(tpot_snap)
                 df_tpot = pd.DataFrame(clean_tpot).dropna()
                 df_tpot = df_tpot[df_tpot["actual_tpot_ms"] > 0]
                 if settings.TPOT_ZERO_TOKEN_COUNT:
@@ -1043,7 +1043,7 @@ def train(self):
                             ("tpot_queued", tpot_queued, "tpot", "actual_tpot_ms", "queued"),
                         ]:
                             try:
-                                clean = self._drop_timestamp(samples)
+                                clean = _drop_timestamp(samples)
                                 raw = pd.DataFrame(clean).dropna()
                                 raw = raw[raw[target_col] > 0]
                                 X = self._prepare_features_for_ensemble(raw.copy(), model_name, regime)