Upload src/hw_in_loop.py with huggingface_hub

src/hw_in_loop.py

@@ -0,0 +1,202 @@
+"""
+hw_in_loop.py
+=============
+Hardware-in-the-loop (HIL) stub interface (G-5).
+Provides a generic socket-based bridge to real PIM/CPU/GPU hardware measurements.
+
+Architecture:
+  - HILClient: connects to a local HIL daemon (TCP/Unix socket)
+  - HILDaemon: stub server that replays pre-recorded traces or forwards to real driver
+  - Measures: latency_ms, energy_mJ, power_mW, temperature_C, accuracy_drop
+"""
+
+import json
+import time
+import warnings
+from dataclasses import dataclass, asdict
+from typing import Dict, Optional, Callable
+import numpy as np
+
+try:
+    import torch
+    import torch.nn as nn
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+
+@dataclass
+class HILMeasurement:
+    target: str
+    latency_ms: float
+    energy_mJ: float
+    power_mW: float
+    temperature_C: float
+    accuracy_drop: float = 0.0
+    timestamp: float = 0.0
+    raw: Dict = None
+
+    def __post_init__(self):
+        if self.raw is None:
+            self.raw = {}
+
+
+class HILClient:
+    """
+    G-5: Hardware-in-the-loop measurement client.
+    Connects to a HIL daemon to get real (or trace-replayed) measurements.
+    """
+
+    def __init__(self, host: str = "127.0.0.1", port: int = 9876,
+                 timeout_s: float = 5.0,
+                 trace_file: Optional[str] = None):
+        self.host = host
+        self.port = port
+        self.timeout_s = timeout_s
+        self.sock = None
+        self._connected = False
+        self._trace: Dict[str, list] = {}  # target -> list of HILMeasurement
+        if trace_file:
+            self._load_trace(trace_file)
+
+    def _load_trace(self, path: str):
+        try:
+            with open(path) as f:
+                data = json.load(f)
+            for entry in data:
+                t = entry["target"]
+                self._trace.setdefault(t, []).append(HILMeasurement(**entry))
+        except FileNotFoundError:
+            warnings.warn(f"Trace file {path} not found — using dummy traces.")
+
+    def connect(self) -> bool:
+        """Attempt TCP connection to HIL daemon."""
+        try:
+            import socket
+            self.sock = socket.create_connection((self.host, self.port),
+                                                  timeout=self.timeout_s)
+            self._connected = True
+            return True
+        except Exception as exc:
+            warnings.warn(f"HIL connection failed: {exc} — falling back to trace/dummy.")
+            self._connected = False
+            return False
+
+    def measure(self, model: "nn.Module", target: str,
+                input_shape: tuple, timesteps: int = 1) -> HILMeasurement:
+        """
+        G-5: Execute model on target hardware (or replay trace) and return measurement.
+        """
+        if self._connected and self.sock:
+            req = {
+                "target": target,
+                "model_hash": str(hash(str(model))),
+                "input_shape": input_shape,
+                "timesteps": timesteps,
+            }
+            try:
+                import socket
+                self.sock.sendall((json.dumps(req) + "\n").encode())
+                resp = self.sock.recv(4096).decode()
+                data = json.loads(resp)
+                return HILMeasurement(**data)
+            except Exception as exc:
+                warnings.warn(f"HIL measure failed: {exc} — using fallback.")
+
+        # Fallback 1: trace replay
+        if target in self._trace and self._trace[target]:
+            return self._trace[target].pop(0)
+
+        # Fallback 2: dummy deterministic measurement
+        return self._dummy_measurement(target, input_shape, timesteps)
+
+    def _dummy_measurement(self, target: str, input_shape: tuple, timesteps: int) -> HILMeasurement:
+        # Deterministic dummy based on target + workload size
+        base = np.prod(input_shape) * timesteps
+        if target == "PIM":
+            lat = 0.5 + base / 1e6 * 0.1
+            pwr = 30.0
+        elif target == "CPU":
+            lat = 2.0 + base / 1e6 * 0.5
+            pwr = 8000.0
+        else:
+            lat = 0.8 + base / 1e6 * 0.05
+            pwr = 150000.0
+        return HILMeasurement(
+            target=target,
+            latency_ms=lat,
+            energy_mJ=pwr * lat / 1000.0,
+            power_mW=pwr,
+            temperature_C=35.0 + np.random.normal(0, 2),
+            accuracy_drop=0.0,
+            timestamp=time.time(),
+        )
+
+    def close(self):
+        if self.sock:
+            self.sock.close()
+            self.sock = None
+            self._connected = False
+
+    def __enter__(self):
+        self.connect()
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
+class HILDaemon:
+    """
+    G-5: Stub HIL daemon that accepts TCP connections and either
+    (a) forwards to a real driver shim, or (b) replays traces.
+    This is a server — run in a separate process or thread.
+    """
+
+    def __init__(self, host: str = "127.0.0.1", port: int = 9876,
+                 driver_fn: Optional[Callable] = None,
+                 trace_file: Optional[str] = None):
+        self.host = host
+        self.port = port
+        self.driver_fn = driver_fn
+        self._trace: Dict[str, list] = {}
+        if trace_file:
+            self._load_trace(trace_file)
+
+    def _load_trace(self, path: str):
+        with open(path) as f:
+            data = json.load(f)
+        for entry in data:
+            t = entry["target"]
+            self._trace.setdefault(t, []).append(HILMeasurement(**entry))
+
+    def _handle_request(self, req: dict) -> dict:
+        target = req.get("target", "CPU")
+        if self.driver_fn:
+            return asdict(self.driver_fn(req))
+        if target in self._trace and self._trace[target]:
+            return asdict(self._trace[target].pop(0))
+        # Dummy
+        m = HILMeasurement(
+            target=target, latency_ms=1.0, energy_mJ=1.0,
+            power_mW=1000.0, temperature_C=40.0)
+        return asdict(m)
+
+    def run(self):
+        import socket
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind((self.host, self.port))
+            s.listen(1)
+            print(f"[HILDaemon] listening on {self.host}:{self.port}")
+            while True:
+                conn, addr = s.accept()
+                with conn:
+                    data = conn.recv(4096)
+                    if not data:
+                        continue
+                    try:
+                        req = json.loads(data.decode())
+                        resp = self._handle_request(req)
+                        conn.sendall((json.dumps(resp) + "\n").encode())
+                    except Exception as exc:
+                        conn.sendall((json.dumps({"error": str(exc)}) + "\n").encode())