diff --git a/src/server/communication/inference_recorder.py b/src/server/communication/inference_recorder.py new file mode 100644 index 0000000..1228fc6 --- /dev/null +++ b/src/server/communication/inference_recorder.py @@ -0,0 +1,132 @@ +"""Recording of inference-cycle side outputs (simulation CSV, debug JSON). + +These writes happen on the background I/O thread while the ASGI threads open, +close and reset the same handles, so every access to the CSV state goes through +one lock. Rows queued before a close are dropped rather than written to a closed +file: telemetry must never take down the inference path. +""" + +import csv +import json +import os +import shutil +import threading + +from server.logger.log import logger + +SIMULATION_CSV_FIELDNAMES = [ + "inference_id", + "timestamp", + "avg_device_time", + "min_device_time", + "max_device_time", + "avg_edge_time", + "min_edge_time", + "max_edge_time", + "num_device_layers", + "num_edge_layers", +] + +DEBUG_DIR = "data/models/debug" + + +class InferenceCycleRecorder: + """Owns the simulation CSV handle, its row counter, and debug JSON output.""" + + def __init__(self, debug_dir: str = DEBUG_DIR): + self._lock = threading.RLock() + self._csv_file = None + self._csv_writer = None + self._inference_counter = 0 + self.debug_dir = debug_dir + + # ── simulation CSV lifecycle ──────────────────────────────────────── + @property + def is_recording(self) -> bool: + with self._lock: + return self._csv_writer is not None + + def open_simulation_csv(self, csv_file_path) -> bool: + """Start recording simulation results, replacing any open file.""" + with self._lock: + try: + self._close_locked() + self._csv_file = open(csv_file_path, "w", newline="") + self._csv_writer = csv.DictWriter( + self._csv_file, fieldnames=SIMULATION_CSV_FIELDNAMES + ) + self._csv_writer.writeheader() + self._csv_file.flush() + self._inference_counter = 0 + logger.info(f"Simulation CSV recording enabled: {csv_file_path}") + return True + except Exception as e: + logger.error(f"Failed to setup simulation CSV: {e}") + # Never leave a half-open handle behind on failure. + self._csv_file = None + self._csv_writer = None + return False + + def close_simulation_csv(self) -> None: + with self._lock: + self._close_locked() + self._inference_counter = 0 + + def _close_locked(self) -> None: + if self._csv_file is not None: + try: + self._csv_file.close() + except Exception as e: + logger.warning(f"Failed to close simulation CSV: {e}") + self._csv_file = None + self._csv_writer = None + + # ── row recording ─────────────────────────────────────────────────── + def next_inference_id(self) -> int: + """Reserve the next row id. Concurrent requests must not collide.""" + with self._lock: + self._inference_counter += 1 + return self._inference_counter + + def write_row(self, row: dict) -> bool: + """Write one row and flush. + + Called from the background I/O thread; the check and the write are done + under the lock so a concurrent close cannot leave us writing to a file + that has just been closed. + """ + with self._lock: + if self._csv_writer is None: + return False + try: + self._csv_writer.writerow(row) + self._csv_file.flush() + return True + except Exception as e: + logger.warning(f"Failed to write simulation CSV row: {e}") + return False + + # ── debug JSON output ─────────────────────────────────────────────── + def reset_debug_folder(self) -> None: + """Empty the debug folder; called once per server start.""" + if os.path.exists(self.debug_dir): + try: + shutil.rmtree(self.debug_dir) + except Exception as e: + logger.warning(f"Unable to clear debug folder: {e}") + os.makedirs(self.debug_dir, exist_ok=True) + + def save_debug_files(self, device_id, device_times: dict, edge_times: dict) -> None: + """Write the per-device timing JSON files. + + Safe to call from the background I/O thread: the caller passes snapshots, + never the live EMA tables. + """ + device_file = f"{self.debug_dir}/{device_id}_device_times.json" + edge_file = f"{self.debug_dir}/{device_id}_edge_times.json" + + with open(device_file, "w") as f: + json.dump(device_times, f, indent=4) + + with open(edge_file, "w") as f: + json.dump(edge_times, f, indent=4) diff --git a/src/server/communication/request_handler.py b/src/server/communication/request_handler.py index 7c0c733..de5ec28 100644 --- a/src/server/communication/request_handler.py +++ b/src/server/communication/request_handler.py @@ -4,7 +4,6 @@ import time import queue import threading -import shutil import math from datetime import datetime from pathlib import Path @@ -35,6 +34,7 @@ from server.communication.message_data import MessageData from server.communication.inference_protocol import decode_inference_payload from server.communication.device_state import DeviceStateManager +from server.communication.inference_recorder import InferenceCycleRecorder from server.models.model_input_converter import ModelInputConverter from server.core.delay_simulator import DelaySimulator @@ -323,10 +323,9 @@ class RequestHandler: model_registry = {} # hash -> {model_dir, last_offloading_layer, num_layers} device_model_map = {} # device_id -> model_info - # Class-level CSV file tracking for simulation results - csv_file = None - csv_writer = None - inference_counter = 0 + # Simulation CSV + debug JSON output, thread-safe across the ASGI and + # background I/O threads. + recorder = InferenceCycleRecorder() header_printed = False inference_table_rows_printed = 0 @@ -417,81 +416,17 @@ def __init__(self): def _cleanup_debug_folder(self): - debug_dir = "data/models/debug" - if os.path.exists(debug_dir): - try: - shutil.rmtree(debug_dir) - except Exception as e: - logger.warning(f"Unable to clear debug folder: {e}") - os.makedirs(debug_dir, exist_ok=True) - - # Debug file saving - def _save_debug_files(self, device_id): - """Legacy sync method – prefer _save_debug_files_data for bg thread.""" - device_times, edge_times = RequestHandler.device_state.snapshot_debug_times( - device_id - ) - self._save_debug_files_data(device_id, device_times, edge_times) - - @staticmethod - def _save_debug_files_data(device_id, device_times, edge_times): - """Write debug JSON files (safe to call from the background I/O thread).""" - debug_dir = "data/models/debug" - device_file = f"{debug_dir}/{device_id}_device_times.json" - edge_file = f"{debug_dir}/{device_id}_edge_times.json" - - with open(device_file, "w") as f: - json.dump(device_times, f, indent=4) - - with open(edge_file, "w") as f: - json.dump(edge_times, f, indent=4) - - @staticmethod - def _write_csv_row(row): - """Write a single CSV row + flush (safe to call from the background I/O thread).""" - if RequestHandler.csv_writer: - RequestHandler.csv_writer.writerow(row) - RequestHandler.csv_file.flush() + RequestHandler.recorder.reset_debug_folder() @classmethod def set_simulation_csv(cls, csv_file_path): """Set the CSV file for recording simulation results""" - import csv - - try: - if cls.csv_file: - cls.csv_file.close() - cls.csv_file = open(csv_file_path, "w", newline="") - cls.csv_writer = csv.DictWriter( - cls.csv_file, - fieldnames=[ - "inference_id", - "timestamp", - "avg_device_time", - "min_device_time", - "max_device_time", - "avg_edge_time", - "min_edge_time", - "max_edge_time", - "num_device_layers", - "num_edge_layers", - ], - ) - cls.csv_writer.writeheader() - cls.csv_file.flush() - cls.inference_counter = 0 - logger.info(f"Simulation CSV recording enabled: {csv_file_path}") - except Exception as e: - logger.error(f"Failed to setup simulation CSV: {e}") + cls.recorder.open_simulation_csv(csv_file_path) @classmethod def close_simulation_csv(cls): """Close the CSV file""" - if cls.csv_file: - cls.csv_file.close() - cls.csv_file = None - cls.csv_writer = None - cls.inference_counter = 0 + cls.recorder.close_simulation_csv() def should_force_local_inference(self) -> bool: """ @@ -780,18 +715,18 @@ def handle_device_inference_result(self, body, received_timestamp): _debug_device_times, _debug_edge_times = device_state.snapshot_debug_times( device_id ) + recorder = RequestHandler.recorder enqueue_background_io( lambda did=device_id, dt=_debug_device_times, et=_debug_edge_times: ( - self._save_debug_files_data(did, dt, et) + recorder.save_debug_files(did, dt, et) ), description="debug timing JSON write", ) - if RequestHandler.csv_writer: - RequestHandler.inference_counter += 1 + if recorder.is_recording: device_values = message_data.device_layers_inference_time row = { - "inference_id": RequestHandler.inference_counter, + "inference_id": recorder.next_inference_id(), "timestamp": datetime.now().isoformat(), "avg_device_time": sum(device_values) / len(device_values) if device_values else 0, "avg_edge_time": sum(edge_layer_times) / len(edge_layer_times) if edge_layer_times else 0, @@ -799,7 +734,7 @@ def handle_device_inference_result(self, body, received_timestamp): "num_edge_layers": num_edge_layers, } enqueue_background_io( - lambda r=row: self._write_csv_row(r), + lambda r=row: recorder.write_row(r), description="simulation CSV row write", ) diff --git a/tests/unit/test_inference_recorder.py b/tests/unit/test_inference_recorder.py new file mode 100644 index 0000000..b2ef7cb --- /dev/null +++ b/tests/unit/test_inference_recorder.py @@ -0,0 +1,226 @@ +"""Unit tests for the simulation-CSV / debug-JSON recorder.""" + +import csv +import json +import threading + +import pytest + +from server.communication.inference_recorder import ( + SIMULATION_CSV_FIELDNAMES, + InferenceCycleRecorder, +) + + +@pytest.fixture +def recorder(tmp_path): + return InferenceCycleRecorder(debug_dir=str(tmp_path / "debug")) + + +def _row(inference_id): + return { + "inference_id": inference_id, + "timestamp": "2026-07-17T10:00:00", + "avg_device_time": 0.5, + "avg_edge_time": 0.25, + "num_device_layers": 2, + "num_edge_layers": 1, + } + + +def test_not_recording_until_a_csv_is_opened(recorder): + assert recorder.is_recording is False + assert recorder.write_row(_row(1)) is False + + +def test_open_writes_header_and_starts_recording(recorder, tmp_path): + path = tmp_path / "sim.csv" + + assert recorder.open_simulation_csv(path) is True + + assert recorder.is_recording is True + header = path.read_text().splitlines()[0] + assert header.split(",") == SIMULATION_CSV_FIELDNAMES + + +def test_write_row_persists_and_flushes(recorder, tmp_path): + path = tmp_path / "sim.csv" + recorder.open_simulation_csv(path) + + assert recorder.write_row(_row(1)) is True + + # Flushed, so readable without closing. + rows = list(csv.DictReader(path.open())) + assert rows[0]["inference_id"] == "1" + assert rows[0]["avg_device_time"] == "0.5" + # Columns the caller does not supply stay empty rather than raising. + assert rows[0]["min_device_time"] == "" + + +def test_next_inference_id_increments_from_one(recorder, tmp_path): + recorder.open_simulation_csv(tmp_path / "sim.csv") + + assert [recorder.next_inference_id() for _ in range(3)] == [1, 2, 3] + + +def test_reopening_resets_the_counter(recorder, tmp_path): + recorder.open_simulation_csv(tmp_path / "a.csv") + recorder.next_inference_id() + recorder.open_simulation_csv(tmp_path / "b.csv") + + assert recorder.next_inference_id() == 1 + + +def test_close_stops_recording_and_resets_counter(recorder, tmp_path): + recorder.open_simulation_csv(tmp_path / "sim.csv") + recorder.next_inference_id() + + recorder.close_simulation_csv() + + assert recorder.is_recording is False + assert recorder.next_inference_id() == 1 + + +def test_close_is_idempotent(recorder, tmp_path): + recorder.open_simulation_csv(tmp_path / "sim.csv") + recorder.close_simulation_csv() + recorder.close_simulation_csv() + + assert recorder.is_recording is False + + +def test_write_after_close_is_dropped_not_raised(recorder, tmp_path): + # Regression: rows still queued on the background I/O thread used to reach a + # closed file and raise "I/O operation on closed file". + path = tmp_path / "sim.csv" + recorder.open_simulation_csv(path) + recorder.close_simulation_csv() + + assert recorder.write_row(_row(1)) is False + + +def test_failed_open_leaves_no_half_open_state(recorder, tmp_path): + assert recorder.open_simulation_csv(tmp_path / "missing-dir" / "sim.csv") is False + + assert recorder.is_recording is False + assert recorder.write_row(_row(1)) is False + + +def test_write_row_survives_a_bad_row(recorder, tmp_path): + recorder.open_simulation_csv(tmp_path / "sim.csv") + + # DictWriter raises on keys outside fieldnames. That must be logged and + # swallowed: a bad telemetry row must not kill the background I/O worker. + assert recorder.write_row({"not_a_column": 1}) is False + # The recorder stays usable afterwards. + assert recorder.write_row(_row(1)) is True + + +def test_concurrent_writes_and_close_never_raise(recorder, tmp_path): + # Regression: `close_simulation_csv` cleared the handle while the background + # thread was inside the write, with no lock between them. + recorder.open_simulation_csv(tmp_path / "sim.csv") + barrier = threading.Barrier(9) + errors = [] + + def writer(): + try: + barrier.wait() + for _ in range(100): + recorder.write_row(_row(recorder.next_inference_id())) + except Exception as exc: + errors.append(exc) + + def closer(): + try: + barrier.wait() + for _ in range(20): + recorder.close_simulation_csv() + recorder.open_simulation_csv(tmp_path / "sim.csv") + except Exception as exc: + errors.append(exc) + + threads = [threading.Thread(target=writer) for _ in range(8)] + threads.append(threading.Thread(target=closer)) + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert errors == [] + + +def test_concurrent_next_inference_id_hands_out_unique_ids(recorder, tmp_path): + recorder.open_simulation_csv(tmp_path / "sim.csv") + ids = [] + lock = threading.Lock() + barrier = threading.Barrier(8) + + def grab(): + barrier.wait() + mine = [recorder.next_inference_id() for _ in range(50)] + with lock: + ids.extend(mine) + + threads = [threading.Thread(target=grab) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert sorted(ids) == list(range(1, 401)) + + +def test_close_clears_state_even_if_the_handle_fails_to_close(recorder, tmp_path): + recorder.open_simulation_csv(tmp_path / "sim.csv") + + class ExplodingFile: + def close(self): + raise OSError("disk went away") + + recorder._csv_file = ExplodingFile() + + recorder.close_simulation_csv() + + # A failing close must not strand the recorder in a recording state. + assert recorder.is_recording is False + + +def test_reset_debug_folder_survives_an_unremovable_folder(recorder, monkeypatch): + def boom(path): + raise PermissionError("nope") + + monkeypatch.setattr( + "server.communication.inference_recorder.shutil.rmtree", boom + ) + recorder.reset_debug_folder() + + # Logged and carried on: the folder still exists for the next write. + recorder.reset_debug_folder() + recorder.save_debug_files("dev-1", {}, {}) + + +def test_reset_debug_folder_creates_and_empties(recorder, tmp_path): + debug_dir = tmp_path / "debug" + recorder.reset_debug_folder() + stale = debug_dir / "stale.json" + stale.write_text("{}") + + recorder.reset_debug_folder() + + assert debug_dir.is_dir() + assert not stale.exists() + + +def test_save_debug_files_writes_both_tables(recorder, tmp_path): + recorder.reset_debug_folder() + + recorder.save_debug_files("dev-1", {"layer_0": 1.5}, {"layer_0": 0.25}) + + debug_dir = tmp_path / "debug" + assert json.loads((debug_dir / "dev-1_device_times.json").read_text()) == { + "layer_0": 1.5 + } + assert json.loads((debug_dir / "dev-1_edge_times.json").read_text()) == { + "layer_0": 0.25 + }