feat(infra): Ray VCN bridge — FrameDispatcher over Ray transport

ray_vcn_bridge.py: Ray transport for the VCN-LUPINE bridge.
  Replaces GPUNodeConnection TCP/MKV transport with Ray ObjectRef.
  FrameDispatcher, BraidBackend, CUDABackend are unchanged — only
  the wire between daemon and GPU node changes.

  - RayBraidBackend: compute actor matching VCNBraidBackend pattern
  - RayCUDABackend: GPU actor with /dev/dri (Mesa, no NVIDIA plugin)
  - RayVCNBridge: full bridge as Ray actor (replaces daemon)
  - RayGPUNodeConnection: drop-in for GPUNodeConnection
  - SyncBraidWrapper/SyncCUDAWrapper: bridge Ray actors to sync interface

  STRAND 42B → 63B, CROSSING 42B → 22B, PIST 24B → 57B
  Batch 10: 5ms (0.5ms/frame), 10/10 non-empty
This commit is contained in:
Brandon Schneider 2026-05-30 19:48:34 -05:00
parent c87dfaaea5
commit 1c272eb197

View file

@ -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()