Codex
89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
|
|
from crypto_spot_bot.training_coordination import TrainingCoordinator
|
|
|
|
|
|
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"})
|
|
|
|
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"},
|
|
)
|
|
|
|
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"})
|
|
|
|
assert completed["job"]["status"] == "completed"
|
|
assert coordinator.status()["active_job"] is None
|
|
|
|
|
|
def test_training_coordinator_accepts_chunked_artifact_upload(tmp_path) -> None:
|
|
coordinator = TrainingCoordinator(tmp_path)
|
|
job = coordinator.request_retrain({"source": "test"})["job"]
|
|
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"),
|
|
},
|
|
)
|
|
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"),
|
|
},
|
|
)
|
|
|
|
assert part_1["complete"] is False
|
|
assert part_2["complete"] is True
|
|
assert (tmp_path / "lstm_forecaster.json").read_bytes() == payload
|
|
assert coordinator.status()["latest_job"]["artifacts"][0]["sha256"] == sha256
|
|
|
|
|
|
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
|