318 lines
11 KiB
Python
318 lines
11 KiB
Python
"""Solana contract boundary for the Distributed Inference Network.
|
|
|
|
The prototype uses deterministic local state with Solana-shaped wrapper classes.
|
|
ADR-0007 keeps the public Python boundary stable so real Solana programs can
|
|
replace this adapter later without changing gateway or tracker call sites.
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
__version__ = "0.1.0"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RegistryWallet:
|
|
"""Stake, strike, and ban state for a node operator wallet."""
|
|
|
|
stake_balance: int = 0
|
|
strike_count: int = 0
|
|
banned: bool = False
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ApiKeyBalance:
|
|
"""Client API key payment account balance."""
|
|
|
|
lamports: int = 0
|
|
usdc_micro: int = 0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ComputeAttribution:
|
|
"""On-chain work attribution recorded by the gateway after inference."""
|
|
|
|
session_id: str
|
|
api_key: str
|
|
node_wallet: str
|
|
layer_start: int
|
|
layer_end: int
|
|
tokens: int
|
|
speed_score: float = 1.0
|
|
settled_epoch: int | None = None
|
|
|
|
@property
|
|
def layer_count(self) -> int:
|
|
return self.layer_end - self.layer_start + 1
|
|
|
|
@property
|
|
def work_units(self) -> int:
|
|
return self.layer_count * self.tokens
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SettlementResult:
|
|
"""Summary returned by an epoch settlement transaction."""
|
|
|
|
epoch: int
|
|
client_debits: dict[str, int]
|
|
node_rewards: dict[str, int]
|
|
validator_rewards: dict[str, int]
|
|
|
|
|
|
@dataclass
|
|
class _LocalContractState:
|
|
registry: dict[str, RegistryWallet] = field(default_factory=dict)
|
|
api_keys: dict[str, ApiKeyBalance] = field(default_factory=dict)
|
|
attributions: list[ComputeAttribution] = field(default_factory=list)
|
|
token_balances: dict[str, int] = field(default_factory=dict)
|
|
|
|
|
|
class RegistryContract:
|
|
"""Registry wrapper for node stake, strikes, and bans."""
|
|
|
|
def __init__(self, state: _LocalContractState, cluster: str) -> None:
|
|
self._state = state
|
|
self._cluster = cluster
|
|
|
|
def submit_stake(self, wallet_address: str, amount: int) -> dict:
|
|
if amount <= 0:
|
|
raise ValueError("stake amount must be positive")
|
|
current = self.get_wallet(wallet_address)
|
|
self._state.registry[wallet_address] = RegistryWallet(
|
|
stake_balance=current.stake_balance + amount,
|
|
strike_count=current.strike_count,
|
|
banned=current.banned,
|
|
)
|
|
return {"signature": f"local-stake-{wallet_address}", "cluster": self._cluster}
|
|
|
|
def get_wallet(self, wallet_address: str) -> RegistryWallet:
|
|
return self._state.registry.get(wallet_address, RegistryWallet())
|
|
|
|
def record_strike(self, wallet_address: str, ban_threshold: int = 3) -> dict:
|
|
current = self.get_wallet(wallet_address)
|
|
strike_count = current.strike_count + 1
|
|
self._state.registry[wallet_address] = RegistryWallet(
|
|
stake_balance=current.stake_balance,
|
|
strike_count=strike_count,
|
|
banned=current.banned or strike_count >= ban_threshold,
|
|
)
|
|
return {"signature": f"local-strike-{wallet_address}", "cluster": self._cluster}
|
|
|
|
def ban_wallet(self, wallet_address: str) -> dict:
|
|
current = self.get_wallet(wallet_address)
|
|
self._state.registry[wallet_address] = RegistryWallet(
|
|
stake_balance=current.stake_balance,
|
|
strike_count=current.strike_count,
|
|
banned=True,
|
|
)
|
|
return {"signature": f"local-ban-{wallet_address}", "cluster": self._cluster}
|
|
|
|
def has_minimum_stake(self, wallet_address: str | None, minimum_stake: int) -> bool:
|
|
if not wallet_address:
|
|
return False
|
|
wallet = self.get_wallet(wallet_address)
|
|
return wallet.stake_balance >= minimum_stake
|
|
|
|
|
|
class PaymentContract:
|
|
"""Payment wrapper for funded API keys and compute attribution."""
|
|
|
|
def __init__(self, state: _LocalContractState, cluster: str, starting_credit_lamports: int) -> None:
|
|
self._state = state
|
|
self._cluster = cluster
|
|
self._starting_credit_lamports = starting_credit_lamports
|
|
|
|
def fund_api_key(
|
|
self,
|
|
api_key: str,
|
|
*,
|
|
lamports: int = 0,
|
|
usdc_micro: int = 0,
|
|
) -> dict:
|
|
if lamports < 0 or usdc_micro < 0:
|
|
raise ValueError("funding amounts must be non-negative")
|
|
if lamports == 0 and usdc_micro == 0:
|
|
raise ValueError("funding amount must be positive")
|
|
current = self.get_balance(api_key)
|
|
self._state.api_keys[api_key] = ApiKeyBalance(
|
|
lamports=current.lamports + lamports,
|
|
usdc_micro=current.usdc_micro + usdc_micro,
|
|
)
|
|
return {"signature": f"local-fund-{api_key}", "cluster": self._cluster}
|
|
|
|
def get_balance(self, api_key: str) -> ApiKeyBalance:
|
|
return self._state.api_keys.get(
|
|
api_key,
|
|
ApiKeyBalance(lamports=self._starting_credit_lamports),
|
|
)
|
|
|
|
def record_attribution(
|
|
self,
|
|
*,
|
|
session_id: str,
|
|
api_key: str,
|
|
node_wallet: str,
|
|
layer_start: int,
|
|
layer_end: int,
|
|
tokens: int,
|
|
speed_score: float = 1.0,
|
|
) -> dict:
|
|
if not api_key:
|
|
raise ValueError("api_key is required")
|
|
if not node_wallet:
|
|
raise ValueError("node_wallet is required")
|
|
if layer_start < 0 or layer_end < layer_start:
|
|
raise ValueError("layer range must be non-negative and ordered")
|
|
if tokens <= 0:
|
|
raise ValueError("tokens must be positive")
|
|
if speed_score <= 0:
|
|
raise ValueError("speed_score must be positive")
|
|
self._state.attributions.append(ComputeAttribution(
|
|
session_id=session_id,
|
|
api_key=api_key,
|
|
node_wallet=node_wallet,
|
|
layer_start=layer_start,
|
|
layer_end=layer_end,
|
|
tokens=tokens,
|
|
speed_score=speed_score,
|
|
))
|
|
return {"signature": f"local-attribution-{session_id}", "cluster": self._cluster}
|
|
|
|
def list_attributions(self, *, include_settled: bool = True) -> list[ComputeAttribution]:
|
|
if include_settled:
|
|
return list(self._state.attributions)
|
|
return [a for a in self._state.attributions if a.settled_epoch is None]
|
|
|
|
|
|
class SettlementContract:
|
|
"""Settlement wrapper that debits clients and credits token rewards."""
|
|
|
|
def __init__(
|
|
self,
|
|
state: _LocalContractState,
|
|
cluster: str,
|
|
cost_per_layer_token_lamport: int,
|
|
) -> None:
|
|
self._state = state
|
|
self._cluster = cluster
|
|
self._cost_per_layer_token_lamport = cost_per_layer_token_lamport
|
|
|
|
def settle_epoch(
|
|
self,
|
|
*,
|
|
epoch: int,
|
|
validator_wallet: str | None = None,
|
|
validator_reward_share_bps: int = 0,
|
|
) -> SettlementResult:
|
|
if epoch < 0:
|
|
raise ValueError("epoch must be non-negative")
|
|
if not 0 <= validator_reward_share_bps <= 10_000:
|
|
raise ValueError("validator_reward_share_bps must be between 0 and 10000")
|
|
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] = {}
|
|
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
|
|
)
|
|
|
|
for api_key, debit in client_debits.items():
|
|
current = self._state.api_keys.get(api_key, ApiKeyBalance())
|
|
if current.lamports < debit:
|
|
raise ValueError(f"insufficient SOL balance for API key: {api_key}")
|
|
self._state.api_keys[api_key] = ApiKeyBalance(
|
|
lamports=current.lamports - debit,
|
|
usdc_micro=current.usdc_micro,
|
|
)
|
|
|
|
total_weighted_work = sum(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():
|
|
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] = (
|
|
self._state.token_balances.get(node_wallet, 0) + reward
|
|
)
|
|
|
|
for wallet, reward in validator_rewards.items():
|
|
self._state.token_balances[wallet] = self._state.token_balances.get(wallet, 0) + reward
|
|
|
|
self._state.attributions = [
|
|
ComputeAttribution(
|
|
session_id=a.session_id,
|
|
api_key=a.api_key,
|
|
node_wallet=a.node_wallet,
|
|
layer_start=a.layer_start,
|
|
layer_end=a.layer_end,
|
|
tokens=a.tokens,
|
|
speed_score=a.speed_score,
|
|
settled_epoch=epoch,
|
|
)
|
|
if a.settled_epoch is None else a
|
|
for a in self._state.attributions
|
|
]
|
|
return SettlementResult(
|
|
epoch=epoch,
|
|
client_debits=client_debits,
|
|
node_rewards=node_rewards,
|
|
validator_rewards=validator_rewards,
|
|
)
|
|
|
|
def get_token_balance(self, wallet_address: str) -> int:
|
|
return self._state.token_balances.get(wallet_address, 0)
|
|
|
|
def deployment_plan(self) -> dict:
|
|
"""Return the configured manual testnet deployment targets."""
|
|
return {
|
|
"cluster": self._cluster,
|
|
"programs": ["registry", "payment", "settlement"],
|
|
}
|
|
|
|
|
|
class LocalSolanaContracts:
|
|
"""Facade that exposes all three contract wrappers over local validator state."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
cluster: str = "local-test-validator",
|
|
cost_per_layer_token_lamport: int = 1,
|
|
starting_credit_lamports: int = 1_000,
|
|
) -> 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")
|
|
self.cluster = cluster
|
|
self._state = _LocalContractState()
|
|
self.registry = RegistryContract(self._state, cluster)
|
|
self.payment = PaymentContract(self._state, cluster, starting_credit_lamports)
|
|
self.settlement = SettlementContract(
|
|
self._state,
|
|
cluster,
|
|
cost_per_layer_token_lamport,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"ApiKeyBalance",
|
|
"ComputeAttribution",
|
|
"LocalSolanaContracts",
|
|
"PaymentContract",
|
|
"RegistryContract",
|
|
"RegistryWallet",
|
|
"SettlementContract",
|
|
"SettlementResult",
|
|
]
|