Spaces:
Sleeping
Sleeping
| """SiliconMind — Pydantic Models for the OpenEnv Memory Subsystem Intelligence Gym. | |
| All models inherit from openenv.core.env_server base classes to ensure | |
| full OpenEnv spec compliance. Every field is typed and documented so the | |
| LLM agent (and human reviewers) can reason about the observation space. | |
| Architecture note: | |
| Action → what the agent sends | |
| Observation → what the agent receives (includes reward) | |
| State → internal episode state (for /state endpoint) | |
| """ | |
| from __future__ import annotations | |
| from enum import Enum | |
| from typing import Any, Dict, List, Optional | |
| from pydantic import BaseModel, Field | |
| # --------------------------------------------------------------------------- | |
| # OpenEnv base classes — same import path used by CasualOps (proven pattern) | |
| # --------------------------------------------------------------------------- | |
| from openenv.core.env_server import ( | |
| Action as OEAction, | |
| Observation as OEObservation, | |
| State as OEState, | |
| ) | |
| # ╔══════════════════════════════════════════════════════════════════════════╗ | |
| # ║ ENUMS ║ | |
| # ╚══════════════════════════════════════════════════════════════════════════╝ | |
| class ActionType(str, Enum): | |
| """The five control levers available to the agent. | |
| Easy tasks expose only PREFETCH / NO_PREFETCH. | |
| Medium adds SET_AGGRESSIVENESS. | |
| Hard/Extreme unlock THROTTLE_BANDWIDTH and PARTITION_CACHE. | |
| """ | |
| PREFETCH = "prefetch" | |
| NO_PREFETCH = "no_prefetch" | |
| SET_AGGRESSIVENESS = "set_aggressiveness" | |
| THROTTLE_BANDWIDTH = "throttle_bandwidth" | |
| PARTITION_CACHE = "partition_cache" | |
| class PrefetchOutcome(str, Enum): | |
| """Pythia-inspired 9-level classification for every prefetch decision. | |
| Reference: Pythia (MICRO '21) — Bera et al. | |
| Extensions: POLLUTION and CORRECT_RESTRAINT are SiliconMind originals. | |
| """ | |
| # Positive outcomes | |
| ACCURATE_TIMELY = "RAT" # Prefetch was used AND arrived before demand | |
| ACCURATE_LATE = "RAL" # Prefetch was used BUT arrived after demand started | |
| CORRECT_RESTRAINT = "RESTRAINT" # Agent correctly chose NO_PREFETCH on unpredictable | |
| # Neutral / mild negative | |
| REDUNDANT = "REDUNDANT" # Prefetched a line already in cache | |
| NO_PREFETCH_LOW_BW = "RNP_L" # Did not prefetch, bandwidth was available (missed opportunity) | |
| NO_PREFETCH_HIGH_BW = "RNP_H" # Did not prefetch, bandwidth was congested (acceptable) | |
| # Negative outcomes | |
| INACCURATE_LOW_BW = "RIN_L" # Wrong prefetch, bandwidth was available | |
| INACCURATE_HIGH_BW = "RIN_H" # Wrong prefetch, bandwidth was congested (worse) | |
| POLLUTION = "POLLUTION" # Prefetch evicted a useful line from cache | |
| class CacheLevel(str, Enum): | |
| """Where a demand access resolved in the cache hierarchy.""" | |
| L1_HIT = "L1_HIT" | |
| L2_HIT = "L2_HIT" | |
| LLC_HIT = "LLC_HIT" | |
| MISS = "MISS" | |
| class BandwidthPressure(str, Enum): | |
| """Qualitative bandwidth state for LLM-friendly observation.""" | |
| LOW = "low" # utilization < 50% | |
| MODERATE = "moderate" # 50% - 70% | |
| HIGH = "high" # 70% - 90% | |
| CRITICAL = "critical" # > 90% | |
| class MSHRRisk(str, Enum): | |
| """MSHR stall risk level.""" | |
| NONE = "none" # < 50% utilization | |
| APPROACHING = "approaching" # 50% - 80% | |
| STALLING = "stalling" # > 80% or actual stall | |
| class SLAStatus(str, Enum): | |
| """Per-workload SLA compliance status.""" | |
| MET = "met" | |
| AT_RISK = "at_risk" | |
| VIOLATED = "violated" | |
| # ╔══════════════════════════════════════════════════════════════════════════╗ | |
| # ║ REWARD MODELS ║ | |
| # ╚══════════════════════════════════════════════════════════════════════════╝ | |
| # Pythia-inspired reward values (raw, before normalization) | |
| REWARD_VALUES: Dict[PrefetchOutcome, float] = { | |
| PrefetchOutcome.ACCURATE_TIMELY: +20.0, | |
| PrefetchOutcome.ACCURATE_LATE: +12.0, | |
| PrefetchOutcome.CORRECT_RESTRAINT: +8.0, | |
| PrefetchOutcome.REDUNDANT: -2.0, | |
| PrefetchOutcome.NO_PREFETCH_LOW_BW: -4.0, | |
| PrefetchOutcome.NO_PREFETCH_HIGH_BW: -2.0, | |
| PrefetchOutcome.INACCURATE_LOW_BW: -8.0, | |
| PrefetchOutcome.INACCURATE_HIGH_BW: -14.0, | |
| PrefetchOutcome.POLLUTION: -18.0, | |
| } | |
| # For normalization to [0, 1] | |
| _REWARD_MIN = min(REWARD_VALUES.values()) # -18 | |
| _REWARD_MAX = max(REWARD_VALUES.values()) # +20 | |
| def normalize_reward(raw: float) -> float: | |
| """Map raw Pythia-scale reward to [0.001, 0.999] for OpenEnv compliance.""" | |
| if _REWARD_MAX == _REWARD_MIN: | |
| return 0.5 | |
| normalized = (raw - _REWARD_MIN) / (_REWARD_MAX - _REWARD_MIN) | |
| return round(max(0.001, min(0.999, normalized)), 3) | |
| class RewardComponents(BaseModel): | |
| """Decomposed per-step reward with 6 named components. | |
| All components are in [0.0, 0.25] so they sum to at most 1.0. | |
| The component weights change per-task (defined in task config). | |
| """ | |
| accuracy: float = Field(default=0.0, ge=0.0, le=0.25, | |
| description="Prefetch prediction accuracy") | |
| coverage: float = Field(default=0.0, ge=0.0, le=0.20, | |
| description="Miss elimination rate") | |
| pollution: float = Field(default=0.0, ge=0.0, le=0.20, | |
| description="Cache pollution control") | |
| bandwidth: float = Field(default=0.0, ge=0.0, le=0.15, | |
| description="Memory bus efficiency") | |
| timeliness: float = Field(default=0.0, ge=0.0, le=0.10, | |
| description="Prefetch arrived before demand") | |
| restraint: float = Field(default=0.0, ge=0.0, le=0.10, | |
| description="Correct NO_PREFETCH decisions") | |
| def total(self) -> float: | |
| """Sum of all components, clamped to [0, 1].""" | |
| return min(1.0, ( | |
| self.accuracy + self.coverage + self.pollution | |
| + self.bandwidth + self.timeliness + self.restraint | |
| )) | |
| class Reward(BaseModel): | |
| """Per-step reward returned to the agent. | |
| Contains the total normalized scalar, the component breakdown, | |
| the Pythia classification, and a human-readable explanation. | |
| """ | |
| total: float = Field(default=0.0, ge=0.0, le=1.0, | |
| description="Normalized scalar reward [0, 1]") | |
| components: RewardComponents = Field(default_factory=RewardComponents) | |
| outcome: str = Field(default="", | |
| description="Pythia classification (RAT/RAL/RIN_L/...)") | |
| explanation: str = Field(default="", | |
| description="Human-readable reward explanation") | |
| # ╔══════════════════════════════════════════════════════════════════════════╗ | |
| # ║ OBSERVATION SUB-MODELS ║ | |
| # ╚══════════════════════════════════════════════════════════════════════════╝ | |
| class PatternAnalysis(BaseModel): | |
| """Environment-computed pattern detection shared with the agent. | |
| This is the key LLM-friendly feature: instead of making the LLM | |
| figure out patterns from raw hex, we tell it what we detected. | |
| """ | |
| detected_pattern: str = Field( | |
| default="unknown", | |
| description="Identified pattern: 'stride-1', 'stride-4', 'random', 'unknown'") | |
| confidence: float = Field( | |
| default=0.0, ge=0.0, le=1.0, | |
| description="How confident the detector is (drops on phase change)") | |
| pattern_age: int = Field( | |
| default=0, ge=0, | |
| description="Steps since this pattern was first detected") | |
| previous_pattern: str = Field( | |
| default="", | |
| description="What pattern was active before current one") | |
| previous_confidence: float = Field( | |
| default=0.0, ge=0.0, le=1.0, | |
| description="Confidence of the previous pattern") | |
| class CacheStats(BaseModel): | |
| """Aggregate cache performance counters exposed to the agent.""" | |
| hit_rate: float = Field(default=0.0, ge=0.0, le=1.0) | |
| l1_hit_rate: float = Field(default=0.0, ge=0.0, le=1.0) | |
| l2_hit_rate: float = Field(default=0.0, ge=0.0, le=1.0) | |
| llc_hit_rate: float = Field(default=0.0, ge=0.0, le=1.0) | |
| total_hits: int = Field(default=0, ge=0) | |
| total_misses: int = Field(default=0, ge=0) | |
| total_prefetches_issued: int = Field(default=0, ge=0) | |
| total_prefetches_useful: int = Field(default=0, ge=0) | |
| total_prefetches_wasted: int = Field(default=0, ge=0) | |
| total_pollution_evictions: int = Field(default=0, ge=0) | |
| class BandwidthState(BaseModel): | |
| """Memory bus state for bandwidth-aware decision making.""" | |
| utilization: float = Field(default=0.0, ge=0.0, le=1.0, | |
| description="Current bandwidth utilization [0, 1]") | |
| pressure: BandwidthPressure = Field( | |
| default=BandwidthPressure.LOW, | |
| description="Qualitative pressure level") | |
| outstanding_requests: int = Field(default=0, ge=0) | |
| dropped_requests: int = Field(default=0, ge=0, | |
| description="Prefetches dropped due to congestion") | |
| class MSHRState(BaseModel): | |
| """Miss Status Holding Register state.""" | |
| utilization: float = Field(default=0.0, ge=0.0, le=1.0) | |
| stall_risk: MSHRRisk = Field(default=MSHRRisk.NONE) | |
| total_stalls: int = Field(default=0, ge=0, | |
| description="Pipeline stalls caused by full MSHR") | |
| demand_blocked: int = Field(default=0, ge=0, | |
| description="Demand requests blocked by prefetch MSHR usage") | |
| # ╔══════════════════════════════════════════════════════════════════════════╗ | |
| # ║ ACTION MODEL ║ | |
| # ╚══════════════════════════════════════════════════════════════════════════╝ | |
| class SiliconMindAction(OEAction): | |
| """Agent → Environment: the prefetch decision. | |
| Target format varies by action type: | |
| PREFETCH → "1,2,4" (comma-sep cache-line offsets) | |
| NO_PREFETCH → "skip" | |
| SET_AGGRESSIVENESS → "1" | "2" | "3" | "4" | |
| THROTTLE_BANDWIDTH → "0.6" (max utilization target) | |
| PARTITION_CACHE → "web:6,db:6,video:4" (LLC way allocation) | |
| """ | |
| type: ActionType = Field( | |
| ..., | |
| description="Which control lever to use") | |
| target: str = Field( | |
| ..., | |
| description="Action parameter (format depends on type)") | |
| detail: str = Field( | |
| default="", | |
| description="Agent's reasoning (optional, for interpretability)") | |
| # ╔══════════════════════════════════════════════════════════════════════════╗ | |
| # ║ OBSERVATION MODEL ║ | |
| # ╚══════════════════════════════════════════════════════════════════════════╝ | |
| class SiliconMindObservation(OEObservation): | |
| """Environment → Agent: everything the LLM sees each step. | |
| The `narrative` field is the #1 innovation: a 2-3 sentence natural | |
| language summary designed for LLM reasoning. All other fields provide | |
| structured data for precision. | |
| """ | |
| # ── LLM-OPTIMIZED NARRATIVE ────────────────────────────────────────── | |
| narrative: str = Field( | |
| default="", | |
| description="2-3 sentence natural language summary of current state") | |
| # ── CURRENT MEMORY ACCESS ──────────────────────────────────────────── | |
| step_number: int = Field(default=0, ge=0) | |
| steps_remaining: int = Field(default=0, ge=0) | |
| current_pc: str = Field( | |
| default="0x0", | |
| description="Program counter (hex string)") | |
| current_address: str = Field( | |
| default="0x0", | |
| description="Memory address being accessed (hex string)") | |
| current_cache_result: CacheLevel = Field( | |
| default=CacheLevel.MISS, | |
| description="Where this access resolved") | |
| is_write: bool = Field(default=False) | |
| current_workload: str = Field( | |
| default="default", | |
| description="Which workload generated this access") | |
| current_phase: str = Field( | |
| default="unknown", | |
| description="Current phase within the workload (e.g., 'forward_pass')") | |
| # ── PATTERN CONTEXT ────────────────────────────────────────────────── | |
| recent_deltas: List[int] = Field( | |
| default_factory=list, | |
| description="Last 8 address deltas in cache-line units") | |
| recent_pcs: List[str] = Field( | |
| default_factory=list, | |
| description="Last 8 program counters (hex)") | |
| recent_cache_results: List[str] = Field( | |
| default_factory=list, | |
| description="Last 8 cache results") | |
| pattern_analysis: PatternAnalysis = Field( | |
| default_factory=PatternAnalysis) | |
| # ── PERFORMANCE FEEDBACK ───────────────────────────────────────────── | |
| last_action_outcome: str = Field( | |
| default="", | |
| description="Pythia classification of last action (RAT/WASTED/...)") | |
| last_action_explanation: str = Field( | |
| default="", | |
| description="Why the last action got its outcome") | |
| prefetch_accuracy_recent: float = Field( | |
| default=0.0, ge=0.0, le=1.0, | |
| description="Accuracy over last 10 prefetches") | |
| prefetch_accuracy_total: float = Field( | |
| default=0.0, ge=0.0, le=1.0, | |
| description="Accuracy over entire episode") | |
| # ── SYSTEM STATE ───────────────────────────────────────────────────── | |
| cache_stats: CacheStats = Field(default_factory=CacheStats) | |
| bandwidth: BandwidthState = Field(default_factory=BandwidthState) | |
| mshr: MSHRState = Field(default_factory=MSHRState) | |
| pollution_rate: float = Field( | |
| default=0.0, ge=0.0, le=1.0, | |
| description="pollution_evictions / total_prefetches") | |
| # ── MULTI-TENANT (Tasks 3-4) ───────────────────────────────────────── | |
| active_workloads: List[str] = Field(default_factory=list) | |
| workload_sla_status: Dict[str, str] = Field( | |
| default_factory=dict, | |
| description="workload_id → 'met'|'at_risk'|'violated'") | |
| llc_partition: Dict[str, int] = Field( | |
| default_factory=dict, | |
| description="Current LLC way allocation per workload") | |
| # ── AVAILABLE ACTIONS ──────────────────────────────────────────────── | |
| available_actions: List[str] = Field( | |
| default_factory=list, | |
| description="Which ActionTypes are available in this task") | |
| action_hint: str = Field( | |
| default="", | |
| description="Environment hint (e.g., 'Consider reducing aggressiveness')") | |
| # ── EPISODE METADATA ───────────────────────────────────────────────── | |
| task_id: str = Field(default="") | |
| task_description: str = Field(default="") | |
| done: bool = Field(default=False) | |
| reward: float = Field(default=0.0, | |
| description="Normalized scalar reward [0, 1]") | |
| reward_detail: Reward = Field(default_factory=Reward, | |
| description="Rich reward breakdown (outcome, components, explanation)") | |
| info: Dict[str, Any] = Field( | |
| default_factory=dict, | |
| description="Extra data: forensics, baseline comparison, etc.") | |
| # ╔══════════════════════════════════════════════════════════════════════════╗ | |
| # ║ STATE MODEL ║ | |
| # ╚══════════════════════════════════════════════════════════════════════════╝ | |
| class SiliconMindState(OEState): | |
| """Internal episode state returned by the /state endpoint. | |
| This is NOT exposed to the agent during normal play — it is used | |
| by the OpenEnv framework and for debugging. | |
| """ | |
| task_id: str = Field(default="") | |
| step_number: int = Field(default=0, ge=0) | |
| max_steps: int = Field(default=0, ge=0) | |
| done: bool = Field(default=False) | |
| total_reward: float = Field(default=0.0) | |
| seed: int = Field(default=0) | |
| # Aggregate counters | |
| cache_stats: CacheStats = Field(default_factory=CacheStats) | |
| bandwidth: BandwidthState = Field(default_factory=BandwidthState) | |
| mshr: MSHRState = Field(default_factory=MSHRState) | |
| # Agent config state | |
| prefetch_aggressiveness: int = Field(default=2, ge=1, le=4) | |
| bandwidth_throttle: float = Field(default=1.0, ge=0.0, le=1.0) | |
| llc_partition: Dict[str, int] = Field(default_factory=dict) | |
| # Workload tracking | |
| current_workload: str = Field(default="") | |
| current_phase: str = Field(default="") | |