Spaces:
Sleeping
Sleeping
| """ | |
| environment/node.py — Single worker node in the training cluster. | |
| Finale upgrades: | |
| - False alarm node support (_is_false_alarm flag) | |
| - False alarm memory elevation in tick() | |
| - Intermittent straggler recovery handled by FaultInjector | |
| (node.py just needs to support the state transitions) | |
| """ | |
| import random | |
| from dataclasses import dataclass, field | |
| from typing import Literal | |
| NodeStatus = Literal["healthy", "slow", "oom", "crashed"] | |
| class Node: | |
| """ | |
| Simulates one GPU worker node in a distributed training ring. | |
| All values updated by RingCluster.tick() each timestep. | |
| """ | |
| id: str | |
| status: NodeStatus = "healthy" | |
| in_ring: bool = True | |
| memory: float = 0.3 | |
| throughput: float = 1.0 | |
| latency: float = 5.0 | |
| # Internal fault tracking — NOT directly observable by agent | |
| _oom_progress: float = field(default=0.0, repr=False) | |
| _fault_type: str = field(default="none", repr=False) | |
| _steps_since_fault: int = field(default=0, repr=False) | |
| # Baseline values for reward computation | |
| _baseline_throughput: float = field(default=1.0, repr=False) | |
| _baseline_memory: float = field(default=0.3, repr=False) | |
| # NEW: False alarm support | |
| # Node shows elevated memory but never actually fails | |
| # Set by FaultInjector.tick() when false_alarm config is present | |
| _is_false_alarm: bool = field(default=False, repr=False) | |
| _false_alarm_memory: float = field(default=0.75, repr=False) | |
| def tick(self, noise: bool = True): | |
| """ | |
| Advance node state by one timestep. | |
| Applies fault progression and small random noise. | |
| Called by RingCluster.tick() every step. | |
| """ | |
| if self.status == "crashed": | |
| self.throughput = 0.0 | |
| self.in_ring = False | |
| return | |
| # --- False alarm handling --- | |
| # Memory stays elevated but node NEVER crashes or slows | |
| # Agent that restarts this node takes a penalty | |
| # Must be checked BEFORE OOM progression so it doesn't | |
| # accidentally trigger real OOM logic | |
| if self._is_false_alarm: | |
| # Small oscillation around target memory to look realistic | |
| target = self._false_alarm_memory | |
| self.memory = max(0.65, min(0.88, | |
| target + random.uniform(-0.02, 0.02))) | |
| # Status stays healthy — this is the trap | |
| self.status = "healthy" | |
| self.throughput = max(0.85, min(1.0, | |
| self.throughput + random.uniform(-0.02, 0.02))) | |
| self._steps_since_fault += 1 | |
| return # skip all other fault logic for false alarm nodes | |
| # --- OOM progression --- | |
| # Memory climbs gradually — agent should catch this early | |
| if self._fault_type == "oom": | |
| self._oom_progress = min(1.0, self._oom_progress + 0.06) | |
| self.memory = min(1.0, | |
| self._baseline_memory + self._oom_progress * 0.7) | |
| if self.memory > 0.85 and self.status == "healthy": | |
| self.status = "slow" | |
| self.throughput = max(0.1, self.throughput - 0.3) | |
| if self.memory > 0.95: | |
| self.status = "oom" | |
| self.throughput = max(0.05, self.throughput - 0.2) | |
| if self.memory >= 1.0: | |
| self._crash() | |
| return | |
| # --- Straggler fault --- | |
| if self._fault_type == "straggler": | |
| self.status = "slow" | |
| self.latency = min(200.0, self.latency + 0.5) | |
| # --- Crash fault --- | |
| if self._fault_type == "crash": | |
| self._crash() | |
| return | |
| # --- Small random noise on healthy nodes --- | |
| # Makes observations realistic, agent cant threshold on exact values | |
| if noise and self.status == "healthy": | |
| self.memory = max(0.1, min(0.95, | |
| self.memory + random.uniform(-0.01, 0.01))) | |
| self.throughput = max(0.1, min(1.0, | |
| self.throughput + random.uniform(-0.02, 0.02))) | |
| self.latency = max(1.0, min(50.0, | |
| self.latency + random.uniform(-0.5, 0.5))) | |
| self._steps_since_fault += 1 | |
| def _crash(self): | |
| """Transition node to crashed state.""" | |
| self.status = "crashed" | |
| self.in_ring = False | |
| self.throughput = 0.0 | |
| self.latency = 9999.0 | |
| def restart(self): | |
| """ | |
| Agent action: restart this node. | |
| Starts at 60% throughput — recovers over 2 warmup ticks. | |
| IMPORTANT: restarting a false alarm node clears the false alarm | |
| flag but gives a penalty in reward.py — the node was healthy, | |
| restarting it was wasteful. | |
| """ | |
| self.status = "healthy" | |
| self.in_ring = True | |
| self.memory = self._baseline_memory | |
| self.throughput = 0.6 | |
| self.latency = 5.0 | |
| self._oom_progress = 0.0 | |
| self._fault_type = "none" | |
| self._steps_since_fault = 0 | |
| # Clear false alarm flag if agent restarts it | |
| # (penalty already applied in reward.py) | |
| self._is_false_alarm = False | |
| def remove_from_ring(self): | |
| """ | |
| Agent action: remove node from ring without restarting. | |
| Stops the node from blocking all-reduce. | |
| """ | |
| self.in_ring = False | |
| def reduce_batch(self): | |
| """ | |
| Agent action: halve batch size on this node. | |
| Reduces memory pressure. Only helps before memory >= 0.95. | |
| NOTE: reduce_batch on a false alarm node is a mild waste | |
| (small penalty) but not as bad as restarting it. | |
| """ | |
| if self._is_false_alarm: | |
| # Slightly reduces the false alarm memory display | |
| # but doesnt fix anything — node was never broken | |
| self._false_alarm_memory = max( | |
| 0.65, self._false_alarm_memory - 0.05) | |
| return | |
| if self.memory < 0.95: | |
| self._oom_progress = max(0.0, self._oom_progress - 0.3) | |
| self.memory = max(self._baseline_memory, self.memory - 0.2) | |
| self.throughput = max(0.5, self.throughput - 0.1) | |
| def warmup_tick(self): | |
| """ | |
| Gradually restore throughput after restart. | |
| Simulates node coming back online over 2 steps. | |
| """ | |
| if self.status == "healthy" and self.throughput < 1.0: | |
| self.throughput = min(1.0, self.throughput + 0.2) | |
| def to_state(self) -> dict: | |
| """ | |
| Return observable state as dict. | |
| NOTE: _is_false_alarm is NOT exposed here — agent must | |
| infer from behavior (memory elevated but never crashes). | |
| """ | |
| return { | |
| "id": self.id, | |
| "status": self.status, | |
| "memory": round(self.memory, 3), | |
| "throughput": round(self.throughput, 3), | |
| "latency": round(self.latency, 2), | |
| "in_ring": self.in_ring, | |
| } | |