from __future__ import annotations import base64 import hashlib import json from datetime import UTC, datetime, timedelta import pytest from crypto_spot_bot.training_coordination import TrainingCoordinator, _validate_symbol_models def test_training_coordinator_claims_and_completes_job(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) requested = coordinator.request_retrain({"source": "android"}) job_id = requested["job"]["id"] heartbeat = coordinator.heartbeat({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"}) claimed = coordinator.claim({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"}) lease_token = claimed["lease_token"] assert requested["queued"] is True assert heartbeat["status"]["agent_online"] is True assert claimed["claimed"] is True assert claimed["job"]["id"] == job_id assert coordinator.status()["active_job"]["status"] == "running" progress = coordinator.progress( job_id, { "status": "running", "phase": "training", "progress_percent": 42, "message": "epoch 1", "lease_token": lease_token, }, ) assert progress["job"]["phase"] == "training" assert progress["job"]["progress_percent"] == 42 assert coordinator.status()["active_job"]["message"] == "epoch 1" completed = coordinator.complete( job_id, {"success": True, "message": "ok", "lease_token": lease_token}, ) assert completed["job"]["status"] == "completed" assert coordinator.status()["active_job"] is None def test_training_coordinator_preserves_boolean_resume_candidate_parameter(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) requested = coordinator.request_retrain( {"source": "recovery", "parameters": {"resume_candidate": True}} ) assert requested["job"]["parameters"] == {"resume_candidate": True} def test_training_coordinator_sanitizes_independent_training_parameters(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) requested = coordinator.request_retrain( { "source": "recovery", "parameters": { "pooled": False, "limit": 6000, "validation_window": 720, "ensemble_seeds": "7,19", "selection_folds": 3, "learning_rate": 0.0007, "weight_decay": 0.0005, "horizon": 12, "horizons": "3,6,12,24", "patience": 8, "seed": 7, }, } ) assert requested["job"]["parameters"] == { "pooled": False, "limit": 6000, "validation_window": 720, "ensemble_seeds": "7,19", "selection_folds": 3, "learning_rate": 0.0007, "weight_decay": 0.0005, "horizon": 12, "horizons": "3,6,12,24", "patience": 8, "seed": 7, } def test_training_coordinator_reports_worker_identity_from_heartbeat(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) heartbeat = coordinator.heartbeat( { "worker_id": "SEVENHILL:G:\\Repos\\TradeBot", "name": "SEVENHILL", "path": "G:\\Repos\\TradeBot", } ) assert heartbeat["worker"]["name"] == "SEVENHILL" assert heartbeat["worker"]["path"] == "G:\\Repos\\TradeBot" assert heartbeat["status"]["worker"] == heartbeat["worker"] def test_training_coordinator_records_rejected_candidate_as_completed_training(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) job = coordinator.request_retrain({"source": "android"})["job"] lease_token = coordinator.claim({"worker_id": "worker-1"})["lease_token"] completed = coordinator.complete( job["id"], { "success": True, "message": "training completed; candidate rejected by quality gate", "summary": {"accepted": False, "reason": "candidate_failed_honest_validation"}, "lease_token": lease_token, }, ) assert completed["job"]["status"] == "completed" assert completed["job"]["phase"] == "completed" assert completed["job"]["progress_percent"] == 100 assert completed["job"]["model_decision"] == "rejected" def test_training_coordinator_accepts_chunked_artifact_upload(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) job = coordinator.request_retrain({"source": "test"})["job"] lease_token = coordinator.claim({"worker_id": "test-worker"})["lease_token"] payload = b'{"type":"pytorch_recurrent_forecaster","symbols":{}}\n' sha256 = hashlib.sha256(payload).hexdigest() first = payload[:20] second = payload[20:] part_1 = coordinator.save_artifact_chunk( job["id"], { "name": "lstm_forecaster.json", "index": 0, "total": 2, "sha256": sha256, "data_base64": base64.b64encode(first).decode("ascii"), "lease_token": lease_token, }, ) part_2 = coordinator.save_artifact_chunk( job["id"], { "name": "lstm_forecaster.json", "index": 1, "total": 2, "sha256": sha256, "data_base64": base64.b64encode(second).decode("ascii"), "lease_token": lease_token, }, ) assert part_1["complete"] is False assert part_2["complete"] is True assert not (tmp_path / "lstm_forecaster.json").exists() assert (tmp_path / ".training_uploads" / job["id"] / "ready" / "lstm_forecaster.json").read_bytes() == payload assert coordinator.status()["latest_job"]["artifacts"][0]["sha256"] == sha256 def test_model_validation_accepts_multitask_ensemble_members() -> None: head = { "state_dict": {"weight_ih_l0": [[0.0]]}, "head_hidden_weight": [[0.0]], "head_hidden_bias": [0.0], "return_head_weight": [[0.0]], "return_head_bias": [0.0], "event_head_weight": [[0.0]], "event_head_bias": [0.0], } symbols = { "BTCUSDT": { "model": "torch_gru", "lookback": 8, "input_size": 2, "hidden_size": 4, "multitask_head": True, "ensemble_members": [head, head], } } _validate_symbol_models(symbols) def test_running_claimed_job_keeps_agent_online_when_heartbeat_is_stale(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) coordinator.request_retrain({"source": "android"}) coordinator.claim({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"}) state_path = tmp_path / "training_coordination.json" state = json.loads(state_path.read_text(encoding="utf-8")) state["worker"]["last_seen_at"] = "2026-01-01T00:00:00+00:00" state_path.write_text(json.dumps(state), encoding="utf-8") status = coordinator.status() assert status["agent_recently_seen"] is False assert status["agent_busy"] is True assert status["agent_online"] is True def test_stale_training_lease_is_requeued_and_old_lease_is_rejected(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) job = coordinator.request_retrain({"source": "android"})["job"] first_claim = coordinator.claim({"worker_id": "worker-1"}) state_path = tmp_path / "training_coordination.json" state = json.loads(state_path.read_text(encoding="utf-8")) state["jobs"][0]["updated_at"] = ( datetime.now(UTC) - timedelta(minutes=11) ).isoformat() state_path.write_text(json.dumps(state), encoding="utf-8") second_claim = coordinator.claim({"worker_id": "worker-2"}) assert second_claim["claimed"] is True assert second_claim["job"]["id"] == job["id"] assert second_claim["job"]["attempts"] == 2 assert second_claim["lease_token"] != first_claim["lease_token"] with pytest.raises(ValueError, match="lease"): coordinator.progress( job["id"], { "phase": "training", "progress_percent": 10, "lease_token": first_claim["lease_token"], }, ) def test_training_upload_rejects_unknown_job(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) payload = b"{}" with pytest.raises(ValueError, match="not found"): coordinator.save_artifact_chunk( "11111111-1111-4111-8111-111111111111", { "name": "lstm_forecaster.json", "index": 0, "total": 1, "sha256": hashlib.sha256(payload).hexdigest(), "data_base64": base64.b64encode(payload).decode("ascii"), }, ) def test_training_bundle_promotes_only_after_successful_guard(tmp_path) -> None: coordinator = TrainingCoordinator(tmp_path) job = coordinator.request_retrain({"source": "test"})["job"] lease_token = coordinator.claim({"worker_id": "worker-1"})["lease_token"] model = { "type": "pytorch_recurrent_forecaster", "symbols": { "BTCUSDT": { "model": "torch_gru", "lookback": 4, "input_size": 1, "hidden_size": 1, "state_dict": {"weight_ih_l0": [[0.0]]}, "head_weight": [[0.0]], "head_bias": [0.0], } }, } model_payload = (json.dumps(model) + "\n").encode() model_sha256 = hashlib.sha256(model_payload).hexdigest() artifacts = { "lstm_forecaster.json": model, "torch_retrain_guard.json": { "accepted": True, "candidate_artifact_sha256": model_sha256, }, "torch_threshold_calibration.json": { "artifact_sha256": model_sha256, "validation": { "passed": True, "protocol": "untouched_model_holdout_with_threshold_walk_forward", } }, } for name, data in artifacts.items(): payload = model_payload if name == "lstm_forecaster.json" else (json.dumps(data) + "\n").encode() coordinator.save_artifact_chunk( job["id"], { "name": name, "index": 0, "total": 1, "sha256": hashlib.sha256(payload).hexdigest(), "data_base64": base64.b64encode(payload).decode("ascii"), "lease_token": lease_token, }, ) completed = coordinator.complete( job["id"], {"success": True, "lease_token": lease_token}, ) assert completed["job"]["status"] == "completed" assert json.loads((tmp_path / "lstm_forecaster.json").read_text())["symbols"]["BTCUSDT"]