diff --git a/alembic/versions/2026_07_24_add_disqualified_agents.py b/alembic/versions/2026_07_24_add_disqualified_agents.py new file mode 100644 index 00000000..a1e3bda9 --- /dev/null +++ b/alembic/versions/2026_07_24_add_disqualified_agents.py @@ -0,0 +1,89 @@ +"""Add disqualified_agents table, disqualified_agent_ids view, and disqualification_jobs table. + +Revision ID: e5c8a1f0b942 +Revises: b3f1a9c4d210 +Create Date: 2026-07-24 00:00:00.000000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +revision: str = "e5c8a1f0b942" +down_revision: Union[str, Sequence[str], None] = "b3f1a9c4d210" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +_CREATE_VIEW = """ +CREATE VIEW disqualified_agent_ids AS + SELECT a.agent_id + FROM agents a + JOIN banned_coldkeys bc ON bc.miner_coldkey = a.miner_coldkey + UNION + SELECT agent_id + FROM disqualified_agents; +""" + + +def upgrade() -> None: + op.create_table( + "disqualified_agents", + sa.Column( + "agent_id", + sa.UUID(), + sa.ForeignKey("agents.agent_id", ondelete="CASCADE"), + primary_key=True, + ), + sa.Column("reason", sa.Text(), nullable=False), + sa.Column( + "disqualified_at", + sa.TIMESTAMP(timezone=True), + server_default=sa.text("NOW()"), + nullable=False, + ), + ) + op.execute(_CREATE_VIEW) + + op.create_table( + "disqualification_jobs", + sa.Column( + "id", + sa.UUID(), + server_default=sa.text("gen_random_uuid()"), + primary_key=True, + ), + sa.Column( + "agent_id", + sa.UUID(), + sa.ForeignKey("agents.agent_id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column("set_id", sa.Integer(), nullable=False), + sa.Column( + "created_at", + sa.TIMESTAMP(timezone=True), + server_default=sa.text("NOW()"), + nullable=False, + ), + sa.Column("processed_at", sa.TIMESTAMP(timezone=True), nullable=True), + sa.Column("attempts", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("error", sa.Text(), nullable=True), + ) + op.create_index( + "uq_disqualification_jobs_pending", + "disqualification_jobs", + ["agent_id"], + unique=True, + postgresql_where=sa.text("processed_at IS NULL"), + ) + + +def downgrade() -> None: + op.drop_index("uq_disqualification_jobs_pending", table_name="disqualification_jobs") + op.drop_table("disqualification_jobs") + op.execute("DROP VIEW IF EXISTS disqualified_agent_ids") + op.drop_table("disqualified_agents") diff --git a/api/config.py b/api/config.py index 4bbd136d..56b9daf6 100644 --- a/api/config.py +++ b/api/config.py @@ -207,6 +207,7 @@ PRE_SCREENING_PROJECTOR_POLL_INTERVAL_SECONDS = int(os.getenv("PRE_SCREENING_PROJECTOR_POLL_INTERVAL_SECONDS", "5")) AUTO_APPROVAL_ENABLED = os.getenv("AUTO_APPROVAL_ENABLED", "false").lower() == "true" AUTO_APPROVAL_RUN_LOOP = SHOULD_RUN_LOOPS and AUTO_APPROVAL_ENABLED +DISQUALIFICATION_REAPPROVAL_RUN = SHOULD_RUN_LOOPS AUTO_APPROVAL_POLICY_VERSION = os.getenv("AUTO_APPROVAL_POLICY_VERSION", "approval-v1") APPROVAL_PROJECTOR_POLL_INTERVAL_SECONDS = int(os.getenv("APPROVAL_PROJECTOR_POLL_INTERVAL_SECONDS", "5")) @@ -309,6 +310,7 @@ def _fraction_setting(name: str, default: str) -> float: logger.info(f"Pre-Screening Projector Poll Interval: {PRE_SCREENING_PROJECTOR_POLL_INTERVAL_SECONDS} second(s)") logger.info(f"Auto Approval Enabled: {AUTO_APPROVAL_ENABLED}") logger.info(f"Auto Approval Projector Loop Enabled: {AUTO_APPROVAL_RUN_LOOP}") +logger.info(f"Disqualification Reapproval Enabled: {DISQUALIFICATION_REAPPROVAL_RUN}") logger.info(f"Approval Projector Poll Interval: {APPROVAL_PROJECTOR_POLL_INTERVAL_SECONDS} second(s)") logger.info(f"Earliest SET ID with good data: {EARLIEST_SET_ID_WITH_GOOD_DATA}") logger.info(f"Incentive Start Set ID: {INCENTIVE_START_SET_ID}") diff --git a/api/endpoints/admin.py b/api/endpoints/admin.py index 469f6d92..e7be6004 100644 --- a/api/endpoints/admin.py +++ b/api/endpoints/admin.py @@ -1,5 +1,8 @@ +import asyncio +import logging import secrets from typing import Annotated +from uuid import UUID from bittensor_wallet.keypair import Keypair from fastapi import APIRouter, Depends, HTTPException, Response, status @@ -8,12 +11,24 @@ import api.config as config from models.banned_coldkey import BannedColdkey +from models.disqualified_agent import DisqualifiedAgent +from queries.agent import get_agent_by_id +from queries.approval import process_pending_disqualification_jobs from queries.banned_coldkey import ban_coldkey, unban_coldkey +from queries.disqualification_job import enqueue_disqualification_job +from queries.disqualified_agent import disqualify_agent +from utils.database import DatabaseConnection, db_operation from utils.ttl import clear_all_ttl_caches +logger = logging.getLogger(__name__) + router = APIRouter(tags=["admin"]) admin_bearer = HTTPBearer(auto_error=False) +# Retains references to fire-and-forget drain tasks so they can't be garbage-collected +# before completion (asyncio only holds a weak reference to a bare create_task result). +_background_tasks: set[asyncio.Task[None]] = set() + class ColdkeyBanRequest(BaseModel): reason: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=1000)] @@ -62,3 +77,48 @@ async def delete_banned_coldkey(miner_coldkey: str) -> Response: await unban_coldkey(miner_coldkey) clear_all_ttl_caches() return Response(status_code=status.HTTP_204_NO_CONTENT) + + +@router.put( + "/disqualified-agents/{agent_id}", + response_model=DisqualifiedAgent, + dependencies=[Depends(require_coldkey_ban_admin)], +) +async def put_disqualified_agent(agent_id: UUID, request: ColdkeyBanRequest) -> DisqualifiedAgent: + agent = await get_agent_by_id(agent_id) + if agent is None: + raise HTTPException(status_code=404, detail="Agent not found") + + disqualified = await disqualify_agent(agent_id, request.reason) + + set_id = await _enqueue_disqualification_job_operation(agent_id=agent_id) + if set_id is not None: + _fire_disqualification_drain() + + clear_all_ttl_caches() + return disqualified + + +@db_operation +async def _enqueue_disqualification_job_operation(conn: DatabaseConnection, *, agent_id: UUID) -> int | None: + """Enqueue a reapproval job for the agent's set. Returns the set_id, or None if the agent has none.""" + async with conn.conn.transaction(): + set_id = await conn.fetchval("SELECT set_id FROM agents WHERE agent_id = $1", agent_id) + if set_id is None: + return None + await enqueue_disqualification_job(conn, agent_id=agent_id, set_id=set_id) + return set_id + + +async def _run_disqualification_drain() -> None: + try: + await process_pending_disqualification_jobs() + except Exception as exc: # noqa: BLE001 + logger.error(f"Disqualification drain task failed: {type(exc).__name__}: {exc}") + + +def _fire_disqualification_drain() -> None: + """Fire the drain as a background task, retaining a reference until it completes.""" + task = asyncio.create_task(_run_disqualification_drain()) + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) diff --git a/api/src/main.py b/api/src/main.py index 68677bef..07614b58 100644 --- a/api/src/main.py +++ b/api/src/main.py @@ -29,6 +29,7 @@ from api.src.endpoints.upload import router as upload_router from api.src.middleware.request_interceptor import RequestInterceptorMiddleware from api.src.utils.sentry import initialize_sentry +from queries.approval import process_pending_disqualification_jobs from queries.evaluation import set_all_unfinished_evaluation_runs_to_errored from utils.bittensor import subtensor_client from utils.database import deinitialize_database, initialize_database @@ -96,6 +97,14 @@ async def lifespan(app: FastAPI): error_message="Platform crashed while running this evaluation" ) + if config.DISQUALIFICATION_REAPPROVAL_RUN: + try: + drained = await process_pending_disqualification_jobs() + if drained: + logger.info(f"Drained {drained} pending disqualification job(s) on startup") + except Exception as exc: # noqa: BLE001 + logger.error(f"Startup disqualification drain failed: {type(exc).__name__}: {exc}") + yield tasks_to_cancel = tuple(background_tasks) diff --git a/db/models/agent.py b/db/models/agent.py index 7f1a2151..8acca127 100644 --- a/db/models/agent.py +++ b/db/models/agent.py @@ -68,6 +68,40 @@ class BannedColdkey(Base): ) +class DisqualifiedAgent(Base): + __tablename__ = "disqualified_agents" + + agent_id: Mapped[UUID] = mapped_column( + PG_UUID(as_uuid=True), + sa.ForeignKey("agents.agent_id", ondelete="CASCADE"), + primary_key=True, + ) + reason: Mapped[str] = mapped_column(sa.Text, nullable=False) + disqualified_at: Mapped[datetime] = mapped_column( + sa.TIMESTAMP(timezone=True), + nullable=False, + server_default=sa.text("NOW()"), + ) + + +class DisqualificationJob(Base): + __tablename__ = "disqualification_jobs" + + id: Mapped[UUID] = mapped_column( + PG_UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()") + ) + agent_id: Mapped[UUID] = mapped_column( + PG_UUID(as_uuid=True), sa.ForeignKey("agents.agent_id", ondelete="CASCADE"), nullable=False + ) + set_id: Mapped[int] = mapped_column(sa.Integer, nullable=False) + created_at: Mapped[datetime] = mapped_column( + sa.TIMESTAMP(timezone=True), nullable=False, server_default=sa.text("NOW()") + ) + processed_at: Mapped[Optional[datetime]] = mapped_column(sa.TIMESTAMP(timezone=True)) + attempts: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=sa.text("0")) + error: Mapped[Optional[str]] = mapped_column(sa.Text) + + class BenchmarkAgentId(Base): __tablename__ = "benchmark_agent_ids" diff --git a/models/disqualification_job.py b/models/disqualification_job.py new file mode 100644 index 00000000..79106b82 --- /dev/null +++ b/models/disqualification_job.py @@ -0,0 +1,14 @@ +from datetime import datetime +from uuid import UUID + +from pydantic import BaseModel + + +class DisqualificationJob(BaseModel): + id: UUID + agent_id: UUID + set_id: int + created_at: datetime + processed_at: datetime | None = None + attempts: int = 0 + error: str | None = None diff --git a/models/disqualified_agent.py b/models/disqualified_agent.py new file mode 100644 index 00000000..1f385d5a --- /dev/null +++ b/models/disqualified_agent.py @@ -0,0 +1,10 @@ +from datetime import datetime +from uuid import UUID + +from pydantic import BaseModel + + +class DisqualifiedAgent(BaseModel): + agent_id: UUID + reason: str + disqualified_at: datetime diff --git a/queries/agent.py b/queries/agent.py index 2715c952..51442b89 100644 --- a/queries/agent.py +++ b/queries/agent.py @@ -544,8 +544,8 @@ async def get_top_agents(conn: DatabaseConnection, number_of_agents: int = 10, p and ass.agent_id not in (select agent_id from benchmark_agent_ids) and not exists ( select 1 - from banned_coldkeys bc - where bc.miner_coldkey = a.miner_coldkey + from disqualified_agent_ids dq + where dq.agent_id = a.agent_id ) and ass.status::text <> 'cancelled' and ( @@ -590,8 +590,8 @@ async def get_code_hiding_score_cutoff( AND ass.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND review.approval_review_status IS DISTINCT FROM 'rejected' ) diff --git a/queries/approval.py b/queries/approval.py index 1ca465c1..0171cd90 100644 --- a/queries/approval.py +++ b/queries/approval.py @@ -17,6 +17,11 @@ from models.evaluation import EvaluationStatus from models.evaluation_set import EvaluationSetGroup from queries.banned_coldkey import get_banned_coldkey, lock_coldkey_ban_state +from queries.disqualification_job import ( + claim_next_pending_disqualification_job, + mark_disqualification_job_processed, + record_disqualification_job_error, +) from queries.evaluation import ( AgentRankingProfile, get_approved_leader_ranking_for_set, @@ -196,47 +201,20 @@ async def project_next_approval_job_state(conn: DatabaseConnection) -> bool: return True -async def _insert_incentive_approval( +async def _apply_incentive_decision( conn: DatabaseConnection, + *, agent_id: UUID, set_id: int, + candidate: AgentRankingProfile, + leader: AgentRankingProfile | None, + decision_time: datetime, ) -> str | None: - await conn.execute( - "SELECT pg_advisory_xact_lock($1, $2)", - INCENTIVE_APPROVAL_LOCK_NAMESPACE, - set_id, - ) - - agent = await conn.fetchrow( - """ - SELECT miner_coldkey, status - FROM agents - WHERE agent_id = $1 - FOR UPDATE - """, - agent_id, - ) - if agent is None: - return "Candidate no longer exists" - - if agent["status"] != AgentStatus.finished.value: - return f"Candidate is not finished (status={agent['status']})" - - miner_coldkey = agent["miner_coldkey"] - if miner_coldkey is not None: - await lock_coldkey_ban_state(conn, miner_coldkey) - if await get_banned_coldkey(miner_coldkey) is not None: - return "Candidate coldkey is banned" + """Decide and record one incentive approval against an explicit leader. - candidate = await get_validator_agent_score_for_set(agent_id, set_id, config.NUM_EVALS_PER_AGENT) - if candidate is None: - return "Candidate no longer has a complete validator score" - - leader = await _get_ban_stable_leader( - conn, - set_id, - agent_id, - ) + Returns a rejection reason if the candidate does not qualify (nothing written), + or None on success (approved_agents row inserted). + """ improvement = calculate_relative_improvement( candidate_score=candidate.final_score, candidate_cost=candidate.avg_cost_usd, @@ -248,7 +226,6 @@ async def _insert_incentive_approval( if not improvement.qualified: return "Candidate no longer meets the relative improvement threshold" - decision_time: datetime = await conn.fetchval("SELECT NOW()") last_competition_improvement = None if leader is None else leader.approved_at competition_elapsed_hours = _elapsed_hours(last_competition_improvement, decision_time) time_multiplier = calculate_time_multiplier( @@ -256,7 +233,6 @@ async def _insert_incentive_approval( half_life_hours=config.INCENTIVE_TIME_MULTIPLIER_HALF_LIFE_HOURS, maximum=config.INCENTIVE_TIME_MULTIPLIER_MAX, ) - initial_reward_score = calculate_initial_reward_score( relative_improvement_units=improvement.relative_improvement_units, time_multiplier=time_multiplier, @@ -290,6 +266,324 @@ async def _insert_incentive_approval( return None +@db_operation +async def run_disqualification_reapproval( + conn: DatabaseConnection, + *, + set_id: int, + disqualified_agent_id: UUID, +) -> None: + """Replay downstream incentive decisions after an agent is disqualified. + + Walks every agent whose incentive decision was made after the disqualified agent B's + decision, in decision order, re-running the approve/reject gate against a leader lineage + that no longer includes B. + """ + + async with conn.conn.transaction(): + await conn.execute( + "SELECT pg_advisory_xact_lock($1, $2)", + INCENTIVE_APPROVAL_LOCK_NAMESPACE, + set_id, + ) + + # B's score and its decision moment. final_score comes from agent_scores; the decision + # time is approval_jobs.projected_at of B's latest decision job (write-once, exists for + # approved AND rejected agents). Do NOT use agent_scores.created_at (that is upload time). + b_row = await conn.fetchrow( + """ + SELECT ass.final_score, job.projected_at + FROM agent_scores ass + LEFT JOIN agent_approval_states st + ON st.agent_id = ass.agent_id AND st.set_id = ass.set_id + LEFT JOIN approval_jobs job + ON job.job_id = st.latest_job_id + WHERE ass.agent_id = $1 AND ass.set_id = $2 + """, + disqualified_agent_id, + set_id, + ) + if b_row is None or b_row["projected_at"] is None: + # B was never scored, or never reached an incentive decision: nothing downstream + # was gated against it as leader. + return + + b_approved = await conn.fetchrow( + "SELECT baseline_agent_id FROM approved_agents WHERE agent_id = $1 AND set_id = $2", + disqualified_agent_id, + set_id, + ) + if b_approved is None: + # B never held an approved_agents row, so it never occupied the leader lineage: + # nothing downstream was ever gated against it as leader. Nothing to replay. + return + baseline_agent_id = b_approved["baseline_agent_id"] + + await conn.execute( + "DELETE FROM approved_agents WHERE agent_id = $1 AND set_id = $2", + disqualified_agent_id, + set_id, + ) + + current_leader = await _ranking_profile_for_agent(conn, baseline_agent_id, set_id) + + candidates = await conn.fetch( + """ + SELECT + ass.agent_id, + ass.final_score, + ass.created_at, + ass.approved_at, + job.projected_at, + rt.avg_cost_usd, + (aa.agent_id IS NOT NULL) AS is_approved + FROM agent_scores ass + INNER JOIN agents agent ON agent.agent_id = ass.agent_id + INNER JOIN agent_approval_states st + ON st.agent_id = ass.agent_id AND st.set_id = ass.set_id + INNER JOIN approval_jobs job + ON job.job_id = st.latest_job_id + LEFT JOIN approved_agents aa + ON aa.agent_id = ass.agent_id AND aa.set_id = ass.set_id + LEFT JOIN LATERAL ( + SELECT AVG(eh.avg_cost_usd) AS avg_cost_usd + FROM evaluations_hydrated eh + WHERE eh.agent_id = ass.agent_id + AND eh.set_id = ass.set_id + AND eh.evaluation_set_group = 'validator'::EvaluationSetGroup + AND eh.status = 'success'::EvaluationStatus + ) rt ON true + WHERE ass.set_id = $1 + AND ass.agent_id <> $2 + AND ass.final_score >= $3 + AND job.projected_at > $4 + AND ass.validator_count = $5 + AND st.system_verdict IN ('approved', 'rejected') + AND NOT EXISTS ( + SELECT 1 FROM benchmark_agent_ids b WHERE b.agent_id = ass.agent_id + ) + AND NOT EXISTS ( + SELECT 1 FROM disqualified_agent_ids dq WHERE dq.agent_id = ass.agent_id + ) + ORDER BY job.projected_at ASC, ass.agent_id ASC + """, + set_id, + disqualified_agent_id, + b_row["final_score"], + b_row["projected_at"], + config.NUM_EVALS_PER_AGENT, + ) + + decision_time: datetime = await conn.fetchval("SELECT NOW()") + + for row in candidates: + candidate = AgentRankingProfile( + final_score=row["final_score"], + avg_cost_usd=row["avg_cost_usd"], + created_at=row["created_at"], + agent_id=row["agent_id"], + approved_at=row["approved_at"], + ) + if row["is_approved"]: + improvement = calculate_relative_improvement( + candidate_score=candidate.final_score, + candidate_cost=candidate.avg_cost_usd, + leader_score=None if current_leader is None else current_leader.final_score, + leader_cost=None if current_leader is None else current_leader.avg_cost_usd, + performance_threshold=config.INCENTIVE_PERFORMANCE_THRESHOLD, + cost_threshold=config.INCENTIVE_COST_THRESHOLD, + ) + if improvement.qualified: + last_leader_approved_at = None if current_leader is None else current_leader.approved_at + elapsed_hours = _elapsed_hours(last_leader_approved_at, decision_time) + time_multiplier = calculate_time_multiplier( + elapsed_hours=elapsed_hours, + half_life_hours=config.INCENTIVE_TIME_MULTIPLIER_HALF_LIFE_HOURS, + maximum=config.INCENTIVE_TIME_MULTIPLIER_MAX, + ) + initial_reward_score = calculate_initial_reward_score( + relative_improvement_units=improvement.relative_improvement_units, + time_multiplier=time_multiplier, + ) + await conn.execute( + """ + UPDATE approved_agents + SET baseline_agent_id = $3, + performance_delta = $4, + cost_delta = $5, + relative_improvement_units = $6, + time_multiplier = $7, + initial_reward_score = $8 + WHERE agent_id = $1 AND set_id = $2 + """, + row["agent_id"], + set_id, + None if current_leader is None else current_leader.agent_id, + improvement.performance_delta, + improvement.cost_delta, + improvement.relative_improvement_units, + time_multiplier, + initial_reward_score, + ) + current_leader = candidate + else: + await conn.execute( + "DELETE FROM approved_agents WHERE agent_id = $1 AND set_id = $2", + row["agent_id"], + set_id, + ) + await _set_system_verdict(conn, row["agent_id"], set_id, "rejected") + else: + reason = await _apply_incentive_decision( + conn, + agent_id=row["agent_id"], + set_id=set_id, + candidate=candidate, + leader=current_leader, + decision_time=decision_time, + ) + if reason is None: + await _set_system_verdict(conn, row["agent_id"], set_id, "approved") + current_leader = AgentRankingProfile( + final_score=candidate.final_score, + avg_cost_usd=candidate.avg_cost_usd, + created_at=candidate.created_at, + agent_id=candidate.agent_id, + approved_at=decision_time, + ) + + +async def process_pending_disqualification_jobs() -> int: + """Drain currently-pending disqualification jobs. Safe to call on a task or at startup. + + Not a @db_operation: each nested query acquires its own connection so a job's + claim + replay + mark commit independently and no row lock is held across the replay. + A failure on one job records an error and the drain continues with the next. + + Each invocation processes every DISTINCT currently-pending job at most once, then returns. + A job that fails its replay is left pending (error recorded) for a LATER invocation to retry. + The ids already attempted this invocation are passed to the claim query as an exclusion list, + so the claim advances past a stuck failing job to the next distinct pending job instead of + starving it forever behind the head-of-queue failure (claim orders by created_at, filtering + processed_at IS NULL, so without exclusion it would keep re-returning the same failing job). + """ + processed = 0 + attempted: list = [] + while True: + job = await claim_next_pending_disqualification_job(attempted or None) + if job is None: + return processed + attempted.append(job["id"]) + try: + await run_disqualification_reapproval( + set_id=job["set_id"], + disqualified_agent_id=job["agent_id"], + ) + await mark_disqualification_job_processed(job["id"]) + processed += 1 + except Exception as exc: # noqa: BLE001 - one bad job must not wedge the drain + logger.error(f"Disqualification reapproval job {job['id']} failed: {type(exc).__name__}: {exc}") + await record_disqualification_job_error(job["id"], f"{type(exc).__name__}: {exc}") + + +async def _ranking_profile_for_agent( + conn: DatabaseConnection, agent_id: UUID | None, set_id: int +) -> AgentRankingProfile | None: + if agent_id is None: + return None + row = await conn.fetchrow( + """ + SELECT ass.agent_id, ass.final_score, ass.approved_at, ass.created_at, rt.avg_cost_usd + FROM agent_scores ass + LEFT JOIN LATERAL ( + SELECT AVG(eh.avg_cost_usd) AS avg_cost_usd + FROM evaluations_hydrated eh + WHERE eh.agent_id = ass.agent_id AND eh.set_id = ass.set_id + AND eh.evaluation_set_group = 'validator'::EvaluationSetGroup + AND eh.status = 'success'::EvaluationStatus + ) rt ON true + WHERE ass.agent_id = $1 AND ass.set_id = $2 + """, + agent_id, + set_id, + ) + if row is None: + return None + return AgentRankingProfile( + final_score=row["final_score"], + avg_cost_usd=row["avg_cost_usd"], + created_at=row["created_at"], + agent_id=row["agent_id"], + approved_at=row["approved_at"], + ) + + +async def _set_system_verdict(conn: DatabaseConnection, agent_id: UUID, set_id: int, verdict: str) -> None: + await conn.execute( + """ + UPDATE agent_approval_states + SET system_verdict = $3, updated_at = NOW() + WHERE agent_id = $1 AND set_id = $2 + """, + agent_id, + set_id, + verdict, + ) + + +async def _insert_incentive_approval( + conn: DatabaseConnection, + agent_id: UUID, + set_id: int, +) -> str | None: + await conn.execute( + "SELECT pg_advisory_xact_lock($1, $2)", + INCENTIVE_APPROVAL_LOCK_NAMESPACE, + set_id, + ) + + agent = await conn.fetchrow( + """ + SELECT miner_coldkey, status + FROM agents + WHERE agent_id = $1 + FOR UPDATE + """, + agent_id, + ) + if agent is None: + return "Candidate no longer exists" + + if agent["status"] != AgentStatus.finished.value: + return f"Candidate is not finished (status={agent['status']})" + + miner_coldkey = agent["miner_coldkey"] + if miner_coldkey is not None: + await lock_coldkey_ban_state(conn, miner_coldkey) + if await get_banned_coldkey(miner_coldkey) is not None: + return "Candidate coldkey is banned" + + candidate = await get_validator_agent_score_for_set(agent_id, set_id, config.NUM_EVALS_PER_AGENT) + if candidate is None: + return "Candidate no longer has a complete validator score" + + leader = await _get_ban_stable_leader( + conn, + set_id, + agent_id, + ) + + decision_time: datetime = await conn.fetchval("SELECT NOW()") + return await _apply_incentive_decision( + conn, + agent_id=agent_id, + set_id=set_id, + candidate=candidate, + leader=leader, + decision_time=decision_time, + ) + + async def _get_ban_stable_leader( conn: DatabaseConnection, set_id: int, diff --git a/queries/disqualification_job.py b/queries/disqualification_job.py new file mode 100644 index 00000000..ae54dc91 --- /dev/null +++ b/queries/disqualification_job.py @@ -0,0 +1,79 @@ +from uuid import UUID + +from asyncpg import Record + +from utils.database import DatabaseConnection, db_operation + + +async def enqueue_disqualification_job( + conn: DatabaseConnection, + *, + agent_id: UUID, + set_id: int, +) -> UUID | None: + """Insert a pending disqualification job. Returns None if one is already pending for this agent. + + Called inside the caller's transaction (e.g. the disqualify endpoint), so no @db_operation. + """ + row = await conn.fetchrow( + """ + INSERT INTO disqualification_jobs (agent_id, set_id) + VALUES ($1, $2) + ON CONFLICT (agent_id) WHERE processed_at IS NULL DO NOTHING + RETURNING id + """, + agent_id, + set_id, + ) + return row["id"] if row is not None else None + + +@db_operation +async def claim_next_pending_disqualification_job( + conn: DatabaseConnection, exclude_ids: list[UUID] | None = None +) -> Record | None: + """Claim the oldest pending job, optionally excluding ids already attempted this invocation. + + The exclusion lets a single drain invocation skip past a job it already attempted (and that + failed, leaving it pending) so it advances to the next distinct pending job instead of being + starved behind it. See process_pending_disqualification_jobs in queries/approval.py. + """ + return await conn.fetchrow( + """ + UPDATE disqualification_jobs + SET attempts = attempts + 1 + WHERE id = ( + SELECT id + FROM disqualification_jobs + WHERE processed_at IS NULL + AND ($1::uuid[] IS NULL OR id <> ALL($1::uuid[])) + ORDER BY created_at ASC + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + RETURNING id, agent_id, set_id + """, + exclude_ids, + ) + + +@db_operation +async def mark_disqualification_job_processed(conn: DatabaseConnection, id: UUID) -> None: + await conn.execute( + "UPDATE disqualification_jobs SET processed_at = NOW(), error = NULL WHERE id = $1", + id, + ) + + +@db_operation +async def record_disqualification_job_error(conn: DatabaseConnection, id: UUID, error: str) -> None: + await conn.execute( + "UPDATE disqualification_jobs SET error = $2 WHERE id = $1", + id, + error, + ) + + +@db_operation +async def count_pending_disqualification_jobs(conn: DatabaseConnection) -> int: + return await conn.fetchval("SELECT COUNT(*) FROM disqualification_jobs WHERE processed_at IS NULL") diff --git a/queries/disqualified_agent.py b/queries/disqualified_agent.py new file mode 100644 index 00000000..b7e9852f --- /dev/null +++ b/queries/disqualified_agent.py @@ -0,0 +1,49 @@ +from typing import Optional +from uuid import UUID + +from models.disqualified_agent import DisqualifiedAgent +from utils.database import DatabaseConnection, db_operation + +DISQUALIFIED_AGENT_LOCK_NAMESPACE = -1731 + + +async def lock_disqualified_agent_state(conn: DatabaseConnection, agent_id: UUID) -> None: + await conn.execute( + "SELECT pg_advisory_xact_lock($1, hashtext($2))", + DISQUALIFIED_AGENT_LOCK_NAMESPACE, + str(agent_id), + ) + + +@db_operation +async def get_disqualified_agent( + conn: DatabaseConnection, + agent_id: UUID, +) -> Optional[DisqualifiedAgent]: + row = await conn.fetchrow( + "SELECT * FROM disqualified_agents WHERE agent_id = $1", + agent_id, + ) + return DisqualifiedAgent(**row) if row is not None else None + + +@db_operation +async def disqualify_agent( + conn: DatabaseConnection, + agent_id: UUID, + reason: str, +) -> DisqualifiedAgent: + async with conn.conn.transaction(): + await lock_disqualified_agent_state(conn, agent_id) + row = await conn.fetchrow( + """ + INSERT INTO disqualified_agents (agent_id, reason) + VALUES ($1, $2) + ON CONFLICT (agent_id) DO UPDATE + SET reason = EXCLUDED.reason + RETURNING * + """, + agent_id, + reason, + ) + return DisqualifiedAgent(**row) diff --git a/queries/evaluation.py b/queries/evaluation.py index 4edac0f0..86316730 100644 --- a/queries/evaluation.py +++ b/queries/evaluation.py @@ -224,8 +224,8 @@ async def get_approved_validator_leader_score_for_set( ) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys banned_coldkey - WHERE banned_coldkey.miner_coldkey = agent.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = agent.agent_id ) """, set_id, @@ -329,8 +329,8 @@ async def get_approved_leader_ranking_for_set( ) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys banned_coldkey - WHERE banned_coldkey.miner_coldkey = agent.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = agent.agent_id ) ORDER BY ass.final_score DESC, rt.avg_cost_usd ASC NULLS LAST, ass.created_at ASC LIMIT 1 diff --git a/queries/evaluation_run.py b/queries/evaluation_run.py index 25452fbc..77a076b4 100644 --- a/queries/evaluation_run.py +++ b/queries/evaluation_run.py @@ -169,12 +169,12 @@ async def _get_evaluation_run_metrics_by_ids( JOIN evaluation_runs er2 ON er2.evaluation_id = e2.evaluation_id AND er2.problem_name = dp.problem_name JOIN agents a ON a.agent_id = e2.agent_id - LEFT JOIN banned_coldkeys bc ON bc.miner_coldkey = a.miner_coldkey + LEFT JOIN disqualified_agent_ids dq ON dq.agent_id = a.agent_id LEFT JOIN unapproved_agent_ids uai ON uai.agent_id = a.agent_id LEFT JOIN benchmark_agent_ids bai ON bai.agent_id = a.agent_id WHERE e2.set_id = dp.set_id - AND bc.miner_coldkey IS NULL + AND dq.agent_id IS NULL AND uai.agent_id IS NULL AND bai.agent_id IS NULL ) agg diff --git a/queries/evaluation_set.py b/queries/evaluation_set.py index d1177efa..d26ac98b 100644 --- a/queries/evaluation_set.py +++ b/queries/evaluation_set.py @@ -48,8 +48,8 @@ def _sql_agents_in_window_cte(select_columns: str) -> str: WHEN review.approval_review_status = 'rejected' OR EXISTS ( SELECT 1 - FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) THEN true ELSE false @@ -456,8 +456,8 @@ async def get_evaluation_set_score_stats(conn: DatabaseConnection, set_id: int) ) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys bc - WHERE bc.miner_coldkey = previous_agent.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = previous_agent.agent_id ) ) SELECT @@ -741,8 +741,8 @@ async def get_approved_agents_for_set(conn: DatabaseConnection, set_id: int) -> AND aa.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND ass.status = 'finished' AND review.approval_review_status is distinct from 'rejected' diff --git a/queries/problem_statistics.py b/queries/problem_statistics.py index 03193782..a3625111 100644 --- a/queries/problem_statistics.py +++ b/queries/problem_statistics.py @@ -175,8 +175,8 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ JOIN agents a on e.agent_id = a.agent_id WHERE e.set_id = $1 AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -213,8 +213,8 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ er.status = 'finished' AND e.set_id = $1 AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -275,8 +275,8 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ WHERE er.status = 'error' AND e.set_id = $1 AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -303,8 +303,8 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ JOIN agents a ON e.agent_id = a.agent_id WHERE e.set_id = $1 AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -345,8 +345,8 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ AND erh.solved AND e.set_id = $1 AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) diff --git a/queries/scores.py b/queries/scores.py index e26c5485..d222aa17 100644 --- a/queries/scores.py +++ b/queries/scores.py @@ -37,8 +37,8 @@ async def get_weight_receiving_agent_hotkey(conn: DatabaseConnection) -> Optiona AND ass.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) ORDER BY ass.final_score DESC, rt.avg_cost_usd ASC NULLS LAST, ass.created_at ASC LIMIT 1 @@ -84,8 +84,8 @@ async def get_weight_receiving_agent_info(conn: DatabaseConnection) -> Optional[ AND ass.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) ORDER BY ass.final_score DESC, rt.avg_cost_usd ASC NULLS LAST, ass.created_at ASC LIMIT 1 @@ -142,8 +142,8 @@ async def get_incentive_reward_candidates( ) AND NOT EXISTS ( SELECT 1 - FROM banned_coldkeys banned - WHERE banned.miner_coldkey = agent.miner_coldkey + FROM disqualified_agent_ids dq + WHERE dq.agent_id = agent.agent_id ) """, set_id, diff --git a/queries/statistics.py b/queries/statistics.py index e0d5f584..fb2c999f 100644 --- a/queries/statistics.py +++ b/queries/statistics.py @@ -34,8 +34,8 @@ async def get_top_scores_over_time(conn: DatabaseConnection) -> list[TopScoreOve agent_scores.final_score IS NOT NULL AND agent_scores.set_id = (SELECT set_id FROM max_set) AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -59,8 +59,8 @@ async def get_top_scores_over_time(conn: DatabaseConnection) -> list[TopScoreOve AND agent_scores.created_at <= ts.hour AND agent_scores.set_id = (SELECT set_id FROM max_set) AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -107,8 +107,8 @@ async def get_perfectly_solved_over_time(conn: DatabaseConnection) -> list[Perfe AND erh.benchmark_family <> '' AND erh.benchmark_family <> 'custom' AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND e.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND e.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -167,8 +167,8 @@ async def get_average_score_per_evaluation_set_group( WHERE eh.status = 'success' AND eh.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND eh.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND eh.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -206,8 +206,8 @@ async def get_average_wait_time_per_evaluation_set_group( AND e.evaluation_set_group = '{EvaluationSetGroup.screener_1.value}'::EvaluationSetGroup AND e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -228,8 +228,8 @@ async def get_average_wait_time_per_evaluation_set_group( AND sc1_e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND sc2_e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) @@ -259,8 +259,8 @@ async def get_average_wait_time_per_evaluation_set_group( AND sc2_e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND v_e.validator_count = {config.NUM_EVALS_PER_AGENT} AND NOT EXISTS ( - SELECT 1 FROM banned_coldkeys bc - WHERE bc.miner_coldkey = a.miner_coldkey + SELECT 1 FROM disqualified_agent_ids dq + WHERE dq.agent_id = a.agent_id ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) diff --git a/tests/api/test_admin.py b/tests/api/test_admin.py new file mode 100644 index 00000000..cfcdeeb3 --- /dev/null +++ b/tests/api/test_admin.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from uuid import UUID, uuid4 + +import pytest +from fastapi import HTTPException +from fastapi.security import HTTPAuthorizationCredentials + +import api.config as config +import utils.database as _db +from api.endpoints.admin import ( + ColdkeyBanRequest, # reused body model + put_disqualified_agent, + require_coldkey_ban_admin, +) + + +@pytest.fixture +async def clean_tables(postgres_db): + async with _db.pool.acquire() as conn: + await conn.execute("TRUNCATE disqualification_jobs, disqualified_agents, agents RESTART IDENTITY CASCADE") + yield + async with _db.pool.acquire() as conn: + await conn.execute("TRUNCATE disqualification_jobs, disqualified_agents, agents RESTART IDENTITY CASCADE") + + +async def _insert_agent() -> UUID: + agent_id = uuid4() + async with _db.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO agents ( + agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address + ) + VALUES ($1, $2, $3, 'test-agent', 0, 'evaluating', NOW(), '127.0.0.1') + """, + agent_id, + f"hotkey-{agent_id}", + f"coldkey-{agent_id}", + ) + return agent_id + + +def _admin_creds() -> HTTPAuthorizationCredentials: + return HTTPAuthorizationCredentials(scheme="Bearer", credentials=config.COLDKEY_BAN_ADMIN_API_KEY) + + +def test_auth_rejects_missing_credentials() -> None: + with pytest.raises(HTTPException) as exc: + require_coldkey_ban_admin(None) + assert exc.value.status_code == 401 + + +def test_auth_rejects_wrong_credentials() -> None: + creds = HTTPAuthorizationCredentials(scheme="Bearer", credentials="wrong-key") + with pytest.raises(HTTPException) as exc: + require_coldkey_ban_admin(creds) + assert exc.value.status_code == 401 + + +@pytest.mark.anyio +async def test_disqualify_existing_agent_succeeds(clean_tables) -> None: + agent_id = await _insert_agent() + + result = await put_disqualified_agent(agent_id, ColdkeyBanRequest(reason="cheating")) + + assert result.agent_id == agent_id + assert result.reason == "cheating" + + async with _db.pool.acquire() as conn: + count = await conn.fetchval("SELECT COUNT(*) FROM disqualified_agents WHERE agent_id = $1", agent_id) + assert count == 1 + + +@pytest.mark.anyio +async def test_disqualify_unknown_agent_returns_404(clean_tables) -> None: + with pytest.raises(HTTPException) as exc: + await put_disqualified_agent(uuid4(), ColdkeyBanRequest(reason="cheating")) + assert exc.value.status_code == 404 + + +@pytest.mark.anyio +async def test_disqualify_enqueues_reapproval_job(clean_tables) -> None: + agent_id = uuid4() + async with _db.pool.acquire() as conn: + # seed competition + agent with a set_id + await conn.execute("INSERT INTO competitions (set_id) VALUES (71) ON CONFLICT DO NOTHING") + await conn.execute( + """ + INSERT INTO agents (agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address, set_id) + VALUES ($1, $2, $3, 'test-agent', 0, 'evaluating', NOW(), '127.0.0.1', 71) + """, + agent_id, + f"hotkey-{agent_id}", + f"coldkey-{agent_id}", + ) + + result = await put_disqualified_agent(agent_id, ColdkeyBanRequest(reason="cheating")) + assert result.agent_id == agent_id + + async with _db.pool.acquire() as conn: + row = await conn.fetchrow("SELECT agent_id FROM disqualification_jobs WHERE agent_id = $1", agent_id) + assert row is not None diff --git a/tests/queries/test_disqualification_job.py b/tests/queries/test_disqualification_job.py new file mode 100644 index 00000000..460bd173 --- /dev/null +++ b/tests/queries/test_disqualification_job.py @@ -0,0 +1,86 @@ +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest + +import utils.database as _db +from models.disqualification_job import DisqualificationJob +from queries.disqualification_job import ( + claim_next_pending_disqualification_job, + count_pending_disqualification_jobs, + enqueue_disqualification_job, + mark_disqualification_job_processed, +) + +SET_ID = 71 + + +def test_disqualification_job_model_roundtrips(): + now = datetime.now(timezone.utc) + agent_id = uuid4() + job_id = uuid4() + job = DisqualificationJob( + id=job_id, + agent_id=agent_id, + set_id=71, + created_at=now, + processed_at=None, + attempts=0, + error=None, + ) + assert job.id == job_id + assert job.agent_id == agent_id + assert job.set_id == 71 + assert job.processed_at is None + assert job.attempts == 0 + assert job.error is None + + +@pytest.fixture +async def clean_jobs(postgres_db): + async with _db.pool.acquire() as conn: + await conn.execute("TRUNCATE disqualification_jobs, agents RESTART IDENTITY CASCADE") + yield + + +async def _insert_agent(conn, agent_id): + await conn.execute( + """ + INSERT INTO agents (agent_id, miner_hotkey, name, version_num, status, created_at, ip_address) + VALUES ($1, 'hk', 'hk', 1, 'finished', NOW(), '127.0.0.1') + """, + agent_id, + ) + + +@pytest.mark.anyio +async def test_enqueue_is_deduped_while_pending(clean_jobs): + agent_id = uuid4() + async with _db.pool.acquire() as conn: + await _insert_agent(conn, agent_id) + async with conn.transaction(): + first = await enqueue_disqualification_job(conn, agent_id=agent_id, set_id=SET_ID) + second = await enqueue_disqualification_job(conn, agent_id=agent_id, set_id=SET_ID) + assert first is not None + assert second is None + + +@pytest.mark.anyio +async def test_claim_marks_attempts_and_mark_processed(clean_jobs): + agent_id = uuid4() + async with _db.pool.acquire() as conn: + await _insert_agent(conn, agent_id) + async with conn.transaction(): + await enqueue_disqualification_job(conn, agent_id=agent_id, set_id=SET_ID) + + assert await count_pending_disqualification_jobs() == 1 + + async with _db.pool.acquire() as conn: + async with conn.transaction(): + job = await claim_next_pending_disqualification_job.__wrapped__(conn) + assert job is not None + assert job["agent_id"] == agent_id + assert job["set_id"] == SET_ID + await mark_disqualification_job_processed.__wrapped__(conn, job["id"]) + + assert await count_pending_disqualification_jobs() == 0 diff --git a/tests/queries/test_disqualification_reapproval.py b/tests/queries/test_disqualification_reapproval.py new file mode 100644 index 00000000..8d871515 --- /dev/null +++ b/tests/queries/test_disqualification_reapproval.py @@ -0,0 +1,859 @@ +from datetime import datetime, timedelta, timezone +from uuid import UUID, uuid4 + +import pytest + +import api.config as config +import queries.approval as approval_module +import utils.database as _db +from queries.approval import process_pending_disqualification_jobs, run_disqualification_reapproval +from queries.disqualification_job import count_pending_disqualification_jobs, enqueue_disqualification_job +from utils.incentives import calculate_time_multiplier + +SET_ID = 71 + + +@pytest.fixture(autouse=True) +async def clean_tables(postgres_db, monkeypatch): + monkeypatch.setattr(config, "INCENTIVE_START_SET_ID", SET_ID) + async with _db.pool.acquire() as conn: + await conn.execute( + "TRUNCATE disqualification_jobs, agent_approval_states, approval_jobs, approved_agents, " + "agent_scores, disqualified_agents, benchmark_agent_ids, agents RESTART IDENTITY CASCADE" + ) + yield + + +async def _insert_scored_agent( + conn, + *, + hotkey, + final_score, + created_at, + approved, + approved_at=None, + baseline_agent_id=None, + system_verdict, + projected_at=None, +) -> UUID: + """Insert an agent that has reached an incentive decision. + + `created_at` is the agent's upload time. `projected_at` is the incentive DECISION time + (defaults to created_at) — it drives replay chronology via a linked approval_jobs row. + Pass distinct created_at / projected_at to test the case where upload order != decision order. + """ + if projected_at is None: + projected_at = created_at + agent_id = uuid4() + await conn.execute( + """ + INSERT INTO agents (agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address) + VALUES ($1, $2, $2, $2, 1, 'finished', $3, '127.0.0.1') + """, + agent_id, + hotkey, + created_at, + ) + if approved: + # approved_agents must be inserted BEFORE agent_scores: agent_scores is a derived table + # (refresh_agent_scores_for_agent trigger rebuilds it from evaluations_hydrated whenever + # approved_agents/agents change), so inserting approved_agents afterward would wipe the + # manually-inserted agent_scores row. + await conn.execute( + """ + INSERT INTO approved_agents (agent_id, set_id, approved_at, baseline_agent_id, + relative_improvement_units, time_multiplier, initial_reward_score) + VALUES ($1, $2, $3, $4, 1, 1, 1) + """, + agent_id, + SET_ID, + approved_at, + baseline_agent_id, + ) + await conn.execute( + """ + INSERT INTO agent_scores (agent_id, miner_hotkey, name, version_num, created_at, status, + set_id, approved, approved_at, validator_count, final_score) + VALUES ($1, $2, $2, 1, $3, 'finished', $4, $5, $6, $7, $8) + """, + agent_id, + hotkey, + created_at, + SET_ID, + approved, + approved_at, + config.NUM_EVALS_PER_AGENT, + final_score, + ) + # A completed+projected approval job carries the decision timestamp the replay orders on. + job_id = uuid4() + await conn.execute( + """ + INSERT INTO approval_jobs (job_id, agent_id, set_id, status, policy_version, + input_snapshot, aggregate_verdict, projected_at) + VALUES ($1, $2, $3, 'completed', 'test', '{}'::jsonb, $4, $5) + """, + job_id, + agent_id, + SET_ID, + system_verdict, + projected_at, + ) + await conn.execute( + """ + INSERT INTO agent_approval_states (agent_id, set_id, latest_job_id, processing_status, + system_verdict, updated_at) + VALUES ($1, $2, $3, 'completed', $4, NOW()) + ON CONFLICT (agent_id, set_id) DO UPDATE + SET latest_job_id = EXCLUDED.latest_job_id, system_verdict = EXCLUDED.system_verdict + """, + agent_id, + SET_ID, + job_id, + system_verdict, + ) + return agent_id + + +async def _is_approved(conn, agent_id) -> bool: + row = await conn.fetchrow("SELECT 1 FROM approved_agents WHERE agent_id = $1 AND set_id = $2", agent_id, SET_ID) + return row is not None + + +async def _system_verdict(conn, agent_id) -> str: + return await conn.fetchval( + "SELECT system_verdict FROM agent_approval_states WHERE agent_id = $1 AND set_id = $2", + agent_id, + SET_ID, + ) + + +async def _disqualify(conn, agent_id) -> None: + await conn.execute("INSERT INTO disqualified_agents (agent_id, reason) VALUES ($1, 'test')", agent_id) + + +@pytest.mark.anyio +async def test_case1_promote_b1_keep_b2_demote_c(): + # Threshold check (INCENTIVE_PERFORMANCE_THRESHOLD = 3%): + # B1=0.55 vs seeded leader A=0.50 -> +10% -> qualifies (promoted) + # B2=0.55 vs new leader B1=0.55 -> 0% -> rejects (stays rejected) + # C =0.56 vs leader B1=0.55 -> +1.82% -> below 3% (demoted) + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a, + system_verdict="approved", + ) + b1 = await _insert_scored_agent( + conn, + hotkey="B1", + final_score=0.55, + created_at=base + timedelta(hours=2), + approved=False, + system_verdict="rejected", + ) + b2 = await _insert_scored_agent( + conn, + hotkey="B2", + final_score=0.55, + created_at=base + timedelta(hours=3), + approved=False, + system_verdict="rejected", + ) + c = await _insert_scored_agent( + conn, + hotkey="C", + final_score=0.56, + created_at=base + timedelta(hours=4), + approved=True, + approved_at=base + timedelta(hours=4), + baseline_agent_id=b, + system_verdict="approved", + ) + await _disqualify(conn, b) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, a) is True # untouched, before B + assert await _is_approved(conn, b) is False # removed + assert await _is_approved(conn, b1) is True # promoted + assert await _system_verdict(conn, b1) == "approved" + assert await _is_approved(conn, b2) is False # 0.55 vs 0.55 not an improvement + assert await _system_verdict(conn, b2) == "rejected" + assert await _is_approved(conn, c) is False # demoted: 0.56 vs 0.55 < 3% threshold + assert await _system_verdict(conn, c) == "rejected" + + +@pytest.mark.anyio +async def test_case2_promote_b1_keep_b2_no_c(): + # Same as Case 1 but without C downstream at all. + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a, + system_verdict="approved", + ) + b1 = await _insert_scored_agent( + conn, + hotkey="B1", + final_score=0.55, + created_at=base + timedelta(hours=2), + approved=False, + system_verdict="rejected", + ) + b2 = await _insert_scored_agent( + conn, + hotkey="B2", + final_score=0.55, + created_at=base + timedelta(hours=3), + approved=False, + system_verdict="rejected", + ) + await _disqualify(conn, b) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, a) is True + assert await _is_approved(conn, b) is False + assert await _is_approved(conn, b1) is True + assert await _system_verdict(conn, b1) == "approved" + assert await _is_approved(conn, b2) is False + assert await _system_verdict(conn, b2) == "rejected" + + +@pytest.mark.anyio +async def test_case3_c_demoted_against_b1_baseline(): + # Same shape as Case 1, but C's baseline_agent_id references B1 (the agent that will become + # the new leader after B's removal) rather than B, to exercise the "C's original baseline + # itself gets replaced downstream" path. Outcome (demote C) is identical: after replay the + # leader when we reach C is B1 = 0.55, and C = 0.56 is +1.82%, below the 3% threshold. + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a, + system_verdict="approved", + ) + b1 = await _insert_scored_agent( + conn, + hotkey="B1", + final_score=0.55, + created_at=base + timedelta(hours=2), + approved=False, + system_verdict="rejected", + ) + c = await _insert_scored_agent( + conn, + hotkey="C", + final_score=0.56, + created_at=base + timedelta(hours=3), + approved=True, + approved_at=base + timedelta(hours=3), + baseline_agent_id=b1, + system_verdict="approved", + ) + await _disqualify(conn, b) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, a) is True + assert await _is_approved(conn, b) is False + assert await _is_approved(conn, b1) is True # promoted: 0.55 vs seeded leader A=0.50 -> +10% + assert await _system_verdict(conn, b1) == "approved" + assert await _is_approved(conn, c) is False # demoted: 0.56 vs new leader B1=0.55 -> +1.82% < 3% + assert await _system_verdict(conn, c) == "rejected" + + +@pytest.mark.anyio +async def test_first_approved_disqualified_new_baseline(): + # B was itself the first-approved agent for the set (baseline_agent_id=None). Disqualifying + # it seeds current_leader=None, so the first downstream candidate that qualifies re-bootstraps + # the competition and must get relative_improvement_units == 1.0 (calculate_relative_improvement's + # "first approved agent" branch, triggered whenever leader_score is None). + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + baseline_agent_id=None, + system_verdict="approved", + ) + # D1's score must be >= B's (0.50) to survive the candidate query's floor filter. + d1 = await _insert_scored_agent( + conn, + hotkey="D1", + final_score=0.51, + created_at=base + timedelta(hours=1), + approved=False, + system_verdict="rejected", + ) + await _disqualify(conn, b) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, b) is False + assert await _is_approved(conn, d1) is True + assert await _system_verdict(conn, d1) == "approved" + units = await conn.fetchval( + "SELECT relative_improvement_units FROM approved_agents WHERE agent_id = $1 AND set_id = $2", + d1, + SET_ID, + ) + assert units == pytest.approx(1.0) + + +@pytest.mark.anyio +async def test_b_never_approved_is_noop(): + # B reached a "rejected" incentive decision (never occupied approved_agents), so nothing + # downstream was ever gated against it as leader. Disqualifying it must not touch anything. + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.40, + created_at=base + timedelta(hours=1), + approved=False, + system_verdict="rejected", + ) + d1 = await _insert_scored_agent( + conn, + hotkey="D1", + final_score=0.51, + created_at=base + timedelta(hours=2), + approved=False, + system_verdict="rejected", + ) + await _disqualify(conn, b) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, a) is True + assert await _system_verdict(conn, a) == "approved" + assert await _is_approved(conn, d1) is False + assert await _system_verdict(conn, d1) == "rejected" + + +@pytest.mark.anyio +async def test_orders_by_decision_time_not_upload_time(): + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + projected_at=base, + ) + # X: uploaded AFTER B, but its incentive decision happened BEFORE B's. + x = await _insert_scored_agent( + conn, + hotkey="X", + final_score=0.90, + created_at=base + timedelta(hours=5), # late upload + approved=False, + system_verdict="rejected", + projected_at=base + timedelta(hours=1), # early decision + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=2), + approved=True, + approved_at=base + timedelta(hours=2), + baseline_agent_id=a, + system_verdict="approved", + projected_at=base + timedelta(hours=2), + ) + await _disqualify(conn, b) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + # X decided before B → not in B's downstream set → untouched despite its high score/late upload. + assert await _is_approved(conn, x) is False + assert await _system_verdict(conn, x) == "rejected" + + +@pytest.mark.anyio +async def test_process_pending_runs_enqueued_job(): + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a, + system_verdict="approved", + ) + b1 = await _insert_scored_agent( + conn, + hotkey="B1", + final_score=0.60, + created_at=base + timedelta(hours=2), + approved=False, + system_verdict="rejected", + ) + await _disqualify(conn, b) + async with conn.transaction(): + await enqueue_disqualification_job(conn, agent_id=b, set_id=SET_ID) + + processed = await process_pending_disqualification_jobs() + assert processed == 1 + assert await count_pending_disqualification_jobs() == 0 + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, b) is False + assert await _is_approved(conn, b1) is True # 0.60 vs 0.50 = 20% > 3% + + +@pytest.mark.anyio +async def test_process_pending_stops_after_one_pass_on_permanent_failure(monkeypatch): + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + await _disqualify(conn, b) + async with conn.transaction(): + await enqueue_disqualification_job(conn, agent_id=b, set_id=SET_ID) + + async def _always_raise(*, set_id, disqualified_agent_id): + raise RuntimeError("boom") + + monkeypatch.setattr(approval_module, "run_disqualification_reapproval", _always_raise) + + processed = await process_pending_disqualification_jobs() + assert processed == 0 + + assert await count_pending_disqualification_jobs() == 1 + async with _db.pool.acquire() as conn: + row = await conn.fetchrow( + "SELECT attempts, error, processed_at FROM disqualification_jobs WHERE agent_id = $1", + b, + ) + assert row is not None + assert row["processed_at"] is None + assert row["error"] is not None + assert row["attempts"] <= 2 + + first_attempts = row["attempts"] + + # A second invocation (e.g. a later startup) must retry the still-pending failing job. + processed_again = await process_pending_disqualification_jobs() + assert processed_again == 0 + + async with _db.pool.acquire() as conn: + row_after = await conn.fetchrow( + "SELECT attempts FROM disqualification_jobs WHERE agent_id = $1", + b, + ) + assert row_after["attempts"] > first_attempts + + +@pytest.mark.anyio +async def test_process_pending_skips_failing_job_and_processes_healthy_one(monkeypatch): + # A (older, permanently failing) must not starve B (newer, healthy) out of the same drain + # invocation: the claim query must advance past A once it has been attempted. + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a_leader = await _insert_scored_agent( + conn, + hotkey="A-leader", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + failing_agent = await _insert_scored_agent( + conn, + hotkey="FAILING", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a_leader, + system_verdict="approved", + ) + healthy_agent = await _insert_scored_agent( + conn, + hotkey="HEALTHY", + final_score=0.70, + created_at=base + timedelta(hours=2), + approved=True, + approved_at=base + timedelta(hours=2), + baseline_agent_id=a_leader, + system_verdict="approved", + ) + promotable = await _insert_scored_agent( + conn, + hotkey="PROMOTABLE", + final_score=0.75, + created_at=base + timedelta(hours=3), + approved=False, + system_verdict="rejected", + ) + await _disqualify(conn, failing_agent) + await _disqualify(conn, healthy_agent) + async with conn.transaction(): + # Enqueue the failing job FIRST (older) so it sits at the head of the pending queue. + await enqueue_disqualification_job(conn, agent_id=failing_agent, set_id=SET_ID) + async with conn.transaction(): + await enqueue_disqualification_job(conn, agent_id=healthy_agent, set_id=SET_ID) + + real_run_disqualification_reapproval = run_disqualification_reapproval + + async def _fail_only_for_failing_agent(*, set_id, disqualified_agent_id): + if disqualified_agent_id == failing_agent: + raise RuntimeError("boom") + return await real_run_disqualification_reapproval(set_id=set_id, disqualified_agent_id=disqualified_agent_id) + + monkeypatch.setattr(approval_module, "run_disqualification_reapproval", _fail_only_for_failing_agent) + + processed = await process_pending_disqualification_jobs() + assert processed == 1 + + async with _db.pool.acquire() as conn: + failing_row = await conn.fetchrow( + "SELECT processed_at, error FROM disqualification_jobs WHERE agent_id = $1", + failing_agent, + ) + healthy_row = await conn.fetchrow( + "SELECT processed_at, error FROM disqualification_jobs WHERE agent_id = $1", + healthy_agent, + ) + assert await _is_approved(conn, promotable) is True # proves healthy job's replay actually ran + + assert failing_row["processed_at"] is None + assert failing_row["error"] is not None + + assert healthy_row is not None + assert healthy_row["processed_at"] is not None # B was NOT starved behind A + + +@pytest.mark.anyio +async def test_kept_leader_preserves_approved_at_for_time_multiplier(): + # Regression test for the bug where a kept already-approved leader lost its real + # `approved_at` (rebuilt via a fresh AgentRankingProfile defaulting to None), flooring + # `time_multiplier` to 1.0 for whatever gets promoted against it later. + # + # Chain: A (seed leader) -> B (disqualified) -> D (already-approved, still qualifies, + # kept as the new leader with its REAL approved_at far in the past) -> E (rejected, + # promoted against D). + # + # Threshold check (INCENTIVE_PERFORMANCE_THRESHOLD = 3%): + # D=0.55 vs seeded leader A=0.50 -> +10% -> qualifies (kept as leader) + # E=0.60 vs new leader D=0.55 -> +9.09% -> qualifies (promoted) + base = datetime.now(timezone.utc) - timedelta(days=2) + d_approved_at = datetime.now(timezone.utc) - timedelta(days=5) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a, + system_verdict="approved", + ) + d = await _insert_scored_agent( + conn, + hotkey="D", + final_score=0.55, + created_at=base + timedelta(hours=2), + approved=True, + approved_at=d_approved_at, + baseline_agent_id=b, + system_verdict="approved", + ) + e = await _insert_scored_agent( + conn, + hotkey="E", + final_score=0.60, + created_at=base + timedelta(hours=3), + approved=False, + system_verdict="rejected", + ) + await _disqualify(conn, b) + + before = datetime.now(timezone.utc) + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + after = datetime.now(timezone.utc) + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, a) is True + assert await _is_approved(conn, b) is False + assert await _is_approved(conn, d) is True # kept: still qualifies against A + assert await _system_verdict(conn, d) == "approved" + assert await _is_approved(conn, e) is True # promoted against D + assert await _system_verdict(conn, e) == "approved" + + e_row = await conn.fetchrow( + "SELECT time_multiplier, baseline_agent_id FROM approved_agents WHERE agent_id = $1 AND set_id = $2", + e, + SET_ID, + ) + assert e_row["baseline_agent_id"] == d + + # If the bug were present, current_leader for D would have approved_at=None, so + # E's elapsed_hours would floor to 0.0 and time_multiplier would be exactly 1.0. + floored_multiplier = calculate_time_multiplier( + elapsed_hours=0.0, + half_life_hours=config.INCENTIVE_TIME_MULTIPLIER_HALF_LIFE_HOURS, + maximum=config.INCENTIVE_TIME_MULTIPLIER_MAX, + ) + assert floored_multiplier == pytest.approx(1.0) + assert e_row["time_multiplier"] > floored_multiplier + 0.1 + + # Sanity bound: elapsed_hours must reflect D's real approved_at (~5 days), not 0. + min_elapsed_hours = (before - d_approved_at).total_seconds() / 3600 + max_elapsed_hours = (after - d_approved_at).total_seconds() / 3600 + expected_min = calculate_time_multiplier( + elapsed_hours=min_elapsed_hours, + half_life_hours=config.INCENTIVE_TIME_MULTIPLIER_HALF_LIFE_HOURS, + maximum=config.INCENTIVE_TIME_MULTIPLIER_MAX, + ) + expected_max = calculate_time_multiplier( + elapsed_hours=max_elapsed_hours, + half_life_hours=config.INCENTIVE_TIME_MULTIPLIER_HALF_LIFE_HOURS, + maximum=config.INCENTIVE_TIME_MULTIPLIER_MAX, + ) + assert expected_min - 1e-9 <= e_row["time_multiplier"] <= expected_max + 1e-9 + + +@pytest.mark.anyio +async def test_surviving_approved_agent_snapshot_unchanged(): + # Spec test: byte-for-byte frozen-snapshot proof that a surviving (still-qualifying, + # untouched) approved agent's approved_agents row is left completely alone by the replay, + # not merely that it is still present (_is_approved only checks existence). + base = datetime.now(timezone.utc) - timedelta(days=2) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a, + system_verdict="approved", + ) + b1 = await _insert_scored_agent( + conn, + hotkey="B1", + final_score=0.55, + created_at=base + timedelta(hours=2), + approved=False, + system_verdict="rejected", + ) + await _disqualify(conn, b) + + async with _db.pool.acquire() as conn: + # A is untouched by B's disqualification: it decided before B and never competes again. + before_row = dict( + await conn.fetchrow( + """ + SELECT approved_at, baseline_agent_id, performance_delta, cost_delta, + relative_improvement_units, time_multiplier, initial_reward_score + FROM approved_agents + WHERE agent_id = $1 AND set_id = $2 + """, + a, + SET_ID, + ) + ) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, a) is True + after_row = dict( + await conn.fetchrow( + """ + SELECT approved_at, baseline_agent_id, performance_delta, cost_delta, + relative_improvement_units, time_multiplier, initial_reward_score + FROM approved_agents + WHERE agent_id = $1 AND set_id = $2 + """, + a, + SET_ID, + ) + ) + + assert after_row == before_row + # b1 promoted as an independent sanity check that the replay actually ran. + async with _db.pool.acquire() as conn: + assert await _is_approved(conn, b1) is True + + +@pytest.mark.anyio +async def test_surviving_approved_agent_snapshot_updates_against_new_leader(): + """C stays approved after B is disqualified and B1 becomes the new leader, but C's + relative_improvement_units/time_multiplier/initial_reward_score/baseline_agent_id must be + recomputed against B1 — not left as originally computed against B. C's own approved_at must + NOT change. + """ + base = datetime.now(timezone.utc) - timedelta(days=10) + async with _db.pool.acquire() as conn: + a = await _insert_scored_agent( + conn, + hotkey="A", + final_score=0.50, + created_at=base, + approved=True, + approved_at=base, + system_verdict="approved", + ) + b = await _insert_scored_agent( + conn, + hotkey="B", + final_score=0.54, + created_at=base + timedelta(hours=1), + approved=True, + approved_at=base + timedelta(hours=1), + baseline_agent_id=a, + system_verdict="approved", + ) + b1 = await _insert_scored_agent( + conn, + hotkey="B1", + final_score=0.60, + created_at=base + timedelta(hours=2), + approved=False, + system_verdict="rejected", + ) + c_approved_at = base + timedelta(hours=3) + c = await _insert_scored_agent( + conn, + hotkey="C", + final_score=0.62, + created_at=base + timedelta(hours=3), + approved=True, + approved_at=c_approved_at, + baseline_agent_id=b, + system_verdict="approved", + ) + # capture C's snapshot BEFORE the reapproval, computed against B (per the + # fixture's default relative_improvement_units=1/time_multiplier=1/initial_reward_score=1) + before = await conn.fetchrow("SELECT * FROM approved_agents WHERE agent_id = $1 AND set_id = $2", c, SET_ID) + await _disqualify(conn, b) + + await run_disqualification_reapproval(set_id=SET_ID, disqualified_agent_id=b) + + async with _db.pool.acquire() as conn: + # B1 promoted against A (0.60 vs 0.50 = 20% > 3% threshold), becomes leader. + assert await _is_approved(conn, b1) is True + # C still qualifies against B1 (0.62 vs 0.60 = 3.33% > 3% threshold) -> stays approved. + assert await _is_approved(conn, c) is True + + after = await conn.fetchrow("SELECT * FROM approved_agents WHERE agent_id = $1 AND set_id = $2", c, SET_ID) + # Snapshot fields must have been recomputed against B1, not left as-computed against B. + assert after["baseline_agent_id"] == b1 + assert after["relative_improvement_units"] != before["relative_improvement_units"] + # C's OWN approved_at must be untouched. + assert after["approved_at"] == c_approved_at == before["approved_at"] diff --git a/tests/queries/test_disqualified_agent.py b/tests/queries/test_disqualified_agent.py new file mode 100644 index 00000000..2cc999c1 --- /dev/null +++ b/tests/queries/test_disqualified_agent.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +from uuid import UUID, uuid4 + +import pytest + +import utils.database as _db +from queries.agent import get_top_agents +from queries.disqualified_agent import disqualify_agent, get_disqualified_agent +from queries.evaluation_set import get_evaluation_set_leaderboard_agents +from queries.scores import get_incentive_reward_candidates + + +@pytest.fixture(autouse=True) +async def clean_tables(postgres_db): + async with _db.pool.acquire() as conn: + await conn.execute( + "TRUNCATE disqualified_agents, approved_agents, agent_scores, evaluation_sets, agents " + "RESTART IDENTITY CASCADE" + ) + yield + async with _db.pool.acquire() as conn: + await conn.execute( + "TRUNCATE disqualified_agents, approved_agents, agent_scores, evaluation_sets, agents " + "RESTART IDENTITY CASCADE" + ) + + +async def _insert_agent() -> UUID: + agent_id = uuid4() + async with _db.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO agents ( + agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address + ) + VALUES ($1, $2, $3, 'test-agent', 0, 'evaluating', NOW(), '127.0.0.1') + """, + agent_id, + f"hotkey-{agent_id}", + f"coldkey-{agent_id}", + ) + return agent_id + + +@pytest.mark.anyio +async def test_get_returns_none_when_not_disqualified() -> None: + agent_id = await _insert_agent() + assert await get_disqualified_agent(agent_id) is None + + +@pytest.mark.anyio +async def test_disqualify_inserts_and_get_roundtrips() -> None: + agent_id = await _insert_agent() + + result = await disqualify_agent(agent_id, "cheating") + + assert result.agent_id == agent_id + assert result.reason == "cheating" + assert result.disqualified_at is not None + + fetched = await get_disqualified_agent(agent_id) + assert fetched is not None + assert fetched.agent_id == agent_id + assert fetched.reason == "cheating" + + +@pytest.mark.anyio +async def test_disqualify_is_idempotent_and_updates_reason() -> None: + agent_id = await _insert_agent() + + await disqualify_agent(agent_id, "first reason") + second = await disqualify_agent(agent_id, "second reason") + + assert second.reason == "second reason" + + async with _db.pool.acquire() as conn: + count = await conn.fetchval("SELECT COUNT(*) FROM disqualified_agents WHERE agent_id = $1", agent_id) + assert count == 1 + + +async def _insert_scored_agent(*, final_score: float, coldkey: str) -> UUID: + agent_id = uuid4() + async with _db.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO agents ( + agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address + ) + VALUES ($1, $2, $3, 'scored-agent', 0, 'finished', NOW(), '127.0.0.1') + """, + agent_id, + f"hotkey-{agent_id}", + coldkey, + ) + max_set_id = await conn.fetchval("SELECT MAX(set_id) FROM evaluation_sets") + if max_set_id is None: + max_set_id = 1 + await conn.execute( + """ + INSERT INTO evaluation_sets (set_id, set_group, problem_name) + VALUES ($1, 'validator', 'p1') + ON CONFLICT (set_id, set_group, problem_name) DO NOTHING + """, + max_set_id, + ) + await conn.execute( + """ + INSERT INTO agent_scores ( + agent_id, miner_hotkey, name, version_num, created_at, status, + set_id, approved, approved_at, validator_count, final_score + ) + VALUES ($1, $2, 'scored-agent', 0, NOW(), 'finished', + $3, TRUE, NOW(), 3, $4) + """, + agent_id, + f"hotkey-{agent_id}", + max_set_id, + final_score, + ) + return agent_id + + +@pytest.mark.anyio +async def test_disqualified_agent_excluded_from_top_agents() -> None: + kept = await _insert_scored_agent(final_score=0.9, coldkey="ck-kept") + stopped = await _insert_scored_agent(final_score=0.8, coldkey="ck-stopped") + + before = {agent.agent_id for agent in await get_top_agents(number_of_agents=10)} + assert kept in before + assert stopped in before + + await disqualify_agent(stopped, "cheating") + + after = {agent.agent_id for agent in await get_top_agents(number_of_agents=10)} + assert kept in after + assert stopped not in after + + +@pytest.mark.anyio +async def test_stopped_agent_marked_disqualified_on_leaderboard() -> None: + # Reuse the scored-agent helper; the agent must fall in the latest set window. + stopped = await _insert_scored_agent(final_score=0.8, coldkey="ck-lb-stopped") + max_set_id = None + async with _db.pool.acquire() as conn: + max_set_id = await conn.fetchval("SELECT MAX(set_id) FROM evaluation_sets") + # Pin the agent to the set explicitly so it falls in the leaderboard's window + # regardless of created_at ordering versus the set/competition row. + await conn.execute("UPDATE agents SET set_id = $1 WHERE agent_id = $2", max_set_id, stopped) + + await disqualify_agent(stopped, "cheating") + + rows = await get_evaluation_set_leaderboard_agents(max_set_id) + match = [r for r in rows if r["agent_id"] == stopped] + assert match, "stopped agent should still appear on the leaderboard" + assert match[0]["disqualified"] is True + + +async def _insert_approved_incentive_agent(*, coldkey: str, initial_reward_score: float, set_id: int) -> UUID: + """Insert an agent that qualifies as an incentive reward candidate for set_id.""" + agent_id = uuid4() + async with _db.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO agents ( + agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address + ) + VALUES ($1, $2, $3, 'inc-agent', 0, 'finished', NOW(), '127.0.0.1') + """, + agent_id, + f"hotkey-{agent_id}", + coldkey, + ) + # `agent_scores` is a derived table, rebuilt by triggers on `agents`/`approved_agents`/ + # `evaluations`/`banned_hotkeys`/`unapproved_agent_ids` writes (see refresh_agent_scores() + # in the initial schema migration). Insert `approved_agents` BEFORE the manual `agent_scores` + # row so no later trigger wipes it out; nothing must write to `agents`/`approved_agents`/etc. + # for this agent_id after this point. + await conn.execute( + """ + INSERT INTO approved_agents ( + agent_id, set_id, approved_at, + relative_improvement_units, time_multiplier, initial_reward_score + ) + VALUES ($1, $2, NOW(), 1.0, 1.0, $3) + """, + agent_id, + set_id, + initial_reward_score, + ) + await conn.execute( + """ + INSERT INTO agent_scores ( + agent_id, miner_hotkey, name, version_num, created_at, status, + set_id, approved, approved_at, validator_count, final_score + ) + VALUES ($1, $2, 'inc-agent', 0, NOW(), 'finished', $3, TRUE, NOW(), 3, 0.9) + """, + agent_id, + f"hotkey-{agent_id}", + set_id, + ) + return agent_id + + +@pytest.mark.anyio +async def test_disqualified_agent_excluded_from_incentive_reward_candidates() -> None: + async with _db.pool.acquire() as conn: + set_id = await conn.fetchval("SELECT MAX(set_id) FROM evaluation_sets") + if set_id is None: + set_id = 1 + await conn.execute( + """ + INSERT INTO evaluation_sets (set_id, set_group, problem_name) + VALUES ($1, 'validator', 'p1') + ON CONFLICT (set_id, set_group, problem_name) DO NOTHING + """, + set_id, + ) + + kept = await _insert_approved_incentive_agent(coldkey="ck-inc-kept", initial_reward_score=1.0, set_id=set_id) + stopped = await _insert_approved_incentive_agent(coldkey="ck-inc-stopped", initial_reward_score=2.0, set_id=set_id) + + await disqualify_agent(stopped, "cheating") + + candidates, _observed_at = await get_incentive_reward_candidates(set_id, required_validator_count=3) + candidate_ids = {c.agent_id for c in candidates} + + assert kept in candidate_ids + assert stopped not in candidate_ids