317 lines
11 KiB
Python
317 lines
11 KiB
Python
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"]
|