diff --git a/packages/challenges/agent-challenge/src/agent_challenge/evaluation/raw_weight_push.py b/packages/challenges/agent-challenge/src/agent_challenge/evaluation/raw_weight_push.py index c8aecc7bd..885387472 100644 --- a/packages/challenges/agent-challenge/src/agent_challenge/evaluation/raw_weight_push.py +++ b/packages/challenges/agent-challenge/src/agent_challenge/evaluation/raw_weight_push.py @@ -141,6 +141,30 @@ def _is_all_zero_weights(weights: Mapping[str, float]) -> bool: NON_RETRYABLE_PUSH_STATUSES = frozenset({409, 413, 415, 422}) +MAX_IDENTITY_ADVANCES = 5 + + +def _is_epoch_advance_detail(detail: str) -> bool: + """True when a 409 detail means the client must change epoch/revision identity.""" + + text = (detail or "").lower() + if not text: + return True + return any( + token in text + for token in ( + "sealed", + "stale", + "conflicting raw weight payload", + "conflict", + ) + ) + + +def _is_revision_conflict_detail(detail: str) -> bool: + text = (detail or "").lower() + return "conflicting raw weight payload" in text or text == "conflict" + class RawWeightPushLedger(Base): """Singleton SQLAlchemy row for durable attempt/ack checkpointing. @@ -419,6 +443,45 @@ def _ack_matches( and bool(ack.snapshot_id) ) + async def _fetch_master_target_epoch(self, *, client: httpx.AsyncClient) -> int | None: + """Ask master for the next unsealed aggregation epoch (authoritative).""" + + path = "/internal/v1/aggregation/target-epoch" + url = f"{self.master_base_url}{path}" + headers = { + "Authorization": f"Bearer {self.shared_token}", + "X-Base-Challenge-Slug": self.challenge_slug, + "Accept": "application/json", + } + try: + response = await client.get(url, headers=headers) + if response.status_code != 200: + return None + data = response.json() + if not isinstance(data, dict) or "target_epoch" not in data: + return None + return int(data["target_epoch"]) + except Exception: # noqa: BLE001 - transport, mock, or schema failures fall back + return None + + async def _resolve_target_epoch( + self, + *, + client: httpx.AsyncClient, + explicit: int | None, + now: datetime, + ) -> int: + """Prefer master target-epoch; fall back to epoch_fn / wall bucket only if needed.""" + + if explicit is not None: + return int(explicit) + master_target = await self._fetch_master_target_epoch(client=client) + if master_target is not None: + return int(master_target) + if self.epoch_fn is not None: + return int(self.epoch_fn()) + return int(now.timestamp()) // 360 + @role_contract(role=Role.CHALLENGE, capability=Capability.CHALLENGE_RAW_WEIGHT_PUSH) async def push_once( self, @@ -427,8 +490,14 @@ async def push_once( epoch: int | None = None, force_revision: int | None = None, reuse_pending: bool = True, + _advance_depth: int = 0, ) -> PushAttemptResult: - """Push one snapshot. Cursor advances only on exact durable acknowledgement.""" + """Push one snapshot. Cursor advances only on exact durable acknowledgement. + + On 409 sealed/stale/conflict, clears pending and retries with an advanced + epoch/revision identity (bounded) so a sealed wall-clock bucket cannot + loop forever. + """ await self.store.init() now = self._now_fn() @@ -436,203 +505,234 @@ async def push_once( pending = await self.store.get_pending() if reuse_pending else None payload: RawWeightPushRequest | None = None raw_bytes: bytes | None = None + cleaned_for_retry: dict[str, float] | None = None - if pending is not None: - # Retry exact previous bytes after timeout/restart (no new revision). - try: - pending_bytes = str(pending["canonical_payload"]).encode("utf-8") - payload = RawWeightPushRequest.model_validate_json(pending_bytes) - raw_bytes = payload.canonical_bytes() - if not _pending_still_fresh(payload, now=now): - logger.info( - "raw weight pending expired; rebuilding", - extra={ - "epoch": payload.epoch, - "revision": payload.revision, - "digest": payload.payload_digest[:12], - }, - ) - await self.store.clear_pending() + http = self._http + owns_client = http is None + if owns_client: + http = httpx.AsyncClient(timeout=self.timeout_seconds) + assert http is not None + + try: + if pending is not None: + # Retry exact previous bytes after timeout/restart (no new revision). + try: + pending_bytes = str(pending["canonical_payload"]).encode("utf-8") + payload = RawWeightPushRequest.model_validate_json(pending_bytes) + raw_bytes = payload.canonical_bytes() + if not _pending_still_fresh(payload, now=now): + logger.info( + "raw weight pending expired; rebuilding", + extra={ + "epoch": payload.epoch, + "revision": payload.revision, + "digest": payload.payload_digest[:12], + }, + ) + await self.store.clear_pending() + pending = None + payload = None + raw_bytes = None + except Exception: # noqa: BLE001 - corrupt pending is rebuilt pending = None payload = None raw_bytes = None - except Exception: # noqa: BLE001 - corrupt pending is rebuilt - pending = None - payload = None - raw_bytes = None - - if payload is None or raw_bytes is None: - resolved_weights = await _resolve_weights(weights, self.weights_fn) - # Positive hotkey weights only when synthesizing from get_weights; - # explicit zero maps (caller-supplied zeros) are preserved as zero-contribution. - if weights is not None: - cleaned = {str(hotkey): float(value) for hotkey, value in resolved_weights.items()} - else: - cleaned = { - str(hotkey): float(value) - for hotkey, value in resolved_weights.items() - if float(value) > 0.0 - } - if not cleaned: + + if payload is None or raw_bytes is None: + resolved_weights = await _resolve_weights(weights, self.weights_fn) + # Positive hotkey weights only when synthesizing from get_weights; + # explicit zero maps (caller-supplied zeros) are preserved as zero-contribution. + if weights is not None: + cleaned = { + str(hotkey): float(value) for hotkey, value in resolved_weights.items() + } + else: + cleaned = { + str(hotkey): float(value) + for hotkey, value in resolved_weights.items() + if float(value) > 0.0 + } + if not cleaned: + return PushAttemptResult( + status="skipped_empty", + epoch=0, + revision=0, + payload_digest="", + snapshot_id=None, + cursor_advanced=False, + error="empty weights", + ) + if _is_all_zero_weights(cleaned): + return PushAttemptResult( + status="skipped_zero", + epoch=0, + revision=0, + payload_digest="", + snapshot_id=None, + cursor_advanced=False, + error="all-zero weight map", + ) + # Host-trust path: never silently push unattested scores as if TEE-backed. + _log_unattested_weight_push_if_no_phala(cleaned) + cleaned_for_retry = dict(cleaned) + resolved_epoch = await self._resolve_target_epoch( + client=http, explicit=epoch, now=now + ) + revision = ( + int(force_revision) + if force_revision is not None + else self._next_revision(cursor, resolved_epoch) + ) + nonce = f"agent-challenge-{uuid.uuid4().hex}" + payload, raw_bytes = self._build_payload( + weights=cleaned, + epoch=resolved_epoch, + revision=revision, + nonce=nonce, + now=now, + ) + await self.store.record_pending( + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + canonical_payload=raw_bytes.decode("utf-8"), + nonce=payload.nonce, + attempted_at=now.isoformat(), + ) + + path = self._path_for() + url = f"{self.master_base_url}{path}" + headers = self._headers(path=path, body=raw_bytes, timestamp=int(now.timestamp())) + try: + response = await http.post(url, content=raw_bytes, headers=headers) + except (httpx.TimeoutException, httpx.TransportError) as exc: + # Never log Authorization / token material — digest prefix only. + logger.info( + "raw weight push transport failure", + extra={ + "epoch": payload.epoch, + "revision": payload.revision, + "digest": payload.payload_digest[:12], + }, + ) return PushAttemptResult( - status="skipped_empty", - epoch=0, - revision=0, - payload_digest="", + status="transport_error", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, snapshot_id=None, cursor_advanced=False, - error="empty weights", + error=str(exc), ) - if _is_all_zero_weights(cleaned): + + if response.status_code >= 500: return PushAttemptResult( - status="skipped_zero", - epoch=0, - revision=0, - payload_digest="", + status="server_error", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, snapshot_id=None, cursor_advanced=False, - error="all-zero weight map", + error=f"status={response.status_code}", ) - # Host-trust path: never silently push unattested scores as if TEE-backed. - _log_unattested_weight_push_if_no_phala(cleaned) - resolved_epoch = ( - int(epoch) - if epoch is not None - else int(self.epoch_fn()) - if self.epoch_fn is not None - else int(now.timestamp()) // 3600 - ) - revision = ( - int(force_revision) - if force_revision is not None - else self._next_revision(cursor, resolved_epoch) - ) - nonce = f"agent-challenge-{uuid.uuid4().hex}" - payload, raw_bytes = self._build_payload( - weights=cleaned, - epoch=resolved_epoch, - revision=revision, - nonce=nonce, - now=now, - ) - await self.store.record_pending( + if response.status_code not in {200, 201}: + detail = _safe_response_detail(response) + logger.info( + "raw weight push rejected", + extra={ + "status": response.status_code, + "detail": detail, + "epoch": payload.epoch, + "revision": payload.revision, + "digest": payload.payload_digest[:12], + }, + ) + if response.status_code in NON_RETRYABLE_PUSH_STATUSES: + await self.store.clear_pending() + err = f"status={response.status_code}" + if detail: + err = f"{err} detail={detail}" + # Bounded advance: sealed/stale/conflict must not retry same identity. + if ( + response.status_code == 409 + and _is_epoch_advance_detail(detail) + and _advance_depth < MAX_IDENTITY_ADVANCES + ): + if _is_revision_conflict_detail(detail): + next_epoch: int | None = payload.epoch + next_revision = payload.revision + 1 + else: + # Sealed/stale: re-query master; fall back to epoch+1. + master_next = await self._fetch_master_target_epoch(client=http) + if master_next is not None and int(master_next) > payload.epoch: + next_epoch = int(master_next) + else: + next_epoch = payload.epoch + 1 + next_revision = 1 + retry_weights = ( + cleaned_for_retry + if cleaned_for_retry is not None + else dict(payload.weights) + ) + return await self.push_once( + weights=retry_weights, + epoch=next_epoch, + force_revision=next_revision, + reuse_pending=False, + _advance_depth=_advance_depth + 1, + ) + return PushAttemptResult( + status="rejected", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + snapshot_id=None, + cursor_advanced=False, + error=err, + ) + try: + ack = RawWeightPushAcknowledgement.model_validate(response.json()) + except Exception as exc: # noqa: BLE001 + return PushAttemptResult( + status="malformed_ack", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + snapshot_id=None, + cursor_advanced=False, + error=str(exc), + ) + if not self._ack_matches(ack, payload=payload): + return PushAttemptResult( + status="ack_mismatch", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + snapshot_id=ack.snapshot_id if hasattr(ack, "snapshot_id") else None, + cursor_advanced=False, + error="acknowledgement identity mismatch", + ) + ack_time = self._now_fn().isoformat() + await self.store.acknowledge( epoch=payload.epoch, revision=payload.revision, payload_digest=payload.payload_digest, + snapshot_id=ack.snapshot_id, canonical_payload=raw_bytes.decode("utf-8"), nonce=payload.nonce, - attempted_at=now.isoformat(), - ) - - path = self._path_for() - url = f"{self.master_base_url}{path}" - headers = self._headers(path=path, body=raw_bytes, timestamp=int(now.timestamp())) - client = self._http - owns_client = client is None - if owns_client: - client = httpx.AsyncClient(timeout=self.timeout_seconds) - assert client is not None - try: - response = await client.post(url, content=raw_bytes, headers=headers) - except (httpx.TimeoutException, httpx.TransportError) as exc: - # Never log Authorization / token material — digest prefix only. - logger.info( - "raw weight push transport failure", - extra={ - "epoch": payload.epoch, - "revision": payload.revision, - "digest": payload.payload_digest[:12], - }, + acknowledged_at=ack_time, ) return PushAttemptResult( - status="transport_error", + status="acknowledged", epoch=payload.epoch, revision=payload.revision, payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=str(exc), + snapshot_id=ack.snapshot_id, + cursor_advanced=True, ) finally: if owns_client: - await client.aclose() - - if response.status_code >= 500: - return PushAttemptResult( - status="server_error", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=f"status={response.status_code}", - ) - if response.status_code not in {200, 201}: - detail = _safe_response_detail(response) - logger.info( - "raw weight push rejected", - extra={ - "status": response.status_code, - "detail": detail, - "epoch": payload.epoch, - "revision": payload.revision, - "digest": payload.payload_digest[:12], - }, - ) - if response.status_code in NON_RETRYABLE_PUSH_STATUSES: - await self.store.clear_pending() - err = f"status={response.status_code}" - if detail: - err = f"{err} detail={detail}" - return PushAttemptResult( - status="rejected", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=err, - ) - try: - ack = RawWeightPushAcknowledgement.model_validate(response.json()) - except Exception as exc: # noqa: BLE001 - return PushAttemptResult( - status="malformed_ack", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=str(exc), - ) - if not self._ack_matches(ack, payload=payload): - return PushAttemptResult( - status="ack_mismatch", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=ack.snapshot_id if hasattr(ack, "snapshot_id") else None, - cursor_advanced=False, - error="acknowledgement identity mismatch", - ) - ack_time = self._now_fn().isoformat() - await self.store.acknowledge( - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=ack.snapshot_id, - canonical_payload=raw_bytes.decode("utf-8"), - nonce=payload.nonce, - acknowledged_at=ack_time, - ) - return PushAttemptResult( - status="acknowledged", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=ack.snapshot_id, - cursor_advanced=True, - ) + await http.aclose() def build_weights_loader( @@ -768,6 +868,7 @@ def maybe_build_push_client_from_settings( or DEFAULT_CHALLENGE_SLUG ) + # Fallback only: push_once prefers GET /internal/v1/aggregation/target-epoch. def _epoch() -> int: return int(datetime.now(UTC).timestamp()) // max(epoch_seconds, 1) diff --git a/packages/challenges/agent-challenge/tests/test_ac_raw_weight_push.py b/packages/challenges/agent-challenge/tests/test_ac_raw_weight_push.py index 64141eb29..f3bcd083b 100644 --- a/packages/challenges/agent-challenge/tests/test_ac_raw_weight_push.py +++ b/packages/challenges/agent-challenge/tests/test_ac_raw_weight_push.py @@ -373,3 +373,251 @@ def handler(request: httpx.Request) -> httpx.Response: assert result.status in {"skipped_empty", "skipped_zero"} assert calls == [] await http.aclose() + + +@pytest.mark.asyncio +async def test_successful_push_writes_ack_cursor_fields(database: Database) -> None: + """ACK path must set last_epoch, last_snapshot_id, and acknowledged_at.""" + + clock = FakeClock() + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "target_epoch": 100, + "highest_sealed_epoch": 99, + "max_future_epoch_ahead": 2, + }, + ) + parsed = RawWeightPushRequest.model_validate_json(request.content) + return httpx.Response( + 200, + json={ + "protocol_version": "1.0", + "challenge_slug": SLUG, + "epoch": parsed.epoch, + "revision": parsed.revision, + "snapshot_id": "snap-ack-fields", + "payload_digest": parsed.payload_digest, + "accepted": True, + "idempotent": False, + }, + ) + + http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://master.test", + ) + client = RawWeightPushClient( + database=database, + challenge_slug=SLUG, + master_base_url="http://master.test", + shared_token=TOKEN, + now_fn=clock.now, + http_client=http, + # No epoch_fn: must use master target-epoch. + epoch_fn=None, + ) + await client.init() + with activate_role(Role.CHALLENGE): + result = await client.push_once(weights={WINNER: 1.0}) + assert result.status == "acknowledged" + assert result.cursor_advanced is True + assert result.epoch == 100 + assert result.snapshot_id == "snap-ack-fields" + cursor = await client.store.get_cursor() + assert cursor is not None + assert cursor.epoch == 100 + assert cursor.snapshot_id == "snap-ack-fields" + assert cursor.acknowledged_at is not None + assert cursor.acknowledged_at != "" + await http.aclose() + + +@pytest.mark.asyncio +async def test_sealed_epoch_409_advances_to_master_target(database: Database) -> None: + """Sealed-epoch 409 must not infinite-retry the same epoch; advance via master.""" + + clock = FakeClock() + post_epochs: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + # Master says next unsealed is 50 after first sealed rejection path. + return httpx.Response( + 200, + json={ + "target_epoch": 50, + "highest_sealed_epoch": 49, + "max_future_epoch_ahead": 2, + }, + ) + parsed = RawWeightPushRequest.model_validate_json(request.content) + post_epochs.append(parsed.epoch) + if parsed.epoch == 49: + return httpx.Response( + 409, json={"detail": "epoch is sealed; revision rejected"} + ) + return httpx.Response( + 200, + json={ + "protocol_version": "1.0", + "challenge_slug": SLUG, + "epoch": parsed.epoch, + "revision": parsed.revision, + "snapshot_id": f"snap-{parsed.epoch}", + "payload_digest": parsed.payload_digest, + "accepted": True, + "idempotent": False, + }, + ) + + http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://master.test", + ) + client = RawWeightPushClient( + database=database, + challenge_slug=SLUG, + master_base_url="http://master.test", + shared_token=TOKEN, + now_fn=clock.now, + http_client=http, + # Stale local clock would stick on sealed 49 forever without master. + epoch_fn=lambda: 49, + ) + await client.init() + with activate_role(Role.CHALLENGE): + result = await client.push_once(weights={WINNER: 1.0}) + assert result.status == "acknowledged" + assert result.epoch == 50 + assert 49 in post_epochs or post_epochs == [50] + # Must not hammer the same sealed epoch repeatedly in one push_once. + assert post_epochs.count(49) <= 1 + assert await client.store.get_cursor() is not None + assert (await client.store.get_cursor()).epoch == 50 + await http.aclose() + + +@pytest.mark.asyncio +async def test_conflict_409_does_not_retry_same_epoch_revision( + database: Database, +) -> None: + """Conflicting payload at (epoch, rev) must bump revision, not loop forever.""" + + clock = FakeClock() + seen: list[tuple[int, int]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "target_epoch": 7, + "highest_sealed_epoch": 6, + "max_future_epoch_ahead": 2, + }, + ) + parsed = RawWeightPushRequest.model_validate_json(request.content) + seen.append((parsed.epoch, parsed.revision)) + if parsed.revision == 1: + return httpx.Response( + 409, json={"detail": "conflicting raw weight payload"} + ) + return httpx.Response( + 200, + json={ + "protocol_version": "1.0", + "challenge_slug": SLUG, + "epoch": parsed.epoch, + "revision": parsed.revision, + "snapshot_id": "snap-rev2", + "payload_digest": parsed.payload_digest, + "accepted": True, + "idempotent": False, + }, + ) + + http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://master.test", + ) + client = RawWeightPushClient( + database=database, + challenge_slug=SLUG, + master_base_url="http://master.test", + shared_token=TOKEN, + now_fn=clock.now, + http_client=http, + epoch_fn=None, + ) + await client.init() + with activate_role(Role.CHALLENGE): + result = await client.push_once(weights={WINNER: 1.0}) + assert result.status == "acknowledged" + assert result.revision == 2 + assert (7, 1) in seen + assert (7, 2) in seen + assert seen.count((7, 1)) == 1 + cursor = await client.store.get_cursor() + assert cursor is not None + assert cursor.epoch == 7 + assert cursor.revision == 2 + assert cursor.acknowledged_at is not None + await http.aclose() + + +@pytest.mark.asyncio +async def test_epoch_source_prefers_master_over_local_clock(database: Database) -> None: + """When master target-epoch is available, do not use local wall-clock epoch_fn.""" + + clock = FakeClock() + posted: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "target_epoch": 9001, + "highest_sealed_epoch": 9000, + "max_future_epoch_ahead": 2, + }, + ) + parsed = RawWeightPushRequest.model_validate_json(request.content) + posted.append(parsed.epoch) + return httpx.Response( + 200, + json={ + "protocol_version": "1.0", + "challenge_slug": SLUG, + "epoch": parsed.epoch, + "revision": parsed.revision, + "snapshot_id": "snap-master-epoch", + "payload_digest": parsed.payload_digest, + "accepted": True, + "idempotent": False, + }, + ) + + http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://master.test", + ) + client = RawWeightPushClient( + database=database, + challenge_slug=SLUG, + master_base_url="http://master.test", + shared_token=TOKEN, + now_fn=clock.now, + http_client=http, + epoch_fn=lambda: 1, # would be wrong if preferred + ) + await client.init() + with activate_role(Role.CHALLENGE): + result = await client.push_once(weights={WINNER: 1.0}) + assert result.status == "acknowledged" + assert posted == [9001] + await http.aclose() diff --git a/packages/challenges/prism/src/prism_challenge/raw_weight_push.py b/packages/challenges/prism/src/prism_challenge/raw_weight_push.py index bc1e84226..2c884585d 100644 --- a/packages/challenges/prism/src/prism_challenge/raw_weight_push.py +++ b/packages/challenges/prism/src/prism_challenge/raw_weight_push.py @@ -133,6 +133,31 @@ def _is_all_zero_weights(weights: Mapping[str, float]) -> bool: NON_RETRYABLE_PUSH_STATUSES = frozenset({409, 413, 415, 422}) +MAX_IDENTITY_ADVANCES = 5 + + +def _is_epoch_advance_detail(detail: str) -> bool: + """True when a 409 detail means the client must change epoch/revision identity.""" + + text = (detail or "").lower() + if not text: + return True + return any( + token in text + for token in ( + "sealed", + "stale", + "conflicting raw weight payload", + "conflict", + ) + ) + + +def _is_revision_conflict_detail(detail: str) -> bool: + text = (detail or "").lower() + return "conflicting raw weight payload" in text or text == "conflict" + + RAW_WEIGHT_PUSH_SCHEMA = ( "CREATE TABLE IF NOT EXISTS raw_weight_push_ledger (" "id INTEGER PRIMARY KEY CHECK (id = 1)," @@ -428,6 +453,45 @@ def _ack_matches( and bool(ack.snapshot_id) ) + async def _fetch_master_target_epoch(self, *, client: httpx.AsyncClient) -> int | None: + """Ask master for the next unsealed aggregation epoch (authoritative).""" + + path = "/internal/v1/aggregation/target-epoch" + url = f"{self.master_base_url}{path}" + headers = { + "Authorization": f"Bearer {self.shared_token}", + "X-Base-Challenge-Slug": self.challenge_slug, + "Accept": "application/json", + } + try: + response = await client.get(url, headers=headers) + if response.status_code != 200: + return None + data = response.json() + if not isinstance(data, dict) or "target_epoch" not in data: + return None + return int(data["target_epoch"]) + except Exception: # noqa: BLE001 - transport, mock, or schema failures fall back + return None + + async def _resolve_target_epoch( + self, + *, + client: httpx.AsyncClient, + explicit: int | None, + now: datetime, + ) -> int: + """Prefer master target-epoch; fall back to epoch_fn / wall bucket only if needed.""" + + if explicit is not None: + return int(explicit) + master_target = await self._fetch_master_target_epoch(client=client) + if master_target is not None: + return int(master_target) + if self.epoch_fn is not None: + return int(self.epoch_fn()) + return int(now.timestamp()) // 360 + @role_contract(role=Role.CHALLENGE, capability=Capability.CHALLENGE_RAW_WEIGHT_PUSH) async def push_once( self, @@ -436,8 +500,14 @@ async def push_once( epoch: int | None = None, force_revision: int | None = None, reuse_pending: bool = True, + _advance_depth: int = 0, ) -> PushAttemptResult: - """Push one snapshot. Cursor advances only on exact durable acknowledgement.""" + """Push one snapshot. Cursor advances only on exact durable acknowledgement. + + On 409 sealed/stale/conflict, clears pending and retries with an advanced + epoch/revision identity (bounded) so a sealed wall-clock bucket cannot + loop forever. + """ await self.store.init() now = self._now_fn() @@ -445,206 +515,238 @@ async def push_once( pending = await self.store.get_pending() if reuse_pending else None payload: RawWeightPushRequest | None = None raw_bytes: bytes | None = None + cleaned_for_retry: dict[str, float] | None = None - if pending is not None: - # Retry exact previous bytes after timeout/restart (no new revision). - try: - pending_bytes = str(pending["canonical_payload"]).encode("utf-8") - payload = RawWeightPushRequest.model_validate_json(pending_bytes) - raw_bytes = payload.canonical_bytes() - if not _pending_still_fresh(payload, now=now): - logger.info( - "raw weight pending expired; rebuilding", - extra={ - "epoch": payload.epoch, - "revision": payload.revision, - "digest": payload.payload_digest[:12], - }, - ) - await self.store.clear_pending() + http = self._http + owns_client = http is None + if owns_client: + http = httpx.AsyncClient(timeout=self.timeout_seconds) + assert http is not None + + try: + if pending is not None: + # Retry exact previous bytes after timeout/restart (no new revision). + try: + pending_bytes = str(pending["canonical_payload"]).encode("utf-8") + payload = RawWeightPushRequest.model_validate_json(pending_bytes) + raw_bytes = payload.canonical_bytes() + if not _pending_still_fresh(payload, now=now): + logger.info( + "raw weight pending expired; rebuilding", + extra={ + "epoch": payload.epoch, + "revision": payload.revision, + "digest": payload.payload_digest[:12], + }, + ) + await self.store.clear_pending() + pending = None + payload = None + raw_bytes = None + except Exception: # noqa: BLE001 - corrupt pending is rebuilt pending = None payload = None raw_bytes = None - except Exception: # noqa: BLE001 - corrupt pending is rebuilt - pending = None - payload = None - raw_bytes = None - - if payload is None or raw_bytes is None: - resolved_weights = ( - dict(weights) - if weights is not None - else dict(await self.weights_fn()) - if self.weights_fn is not None - else {} - ) - # Positive hotkey weights only when synthesizing from get_weights; - # explicit zero maps (caller-supplied zeros) are preserved as zero-contribution. - if weights is not None: - cleaned = {str(hotkey): float(value) for hotkey, value in resolved_weights.items()} - else: - cleaned = { - str(hotkey): float(value) - for hotkey, value in resolved_weights.items() - if float(value) > 0.0 - } - if not cleaned: + + if payload is None or raw_bytes is None: + resolved_weights = ( + dict(weights) + if weights is not None + else dict(await self.weights_fn()) + if self.weights_fn is not None + else {} + ) + # Positive hotkey weights only when synthesizing from get_weights; + # explicit zero maps (caller-supplied zeros) are preserved as zero-contribution. + if weights is not None: + cleaned = { + str(hotkey): float(value) for hotkey, value in resolved_weights.items() + } + else: + cleaned = { + str(hotkey): float(value) + for hotkey, value in resolved_weights.items() + if float(value) > 0.0 + } + if not cleaned: + return PushAttemptResult( + status="skipped_empty", + epoch=0, + revision=0, + payload_digest="", + snapshot_id=None, + cursor_advanced=False, + error="empty weights", + ) + if _is_all_zero_weights(cleaned): + return PushAttemptResult( + status="skipped_zero", + epoch=0, + revision=0, + payload_digest="", + snapshot_id=None, + cursor_advanced=False, + error="all-zero weight map", + ) + cleaned_for_retry = dict(cleaned) + resolved_epoch = await self._resolve_target_epoch( + client=http, explicit=epoch, now=now + ) + revision = ( + int(force_revision) + if force_revision is not None + else self._next_revision(cursor, resolved_epoch) + ) + nonce = f"prism-{uuid.uuid4().hex}" + payload, raw_bytes = self._build_payload( + weights=cleaned, + epoch=resolved_epoch, + revision=revision, + nonce=nonce, + now=now, + ) + await self.store.record_pending( + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + canonical_payload=raw_bytes.decode("utf-8"), + nonce=payload.nonce, + attempted_at=now.isoformat(), + ) + + path = self._path_for() + url = f"{self.master_base_url}{path}" + headers = self._headers(path=path, body=raw_bytes, timestamp=int(now.timestamp())) + try: + response = await http.post(url, content=raw_bytes, headers=headers) + except (httpx.TimeoutException, httpx.TransportError) as exc: + # Never log Authorization / token material — digest prefix only. + logger.info( + "raw weight push transport failure", + extra={ + "epoch": payload.epoch, + "revision": payload.revision, + "digest": payload.payload_digest[:12], + }, + ) return PushAttemptResult( - status="skipped_empty", - epoch=0, - revision=0, - payload_digest="", + status="transport_error", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, snapshot_id=None, cursor_advanced=False, - error="empty weights", + error=str(exc), ) - if _is_all_zero_weights(cleaned): + + if response.status_code >= 500: return PushAttemptResult( - status="skipped_zero", - epoch=0, - revision=0, - payload_digest="", + status="server_error", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, snapshot_id=None, cursor_advanced=False, - error="all-zero weight map", + error=f"status={response.status_code}", ) - resolved_epoch = ( - int(epoch) - if epoch is not None - else int(self.epoch_fn()) - if self.epoch_fn is not None - else int(now.timestamp()) // 360 - ) - revision = ( - int(force_revision) - if force_revision is not None - else self._next_revision(cursor, resolved_epoch) - ) - nonce = f"prism-{uuid.uuid4().hex}" - payload, raw_bytes = self._build_payload( - weights=cleaned, - epoch=resolved_epoch, - revision=revision, - nonce=nonce, - now=now, - ) - await self.store.record_pending( + if response.status_code not in {200, 201}: + detail = _safe_response_detail(response) + logger.info( + "raw weight push rejected", + extra={ + "status": response.status_code, + "detail": detail, + "epoch": payload.epoch, + "revision": payload.revision, + "digest": payload.payload_digest[:12], + }, + ) + if response.status_code in NON_RETRYABLE_PUSH_STATUSES: + await self.store.clear_pending() + err = f"status={response.status_code}" + if detail: + err = f"{err} detail={detail}" + # Bounded advance: sealed/stale/conflict must not retry same identity. + if ( + response.status_code == 409 + and _is_epoch_advance_detail(detail) + and _advance_depth < MAX_IDENTITY_ADVANCES + ): + if _is_revision_conflict_detail(detail): + next_epoch: int | None = payload.epoch + next_revision = payload.revision + 1 + else: + # Sealed/stale: re-query master; fall back to epoch+1. + master_next = await self._fetch_master_target_epoch(client=http) + if master_next is not None and int(master_next) > payload.epoch: + next_epoch = int(master_next) + else: + next_epoch = payload.epoch + 1 + next_revision = 1 + retry_weights = ( + cleaned_for_retry + if cleaned_for_retry is not None + else dict(payload.weights) + ) + return await self.push_once( + weights=retry_weights, + epoch=next_epoch, + force_revision=next_revision, + reuse_pending=False, + _advance_depth=_advance_depth + 1, + ) + return PushAttemptResult( + status="rejected", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + snapshot_id=None, + cursor_advanced=False, + error=err, + ) + try: + ack = RawWeightPushAcknowledgement.model_validate(response.json()) + except Exception as exc: # noqa: BLE001 + return PushAttemptResult( + status="malformed_ack", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + snapshot_id=None, + cursor_advanced=False, + error=str(exc), + ) + if not self._ack_matches(ack, payload=payload): + return PushAttemptResult( + status="ack_mismatch", + epoch=payload.epoch, + revision=payload.revision, + payload_digest=payload.payload_digest, + snapshot_id=ack.snapshot_id if hasattr(ack, "snapshot_id") else None, + cursor_advanced=False, + error="acknowledgement identity mismatch", + ) + ack_time = self._now_fn().isoformat() + await self.store.acknowledge( epoch=payload.epoch, revision=payload.revision, payload_digest=payload.payload_digest, + snapshot_id=ack.snapshot_id, canonical_payload=raw_bytes.decode("utf-8"), nonce=payload.nonce, - attempted_at=now.isoformat(), - ) - - path = self._path_for() - url = f"{self.master_base_url}{path}" - headers = self._headers(path=path, body=raw_bytes, timestamp=int(now.timestamp())) - client = self._http - owns_client = client is None - if owns_client: - client = httpx.AsyncClient(timeout=self.timeout_seconds) - assert client is not None - try: - response = await client.post(url, content=raw_bytes, headers=headers) - except (httpx.TimeoutException, httpx.TransportError) as exc: - logger.info( - "raw weight push transport failure", - extra={ - "epoch": payload.epoch, - "revision": payload.revision, - "digest": payload.payload_digest[:12], - }, + acknowledged_at=ack_time, ) return PushAttemptResult( - status="transport_error", + status="acknowledged", epoch=payload.epoch, revision=payload.revision, payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=str(exc), + snapshot_id=ack.snapshot_id, + cursor_advanced=True, ) finally: if owns_client: - await client.aclose() - - if response.status_code >= 500: - return PushAttemptResult( - status="server_error", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=f"status={response.status_code}", - ) - if response.status_code not in {200, 201}: - detail = _safe_response_detail(response) - logger.info( - "raw weight push rejected", - extra={ - "status": response.status_code, - "detail": detail, - "epoch": payload.epoch, - "revision": payload.revision, - "digest": payload.payload_digest[:12], - }, - ) - if response.status_code in NON_RETRYABLE_PUSH_STATUSES: - await self.store.clear_pending() - err = f"status={response.status_code}" - if detail: - err = f"{err} detail={detail}" - return PushAttemptResult( - status="rejected", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=err, - ) - try: - ack = RawWeightPushAcknowledgement.model_validate(response.json()) - except Exception as exc: # noqa: BLE001 - return PushAttemptResult( - status="malformed_ack", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=None, - cursor_advanced=False, - error=str(exc), - ) - if not self._ack_matches(ack, payload=payload): - return PushAttemptResult( - status="ack_mismatch", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=ack.snapshot_id if hasattr(ack, "snapshot_id") else None, - cursor_advanced=False, - error="acknowledgement identity mismatch", - ) - ack_time = self._now_fn().isoformat() - await self.store.acknowledge( - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=ack.snapshot_id, - canonical_payload=raw_bytes.decode("utf-8"), - nonce=payload.nonce, - acknowledged_at=ack_time, - ) - return PushAttemptResult( - status="acknowledged", - epoch=payload.epoch, - revision=payload.revision, - payload_digest=payload.payload_digest, - snapshot_id=ack.snapshot_id, - cursor_advanced=True, - ) + await http.aclose() def build_weights_loader( @@ -730,9 +832,9 @@ def maybe_build_push_client_from_settings( return None # Challenge scoring epoch (architecture crowns) stays on settings.epoch_seconds. epoch_seconds = int(getattr(settings, "epoch_seconds", 3600) or 3600) - # Master weight-seal identity uses BASE_MASTER epoch interval (default 360s), - # NOT the challenge scoring epoch. Wall-clock can lag force-advanced open - # epochs, so advance past last acknowledged push when behind. + # Fallback only when master target-epoch is unreachable. push_once prefers + # GET /internal/v1/aggregation/target-epoch (max sealed + 1). Wall-clock can + # lag force-advanced seals; ledger last_epoch+1 is a secondary catch-up. master_epoch_seconds = int( getattr(settings, "raw_weight_master_epoch_seconds", 0) or __import__("os").environ.get("PRISM_RAW_WEIGHT_MASTER_EPOCH_SECONDS", 0) diff --git a/packages/challenges/prism/tests/test_raw_weight_push.py b/packages/challenges/prism/tests/test_raw_weight_push.py index 116d152ba..2e39205ab 100644 --- a/packages/challenges/prism/tests/test_raw_weight_push.py +++ b/packages/challenges/prism/tests/test_raw_weight_push.py @@ -502,3 +502,166 @@ async def test_all_zero_weights_skipped_locally(database: Database) -> None: assert result.status in {"skipped_empty", "skipped_zero"} assert transport.requests == [] await http.aclose() + + +@pytest.mark.asyncio +async def test_successful_push_writes_ack_cursor_fields(database: Database) -> None: + clock = FakeClock() + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "target_epoch": 100, + "highest_sealed_epoch": 99, + "max_future_epoch_ahead": 2, + }, + ) + parsed = RawWeightPushRequest.model_validate_json(request.content) + return httpx.Response( + 200, + json={ + "protocol_version": "1.0", + "challenge_slug": SLUG, + "epoch": parsed.epoch, + "revision": parsed.revision, + "snapshot_id": "snap-ack-fields", + "payload_digest": parsed.payload_digest, + "accepted": True, + "idempotent": False, + }, + ) + + http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://master.test", + ) + client = RawWeightPushClient( + database=database, + challenge_slug=SLUG, + master_base_url="http://master.test", + shared_token=TOKEN, + now_fn=clock.now, + http_client=http, + epoch_fn=None, + ) + await client.init() + with activate_role(Role.CHALLENGE): + result = await client.push_once(weights={HOTKEY: 1.0}) + assert result.status == "acknowledged" + cursor = await client.store.get_cursor() + assert cursor is not None + assert cursor.epoch == 100 + assert cursor.snapshot_id == "snap-ack-fields" + assert cursor.acknowledged_at is not None + await http.aclose() + + +@pytest.mark.asyncio +async def test_sealed_epoch_409_advances_via_master_target(database: Database) -> None: + clock = FakeClock() + post_epochs: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "target_epoch": 50, + "highest_sealed_epoch": 49, + "max_future_epoch_ahead": 2, + }, + ) + parsed = RawWeightPushRequest.model_validate_json(request.content) + post_epochs.append(parsed.epoch) + if parsed.epoch < 50: + return httpx.Response( + 409, json={"detail": "epoch is sealed; revision rejected"} + ) + return httpx.Response( + 200, + json={ + "protocol_version": "1.0", + "challenge_slug": SLUG, + "epoch": parsed.epoch, + "revision": parsed.revision, + "snapshot_id": f"snap-{parsed.epoch}", + "payload_digest": parsed.payload_digest, + "accepted": True, + "idempotent": False, + }, + ) + + http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://master.test", + ) + client = RawWeightPushClient( + database=database, + challenge_slug=SLUG, + master_base_url="http://master.test", + shared_token=TOKEN, + now_fn=clock.now, + http_client=http, + epoch_fn=lambda: 49, + ) + await client.init() + with activate_role(Role.CHALLENGE): + result = await client.push_once(weights={HOTKEY: 1.0}) + assert result.status == "acknowledged" + assert result.epoch == 50 + assert post_epochs.count(49) <= 1 + await http.aclose() + + +@pytest.mark.asyncio +async def test_epoch_source_prefers_master_over_local_clock(database: Database) -> None: + clock = FakeClock() + posted: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "target_epoch": 9001, + "highest_sealed_epoch": 9000, + "max_future_epoch_ahead": 2, + }, + ) + parsed = RawWeightPushRequest.model_validate_json(request.content) + posted.append(parsed.epoch) + return httpx.Response( + 200, + json={ + "protocol_version": "1.0", + "challenge_slug": SLUG, + "epoch": parsed.epoch, + "revision": parsed.revision, + "snapshot_id": "snap-master-epoch", + "payload_digest": parsed.payload_digest, + "accepted": True, + "idempotent": False, + }, + ) + + http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://master.test", + ) + client = RawWeightPushClient( + database=database, + challenge_slug=SLUG, + master_base_url="http://master.test", + shared_token=TOKEN, + now_fn=clock.now, + http_client=http, + epoch_fn=lambda: 1, + ) + await client.init() + with activate_role(Role.CHALLENGE): + result = await client.push_once(weights={HOTKEY: 1.0}) + assert result.status == "acknowledged" + assert posted == [9001] + await http.aclose() diff --git a/src/base/master/raw_weight_ingress.py b/src/base/master/raw_weight_ingress.py index 71097fed2..60bc72d14 100644 --- a/src/base/master/raw_weight_ingress.py +++ b/src/base/master/raw_weight_ingress.py @@ -116,6 +116,15 @@ def _log_rejection( ) +@dataclass(frozen=True) +class TargetEpochInfo: + """Next unsealed aggregation epoch clients should target for raw-weight push.""" + + target_epoch: int + highest_sealed_epoch: int | None + max_future_epoch_ahead: int + + class RawWeightAuthError(PermissionError): """Missing, malformed, or unknown challenge credential (HTTP 401).""" @@ -406,6 +415,51 @@ async def seal_epoch(self, epoch: int) -> AggregationEpoch: await session.refresh(row) return row + async def get_target_epoch(self) -> TargetEpochInfo: + """Return the next unsealed aggregation epoch (max sealed + 1). + + Matches the production hotpatch semantic: + ``COALESCE(MAX(sealed epoch), 0) + 1`` from ``aggregation_epochs``. + Challenge clients must use this rather than a local wall-clock guess so + pushes land on an epoch master will still accept. + """ + + async with session_scope(self._session_factory) as session: + highest_sealed = ( + await session.execute( + select(AggregationEpoch.epoch) + .where(AggregationEpoch.status == AggregationEpochStatus.SEALED) + .order_by(AggregationEpoch.epoch.desc()) + .limit(1) + ) + ).scalar_one_or_none() + highest = int(highest_sealed) if highest_sealed is not None else None + target = (highest if highest is not None else 0) + 1 + return TargetEpochInfo( + target_epoch=target, + highest_sealed_epoch=highest, + max_future_epoch_ahead=int(self.max_future_epoch_ahead), + ) + + async def authorize_challenge_bearer( + self, + *, + authorization: str | None, + challenge_slug: str | None, + ) -> None: + """Accept a bearer token bound to a known challenge slug (read paths).""" + + if not challenge_slug: + raise RawWeightAuthError(UNAUTHORIZED_DETAIL) + if not authorization or not authorization.lower().startswith("bearer "): + raise RawWeightAuthError(UNAUTHORIZED_DETAIL) + token = authorization.split(" ", 1)[1].strip() + if not token: + raise RawWeightAuthError(UNAUTHORIZED_DETAIL) + expected = await self._credentials.token_for(str(challenge_slug)) + if expected is None or not hmac.compare_digest(token, expected): + raise RawWeightAuthError(UNAUTHORIZED_DETAIL) + def _validate_content_type(self, content_type: str | None) -> None: if content_type is None: raise RawWeightSchemaError(UNSUPPORTED_MEDIA_DETAIL) @@ -1043,6 +1097,37 @@ async def push_raw_weights( idempotent=outcome.idempotent, ) + @router.get( + "/internal/v1/aggregation/target-epoch", + responses={ + 401: {"description": "unauthorized"}, + }, + ) + async def get_aggregation_target_epoch( + authorization: str | None = Header(default=None), + x_base_challenge_slug: str | None = Header( + default=None, alias="X-Base-Challenge-Slug" + ), + ) -> dict[str, Any]: + """Read-only next unsealed epoch for challenge raw-weight push clients.""" + + try: + await service.authorize_challenge_bearer( + authorization=authorization, + challenge_slug=x_base_challenge_slug, + ) + except RawWeightAuthError as exc: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=UNAUTHORIZED_DETAIL, + ) from exc + info = await service.get_target_epoch() + return { + "target_epoch": info.target_epoch, + "highest_sealed_epoch": info.highest_sealed_epoch, + "max_future_epoch_ahead": info.max_future_epoch_ahead, + } + return router @@ -1054,6 +1139,7 @@ async def push_raw_weights( "DEFAULT_MAX_WEIGHT_KEYS", "ZERO_WEIGHT_DETAIL", "PushOutcome", + "TargetEpochInfo", "RAW_WEIGHT_FRESHNESS_POLICY_VERSION", "RawWeightAuthError", "RawWeightConflictError", diff --git a/tests/unit/test_raw_weight_ingress.py b/tests/unit/test_raw_weight_ingress.py index 6bb48e718..702ff0194 100644 --- a/tests/unit/test_raw_weight_ingress.py +++ b/tests/unit/test_raw_weight_ingress.py @@ -798,3 +798,50 @@ async def test_freshness_detail_distinct_from_schema( assert r.status_code == 422 assert r.json()["detail"] == "snapshot outside freshness window" assert r.json()["detail"] != "invalid raw weight payload" + + +@pytest.mark.asyncio +async def test_target_epoch_is_one_past_highest_sealed(harness: dict[str, Any]) -> None: + """Master-authoritative next unsealed epoch: COALESCE(MAX(sealed),0)+1.""" + + client = harness["client"] + service: RawWeightIngressService = harness["service"] + path = "/internal/v1/aggregation/target-epoch" + headers = { + "Authorization": f"Bearer {TOKEN}", + "X-Base-Challenge-Slug": SLUG, + "Accept": "application/json", + } + + empty = await client.get(path, headers=headers) + assert empty.status_code == 200, empty.text + body = empty.json() + assert body["target_epoch"] == 1 + assert body["highest_sealed_epoch"] is None + + caps = (Capability.MASTER_RAW_WEIGHT_INGRESS,) + with activate_role(Role.MASTER, capabilities=caps): + await service.seal_epoch(4959497) + await service.seal_epoch(4959495) + + sealed = await client.get(path, headers=headers) + assert sealed.status_code == 200, sealed.text + body = sealed.json() + assert body["highest_sealed_epoch"] == 4959497 + assert body["target_epoch"] == 4959498 + assert body["max_future_epoch_ahead"] == 2 + + +@pytest.mark.asyncio +async def test_target_epoch_requires_challenge_bearer(harness: dict[str, Any]) -> None: + client = harness["client"] + path = "/internal/v1/aggregation/target-epoch" + assert (await client.get(path)).status_code == 401 + bad = await client.get( + path, + headers={ + "Authorization": "Bearer wrong-token", + "X-Base-Challenge-Slug": SLUG, + }, + ) + assert bad.status_code == 401