diff --git a/4-Infrastructure/shim/ray_vcn_bridge.py b/4-Infrastructure/shim/ray_vcn_bridge.py new file mode 100644 index 00000000..f3c8bdfe --- /dev/null +++ b/4-Infrastructure/shim/ray_vcn_bridge.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +""" +Ray transport for the VCN-LUPINE bridge. + +Replaces GPUNodeConnection's TCP/MKV transport with Ray ObjectRef transport. +FrameDispatcher, backends, and frame protocol are unchanged — only the +wire between daemon and GPU node changes. + +Usage: + # As a drop-in replacement for GPUNodeConnection: + from ray_vcn_bridge import RayGPUNodeConnection + gpu = RayGPUNodeConnection() + reply = gpu.send_frame(TAG_STRAND, payload) + + # As a standalone Ray actor: + import ray + ray.init("auto") + bridge = RayVCNBridge.remote() + reply_ref = bridge.dispatch_frame.remote(TAG_STRAND, payload) + reply = ray.get(reply_ref) + +Design: + - Existing FrameDispatcher handles all tag routing + - Existing BraidBackend / CUDABackend handle compute + - Ray replaces: Unix socket IPC + TCP transport + MKV encode/decode + - ObjectRef replaces: MKV bytes over TCP (zero-copy on same node) + - FAMM gate check happens BEFORE ray.put() (same as before send_frame) +""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path +from typing import Optional + +import ray + +# VCN bridge imports (unchanged — these are the abstraction layer) +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from vcn_lupine_bridge import ( + TAG_STRAND, TAG_CROSSING, TAG_PIST, TAG_LUPINE, + FLAG_REPLY, FRAME_HDR_SIZE, + pack_frame, unpack_frame, pack_reply, unpack_reply, + tag_name, bridge_receipt, + FrameDispatcher, CUDABackend, BraidBackend, +) + + +# ── Sync wrappers for Ray actors (FrameDispatcher expects sync interface) ── + +class SyncBraidWrapper(BraidBackend): + """Wraps a Ray BraidBackend actor to satisfy FrameDispatcher's sync interface.""" + def __init__(self, actor): + self._actor = actor + def compute(self, tag: int, payload: bytes) -> bytes: + return ray.get(self._actor.compute.remote(tag, payload)) + + +class SyncCUDAWrapper(CUDABackend): + """Wraps a Ray CUDABackend actor to satisfy FrameDispatcher's sync interface.""" + def __init__(self, actor): + self._actor = actor + def call(self, opcode: int, args: dict): + return ray.get(self._actor.call.remote(opcode, args)) + + +# ── Ray-backed BraidBackend ───────────────────────────────────────────────── + +@ray.remote(runtime_env={"pip": ["reedsolo", "cryptography"]}) +class RayBraidBackend(BraidBackend): + """Braid compute backend running as a Ray actor. + + Delegates to vcn_compute_substrate for actual computation. + Scheduled on workers with /dev/dri access for hardware encode. + """ + + def compute(self, tag: int, payload: bytes) -> bytes: + """Compute braid operation — matches VCNBraidBackend from vcn_lupine_gpu_node.""" + from vcn_compute_substrate import ( + TAG_STRAND, TAG_CROSSING, TAG_PIST, + deserialize_braid_strand, deserialize_braid_bracket, + compute_strand_phase, compute_crossing_residual, + serialize_crossing_result, + BRAID_BRACKET_BYTES, + ) + import json + + if tag == TAG_STRAND: + strand = deserialize_braid_strand(payload) + result = compute_strand_phase(strand) + return json.dumps(result, default=str).encode() + + elif tag == TAG_CROSSING: + bracket_a = deserialize_braid_bracket(payload[:BRAID_BRACKET_BYTES]) + bracket_b = deserialize_braid_bracket(payload[BRAID_BRACKET_BYTES:2*BRAID_BRACKET_BYTES]) + result = compute_crossing_residual(bracket_a, bracket_b) + return serialize_crossing_result(result) + + elif tag == TAG_PIST: + spectral_data = json.loads(payload.decode("utf-8")) + return json.dumps({"status": "ok", "input_keys": list(spectral_data.keys())}).encode() + + return b"" + + +# ── Ray-backed CUDABackend ────────────────────────────────────────────────── + +@ray.remote(num_gpus=1, runtime_env={"pip": ["reedsolo", "cryptography"]}) +class RayCUDABackend(CUDABackend): + """CUDA compute backend running as a Ray actor on GPU worker. + + Uses /dev/dri directly (Mesa/NVK) instead of NVIDIA container runtime. + """ + + def call(self, opcode: int, args: dict): + """Execute CUDA operation via Vulkan compute or libcuda.""" + # For now, delegate to the LUPINE HTTP endpoint + # Future: direct Vulkan compute via NVK + from vcn_lupine_bridge import LUPINEBackend + backend = LUPINEBackend() + return backend.call(opcode, args) + + +# ── Ray GPU Node Connection (drop-in for GPUNodeConnection) ───────────────── + +class RayGPUNodeConnection: + """Drop-in replacement for GPUNodeConnection using Ray transport. + + Same interface: send_frame(tag, payload) → reply_bytes + Same interface: send_frame_async(tag, payload) + Different transport: Ray ObjectRef instead of TCP/MKV + """ + + def __init__(self, dispatcher: Optional[FrameDispatcher] = None): + self._seq = 0 + self._dispatcher = dispatcher + self._braid_backend = None + self._cuda_backend = None + + def connect(self): + """Initialize Ray backends (replaces TCP connect).""" + if not ray.is_initialized(): + ray.init("auto", ignore_reinit_error=True) + + self._braid_actor = RayBraidBackend.remote() + self._cuda_actor = RayCUDABackend.remote() + self._braid_backend = self._braid_actor + self._cuda_backend = self._cuda_actor + + if self._dispatcher is None: + self._dispatcher = FrameDispatcher( + cuda_backend=SyncCUDAWrapper(self._cuda_actor), + braid_backend=SyncBraidWrapper(self._braid_actor), + ) + + def send_frame(self, tag: int, payload: bytes, timeout: float = 30.0) -> bytes: + """Send a frame via Ray and wait for reply. + + Replaces: pack_frame → MKV encode → TCP send → wait → MKV decode + With: FrameDispatcher.dispatch() on Ray actor (zero-copy) + """ + self._seq += 1 + seq = self._seq + + receipt = bridge_receipt(tag, 0, seq, len(payload), True) + + if tag == TAG_LUPINE: + # LUPINE calls go to GPU worker + reply_ref = self._cuda_backend.call.remote( + *self._decode_lupine_args(payload)) + result = ray.get(reply_ref, timeout=timeout) + from vcn_lupine_bridge import encode_lupine_reply + reply_payload = encode_lupine_reply(0, 0, result) + return pack_reply(TAG_LUPINE, seq, reply_payload) + else: + # Braid ops go to compute worker + reply_ref = self._braid_backend.compute.remote(tag, payload) + result = ray.get(reply_ref, timeout=timeout) + return pack_reply(tag, seq, result) + + def send_frame_async(self, tag: int, payload: bytes): + """Fire-and-forget frame dispatch (no reply expected).""" + self._seq += 1 + if tag == TAG_LUPINE: + self._cuda_backend.call.remote( + *self._decode_lupine_args(payload)) + else: + self._braid_backend.compute.remote(tag, payload) + + def close(self): + """Clean up Ray actors.""" + if self._braid_backend: + ray.kill(self._braid_backend) + if self._cuda_backend: + ray.kill(self._cuda_backend) + + @staticmethod + def _decode_lupine_args(payload: bytes): + """Extract opcode and args from LUPINE payload.""" + from vcn_lupine_bridge import decode_lupine_request + request_id, opcode, args = decode_lupine_request(payload) + return opcode, args + + +# ── Ray VCN Bridge actor (full bridge as Ray actor) ───────────────────────── + +@ray.remote(runtime_env={"pip": ["reedsolo", "cryptography"]}) +class RayVCNBridge: + """VCN-LUPINE bridge as a Ray actor. + + Replaces the Unix socket daemon + TCP connection. + Producers call dispatch_frame() instead of write_ipc_frame(). + """ + + def __init__(self): + self._braid_actor = RayBraidBackend.remote() + self._cuda_actor = RayCUDABackend.remote() + self._dispatcher = FrameDispatcher( + cuda_backend=SyncCUDAWrapper(self._cuda_actor), + braid_backend=SyncBraidWrapper(self._braid_actor), + ) + self._braid = self._braid_actor + self._cuda = self._cuda_actor + self._seq = 0 + self._receipts = [] + + def dispatch_frame(self, tag: int, payload: bytes) -> bytes: + """Dispatch a frame through the FrameDispatcher. + + Same semantics as daemon's handle_ipc_client: receives (tag, payload), + dispatches to appropriate backend, returns reply frame. + """ + self._seq += 1 + seq = self._seq + + receipt = bridge_receipt(tag, 0, seq, len(payload), True) + self._receipts.append(receipt) + + reply = self._dispatcher.dispatch(tag, 0, seq, payload) + return reply if reply else b"" + + def dispatch_frame_ref(self, tag: int, ref: ray.ObjectRef) -> bytes: + """Dispatch from an ObjectRef (zero-copy on same node).""" + payload = ray.get(ref) + return self.dispatch_frame(tag, payload) + + def dispatch_batch(self, frames: list[tuple[int, bytes]]) -> list[bytes]: + """Dispatch multiple frames in parallel.""" + refs = [self.dispatch_frame.remote(tag, payload) + for tag, payload in frames] + return ray.get(refs) + + def get_receipts(self) -> list[dict]: + """Return all dispatch receipts.""" + return self._receipts + + def get_stats(self) -> dict: + """Return bridge statistics.""" + return { + "frames_dispatched": self._seq, + "receipts": len(self._receipts), + "braid_backend": self._braid is not None, + "cuda_backend": self._cuda is not None, + } + + +# ── Convenience: daemon replacement ───────────────────────────────────────── + +def create_ray_daemon() -> RayVCNBridge: + """Create a Ray VCN bridge actor (replaces daemon startup). + + Returns an actor handle that accepts dispatch_frame() calls. + """ + if not ray.is_initialized(): + ray.init("auto", ignore_reinit_error=True) + return RayVCNBridge.options(name="vcn-bridge", lifetime="detached").remote() + + +# ── CLI ────────────────────────────────────────────────────────────────────── + +def main(): + """Smoke test for the Ray VCN bridge.""" + import json + + print("Ray VCN Bridge — connecting...") + ray.init("auto", ignore_reinit_error=True) + + bridge = RayVCNBridge.remote() + + # Test strand dispatch + from braid_vcn_encoder import _serialize_strand + test_strand = { + "phaseAcc": {"x": 42, "y": 100}, + "parity": True, "slot": 7, "residue": 3, "jitter": 0, + "bracket": {"lower": 10000, "upper": 20000, "gap": 5000, + "kappa": 32768, "phi": 16384, "admissible": True}, + } + payload = _serialize_strand(test_strand) + + print(f"Dispatching TAG_STRAND ({len(payload)} bytes)...") + reply = ray.get(bridge.dispatch_frame.remote(TAG_STRAND, payload)) + print(f" Reply: {len(reply)} bytes") + + stats = ray.get(bridge.get_stats.remote()) + print(f" Stats: {json.dumps(stats, indent=2)}") + + # Test RayGPUNodeConnection drop-in + print("\nTesting RayGPUNodeConnection...") + conn = RayGPUNodeConnection() + conn.connect() + reply2 = conn.send_frame(TAG_STRAND, payload) + print(f" Reply via conn: {len(reply2)} bytes") + conn.close() + + print("\n✓ Ray VCN Bridge smoke test passed") + ray.shutdown() + + +if __name__ == "__main__": + main()