Feat/add ck ban mechanism

alembic/versions/2026_07_10_add_coldkey_bans.py

@@ -0,0 +1,267 @@
+"""Add coldkey bans and agent coldkey provenance.
+
+Revision ID: b7e4d2c9a106
+Revises: c8f41e9a2b37
+Create Date: 2026-07-10 00:00:00.000000
+
+"""
+
+import re
+from typing import Match, Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision: str = "b7e4d2c9a106"
+down_revision: Union[str, Sequence[str], None] = "c8f41e9a2b37"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+_HOTKEY_BAN_CONDITION = "agents.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys)"
+_COLDKEY_BAN_CONDITION = """NOT EXISTS (
+            SELECT 1
+            FROM banned_coldkeys
+            WHERE banned_coldkeys.miner_coldkey = agents.miner_coldkey
+          )"""
+
+_REFRESH_AGENT_HOTKEY_FILTER = "          AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys)\n"
+_REFRESH_AGENT_UNAPPROVED_FILTER = "          AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)\n"
+_POPULATE_HOTKEY_FILTER = (
+    "        WHERE miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys)\n"
+    "          AND agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)"
+)
+_POPULATE_UNAPPROVED_FILTER = "        WHERE agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)"
+
+
+def _replace_queue_views(ban_condition: str) -> None:
+    op.execute(f"""
+        CREATE OR REPLACE VIEW pre_screening_queue AS
+        SELECT agents.agent_id, agents.status
+        FROM agents
+        WHERE agents.status IN ('pre_screening', 'pre_screening_needs_review')
+          AND agents.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids)
+          AND {ban_condition}
+          AND agents.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)
+        ORDER BY agents.created_at ASC;
+    """)
+
+    for queue_name, status, set_group in (
+        ("screener_1_queue", "screening_1", "screener_1"),
+        ("screener_2_queue", "screening_2", "screener_2"),
+    ):
+        op.execute(f"""
+            CREATE OR REPLACE VIEW {queue_name} AS
+            SELECT agents.agent_id, agents.status
+            FROM agents
+            WHERE agents.status = '{status}'
+              AND NOT EXISTS (
+                SELECT 1 FROM evaluations e
+                WHERE e.agent_id = agents.agent_id
+                  AND e.evaluation_set_group = '{set_group}'::evaluationsetgroup
+                  AND (
+                    SELECT (CASE
+                        WHEN COUNT(*) = 0 THEN NULL
+                        WHEN EVERY(erh.status = 'finished' OR (erh.status = 'error' AND erh.error_code BETWEEN 1000 AND 1999)) THEN 'success'
+                        WHEN EVERY(erh.status IN ('finished', 'error')) THEN 'failure'
+                        ELSE 'running'
+                    END)::evaluationstatus
+                    FROM evaluation_runs_hydrated erh
+                    WHERE erh.evaluation_id = e.evaluation_id
+                  ) IN ('success', 'running')
+              )
+              AND agents.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids)
+              AND {ban_condition}
+              AND agents.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)
+            ORDER BY agents.created_at ASC;
+        """)
+
+    op.execute(f"""
+        CREATE OR REPLACE VIEW validator_queue AS
+        WITH
+            validator_eval_counts AS (
+                SELECT
+                    agent_id,
+                    COUNT(*) FILTER (WHERE status = 'running') AS num_running_evals,
+                    COUNT(*) FILTER (WHERE status = 'success') AS num_finished_evals
+                FROM evaluations_hydrated
+                WHERE evaluations_hydrated.status IN ('success', 'running')
+                  AND evaluations_hydrated.evaluation_set_group = 'validator'::evaluationsetgroup
+                GROUP BY agent_id
+            ),
+            screener_2_scores AS (
+                SELECT agent_id, MAX(score) AS score FROM evaluations_hydrated
+                WHERE evaluations_hydrated.evaluation_set_group = 'screener_2'::evaluationsetgroup
+                  AND evaluations_hydrated.status = 'success'
+                GROUP BY agent_id
+            )
+        SELECT
+            agent_id,
+            status,
+            COALESCE(num_running_evals, 0) as num_running_evals,
+            COALESCE(num_finished_evals, 0) as num_finished_evals
+        FROM agents
+             INNER JOIN screener_2_scores USING (agent_id)
+             LEFT JOIN validator_eval_counts USING (agent_id)
+        WHERE
+            agents.status = 'evaluating'
+            AND COALESCE(num_running_evals, 0) + COALESCE(num_finished_evals, 0) < 3
+            AND agents.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids)
+            AND {ban_condition}
+            AND agents.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)
+        ORDER BY
+            screener_2_scores.score DESC,
+            agents.created_at ASC,
+            num_finished_evals DESC;
+    """)
+
+
+def _replace_function_fragment(signature: str, old: str, new: str) -> None:
+    definition = (
+        op.get_bind()
+        .execute(
+            sa.text("SELECT pg_get_functiondef(to_regprocedure(:signature))"),
+            {"signature": signature},
+        )
+        .scalar_one()
+    )
+    if definition is None or definition.count(old) != 1:
+        raise RuntimeError(f"Unexpected definition for {signature}")
+    op.execute(definition.replace(old, new))
+
+
+def _replace_function_pattern(signature: str, pattern: str, replacement: str) -> None:
+    definition = (
+        op.get_bind()
+        .execute(
+            sa.text("SELECT pg_get_functiondef(to_regprocedure(:signature))"),
+            {"signature": signature},
+        )
+        .scalar_one()
+    )
+    if definition is None:
+        raise RuntimeError(f"Missing function {signature}")
+    updated, replacement_count = re.subn(pattern, replacement, definition, flags=re.MULTILINE)
+    if replacement_count != 1:
+        raise RuntimeError(f"Unexpected definition for {signature}")
+    op.execute(updated)
+
+
+def _hotkey_trigger_branch_pattern(row: str) -> str:
+    return (
+        rf"^(?P<indent>[ \t]*)ELSIF TG_TABLE_NAME = 'banned_hotkeys' THEN\n"
+        rf"(?P=indent)    DELETE FROM agent_scores\n"
+        rf"(?P=indent)    WHERE agent_id IN "
+        rf"\(SELECT agent_id FROM agents WHERE miner_hotkey = {row}\.miner_hotkey\);\n"
+        rf"(?P=indent)    RETURN {row};\n"
+    )
+
+
+def _restore_hotkey_trigger_branch(row: str) -> None:
+    signature = "refresh_agent_scores()"
+    definition = (
+        op.get_bind()
+        .execute(
+            sa.text("SELECT pg_get_functiondef(to_regprocedure(:signature))"),
+            {"signature": signature},
+        )
+        .scalar_one()
+    )
+    if definition is None:
+        raise RuntimeError(f"Missing function {signature}")
+
+    pattern = (
+        rf"^(?P<indent>[ \t]*)ELSIF TG_TABLE_NAME = 'approved_agents' THEN\n"
+        rf"(?P=indent)    affected_agent_id := {row}\.agent_id;\n"
+        rf"(?P=indent)ELSIF TG_TABLE_NAME = 'unapproved_agent_ids' THEN\n"
+    )
+
+    def add_branch(match: Match[str]) -> str:
+        indent = match.group("indent")
+        return (
+            f"{indent}ELSIF TG_TABLE_NAME = 'approved_agents' THEN\n"
+            f"{indent}    affected_agent_id := {row}.agent_id;\n"
+            f"{indent}ELSIF TG_TABLE_NAME = 'banned_hotkeys' THEN\n"
+            f"{indent}    DELETE FROM agent_scores\n"
+            f"{indent}    WHERE agent_id IN "
+            f"(SELECT agent_id FROM agents WHERE miner_hotkey = {row}.miner_hotkey);\n"
+            f"{indent}    RETURN {row};\n"
+            f"{indent}ELSIF TG_TABLE_NAME = 'unapproved_agent_ids' THEN\n"
+        )
+
+    updated, replacement_count = re.subn(pattern, add_branch, definition, flags=re.MULTILINE)
+    if replacement_count != 1:
+        raise RuntimeError(f"Unexpected definition for {signature}")
+    op.execute(updated)
+
+
+def _remove_hotkey_score_behavior() -> None:
+    _replace_function_fragment(
+        "refresh_agent_scores_for_agent(uuid)",
+        _REFRESH_AGENT_HOTKEY_FILTER,
+        "",
+    )
+    _replace_function_fragment(
+        "populate_agent_scores()",
+        _POPULATE_HOTKEY_FILTER,
+        _POPULATE_UNAPPROVED_FILTER,
+    )
+    _replace_function_pattern(
+        "refresh_agent_scores()",
+        _hotkey_trigger_branch_pattern("OLD"),
+        "",
+    )
+    _replace_function_pattern(
+        "refresh_agent_scores()",
+        _hotkey_trigger_branch_pattern("NEW"),
+        "",
+    )
+
+
+def _restore_hotkey_score_behavior() -> None:
+    _replace_function_fragment(
+        "refresh_agent_scores_for_agent(uuid)",
+        _REFRESH_AGENT_UNAPPROVED_FILTER,
+        _REFRESH_AGENT_HOTKEY_FILTER + _REFRESH_AGENT_UNAPPROVED_FILTER,
+    )
+    _replace_function_fragment(
+        "populate_agent_scores()",
+        _POPULATE_UNAPPROVED_FILTER,
+        _POPULATE_HOTKEY_FILTER,
+    )
+    _restore_hotkey_trigger_branch("OLD")
+    _restore_hotkey_trigger_branch("NEW")
+
+
+def upgrade() -> None:
+    op.create_table(
+        "banned_coldkeys",
+        sa.Column("miner_coldkey", sa.Text(), primary_key=True),
+        sa.Column("banned_reason", sa.Text(), nullable=False),
+        sa.Column(
+            "banned_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("NOW()"),
+            nullable=False,
+        ),
+    )
+
+    op.add_column("agents", sa.Column("miner_coldkey", sa.Text(), nullable=True))
+
+    _replace_queue_views(_COLDKEY_BAN_CONDITION)
+    _remove_hotkey_score_behavior()
+    op.execute("DROP TRIGGER IF EXISTS tr_refresh_agent_scores_banned_hotkeys ON banned_hotkeys;")
+
+
+def downgrade() -> None:
+    _restore_hotkey_score_behavior()
+    op.execute("""
+        CREATE TRIGGER tr_refresh_agent_scores_banned_hotkeys
+        AFTER INSERT OR UPDATE OR DELETE ON banned_hotkeys
+        FOR EACH ROW EXECUTE PROCEDURE refresh_agent_scores();
+    """)
+    _replace_queue_views(_HOTKEY_BAN_CONDITION)
+
+    op.drop_column("agents", "miner_coldkey")
+    op.drop_table("banned_coldkeys")

