feat: add probation and ban enforcement

This commit is contained in:
Dobromir Popov
2026-06-29 09:58:32 +03:00
parent 39f6f23c83
commit 792a9fd97f
8 changed files with 266 additions and 20 deletions

View File

@@ -19,6 +19,7 @@ class RegistryWallet:
stake_balance: int = 0
strike_count: int = 0
banned: bool = False
completed_job_count: int = 0
@dataclass(frozen=True)
@@ -96,6 +97,7 @@ class _LocalContractState:
validation_events: list[ValidationEvent] = field(default_factory=list)
slash_receipts: list[SlashReceipt] = field(default_factory=list)
token_balances: dict[str, int] = field(default_factory=dict)
probationary_job_count: int = 50
class RegistryContract:
@@ -113,6 +115,7 @@ class RegistryContract:
stake_balance=current.stake_balance + amount,
strike_count=current.strike_count,
banned=current.banned,
completed_job_count=current.completed_job_count,
)
return {"signature": f"local-stake-{wallet_address}", "cluster": self._cluster}
@@ -126,6 +129,7 @@ class RegistryContract:
stake_balance=current.stake_balance,
strike_count=strike_count,
banned=current.banned or strike_count >= ban_threshold,
completed_job_count=current.completed_job_count,
)
return {"signature": f"local-strike-{wallet_address}", "cluster": self._cluster}
@@ -153,6 +157,7 @@ class RegistryContract:
stake_balance=stake_after,
strike_count=strike_count,
banned=banned,
completed_job_count=current.completed_job_count,
)
receipt = SlashReceipt(
signature=f"local-slash-{wallet_address}-{len(self._state.slash_receipts) + 1}",
@@ -175,6 +180,7 @@ class RegistryContract:
stake_balance=current.stake_balance,
strike_count=current.strike_count,
banned=True,
completed_job_count=current.completed_job_count,
)
return {"signature": f"local-ban-{wallet_address}", "cluster": self._cluster}
@@ -184,6 +190,21 @@ class RegistryContract:
wallet = self.get_wallet(wallet_address)
return wallet.stake_balance >= minimum_stake
def record_completed_job(self, wallet_address: str) -> RegistryWallet:
current = self.get_wallet(wallet_address)
updated = RegistryWallet(
stake_balance=current.stake_balance,
strike_count=current.strike_count,
banned=current.banned,
completed_job_count=current.completed_job_count + 1,
)
self._state.registry[wallet_address] = updated
return updated
def probationary_jobs_remaining(self, wallet_address: str) -> int:
wallet = self.get_wallet(wallet_address)
return max(0, self._state.probationary_job_count - wallet.completed_job_count)
class ValidationContract:
"""Validation event log consumed by the optimistic fraud detector."""
@@ -325,14 +346,23 @@ class SettlementContract:
unsettled = [a for a in self._state.attributions if a.settled_epoch is None]
client_debits: dict[str, int] = {}
work_by_node: dict[str, int] = {}
eligible_work_by_node: dict[str, int] = {}
node_rewards: dict[str, int] = {}
eligible_raw_work = 0
completed_jobs_to_record: dict[str, int] = {}
for attribution in unsettled:
debit = attribution.work_units * self._cost_per_layer_token_lamport
weighted_work = int(attribution.work_units * attribution.speed_score)
client_debits[attribution.api_key] = client_debits.get(attribution.api_key, 0) + debit
work_by_node[attribution.node_wallet] = (
work_by_node.get(attribution.node_wallet, 0) + weighted_work
)
wallet = self._state.registry.get(attribution.node_wallet, RegistryWallet())
completed_before_job = wallet.completed_job_count + completed_jobs_to_record.get(attribution.node_wallet, 0)
if completed_before_job >= self._state.probationary_job_count:
eligible_raw_work += attribution.work_units
weighted_work = int(attribution.work_units * attribution.speed_score)
eligible_work_by_node[attribution.node_wallet] = (
eligible_work_by_node.get(attribution.node_wallet, 0) + weighted_work
)
node_rewards.setdefault(attribution.node_wallet, 0)
completed_jobs_to_record[attribution.node_wallet] = completed_jobs_to_record.get(attribution.node_wallet, 0) + 1
for api_key, debit in client_debits.items():
current = self._state.api_keys.get(api_key, ApiKeyBalance())
@@ -343,16 +373,18 @@ class SettlementContract:
usdc_micro=current.usdc_micro,
)
total_weighted_work = sum(work_by_node.values())
for wallet_address, count in completed_jobs_to_record.items():
self._record_completed_jobs(wallet_address, count)
total_weighted_work = sum(eligible_work_by_node.values())
raw_total_work = sum(a.work_units for a in unsettled)
validator_rewards: dict[str, int] = {}
validator_reward_total = (raw_total_work * validator_reward_share_bps) // 10_000
if validator_wallet and validator_reward_total:
validator_rewards[validator_wallet] = validator_reward_total
node_reward_pool = raw_total_work - validator_reward_total
node_rewards: dict[str, int] = {}
for node_wallet, weighted_work_units in work_by_node.items():
node_reward_pool = max(0, eligible_raw_work - validator_reward_total)
for node_wallet, weighted_work_units in eligible_work_by_node.items():
reward = 0 if total_weighted_work == 0 else (node_reward_pool * weighted_work_units) // total_weighted_work
node_rewards[node_wallet] = reward
self._state.token_balances[node_wallet] = (
@@ -386,6 +418,15 @@ class SettlementContract:
def get_token_balance(self, wallet_address: str) -> int:
return self._state.token_balances.get(wallet_address, 0)
def _record_completed_jobs(self, wallet_address: str, count: int) -> None:
current = self._state.registry.get(wallet_address, RegistryWallet())
self._state.registry[wallet_address] = RegistryWallet(
stake_balance=current.stake_balance,
strike_count=current.strike_count,
banned=current.banned,
completed_job_count=current.completed_job_count + count,
)
def deployment_plan(self) -> dict:
"""Return the configured manual testnet deployment targets."""
return {
@@ -403,13 +444,17 @@ class LocalSolanaContracts:
cluster: str = "local-test-validator",
cost_per_layer_token_lamport: int = 1,
starting_credit_lamports: int = 1_000,
probationary_job_count: int = 50,
) -> None:
if cost_per_layer_token_lamport <= 0:
raise ValueError("cost_per_layer_token_lamport must be positive")
if starting_credit_lamports < 0:
raise ValueError("starting_credit_lamports must be non-negative")
if probationary_job_count < 0:
raise ValueError("probationary_job_count must be non-negative")
self.cluster = cluster
self._state = _LocalContractState()
self._state.probationary_job_count = probationary_job_count
self.registry = RegistryContract(self._state, cluster)
self.validation = ValidationContract(self._state)
self.payment = PaymentContract(self._state, cluster, starting_credit_lamports)

View File

@@ -117,6 +117,14 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
route_nodes = tracker_resp.get("nodes", [])
if route_nodes and not _valid_route_nodes(route_nodes):
raise ValueError("invalid tracker route metadata")
banned_wallet = _banned_route_wallet(server.contracts, route_nodes)
if banned_wallet is not None:
self._send_openai_error(
503,
f"model not available: route includes banned wallet {banned_wallet}",
"model_not_available",
)
return
except urllib.error.HTTPError as exc:
error_body = _safe_error_body(exc)
self._send_openai_error(
@@ -274,6 +282,18 @@ def _valid_route_nodes(route_nodes: object) -> bool:
return True
def _banned_route_wallet(contracts: Any | None, route_nodes: object) -> str | None:
if contracts is None or not isinstance(route_nodes, list):
return None
for node in route_nodes:
if not isinstance(node, dict):
continue
wallet_address = node.get("wallet_address")
if isinstance(wallet_address, str) and contracts.registry.get_wallet(wallet_address).banned:
return wallet_address
return None
def _api_key_from_authorization(header: str | None) -> str | None:
if not header:
return None

View File

@@ -9,6 +9,7 @@ import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any
from .downloader import download_shard
from .hardware import detect_hardware
@@ -38,6 +39,7 @@ def run_startup(
cache_dir: Path | None = None,
host: str = "127.0.0.1",
advertise_host: str | None = None,
contracts: Any | None = None,
) -> StubNodeServer:
"""Execute the full startup sequence and return a running node server.
@@ -73,6 +75,9 @@ def run_startup(
wallet_kwargs["path"] = wallet_path
_, _, address = load_or_create_wallet(**wallet_kwargs)
print(f" Wallet: {address}", flush=True)
probationary_line = _probationary_status_line(contracts, address)
if probationary_line is not None:
print(f" {probationary_line}", flush=True)
# 3. Shard assignment from tracker
print("Querying tracker for shard assignment...", flush=True)
@@ -152,3 +157,13 @@ def run_startup(
)
return node
def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None:
if contracts is None:
return None
remaining = contracts.registry.probationary_jobs_remaining(wallet_address)
if remaining <= 0:
return "Probationary period complete: earning enabled"
suffix = "job" if remaining == 1 else "jobs"
return f"Probationary period: {remaining} {suffix} remaining before earning"

View File

@@ -83,6 +83,14 @@ def _select_route(
return route, ""
def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None:
if contracts is None or not wallet_address:
return None
if contracts.registry.get_wallet(wallet_address).banned:
return "wallet is banned"
return None
class _TrackerHTTPServer(http.server.HTTPServer):
def __init__(
self,
@@ -205,6 +213,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if wallet_address is not None and not isinstance(wallet_address, str):
self._send_json(400, {"error": "wallet_address must be a string"})
return
ban_error = _registration_ban_error(server.contracts, wallet_address)
if ban_error:
self._send_json(403, {"error": ban_error})
return
node_id = str(uuid.uuid4())
entry = _NodeEntry(