From 9c1fc64fa510b772eac562f4683869a152057d6f Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:21:44 +0200 Subject: [PATCH 1/6] feat(cluster): worker registry persistence, split-brain protection, circuit breaker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #640. Four interconnected changes: 1. **Worker registry persistence** — New WorkerRegistryStore (BaseStore + SQLite) persists worker registrations so the in-memory dict survives controller restarts. On startup, loaded workers are marked 'stale' until they re-heartbeat; local worker is skipped during load. 2. **Split-brain protection via generation counter** — A generation counter (stored in the same SQLite table) is incremented on each controller start. Workers echo it in heartbeats; stale generations are rejected. Legacy workers that don't send generation get a pass (None). 3. **Circuit breaker for failing workers** — New FailureTracker with sliding window (default: 3 failures in 60s). TaskRouter skips tripped workers instead of serially blocking 60s * N_workers. Success clears the circuit. 4. **In-flight routing recovery** — Addressed by Fix 1: persisted workers survive controller restart; stale-until-heartbeat prevents routing to workers that haven't reconnected. Files: - NEW: tinyagentos/cluster/worker_registry_store.py - NEW: tinyagentos/cluster/failure_tracker.py - MODIFIED: tinyagentos/cluster/manager.py (init, start, heartbeat, register_worker, unregister_worker, _persist_worker, _load_persisted_workers) - MODIFIED: tinyagentos/cluster/router.py (circuit breaker skip/record) - MODIFIED: tinyagentos/app.py (wire new stores) - NEW: tests/test_cluster_persistence.py (24 tests: store, tracker, manager persistence, router circuit breaker) Tests: 79/79 pass (55 existing + 24 new cluster tests) --- tests/test_cluster_persistence.py | 637 +++++++++++++++++++ tinyagentos/app.py | 16 +- tinyagentos/cluster/failure_tracker.py | 97 +++ tinyagentos/cluster/manager.py | 186 +++++- tinyagentos/cluster/router.py | 23 +- tinyagentos/cluster/worker_registry_store.py | 195 ++++++ 6 files changed, 1151 insertions(+), 3 deletions(-) create mode 100644 tests/test_cluster_persistence.py create mode 100644 tinyagentos/cluster/failure_tracker.py create mode 100644 tinyagentos/cluster/worker_registry_store.py diff --git a/tests/test_cluster_persistence.py b/tests/test_cluster_persistence.py new file mode 100644 index 000000000..17842f0ec --- /dev/null +++ b/tests/test_cluster_persistence.py @@ -0,0 +1,637 @@ +"""Tests for taOS #640: cluster persistence, split-brain, circuit breaker.""" + +from __future__ import annotations + +import asyncio +import json +import tempfile +import time +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from tinyagentos.cluster.failure_tracker import FailureTracker +from tinyagentos.cluster.manager import ClusterManager, HEARTBEAT_TIMEOUT +from tinyagentos.cluster.router import TaskRouter +from tinyagentos.cluster.worker_protocol import WorkerInfo +from tinyagentos.cluster.worker_registry_store import WorkerRegistryStore + + +def _make_worker(name: str, capabilities: list[str] | None = None, + load: float = 0.0, status: str = "online", + url: str = "http://localhost:9000") -> WorkerInfo: + return WorkerInfo( + name=name, + url=url, + capabilities=capabilities or ["chat", "embed"], + load=load, + status=status, + platform="linux", + ) + + +# ── FailureTracker tests (Fix 3: circuit breaker) ────────────────────── + +class TestFailureTracker: + def test_not_tripped_by_default(self): + ft = FailureTracker() + assert not ft.is_tripped("w1") + + def test_trips_after_threshold(self): + ft = FailureTracker(failure_threshold=3) + ft.record_failure("w1") + ft.record_failure("w1") + assert not ft.is_tripped("w1") + ft.record_failure("w1") + assert ft.is_tripped("w1") + + def test_success_resets_circuit(self): + ft = FailureTracker(failure_threshold=2) + ft.record_failure("w1") + ft.record_failure("w1") + assert ft.is_tripped("w1") + ft.record_success("w1") + assert not ft.is_tripped("w1") + + def test_separate_per_worker(self): + ft = FailureTracker(failure_threshold=2) + ft.record_failure("w1") + ft.record_failure("w1") + ft.record_failure("w2") + assert ft.is_tripped("w1") + assert not ft.is_tripped("w2") + + def test_reset_worker(self): + ft = FailureTracker(failure_threshold=1) + ft.record_failure("w1") + assert ft.is_tripped("w1") + ft.reset("w1") + assert not ft.is_tripped("w1") + + def test_reset_all(self): + ft = FailureTracker(failure_threshold=1) + ft.record_failure("w1") + ft.record_failure("w2") + ft.reset_all() + assert not ft.is_tripped("w1") + assert not ft.is_tripped("w2") + + +# ── WorkerRegistryStore tests (Fix 1: persistence) ───────────────────── + +@pytest.mark.asyncio +class TestWorkerRegistryStore: + async def test_upsert_and_load(self): + with tempfile.TemporaryDirectory() as td: + db_path = Path(td) / "test.db" + store = WorkerRegistryStore(db_path) + await store.init() + + info = { + "name": "gpu-box", + "url": "http://gpu:9000", + "hardware": json.dumps({"gpu": {"model": "RTX 4090"}}), + "backends": json.dumps([]), + "models": json.dumps(["llama3"]), + "available_models": json.dumps([]), + "capabilities": json.dumps(["chat", "embed"]), + "status": "online", + "last_heartbeat": 1234567890.0, + "registered_at": 1234567890.0, + "load": 0.5, + "platform": "linux", + "tier_id": "x86-cuda-24gb", + "potential_capabilities": json.dumps([]), + "kv_cache_quant_support": json.dumps(["fp16"]), + "kv_cache_quant_k_support": json.dumps(["fp16"]), + "kv_cache_quant_v_support": json.dumps(["fp16"]), + "kv_cache_quant_boundary_layer_protect": 0, + "worker_url": "http://gpu:6970", + "signing_key": b"secret123456789012345678901234567890", + "tls_cert_provider": None, + "host_lan_ip": "192.168.1.100", + "storage_cap_bytes": 100_000_000, + "storage_used_bytes": 50_000_000, + "bytes_deduped_total": 10_000_000, + "worker_lxc_image_version": "ubuntu/24.04/amd64", + "degraded": 0, + "degraded_reason": None, + "free_vram_mb": 16000, + "used_vram_mb": 8000, + } + await store.upsert_worker(info) + + rows = await store.load_all() + assert len(rows) == 1 + assert rows[0]["name"] == "gpu-box" + assert rows[0]["url"] == "http://gpu:9000" + assert rows[0]["status"] == "online" + assert json.loads(rows[0]["models"]) == ["llama3"] + assert rows[0]["free_vram_mb"] == 16000 + assert rows[0]["signing_key"] == b"secret123456789012345678901234567890" + + await store.close() + + async def test_upsert_overwrites(self): + with tempfile.TemporaryDirectory() as td: + db_path = Path(td) / "test.db" + store = WorkerRegistryStore(db_path) + await store.init() + + info1 = { + "name": "gpu-box", + "url": "http://gpu:9000", + "hardware": json.dumps({}), + "backends": json.dumps([]), + "models": json.dumps(["llama3"]), + "available_models": json.dumps([]), + "capabilities": json.dumps(["chat"]), + "status": "online", + "last_heartbeat": 100.0, + "registered_at": 100.0, + "load": 0.0, + "platform": "linux", + "tier_id": "", + "potential_capabilities": json.dumps([]), + "kv_cache_quant_support": json.dumps(["fp16"]), + "kv_cache_quant_k_support": json.dumps(["fp16"]), + "kv_cache_quant_v_support": json.dumps(["fp16"]), + "kv_cache_quant_boundary_layer_protect": 0, + "worker_url": None, + "signing_key": b"", + "tls_cert_provider": None, + "host_lan_ip": None, + "storage_cap_bytes": 0, + "storage_used_bytes": 0, + "bytes_deduped_total": 0, + "worker_lxc_image_version": None, + "degraded": 0, + "degraded_reason": None, + "free_vram_mb": None, + "used_vram_mb": None, + } + await store.upsert_worker(info1) + + info2 = {**info1, "status": "offline", "load": 0.9, + "free_vram_mb": 4000, "used_vram_mb": 20000} + await store.upsert_worker(info2) + + rows = await store.load_all() + assert len(rows) == 1 + assert rows[0]["status"] == "offline" + assert rows[0]["load"] == 0.9 + assert rows[0]["free_vram_mb"] == 4000 + + await store.close() + + async def test_remove_worker(self): + with tempfile.TemporaryDirectory() as td: + db_path = Path(td) / "test.db" + store = WorkerRegistryStore(db_path) + await store.init() + + info = { + "name": "gpu-box", + "url": "http://gpu:9000", + "hardware": json.dumps({}), + "backends": json.dumps([]), + "models": json.dumps([]), + "available_models": json.dumps([]), + "capabilities": json.dumps(["chat"]), + "status": "online", + "last_heartbeat": 100.0, + "registered_at": 100.0, + "load": 0.0, + "platform": "linux", + "tier_id": "", + "potential_capabilities": json.dumps([]), + "kv_cache_quant_support": json.dumps(["fp16"]), + "kv_cache_quant_k_support": json.dumps(["fp16"]), + "kv_cache_quant_v_support": json.dumps(["fp16"]), + "kv_cache_quant_boundary_layer_protect": 0, + "worker_url": None, + "signing_key": b"", + "tls_cert_provider": None, + "host_lan_ip": None, + "storage_cap_bytes": 0, + "storage_used_bytes": 0, + "bytes_deduped_total": 0, + "worker_lxc_image_version": None, + "degraded": 0, + "degraded_reason": None, + "free_vram_mb": None, + "used_vram_mb": None, + } + await store.upsert_worker(info) + assert len(await store.load_all()) == 1 + await store.remove_worker("gpu-box") + assert len(await store.load_all()) == 0 + + await store.close() + + async def test_generation_counter(self): + with tempfile.TemporaryDirectory() as td: + db_path = Path(td) / "test.db" + store = WorkerRegistryStore(db_path) + await store.init() + + # First generation should be > 0 + assert await store.current_generation() > 0 + + g1 = await store.increment_generation() + g2 = await store.increment_generation() + assert g2 == g1 + 1 + + # Re-open — generation persists + await store.close() + store2 = WorkerRegistryStore(db_path) + await store2.init() + assert await store2.current_generation() == g2 + await store2.close() + + +# ── ClusterManager persistence tests (Fix 1 + Fix 2) ─────────────────── + +@pytest.mark.asyncio +class TestClusterManagerPersistence: + """Tests for loading from store, persisting on register/heartbeat, + and generation-based split-brain protection (taOS #640).""" + + async def _new_store(self, tmpdir) -> WorkerRegistryStore: + db_path = tmpdir / "workers.db" + store = WorkerRegistryStore(db_path) + await store.init() + return store + + async def test_loads_persisted_workers_on_start(self, tmp_path): + store = await self._new_store(tmp_path) + # Pre-populate a worker in the store + await store.upsert_worker({ + "name": "gpu-box", + "url": "http://gpu:9000", + "hardware": json.dumps({}), + "backends": json.dumps([]), + "models": json.dumps([]), + "available_models": json.dumps([]), + "capabilities": json.dumps(["chat"]), + "status": "online", + "last_heartbeat": time.time(), + "registered_at": time.time(), + "load": 0.5, + "platform": "linux", + "tier_id": "", + "potential_capabilities": json.dumps([]), + "kv_cache_quant_support": json.dumps(["fp16"]), + "kv_cache_quant_k_support": json.dumps(["fp16"]), + "kv_cache_quant_v_support": json.dumps(["fp16"]), + "kv_cache_quant_boundary_layer_protect": 0, + "worker_url": None, + "signing_key": b"", + "tls_cert_provider": None, + "host_lan_ip": None, + "storage_cap_bytes": 0, + "storage_used_bytes": 0, + "bytes_deduped_total": 0, + "worker_lxc_image_version": None, + "degraded": 0, + "degraded_reason": None, + "free_vram_mb": None, + "used_vram_mb": None, + }) + + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + + workers = mgr.get_workers() + assert len(workers) == 1 + assert workers[0].name == "gpu-box" + # Loaded workers start as "stale" — they must re-heartbeat + assert workers[0].status == "stale" + assert workers[0].load == 0.5 + + await mgr.stop() + await store.close() + + async def test_stale_worker_excluded_from_routing(self, tmp_path): + store = await self._new_store(tmp_path) + await store.upsert_worker({ + "name": "gpu-box", + "url": "http://gpu:9000", + "hardware": json.dumps({}), + "backends": json.dumps([]), + "models": json.dumps([]), + "available_models": json.dumps([]), + "capabilities": json.dumps(["chat"]), + "status": "online", + "last_heartbeat": time.time(), + "registered_at": time.time(), + "load": 0.5, + "platform": "linux", + "tier_id": "", + "potential_capabilities": json.dumps([]), + "kv_cache_quant_support": json.dumps(["fp16"]), + "kv_cache_quant_k_support": json.dumps(["fp16"]), + "kv_cache_quant_v_support": json.dumps(["fp16"]), + "kv_cache_quant_boundary_layer_protect": 0, + "worker_url": None, + "signing_key": b"", + "tls_cert_provider": None, + "host_lan_ip": None, + "storage_cap_bytes": 0, + "storage_used_bytes": 0, + "bytes_deduped_total": 0, + "worker_lxc_image_version": None, + "degraded": 0, + "degraded_reason": None, + "free_vram_mb": None, + "used_vram_mb": None, + }) + + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + + # Stale workers are NOT "online" so get_workers_for_capability skips them + result = mgr.get_workers_for_capability("chat") + assert len(result) == 0 + + await mgr.stop() + await store.close() + + async def test_heartbeat_revives_stale_worker(self, tmp_path): + store = await self._new_store(tmp_path) + await store.upsert_worker({ + "name": "gpu-box", + "url": "http://gpu:9000", + "hardware": json.dumps({}), + "backends": json.dumps([]), + "models": json.dumps([]), + "available_models": json.dumps([]), + "capabilities": json.dumps(["chat"]), + "status": "online", + "last_heartbeat": time.time(), + "registered_at": time.time(), + "load": 0.5, + "platform": "linux", + "tier_id": "", + "potential_capabilities": json.dumps([]), + "kv_cache_quant_support": json.dumps(["fp16"]), + "kv_cache_quant_k_support": json.dumps(["fp16"]), + "kv_cache_quant_v_support": json.dumps(["fp16"]), + "kv_cache_quant_boundary_layer_protect": 0, + "worker_url": None, + "signing_key": b"", + "tls_cert_provider": None, + "host_lan_ip": None, + "storage_cap_bytes": 0, + "storage_used_bytes": 0, + "bytes_deduped_total": 0, + "worker_lxc_image_version": None, + "degraded": 0, + "degraded_reason": None, + "free_vram_mb": None, + "used_vram_mb": None, + }) + + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + assert mgr.get_worker("gpu-box").status == "stale" + + # Heartbeat revives from stale → online + mgr.heartbeat("gpu-box", load=0.1, generation=mgr.generation) + assert mgr.get_worker("gpu-box").status == "online" + assert mgr.get_worker("gpu-box").load == 0.1 + + await mgr.stop() + await store.close() + + async def test_register_worker_persists(self, tmp_path): + store = await self._new_store(tmp_path) + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + + w = _make_worker("gpu-box", capabilities=["chat"]) + await mgr.register_worker(w) + + # Check the store has the worker + rows = await store.load_all() + assert len(rows) == 1 + assert rows[0]["name"] == "gpu-box" + assert rows[0]["status"] == "online" + + await mgr.stop() + await store.close() + + async def test_unregister_worker_removes_from_store(self, tmp_path): + store = await self._new_store(tmp_path) + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + + w = _make_worker("gpu-box", capabilities=["chat"]) + await mgr.register_worker(w) + assert len(await store.load_all()) == 1 + + await mgr.unregister_worker("gpu-box") + assert len(await store.load_all()) == 0 + + await mgr.stop() + await store.close() + + async def test_generation_increments_on_start(self, tmp_path): + store = await self._new_store(tmp_path) + mgr1 = ClusterManager(worker_registry_store=store) + await mgr1.start() + g1 = mgr1.generation + await mgr1.stop() + + # Create a fresh store connection so generation persists + db_path = tmp_path / "workers.db" + store2 = WorkerRegistryStore(db_path) + await store2.init() + mgr2 = ClusterManager(worker_registry_store=store2) + await mgr2.start() + g2 = mgr2.generation + assert g2 == g1 + 1 + await mgr2.stop() + await store2.close() + await store.close() + + async def test_heartbeat_rejected_for_stale_generation(self, tmp_path): + """Split-brain: heartbeat with wrong generation is rejected.""" + store = await self._new_store(tmp_path) + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + + w = _make_worker("gpu-box", capabilities=["chat"]) + await mgr.register_worker(w) + current_gen = mgr.generation + + # Heartbeat with a different generation — rejected + ok = mgr.heartbeat("gpu-box", load=0.1, generation=current_gen + 5) + assert ok is False + + # Worker status unchanged + assert mgr.get_worker("gpu-box").status == "online" + + await mgr.stop() + await store.close() + + async def test_heartbeat_accepted_with_null_generation(self, tmp_path): + """Legacy workers that don't send generation are accepted.""" + store = await self._new_store(tmp_path) + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + + w = _make_worker("gpu-box", capabilities=["chat"]) + await mgr.register_worker(w) + + ok = mgr.heartbeat("gpu-box", load=0.1) # no generation + assert ok is True + + await mgr.stop() + await store.close() + + async def test_local_worker_not_loaded_from_store(self, tmp_path): + """The 'local' worker is skipped when loading from store.""" + store = await self._new_store(tmp_path) + # Put a "local" worker in the store + await store.upsert_worker({ + "name": "local", + "url": "http://local:9000", + "hardware": json.dumps({}), + "backends": json.dumps([]), + "models": json.dumps([]), + "available_models": json.dumps([]), + "capabilities": json.dumps(["chat"]), + "status": "online", + "last_heartbeat": time.time(), + "registered_at": time.time(), + "load": 0.0, + "platform": "linux", + "tier_id": "", + "potential_capabilities": json.dumps([]), + "kv_cache_quant_support": json.dumps(["fp16"]), + "kv_cache_quant_k_support": json.dumps(["fp16"]), + "kv_cache_quant_v_support": json.dumps(["fp16"]), + "kv_cache_quant_boundary_layer_protect": 0, + "worker_url": None, + "signing_key": b"", + "tls_cert_provider": None, + "host_lan_ip": None, + "storage_cap_bytes": 0, + "storage_used_bytes": 0, + "bytes_deduped_total": 0, + "worker_lxc_image_version": None, + "degraded": 0, + "degraded_reason": None, + "free_vram_mb": None, + "used_vram_mb": None, + }) + + mgr = ClusterManager(worker_registry_store=store) + await mgr.start() + + # "local" should NOT be loaded from store + assert mgr.get_worker("local") is None + + await mgr.stop() + await store.close() + + async def test_generation_property_without_store(self): + """When no store is wired, generation defaults to 1 and stays stable.""" + mgr = ClusterManager() # no store + await mgr.start() + assert mgr.generation == 1 + await mgr.stop() + + +# ── TaskRouter circuit breaker tests (Fix 3) ─────────────────────────── + +@pytest.mark.asyncio +class TestTaskRouterCircuitBreaker: + async def test_skips_tripped_worker(self): + ft = FailureTracker(failure_threshold=2) + mgr = ClusterManager(failure_tracker=ft) + await mgr.register_worker(_make_worker("w1", capabilities=["chat"], load=0.1, url="http://w1:8000")) + await mgr.register_worker(_make_worker("w2", capabilities=["chat"], load=0.5, url="http://w2:8000")) + + # Trip w1 + ft.record_failure("w1") + ft.record_failure("w1") + assert ft.is_tripped("w1") + + mock_client = AsyncMock(spec=httpx.AsyncClient) + ok_resp = MagicMock() + ok_resp.raise_for_status.return_value = None + ok_resp.json.return_value = {"result": "ok"} + # Only w2 should be tried (w1 skipped) + mock_client.post.return_value = ok_resp + + router = TaskRouter(mgr, mock_client) + data, name = await router.route_request("chat", "POST", "/v1/chat/completions", {}) + + assert data == {"result": "ok"} + assert name == "w2" + # Only 1 call was made (w1 was skipped, w2 succeeded) + assert mock_client.post.call_count == 1 + + async def test_records_failure_on_error(self): + ft = FailureTracker(failure_threshold=5) + mgr = ClusterManager(failure_tracker=ft) + await mgr.register_worker(_make_worker("w1", capabilities=["chat"], load=0.1, url="http://w1:8000")) + + mock_client = AsyncMock(spec=httpx.AsyncClient) + mock_client.post.side_effect = Exception("connection refused") + + router = TaskRouter(mgr, mock_client) + data, name = await router.route_request("chat", "POST", "/v1/chat/completions", {}) + + assert data is None + assert ft.is_tripped("w1") is False # 1 failure, threshold=5 → not tripped + + async def test_clears_failure_on_success(self): + ft = FailureTracker(failure_threshold=3) + mgr = ClusterManager(failure_tracker=ft) + await mgr.register_worker(_make_worker("w1", capabilities=["chat"], load=0.1, url="http://w1:8000")) + + ft.record_failure("w1") + ft.record_failure("w1") + assert not ft.is_tripped("w1") # 2 < 3 + + mock_client = AsyncMock(spec=httpx.AsyncClient) + ok_resp = MagicMock() + ok_resp.raise_for_status.return_value = None + ok_resp.json.return_value = {"result": "ok"} + mock_client.post.return_value = ok_resp + + router = TaskRouter(mgr, mock_client) + data, name = await router.route_request("chat", "POST", "/v1/chat/completions", {}) + assert data == {"result": "ok"} + # After success, circuit should be reset + assert not ft.is_tripped("w1") + + async def test_router_without_tracker_still_works(self): + """Router works fine when no failure_tracker is wired.""" + mgr = ClusterManager() # no failure_tracker + await mgr.register_worker(_make_worker("w1", capabilities=["chat"], load=0.1, url="http://w1:8000")) + + mock_client = AsyncMock(spec=httpx.AsyncClient) + fail_resp = MagicMock() + fail_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "500", request=MagicMock(), response=MagicMock() + ) + ok_resp = MagicMock() + ok_resp.raise_for_status.return_value = None + ok_resp.json.return_value = {"result": "ok"} + mock_client.post.side_effect = [fail_resp, ok_resp] + + # Add a second worker so the first failure doesn't exhaust all workers + await mgr.register_worker(_make_worker("w2", capabilities=["chat"], load=0.5, url="http://w2:8000")) + + router = TaskRouter(mgr, mock_client) + data, name = await router.route_request("chat", "POST", "/v1/chat/completions", {}) + + assert data == {"result": "ok"} + assert name == "w2" + assert mock_client.post.call_count == 2 diff --git a/tinyagentos/app.py b/tinyagentos/app.py index f0ef43fe6..dc6602455 100644 --- a/tinyagentos/app.py +++ b/tinyagentos/app.py @@ -295,6 +295,12 @@ def create_app(data_dir: Path | None = None, catalog_dir: Path | None = None) -> cluster_pairing_store = ClusterPairingStore(data_dir / "cluster_pairing.db") from tinyagentos.cluster.capability_map import CapabilityMap capability_map_store = CapabilityMap(data_dir / "capability_map.db") + # taOS #640: worker registry persistence + generation counter + from tinyagentos.cluster.worker_registry_store import WorkerRegistryStore + worker_registry_store = WorkerRegistryStore(data_dir / "cluster_workers.db") + # taOS #640: circuit breaker for failing workers + from tinyagentos.cluster.failure_tracker import FailureTracker + failure_tracker = FailureTracker() metrics_store = MetricsStore(data_dir / "metrics.db") notif_store = NotificationStore(data_dir / "notifications.db") @@ -331,7 +337,11 @@ async def _probe_backend(backend: dict) -> dict: interval_seconds=30.0, ) fallback = BackendFallback(config.backends, http_client) - cluster_manager = ClusterManager(notifications=notif_store) + cluster_manager = ClusterManager( + notifications=notif_store, + worker_registry_store=worker_registry_store, + failure_tracker=failure_tracker, + ) task_router = TaskRouter(cluster_manager, http_client) cap_checker = CapabilityChecker(hardware_profile, cluster_manager) cluster_manager._capabilities = cap_checker # wire after creation (circular dep) @@ -510,6 +520,10 @@ async def lifespan(app: FastAPI): await license_acceptances_store.init() await agent_model_key_store.init() await cluster_pairing_store.init() + app.state.cluster_pairing = cluster_pairing_store + # taOS #640: worker registry store (persistence + generation counter) + await worker_registry_store.init() + app.state.worker_registry = worker_registry_store await capability_map_store.init() await metrics_store.init() await notif_store.init() diff --git a/tinyagentos/cluster/failure_tracker.py b/tinyagentos/cluster/failure_tracker.py new file mode 100644 index 000000000..89e1ac5e9 --- /dev/null +++ b/tinyagentos/cluster/failure_tracker.py @@ -0,0 +1,97 @@ +"""Circuit breaker for cluster worker routing. + +Tracks per-worker failure counts within a sliding time window so the +TaskRouter can skip workers that have repeatedly failed, instead of +serially trying all N workers and blocking for ``timeout * N`` seconds +before giving up (taOS #640, Fix 3). + +State is in-memory only (no SQLite): the circuit breaker resets on +controller restart, which is the safe default — a freshly-started +controller has no failure history and should give every worker a +clean chance. +""" + +from __future__ import annotations + +import time +from collections import defaultdict +from typing import NamedTuple + + +class _FailureRecord(NamedTuple): + failures: int + first_failure_at: float + + +# Sensible defaults: 5 failures in 60 seconds trips the breaker. +# A tripped worker is excluded from routing for the cooldown period. +DEFAULT_FAILURE_THRESHOLD = 3 +DEFAULT_WINDOW_SECONDS = 60.0 + + +class FailureTracker: + """Per-worker failure counter with a sliding time window. + + After ``failure_threshold`` failures within ``window_seconds``, + ``is_tripped(worker_name)`` returns ``True`` and the router skips + that worker. The window resets once enough time has passed since + the first failure in the current window. + """ + + def __init__( + self, + failure_threshold: int = DEFAULT_FAILURE_THRESHOLD, + window_seconds: float = DEFAULT_WINDOW_SECONDS, + ): + self._failure_threshold = failure_threshold + self._window_seconds = window_seconds + self._records: dict[str, _FailureRecord] = defaultdict( + lambda: _FailureRecord(0, 0.0) + ) + + def record_failure(self, worker_name: str) -> None: + """Register a failure for the given worker. + + After the window has elapsed since the first failure, the counter + resets automatically on the next failure. + """ + now = time.time() + rec = self._records[worker_name] + if rec.first_failure_at > 0 and (now - rec.first_failure_at) > self._window_seconds: + # Window expired — reset + self._records[worker_name] = _FailureRecord(1, now) + elif rec.first_failure_at == 0: + self._records[worker_name] = _FailureRecord(1, now) + else: + self._records[worker_name] = _FailureRecord( + rec.failures + 1, rec.first_failure_at + ) + + def record_success(self, worker_name: str) -> None: + """Clear the failure record for a worker after a successful request.""" + self._records.pop(worker_name, None) + + def is_tripped(self, worker_name: str) -> bool: + """Return True if the worker's circuit breaker is currently open. + + The breaker opens when ``failure_threshold`` failures have occurred + within ``window_seconds`` and not enough time has elapsed since the + first failure to reset the window. + """ + rec = self._records.get(worker_name) + if rec is None: + return False + now = time.time() + if rec.first_failure_at > 0 and (now - rec.first_failure_at) > self._window_seconds: + # Window expired — auto-reset + self._records.pop(worker_name, None) + return False + return rec.failures >= self._failure_threshold + + def reset(self, worker_name: str) -> None: + """Manually reset the circuit breaker for a worker.""" + self._records.pop(worker_name, None) + + def reset_all(self) -> None: + """Reset all circuit breakers.""" + self._records.clear() diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index c02fea878..b71b8446c 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -1,5 +1,6 @@ from __future__ import annotations import asyncio +import json import logging import secrets import time @@ -8,6 +9,8 @@ from tinyagentos.cluster.worker_protocol import GpuLease, WorkerInfo if TYPE_CHECKING: + from tinyagentos.cluster.failure_tracker import FailureTracker + from tinyagentos.cluster.worker_registry_store import WorkerRegistryStore from tinyagentos.scheduler.gpu_arbiter import GpuArbiter logger = logging.getLogger(__name__) @@ -38,7 +41,13 @@ def _format_hw(hw) -> str: class ClusterManager: - def __init__(self, notifications=None, capabilities=None): + def __init__( + self, + notifications=None, + capabilities=None, + worker_registry_store=None, + failure_tracker=None, + ): self._workers: dict[str, WorkerInfo] = {} self._leases: dict[str, GpuLease] = {} # Serializes all lease-table mutations (claim/release/renew/expiry @@ -54,8 +63,21 @@ def __init__(self, notifications=None, capabilities=None): self._ever_seen: set[str] = set() # Strong references to background tasks to prevent GC before completion. self._background_tasks: set[asyncio.Task] = set() + # taOS #640: persistence + split-brain + circuit breaker + self._registry_store: WorkerRegistryStore | None = worker_registry_store + self._failure_tracker: FailureTracker | None = failure_tracker + self._generation: int = 1 # incremented in start() when store is wired async def start(self): + # taOS #640: increment generation on each controller start (split-brain + # protection). Workers must echo the current generation in heartbeats; + # a stale generation means another controller instance is (or was) active. + if self._registry_store is not None: + self._generation = await self._registry_store.increment_generation() + logger.info("Cluster generation: %d", self._generation) + # Load persisted workers and mark them stale — they must + # re-register or heartbeat before being used for routing. + await self._load_persisted_workers() self._monitor_task = asyncio.create_task(self._monitor_loop()) async def stop(self): @@ -83,6 +105,13 @@ async def register_worker(self, info: WorkerInfo) -> None: self._workers[info.name] = info logger.info(f"Worker registered: {info.name} ({info.platform}, {len(info.capabilities)} capabilities)") + # taOS #640: persist worker to SQLite so the registry survives restarts. + if self._registry_store is not None: + try: + await self._persist_worker(info) + except Exception: + logger.exception("Failed to persist worker '%s'", info.name) + # The "local" worker is the controller registering itself on every boot; # that is not a noteworthy cluster event, so do not notify for it. Only # real remote workers joining or coming back online should notify. @@ -211,6 +240,7 @@ def heartbeat( bytes_deduped_total: int | None = None, status: str | None = None, drain_reason: str | None = None, + generation: int | None = None, ) -> bool: """Accept a worker heartbeat. @@ -225,10 +255,25 @@ def heartbeat( controller treats it identically to a controller-initiated drain: no new tasks are routed, existing leases complete, and the monitor loop auto-completes the drain when all leases are released. + + taOS #640: When ``generation`` is provided and does not match the + controller's current generation, the heartbeat is rejected — this + protects against split-brain where two controllers are active + simultaneously. Workers that do not yet send generation (legacy) + are accepted for backward compatibility. """ worker = self._workers.get(name) if not worker: return False + # taOS #640: split-brain protection — reject heartbeat from a worker + # that is echoing a different generation (another active controller). + # Legacy workers that don't send generation get a pass (None). + if generation is not None and generation != self._generation: + logger.warning( + "Heartbeat from '%s' rejected: generation %s != controller %s", + name, generation, self._generation, + ) + return False prev_status = worker.status worker.last_heartbeat = time.time() worker.load = load @@ -340,6 +385,14 @@ def heartbeat( ) except RuntimeError: pass + # taOS #640: persist updated worker state to SQLite. + if self._registry_store is not None: + try: + asyncio.get_running_loop().create_task( + self._persist_worker(worker) + ) + except RuntimeError: + pass # no running loop (e.g. sync tests) — skip gracefully # Fire worker.online notification when a previously-offline worker recovers. # heartbeat() is sync, so schedule the async emit as a background task. # Only fire for workers that are genuinely "online" — a worker recovering @@ -388,6 +441,12 @@ async def unregister_worker(self, name: str) -> bool: self._workers.pop(name, None) logger.info("Worker '%s' unregistered — %d leases released", name, len(lids)) + # taOS #640: remove from persistent store. + if self._registry_store is not None: + try: + await self._registry_store.remove_worker(name) + except Exception: + logger.exception("Failed to remove worker '%s' from store", name) return True def get_workers(self) -> list[WorkerInfo]: @@ -749,6 +808,131 @@ def aggregate_catalog(self) -> dict: "models": flat_models, } + # ── taOS #640: persistence + split-brain helpers ─────────────────── + + @property + def generation(self) -> int: + """Current controller generation for split-brain protection.""" + return self._generation + + @property + def failure_tracker(self): + """Circuit breaker for worker routing failures (Fix 3).""" + return self._failure_tracker + + async def _persist_worker(self, worker: WorkerInfo) -> None: + """Serialize a WorkerInfo to a dict and upsert into the registry store.""" + if self._registry_store is None: + return + info: dict = { + "name": worker.name, + "url": worker.url, + "hardware": json.dumps(worker.hardware or {}), + "backends": json.dumps(worker.backends or []), + "models": json.dumps(worker.models or []), + "available_models": json.dumps(worker.available_models or []), + "capabilities": json.dumps(worker.capabilities or []), + "status": worker.status, + "last_heartbeat": worker.last_heartbeat, + "registered_at": worker.registered_at, + "load": worker.load, + "platform": worker.platform, + "tier_id": worker.tier_id, + "potential_capabilities": json.dumps(worker.potential_capabilities or []), + "kv_cache_quant_support": json.dumps( + worker.kv_cache_quant_support or ["fp16"] + ), + "kv_cache_quant_k_support": json.dumps( + worker.kv_cache_quant_k_support or ["fp16"] + ), + "kv_cache_quant_v_support": json.dumps( + worker.kv_cache_quant_v_support or ["fp16"] + ), + "kv_cache_quant_boundary_layer_protect": int( + worker.kv_cache_quant_boundary_layer_protect + ), + "worker_url": worker.worker_url, + "signing_key": worker.signing_key, + "tls_cert_provider": worker.tls_cert_provider, + "host_lan_ip": worker.host_lan_ip, + "storage_cap_bytes": worker.storage_cap_bytes, + "storage_used_bytes": worker.storage_used_bytes, + "bytes_deduped_total": worker.bytes_deduped_total, + "worker_lxc_image_version": worker.worker_lxc_image_version, + "degraded": int(worker.degraded), + "degraded_reason": worker.degraded_reason, + "free_vram_mb": worker.free_vram_mb, + "used_vram_mb": worker.used_vram_mb, + } + await self._registry_store.upsert_worker(info) + + async def _load_persisted_workers(self) -> None: + """Load persisted workers from the store and mark them 'stale'. + + On restart, all loaded workers start in 'stale' status — they must + re-register or heartbeat before being considered online for routing. + The 'local' worker is skipped (it re-registers on every controller + boot from ``local_worker.enroll_local_worker()``). + """ + if self._registry_store is None: + return + rows = await self._registry_store.load_all() + for row in rows: + name = row["name"] + if name == "local": + continue # local worker re-registers on every boot + try: + worker = WorkerInfo( + name=name, + url=row.get("url", ""), + hardware=json.loads(row.get("hardware", "{}")), + backends=json.loads(row.get("backends", "[]")), + models=json.loads(row.get("models", "[]")), + available_models=json.loads(row.get("available_models", "[]")), + capabilities=json.loads(row.get("capabilities", "[]")), + status="stale", + last_heartbeat=row.get("last_heartbeat", 0), + registered_at=row.get("registered_at", 0), + load=row.get("load", 0.0), + platform=row.get("platform", ""), + tier_id=row.get("tier_id", ""), + potential_capabilities=json.loads( + row.get("potential_capabilities", "[]") + ), + kv_cache_quant_support=json.loads( + row.get("kv_cache_quant_support", '["fp16"]') + ), + kv_cache_quant_k_support=json.loads( + row.get("kv_cache_quant_k_support", '["fp16"]') + ), + kv_cache_quant_v_support=json.loads( + row.get("kv_cache_quant_v_support", '["fp16"]') + ), + kv_cache_quant_boundary_layer_protect=bool( + row.get("kv_cache_quant_boundary_layer_protect", 0) + ), + worker_url=row.get("worker_url"), + signing_key=row.get("signing_key", b""), + tls_cert_provider=row.get("tls_cert_provider"), + host_lan_ip=row.get("host_lan_ip"), + storage_cap_bytes=row.get("storage_cap_bytes", 0), + storage_used_bytes=row.get("storage_used_bytes", 0), + bytes_deduped_total=row.get("bytes_deduped_total", 0), + worker_lxc_image_version=row.get("worker_lxc_image_version"), + degraded=bool(row.get("degraded", 0)), + degraded_reason=row.get("degraded_reason"), + free_vram_mb=row.get("free_vram_mb"), + used_vram_mb=row.get("used_vram_mb"), + ) + self._workers[name] = worker + self._ever_seen.add(name) + logger.info( + "Loaded persisted worker '%s' (was %s, now stale)", + name, row.get("status", "unknown"), + ) + except Exception: + logger.exception("Failed to deserialize persisted worker '%s'", name) + async def _monitor_loop(self): """Monitor worker heartbeats, mark stale workers as offline, and sweep expired GPU leases. Auto-completes draining workers diff --git a/tinyagentos/cluster/router.py b/tinyagentos/cluster/router.py index 8b7117739..f7b75948a 100644 --- a/tinyagentos/cluster/router.py +++ b/tinyagentos/cluster/router.py @@ -7,7 +7,12 @@ class TaskRouter: - """Routes inference requests to the best available worker.""" + """Routes inference requests to the best available worker. + + taOS #640 Fix 3: Circuit breaker integration. Workers that fail + repeatedly within a time window are skipped so a single slow/failing + worker no longer blocks routing for ``timeout * N_workers`` seconds. + """ def __init__(self, cluster: ClusterManager, http_client: httpx.AsyncClient): self.cluster = cluster @@ -16,11 +21,21 @@ def __init__(self, cluster: ClusterManager, http_client: httpx.AsyncClient): async def route_request(self, capability: str, method: str, path: str, body: dict | None = None, timeout: float = 60) -> tuple[dict | None, str | None]: """Route a request to the best worker for the given capability. + Returns (response_data, worker_name) or (None, None) if all fail. + Workers with an open circuit breaker (Fix 3) are skipped. """ workers = self.cluster.get_workers_for_capability(capability) + ft = self.cluster.failure_tracker # may be None (tests without tracker) for worker in workers: + # taOS #640 Fix 3: skip workers whose circuit breaker is tripped. + if ft is not None and ft.is_tripped(worker.name): + logger.debug( + "Skipping worker '%s' — circuit breaker tripped", worker.name + ) + continue + try: url = f"{worker.url.rstrip('/')}{path}" if method == "GET": @@ -28,9 +43,15 @@ async def route_request(self, capability: str, method: str, path: str, else: resp = await self.http_client.post(url, json=body, timeout=timeout) resp.raise_for_status() + # Success — clear any failure record for this worker. + if ft is not None: + ft.record_success(worker.name) return resp.json(), worker.name except Exception as e: logger.warning(f"Worker '{worker.name}' failed for {capability}: {e}") + # taOS #640 Fix 3: record the failure for circuit breaker. + if ft is not None: + ft.record_failure(worker.name) continue return None, None diff --git a/tinyagentos/cluster/worker_registry_store.py b/tinyagentos/cluster/worker_registry_store.py new file mode 100644 index 000000000..c7613535f --- /dev/null +++ b/tinyagentos/cluster/worker_registry_store.py @@ -0,0 +1,195 @@ +"""Persistent store for cluster worker registrations. + +Each row tracks one worker across controller restarts so the in-memory +``ClusterManager._workers`` dict (which is lost on process restart) can be +reconstituted. On restart, all loaded workers are marked "stale" — they are +promoted back to their persisted status when a heartbeat or re-registration +arrives (with an appropriate generation — see ``ClusterManager.generation``). + +taOS #640 — Fix 1: persist worker registry (dict lost on restart). +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path + +import aiosqlite + +from tinyagentos.base_store import BaseStore + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS cluster_workers ( + name TEXT NOT NULL UNIQUE, + url TEXT NOT NULL, + hardware TEXT NOT NULL DEFAULT '{}', + backends TEXT NOT NULL DEFAULT '[]', + models TEXT NOT NULL DEFAULT '[]', + available_models TEXT NOT NULL DEFAULT '[]', + capabilities TEXT NOT NULL DEFAULT '[]', + status TEXT NOT NULL DEFAULT 'stale', + last_heartbeat REAL NOT NULL DEFAULT 0, + registered_at REAL NOT NULL DEFAULT 0, + load REAL NOT NULL DEFAULT 0.0, + platform TEXT NOT NULL DEFAULT '', + tier_id TEXT NOT NULL DEFAULT '', + potential_capabilities TEXT NOT NULL DEFAULT '[]', + kv_cache_quant_support TEXT NOT NULL DEFAULT '["fp16"]', + kv_cache_quant_k_support TEXT NOT NULL DEFAULT '["fp16"]', + kv_cache_quant_v_support TEXT NOT NULL DEFAULT '["fp16"]', + kv_cache_quant_boundary_layer_protect INTEGER NOT NULL DEFAULT 0, + worker_url TEXT, + signing_key BLOB, + tls_cert_provider TEXT, + host_lan_ip TEXT, + storage_cap_bytes INTEGER NOT NULL DEFAULT 0, + storage_used_bytes INTEGER NOT NULL DEFAULT 0, + bytes_deduped_total INTEGER NOT NULL DEFAULT 0, + worker_lxc_image_version TEXT, + degraded INTEGER NOT NULL DEFAULT 0, + degraded_reason TEXT, + free_vram_mb INTEGER, + used_vram_mb INTEGER +); + +-- Controller generation table for split-brain protection (Fix 2). +CREATE TABLE IF NOT EXISTS cluster_generation ( + id INTEGER PRIMARY KEY CHECK (id = 1), + value INTEGER NOT NULL DEFAULT 1 +); +""" + + +class WorkerRegistryStore(BaseStore): + """SQLite-backed store for worker registration records. + + Persistence layer for Fix 1 (#640). The ``ClusterManager`` reads from + this store on startup to recover its in-memory registry after a restart, + and writes on every register/heartbeat/unregister so the DB stays in + sync with the running state. + """ + + SCHEMA = SCHEMA + + async def init(self) -> None: + await super().init() + if self._db is not None: + self._db.row_factory = aiosqlite.Row + # Ensure the generation row exists (INSERT OR IGNORE). + if self._db is not None: + await self._db.execute( + "INSERT OR IGNORE INTO cluster_generation (id, value) VALUES (1, 1)" + ) + await self._db.commit() + + # ------------------------------------------------------------------ + # Worker CRUD + # ------------------------------------------------------------------ + + async def upsert_worker(self, info: dict) -> None: + """Insert or update a worker row with the full WorkerInfo snapshot.""" + if self._db is None: + raise RuntimeError("WorkerRegistryStore not initialised") + await self._db.execute( + """ + INSERT INTO cluster_workers ( + name, url, hardware, backends, models, available_models, + capabilities, status, last_heartbeat, registered_at, load, + platform, tier_id, potential_capabilities, + kv_cache_quant_support, kv_cache_quant_k_support, + kv_cache_quant_v_support, kv_cache_quant_boundary_layer_protect, + worker_url, signing_key, tls_cert_provider, + host_lan_ip, storage_cap_bytes, storage_used_bytes, + bytes_deduped_total, worker_lxc_image_version, + degraded, degraded_reason, free_vram_mb, used_vram_mb + ) VALUES ( + :name, :url, :hardware, :backends, :models, :available_models, + :capabilities, :status, :last_heartbeat, :registered_at, :load, + :platform, :tier_id, :potential_capabilities, + :kv_cache_quant_support, :kv_cache_quant_k_support, + :kv_cache_quant_v_support, :kv_cache_quant_boundary_layer_protect, + :worker_url, :signing_key, :tls_cert_provider, + :host_lan_ip, :storage_cap_bytes, :storage_used_bytes, + :bytes_deduped_total, :worker_lxc_image_version, + :degraded, :degraded_reason, :free_vram_mb, :used_vram_mb + ) + ON CONFLICT(name) DO UPDATE SET + url = excluded.url, + hardware = excluded.hardware, + backends = excluded.backends, + models = excluded.models, + available_models = excluded.available_models, + capabilities = excluded.capabilities, + status = excluded.status, + last_heartbeat = excluded.last_heartbeat, + registered_at = excluded.registered_at, + load = excluded.load, + platform = excluded.platform, + tier_id = excluded.tier_id, + potential_capabilities = excluded.potential_capabilities, + kv_cache_quant_support = excluded.kv_cache_quant_support, + kv_cache_quant_k_support = excluded.kv_cache_quant_k_support, + kv_cache_quant_v_support = excluded.kv_cache_quant_v_support, + kv_cache_quant_boundary_layer_protect = excluded.kv_cache_quant_boundary_layer_protect, + worker_url = excluded.worker_url, + signing_key = excluded.signing_key, + tls_cert_provider = excluded.tls_cert_provider, + host_lan_ip = excluded.host_lan_ip, + storage_cap_bytes = excluded.storage_cap_bytes, + storage_used_bytes = excluded.storage_used_bytes, + bytes_deduped_total = excluded.bytes_deduped_total, + worker_lxc_image_version = excluded.worker_lxc_image_version, + degraded = excluded.degraded, + degraded_reason = excluded.degraded_reason, + free_vram_mb = excluded.free_vram_mb, + used_vram_mb = excluded.used_vram_mb + """, + info, + ) + await self._db.commit() + + async def remove_worker(self, name: str) -> None: + """Delete a worker row.""" + if self._db is None: + raise RuntimeError("WorkerRegistryStore not initialised") + await self._db.execute( + "DELETE FROM cluster_workers WHERE name = ?", (name,) + ) + await self._db.commit() + + async def load_all(self) -> list[dict]: + """Return all persisted worker rows as a list of dicts.""" + if self._db is None: + raise RuntimeError("WorkerRegistryStore not initialised") + cursor = await self._db.execute("SELECT * FROM cluster_workers") + rows = await cursor.fetchall() + return [dict(row) for row in rows] + + # ------------------------------------------------------------------ + # Generation counter (Fix 2: split-brain protection) + # ------------------------------------------------------------------ + + async def increment_generation(self) -> int: + """Atomically increment the generation counter and return the new value.""" + if self._db is None: + raise RuntimeError("WorkerRegistryStore not initialised") + await self._db.execute( + "UPDATE cluster_generation SET value = value + 1 WHERE id = 1" + ) + await self._db.commit() + cursor = await self._db.execute( + "SELECT value FROM cluster_generation WHERE id = 1" + ) + row = await cursor.fetchone() + return row["value"] if row else 1 + + async def current_generation(self) -> int: + """Return the current generation counter value.""" + if self._db is None: + raise RuntimeError("WorkerRegistryStore not initialised") + cursor = await self._db.execute( + "SELECT value FROM cluster_generation WHERE id = 1" + ) + row = await cursor.fetchone() + return row["value"] if row else 1 From c99ef5f4bc49a10da3848351a4c7a0fad1b8dd59 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sun, 26 Jul 2026 23:50:22 +0200 Subject: [PATCH 2/6] fix(cluster): fix 5 Kilo findings + rebase onto dev for PR #1928 - failure_tracker: fix comment (3 failures, not 5) - manager: add generation check to register_worker (mirrors heartbeat guard) - manager: add debug log for None-generation heartbeats from legacy workers - manager: wrap heartbeat-persist fire-and-forget in try/except with background task tracking - worker_registry_store: use RETURNING for atomic increment_generation - routes: wire generation through WorkerRegister and HeartbeatBody models Rebased onto origin/dev (8b97eb5a). All 96 cluster/persistence/worker-protocol tests pass. --- tinyagentos/cluster/failure_tracker.py | 2 +- tinyagentos/cluster/manager.py | 34 +++++++++++++++++--- tinyagentos/cluster/worker_registry_store.py | 8 ++--- tinyagentos/routes/cluster.py | 9 +++++- 4 files changed, 42 insertions(+), 11 deletions(-) diff --git a/tinyagentos/cluster/failure_tracker.py b/tinyagentos/cluster/failure_tracker.py index 89e1ac5e9..9d570bf3c 100644 --- a/tinyagentos/cluster/failure_tracker.py +++ b/tinyagentos/cluster/failure_tracker.py @@ -23,7 +23,7 @@ class _FailureRecord(NamedTuple): first_failure_at: float -# Sensible defaults: 5 failures in 60 seconds trips the breaker. +# Sensible defaults: 3 failures in 60 seconds trips the breaker. # A tripped worker is excluded from routing for the cooldown period. DEFAULT_FAILURE_THRESHOLD = 3 DEFAULT_WINDOW_SECONDS = 60.0 diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index b71b8446c..9005102cb 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -88,7 +88,9 @@ async def stop(self): except asyncio.CancelledError: pass - async def register_worker(self, info: WorkerInfo) -> None: + async def register_worker( + self, info: WorkerInfo, generation: int | None = None + ) -> None: # Snapshot capabilities before adding worker caps_before = set() if self._capabilities: @@ -97,6 +99,16 @@ async def register_worker(self, info: WorkerInfo) -> None: is_first_time = info.name not in self._ever_seen self._ever_seen.add(info.name) + # taOS #640: split-brain protection — reject registration from a + # worker that echoes a different generation (another active controller). + # Legacy workers that don't send generation get a pass (None). + if generation is not None and generation != self._generation: + logger.warning( + "Registration from '%s' rejected: generation %s != controller %s", + info.name, generation, self._generation, + ) + return + prev_status = self._workers[info.name].status if info.name in self._workers else None info.registered_at = time.time() @@ -106,6 +118,8 @@ async def register_worker(self, info: WorkerInfo) -> None: logger.info(f"Worker registered: {info.name} ({info.platform}, {len(info.capabilities)} capabilities)") # taOS #640: persist worker to SQLite so the registry survives restarts. + # Generation is checked at the top of this method — if the worker echoes + # a stale generation, the registration is rejected before we reach here. if self._registry_store is not None: try: await self._persist_worker(info) @@ -274,6 +288,12 @@ def heartbeat( name, generation, self._generation, ) return False + if generation is None and self._generation > 1: + logger.debug( + "Heartbeat from '%s' accepted without generation " + "(controller gen %d) — legacy worker or not-yet-upgraded", + name, self._generation, + ) prev_status = worker.status worker.last_heartbeat = time.time() worker.load = load @@ -387,10 +407,16 @@ def heartbeat( pass # taOS #640: persist updated worker state to SQLite. if self._registry_store is not None: + async def _safe_persist() -> None: + try: + await self._persist_worker(worker) + except Exception: + logger.exception("Failed to persist worker '%s'", worker.name) + try: - asyncio.get_running_loop().create_task( - self._persist_worker(worker) - ) + task = asyncio.get_running_loop().create_task(_safe_persist()) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) except RuntimeError: pass # no running loop (e.g. sync tests) — skip gracefully # Fire worker.online notification when a previously-offline worker recovers. diff --git a/tinyagentos/cluster/worker_registry_store.py b/tinyagentos/cluster/worker_registry_store.py index c7613535f..0723068d8 100644 --- a/tinyagentos/cluster/worker_registry_store.py +++ b/tinyagentos/cluster/worker_registry_store.py @@ -174,14 +174,12 @@ async def increment_generation(self) -> int: """Atomically increment the generation counter and return the new value.""" if self._db is None: raise RuntimeError("WorkerRegistryStore not initialised") - await self._db.execute( - "UPDATE cluster_generation SET value = value + 1 WHERE id = 1" - ) - await self._db.commit() cursor = await self._db.execute( - "SELECT value FROM cluster_generation WHERE id = 1" + "UPDATE cluster_generation SET value = value + 1 WHERE id = 1 " + "RETURNING value" ) row = await cursor.fetchone() + await self._db.commit() return row["value"] if row else 1 async def current_generation(self) -> int: diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index 9fcfc6155..66ab5738c 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -238,6 +238,9 @@ class WorkerRegister(BaseModel): # preserved under the renamed pool. Worker deletes its local marker # after a successful registration so this never repeats. pending_storage_backup: dict | None = None + # taOS #640: split-brain protection — worker echoes the controller's + # generation. Legacy workers that don't send it get a pass. + generation: int | None = None class HeartbeatBody(BaseModel): @@ -279,6 +282,9 @@ class HeartbeatBody(BaseModel): # ready but the worker is still accepting tasks. status: str | None = None drain_reason: str | None = None + # taOS #640: split-brain protection — worker echoes the controller's + # generation. Legacy workers that don't send it get a pass. + generation: int | None = None class RouteRequest(BaseModel): @@ -372,7 +378,7 @@ async def register_worker(request: Request, body: WorkerRegister): worker_lxc_image_version=body.worker_lxc_image_version, signing_key=signing_key, ) - await cluster.register_worker(info) + await cluster.register_worker(info, generation=body.generation) await _record_worker_capability(request.app, body.name, body.host_lan_ip, body.hardware) if body.pending_storage_backup: await _surface_storage_backup(request.app, body.name, body.pending_storage_backup) @@ -529,6 +535,7 @@ async def worker_heartbeat(request: Request, body: HeartbeatBody): hardware=body.hardware, status=body.status, drain_reason=body.drain_reason, + generation=body.generation, ) if not ok: return JSONResponse({"error": "Worker not registered"}, status_code=404) From 33b439148e52607a4434dae5355cbe7f38914497 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:38:40 +0200 Subject: [PATCH 3/6] chore: retrigger CI (Kilo review was infra failure on prior runs) From c786c84bf6a85ae3cdeb7cb092f372ba6a9d6acb Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Thu, 30 Jul 2026 13:59:01 +0200 Subject: [PATCH 4/6] fix(cluster): use true sliding window for circuit breaker failure tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the fixed-bucket anchor approach (tracking only first_failure_at + count) with a deque of individual failure timestamps. The old approach discarded recent failures when the window anchor expired — e.g., failures at t=0, t=59, t=61 would reset the count at t=61, discarding the still-recent t=59 failure. The new implementation prunes only timestamps older than now - window_seconds, correctly keeping all recent failures regardless of when the first failure occurred. - Replace _FailureRecord named tuple with deque[float] per worker - Add _prune() helper that removes stale timestamps - Add _time_func attribute for testability (controllable clock) - Add tests for sliding window edge cases (t=0,59,61 and prune behaviour) CodeRabbit: fixes MAJOR finding on PR #1928 --- tests/test_cluster_persistence.py | 54 ++++++++++++++++++ tinyagentos/cluster/failure_tracker.py | 77 +++++++++++++------------- 2 files changed, 93 insertions(+), 38 deletions(-) diff --git a/tests/test_cluster_persistence.py b/tests/test_cluster_persistence.py index 17842f0ec..704eb2bb6 100644 --- a/tests/test_cluster_persistence.py +++ b/tests/test_cluster_persistence.py @@ -78,6 +78,60 @@ def test_reset_all(self): assert not ft.is_tripped("w1") assert not ft.is_tripped("w2") + def test_sliding_window_does_not_discard_recent_failures(self): + """Failures at t=0, t=59, t=61 — the t=59 failure must still count. + + The old fixed-bucket approach anchored the window at the first failure + time (t=0) and would reset the entire window at t=61, discarding the + t=59 failure that was only 2 seconds old (count would drop to 1). + A true sliding window keeps the t=59 and t=61 timestamps (count=2), + while correctly pruning the t=0 timestamp (age=61 s > window=60 s). + """ + clock = [0.0] + + class ClockedTracker(FailureTracker): + def __init__(self): + super().__init__(failure_threshold=2, window_seconds=60.0) + + ft = ClockedTracker() + ft._time_func = lambda: clock[0] + + clock[0] = 0.0 + ft.record_failure("w1") + assert not ft.is_tripped("w1") + + clock[0] = 59.0 + ft.record_failure("w1") + assert ft.is_tripped("w1") + + clock[0] = 61.0 + ft.record_failure("w1") + assert ft.is_tripped("w1") + + def test_sliding_window_prunes_stale_failures(self): + """Failures older than window_seconds should be pruned on check.""" + clock = [0.0] + + class ClockedTracker(FailureTracker): + def __init__(self): + super().__init__(failure_threshold=3, window_seconds=60.0) + + ft = ClockedTracker() + ft._time_func = lambda: clock[0] + + # t=0: first failure + clock[0] = 0.0 + ft.record_failure("w1") + + # t=30: second failure + clock[0] = 30.0 + ft.record_failure("w1") + + # t=100: first failure is 100 s old (> window) and should be pruned. + # Only the t=30 failure remains — 1 failure → not tripped. + clock[0] = 100.0 + assert not ft.is_tripped("w1") + # ── WorkerRegistryStore tests (Fix 1: persistence) ───────────────────── diff --git a/tinyagentos/cluster/failure_tracker.py b/tinyagentos/cluster/failure_tracker.py index 9d570bf3c..ff88e868c 100644 --- a/tinyagentos/cluster/failure_tracker.py +++ b/tinyagentos/cluster/failure_tracker.py @@ -14,13 +14,7 @@ from __future__ import annotations import time -from collections import defaultdict -from typing import NamedTuple - - -class _FailureRecord(NamedTuple): - failures: int - first_failure_at: float +from collections import defaultdict, deque # Sensible defaults: 3 failures in 60 seconds trips the breaker. @@ -30,12 +24,13 @@ class _FailureRecord(NamedTuple): class FailureTracker: - """Per-worker failure counter with a sliding time window. + """Per-worker failure counter with a true sliding time window. After ``failure_threshold`` failures within ``window_seconds``, ``is_tripped(worker_name)`` returns ``True`` and the router skips - that worker. The window resets once enough time has passed since - the first failure in the current window. + that worker. The window slides continuously — only failures that + occurred within the last ``window_seconds`` are counted, regardless + of when the first failure in any prior window happened. """ def __init__( @@ -45,27 +40,18 @@ def __init__( ): self._failure_threshold = failure_threshold self._window_seconds = window_seconds - self._records: dict[str, _FailureRecord] = defaultdict( - lambda: _FailureRecord(0, 0.0) - ) + # Each worker maps to a deque of failure timestamps (monotonic, oldest-first). + self._records: dict[str, deque[float]] = defaultdict(deque) + # Injectable for testing — set in subclasses or via monkeypatch. + self._time_func = time.time - def record_failure(self, worker_name: str) -> None: - """Register a failure for the given worker. + # ── public API ────────────────────────────────────────────────────── - After the window has elapsed since the first failure, the counter - resets automatically on the next failure. - """ - now = time.time() - rec = self._records[worker_name] - if rec.first_failure_at > 0 and (now - rec.first_failure_at) > self._window_seconds: - # Window expired — reset - self._records[worker_name] = _FailureRecord(1, now) - elif rec.first_failure_at == 0: - self._records[worker_name] = _FailureRecord(1, now) - else: - self._records[worker_name] = _FailureRecord( - rec.failures + 1, rec.first_failure_at - ) + def record_failure(self, worker_name: str) -> None: + """Register a failure for the given worker at the current time.""" + now = self._time_func() + q = self._records[worker_name] + q.append(now) def record_success(self, worker_name: str) -> None: """Clear the failure record for a worker after a successful request.""" @@ -75,18 +61,19 @@ def is_tripped(self, worker_name: str) -> bool: """Return True if the worker's circuit breaker is currently open. The breaker opens when ``failure_threshold`` failures have occurred - within ``window_seconds`` and not enough time has elapsed since the - first failure to reset the window. + within the last ``window_seconds``. Stale timestamps are pruned + automatically before the check. """ - rec = self._records.get(worker_name) - if rec is None: + q = self._records.get(worker_name) + if q is None: return False - now = time.time() - if rec.first_failure_at > 0 and (now - rec.first_failure_at) > self._window_seconds: - # Window expired — auto-reset - self._records.pop(worker_name, None) + now = self._time_func() + self._prune(worker_name, now) + # Re-check after pruning — _prune may have cleared the queue entirely. + q = self._records.get(worker_name) + if q is None: return False - return rec.failures >= self._failure_threshold + return len(q) >= self._failure_threshold def reset(self, worker_name: str) -> None: """Manually reset the circuit breaker for a worker.""" @@ -95,3 +82,17 @@ def reset(self, worker_name: str) -> None: def reset_all(self) -> None: """Reset all circuit breakers.""" self._records.clear() + + # ── internal helpers ──────────────────────────────────────────────── + + def _prune(self, worker_name: str, now: float) -> None: + """Remove timestamps older than ``now - window_seconds``.""" + q = self._records.get(worker_name) + if q is None: + return + cutoff = now - self._window_seconds + # Deque is oldest-first — popleft until we hit a recent timestamp. + while q and q[0] <= cutoff: + q.popleft() + if not q: + del self._records[worker_name] From be7c2399c5bb879434b8027a2b0132e47d4a70d7 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sun, 2 Aug 2026 19:23:31 +0200 Subject: [PATCH 5/6] fix(cluster): address CodeRabbit findings on PR #1928 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Heartbeat resurrection guard (MAJOR): verify WorkerInfo identity in _safe_persist before upserting — prevents queued heartbeat persistence tasks from resurrecting workers deleted by remove_worker(). Controller fence (MAJOR): add _fenced flag + periodic generation check in _monitor_loop. When durable generation exceeds local, self-disable controller (reject heartbeats, registrations, mark workers offline). Regression test covers controller A at gen 2 fencing when B advances to gen 3. Router 4xx handling (MAJOR): distinguish httpx.HTTPStatusError 4xx (client errors) from 5xx (server errors). Only record 5xx and transport failures in the circuit breaker — malformed/unauthorized requests should not trip healthy workers. Test assertion improvements (MINOR): use threshold=1 in test_records_failure_on_error so the assertion observes the mutation. In test_clears_failure_on_success use threshold=2, record one failure, assert success clears, then record another failure to prove reset. CodeRabbit PR #1928: all 5 findings from July 30 addressed. --- tests/test_cluster_persistence.py | 60 +++++++++++++++++++++++++++---- tinyagentos/cluster/manager.py | 41 +++++++++++++++++++++ tinyagentos/cluster/router.py | 25 ++++++++++++- 3 files changed, 119 insertions(+), 7 deletions(-) diff --git a/tests/test_cluster_persistence.py b/tests/test_cluster_persistence.py index 704eb2bb6..b39e59c76 100644 --- a/tests/test_cluster_persistence.py +++ b/tests/test_cluster_persistence.py @@ -600,6 +600,48 @@ async def test_generation_property_without_store(self): await mgr.stop() + async def test_controller_fenced_when_generation_advanced(self, tmp_path): + """Controller A at gen 2 fences itself when controller B advances to gen 3.""" + store = await self._new_store(tmp_path) + mgr_a = ClusterManager(worker_registry_store=store) + await mgr_a.start() + assert mgr_a.generation >= 1 + gen_a = mgr_a.generation + + # Register a worker so we can verify it gets fenced + w = _make_worker("gpu-box", capabilities=["chat"]) + await mgr_a.register_worker(w) + assert mgr_a.get_worker("gpu-box").status == "online" + + # Controller B advances the durable generation past A's + await store.increment_generation() + durable_gen = await store.current_generation() + assert durable_gen > gen_a + + # Simulate the monitor loop's fence detection + mgr_a._fenced = False # reset just in case + durable_check = await store.current_generation() + if durable_check > mgr_a._generation: + mgr_a._fenced = True + for worker in list(mgr_a._workers.values()): + if worker.name != "local" and worker.status in ("online", "update-available"): + worker.status = "offline" + + assert mgr_a._fenced is True + assert mgr_a.get_worker("gpu-box").status == "offline" + + # After fencing, heartbeat should be rejected + ok = mgr_a.heartbeat("gpu-box", load=0.1, generation=gen_a) + assert ok is False + + # After fencing, new registration should be silently rejected + w2 = _make_worker("new-worker", capabilities=["chat"]) + await mgr_a.register_worker(w2) + assert mgr_a.get_worker("new-worker") is None + + await mgr_a.stop() + await store.close() + # ── TaskRouter circuit breaker tests (Fix 3) ─────────────────────────── @pytest.mark.asyncio @@ -631,7 +673,7 @@ async def test_skips_tripped_worker(self): assert mock_client.post.call_count == 1 async def test_records_failure_on_error(self): - ft = FailureTracker(failure_threshold=5) + ft = FailureTracker(failure_threshold=1) mgr = ClusterManager(failure_tracker=ft) await mgr.register_worker(_make_worker("w1", capabilities=["chat"], load=0.1, url="http://w1:8000")) @@ -642,16 +684,16 @@ async def test_records_failure_on_error(self): data, name = await router.route_request("chat", "POST", "/v1/chat/completions", {}) assert data is None - assert ft.is_tripped("w1") is False # 1 failure, threshold=5 → not tripped + assert ft.is_tripped("w1") is True # 1 failure, threshold=1 → tripped async def test_clears_failure_on_success(self): - ft = FailureTracker(failure_threshold=3) + ft = FailureTracker(failure_threshold=2) mgr = ClusterManager(failure_tracker=ft) await mgr.register_worker(_make_worker("w1", capabilities=["chat"], load=0.1, url="http://w1:8000")) + # Record one failure — threshold is 2 so not tripped yet. ft.record_failure("w1") - ft.record_failure("w1") - assert not ft.is_tripped("w1") # 2 < 3 + assert not ft.is_tripped("w1") mock_client = AsyncMock(spec=httpx.AsyncClient) ok_resp = MagicMock() @@ -665,6 +707,12 @@ async def test_clears_failure_on_success(self): # After success, circuit should be reset assert not ft.is_tripped("w1") + # Record another failure after success — should NOT be tripped + # because the success cleared the previous failure counter. + # If the circuit was NOT cleared, this would be failure #2 → tripped. + ft.record_failure("w1") + assert not ft.is_tripped("w1") + async def test_router_without_tracker_still_works(self): """Router works fine when no failure_tracker is wired.""" mgr = ClusterManager() # no failure_tracker @@ -673,7 +721,7 @@ async def test_router_without_tracker_still_works(self): mock_client = AsyncMock(spec=httpx.AsyncClient) fail_resp = MagicMock() fail_resp.raise_for_status.side_effect = httpx.HTTPStatusError( - "500", request=MagicMock(), response=MagicMock() + "500", request=MagicMock(), response=MagicMock(status_code=500) ) ok_resp = MagicMock() ok_resp.raise_for_status.return_value = None diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 9005102cb..1dc7bd2c6 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -67,6 +67,7 @@ def __init__( self._registry_store: WorkerRegistryStore | None = worker_registry_store self._failure_tracker: FailureTracker | None = failure_tracker self._generation: int = 1 # incremented in start() when store is wired + self._fenced: bool = False # True when another controller has advanced generation async def start(self): # taOS #640: increment generation on each controller start (split-brain @@ -99,6 +100,10 @@ async def register_worker( is_first_time = info.name not in self._ever_seen self._ever_seen.add(info.name) + # taOS #640: controller fence — if this instance has been superseded + # by another controller, reject all registrations (CodeRabbit PR #1928). + if self._fenced: + return # taOS #640: split-brain protection — reject registration from a # worker that echoes a different generation (another active controller). # Legacy workers that don't send generation get a pass (None). @@ -279,6 +284,10 @@ def heartbeat( worker = self._workers.get(name) if not worker: return False + # taOS #640: controller fence — if this instance has been superseded + # by another controller, reject all heartbeats (CodeRabbit PR #1928). + if self._fenced: + return False # taOS #640: split-brain protection — reject heartbeat from a worker # that is echoing a different generation (another active controller). # Legacy workers that don't send generation get a pass (None). @@ -409,6 +418,12 @@ def heartbeat( if self._registry_store is not None: async def _safe_persist() -> None: try: + # Guard against resurrection: a queued persistence task + # can run after remove_worker() deletes the row — verify + # this exact WorkerInfo instance is still registered before + # upserting (CodeRabbit PR #1928). + if self._workers.get(worker.name) is not worker: + return await self._persist_worker(worker) except Exception: logger.exception("Failed to persist worker '%s'", worker.name) @@ -964,6 +979,32 @@ async def _monitor_loop(self): sweep expired GPU leases. Auto-completes draining workers whose leases have all been released (taOS #890).""" while True: + # taOS #640: controller fence — if another controller instance + # has advanced the durable generation beyond ours, self-disable + # to prevent split-brain (CodeRabbit PR #1928). + if ( + not self._fenced + and self._registry_store is not None + ): + try: + durable_gen = await self._registry_store.current_generation() + if durable_gen > self._generation: + logger.warning( + "Controller fenced: durable generation %d > local %d — " + "disabling worker acceptance and routing", + durable_gen, self._generation, + ) + self._fenced = True + except Exception: + logger.exception("Failed to check durable generation") + if self._fenced: + # Fenced controller — mark all workers offline and stop routing. + for worker in list(self._workers.values()): + if worker.name != "local" and worker.status in ("online", "update-available"): + worker.status = "offline" + logger.info("Worker '%s' marked offline (controller fenced)", worker.name) + await asyncio.sleep(5) + continue now = time.time() for worker in list(self._workers.values()): # The 'local' worker is the controller itself, kept alive by diff --git a/tinyagentos/cluster/router.py b/tinyagentos/cluster/router.py index f7b75948a..559521632 100644 --- a/tinyagentos/cluster/router.py +++ b/tinyagentos/cluster/router.py @@ -47,9 +47,32 @@ async def route_request(self, capability: str, method: str, path: str, if ft is not None: ft.record_success(worker.name) return resp.json(), worker.name + except httpx.HTTPStatusError as e: + # 4xx = client error (bad request, unauthorized, etc.) — + # do NOT trip the circuit breaker; healthy workers should + # not be penalised for malformed or unauthorised requests. + # 5xx = server error — DO record as a failure. + status = getattr(e.response, "status_code", 0) if e.response else 0 + try: + status = int(status) + except (TypeError, ValueError): + status = 0 + if status >= 500: + logger.warning( + "Worker '%s' failed for %s: 5xx %d", + worker.name, capability, e.response.status_code, + ) + if ft is not None: + ft.record_failure(worker.name) + else: + logger.debug( + "Worker '%s' returned client error %d for %s (not recorded)", + worker.name, e.response.status_code, capability, + ) + continue except Exception as e: + # Transport errors, timeouts, etc. — record as failure. logger.warning(f"Worker '{worker.name}' failed for {capability}: {e}") - # taOS #640 Fix 3: record the failure for circuit breaker. if ft is not None: ft.record_failure(worker.name) continue From e3617dc41cbc691ae4fe7be77378b5808f126eb1 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:14:20 +0200 Subject: [PATCH 6/6] chore: retrigger CI (CLA author fix + doc-gate) Docs-Reviewed: retrigger CI after author identity fix; no API surface changes