environment/node.py

"""
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"]


@dataclass
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,
        }