diff --git a/tests/test_dual_server_client.py b/tests/test_dual_server_client.py index daa937b1..3a2fe8f5 100644 --- a/tests/test_dual_server_client.py +++ b/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): diff --git a/tests/test_training_server.py b/tests/test_training_server.py new file mode 100644 index 00000000..095b988e --- /dev/null +++ b/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 diff --git a/training/training_server.py b/training/training_server.py index 109cf66d..a674708e 100644 --- a/training/training_server.py +++ b/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)