api/.env.example

@@ -8,6 +8,7 @@ SUBTENSOR_ADDRESS=wss://entrypoint-finney.opentensor.ai
 SUBTENSOR_NETWORK=finney
 
 OWNER_HOTKEY=
+COLDKEY_BAN_ADMIN_API_KEY=
 UPLOAD_SEND_ADDRESS=
 BURN=False
 DISALLOW_UPLOADS=False

api/config.py

@@ -40,6 +40,10 @@
 if not OWNER_HOTKEY:
     logger.fatal("OWNER_HOTKEY is not set in .env")
 
+COLDKEY_BAN_ADMIN_API_KEY = os.getenv("COLDKEY_BAN_ADMIN_API_KEY")
+if not COLDKEY_BAN_ADMIN_API_KEY:
+    logger.warning("COLDKEY_BAN_ADMIN_API_KEY is not set; coldkey ban administration will be unavailable")
+
 UPLOAD_SEND_ADDRESS = os.getenv("UPLOAD_SEND_ADDRESS")
 if not UPLOAD_SEND_ADDRESS:
     logger.fatal("UPLOAD_SEND_ADDRESS is not set in .env")

api/endpoints/admin.py

@@ -0,0 +1,64 @@
+import secrets
+from typing import Annotated
+
+from bittensor_wallet.keypair import Keypair
+from fastapi import APIRouter, Depends, HTTPException, Response, status
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from pydantic import BaseModel, StringConstraints
+
+import api.config as config
+from models.banned_coldkey import BannedColdkey
+from queries.banned_coldkey import ban_coldkey, unban_coldkey
+from utils.ttl import clear_all_ttl_caches
+
+router = APIRouter(tags=["admin"])
+admin_bearer = HTTPBearer(auto_error=False)
+
+
+class ColdkeyBanRequest(BaseModel):
+    reason: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=1000)]
+
+
+def require_coldkey_ban_admin(
+    credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(admin_bearer)],
+) -> None:
+    expected = config.COLDKEY_BAN_ADMIN_API_KEY
+    if not expected:
+        raise HTTPException(status_code=503, detail="Coldkey ban administration is not configured")
+    if credentials is None or not secrets.compare_digest(credentials.credentials, expected):
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid admin credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+
+def validate_coldkey(miner_coldkey: str) -> None:
+    try:
+        Keypair(ss58_address=miner_coldkey)
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid coldkey SS58 address") from None
+
+
+@router.put(
+    "/banned-coldkeys/{miner_coldkey}",
+    response_model=BannedColdkey,
+    dependencies=[Depends(require_coldkey_ban_admin)],
+)
+async def put_banned_coldkey(miner_coldkey: str, request: ColdkeyBanRequest) -> BannedColdkey:
+    validate_coldkey(miner_coldkey)
+    banned_coldkey = await ban_coldkey(miner_coldkey, request.reason)
+    clear_all_ttl_caches()
+    return banned_coldkey
+
+
+@router.delete(
+    "/banned-coldkeys/{miner_coldkey}",
+    status_code=status.HTTP_204_NO_CONTENT,
+    dependencies=[Depends(require_coldkey_ban_admin)],
+)
+async def delete_banned_coldkey(miner_coldkey: str) -> Response:
+    validate_coldkey(miner_coldkey)
+    await unban_coldkey(miner_coldkey)
+    clear_all_ttl_caches()
+    return Response(status_code=status.HTTP_204_NO_CONTENT)

api/endpoints/validator.py

@@ -54,6 +54,7 @@
     update_agent_status,
 )
 from queries.approval import finish_agent_and_enqueue_approval
+from queries.banned_coldkey import is_agent_coldkey_banned
 from queries.evaluation import (
     create_new_evaluation_and_evaluation_runs,
     get_approved_leader_ranking_for_set,
@@ -1193,6 +1194,10 @@ async def handle_evaluation_if_finished(evaluation_id: UUID) -> None:
 
 
 async def _should_run_auto_approval_judge(*, agent_id: UUID, set_id: int) -> bool:
+    if await is_agent_coldkey_banned(agent_id):
+        logger.info(f"Skipping auto approval for agent_id={agent_id}: miner coldkey is banned")
+        return False
+
     candidate = await get_validator_agent_score_for_set(
         agent_id,
         set_id,