fix(perception): rendezvous completed gpu handoff

This commit is contained in:
DCCONSTRUCTIONS
2026-09-02 01:48:25 +03:00
parent 380a13688c
commit 294585936f
3 changed files with 82 additions and 12 deletions
@@ -4,7 +4,6 @@ No GPU model concurrency. Ingress and completed-GPU results share two pending
slots. Active payloads remain in the shared mailbox byte budget.
"""
import queue
import threading
import traceback
@@ -14,13 +13,18 @@ class GpuStage:
self.mailbox = mailbox
self.compute = compute
self.stop = stop
self.output = queue.Queue(maxsize=1)
self.output_condition = threading.Condition()
self.completed = None
self.completed_reserved = False
self.consumer_waiting = False
# Reserve the output slot BEFORE starting another GPU call. Otherwise a
# blocked put would hide a third pending frame outside the two queues.
self.output_slot = threading.Semaphore(1)
self.finished = threading.Event()
self.error = None
self.peak_pending = 0
self.direct_handoffs = 0
self.buffered_handoffs = 0
self.thread = threading.Thread(target=self._run, daemon=True)
self.thread.start()
@@ -33,26 +37,52 @@ class GpuStage:
if bundle is None:
break
result = self.compute(bundle)
self.mailbox.reserve_completed()
self.output.put_nowait((bundle, result))
self.peak_pending = max(self.peak_pending, self.output.qsize())
with self.output_condition:
if self.completed is not None:
raise ValueError("completed GPU slot is already occupied")
# A receiver already blocked in take() accepts the result as
# a synchronous rendezvous. It never becomes pending and
# therefore must not evict a newer ingress frame merely for
# the few instructions between publish and receive.
reserved = not self.consumer_waiting
if reserved:
self.mailbox.reserve_completed()
self.buffered_handoffs += 1
else:
self.direct_handoffs += 1
self.completed = (bundle, result)
self.completed_reserved = reserved
self.peak_pending = max(self.peak_pending, int(reserved))
self.output_condition.notify_all()
except Exception:
self.error = traceback.format_exc()
finally:
self.finished.set()
with self.output_condition:
self.output_condition.notify_all()
def take(self):
while True:
try:
result = self.output.get(timeout=0.05)
self.mailbox.take_completed()
self.output_slot.release()
return result
except queue.Empty:
if self.finished.is_set():
with self.output_condition:
if self.completed is not None:
result = self.completed
reserved = self.completed_reserved
self.completed = None
self.completed_reserved = False
self.consumer_waiting = False
elif self.finished.is_set():
self.consumer_waiting = False
if self.error:
raise RuntimeError(self.error) from None
return None
else:
self.consumer_waiting = True
self.output_condition.wait(timeout=0.05)
continue
if reserved:
self.mailbox.take_completed()
self.output_slot.release()
return result
def close(self):
self.stop.set()
@@ -487,6 +487,8 @@ def run(args):
"completed_shares_pending_limit": bool(gpu_stage),
"global_pending_limit": 2,
"gpu_peak_pending": gpu_stage.peak_pending if gpu_stage else 0,
"gpu_direct_handoffs": gpu_stage.direct_handoffs if gpu_stage else 0,
"gpu_buffered_handoffs": gpu_stage.buffered_handoffs if gpu_stage else 0,
"byte_limit": mailbox.byte_limit,
"peak_pending": mailbox.peak_pending,
"peak_bytes": mailbox.peak_bytes,
+38
View File
@@ -176,6 +176,44 @@ def test_gpu_stage_preserves_order_and_reserves_bounded_output_before_compute(pi
assert stage.close()
def test_waiting_consumer_uses_direct_handoff_without_evicting_ingress(pilot):
mailbox = pilot("pilot_queue").Mailbox(capacity=2, byte_limit=100)
compute_started = threading.Event()
release_compute = threading.Event()
def compute(bundle):
if bundle["sequence"] == 0:
compute_started.set()
assert release_compute.wait(1)
return bundle["sequence"] * 10
stage = pilot("pilot_scheduler").GpuStage(mailbox, compute, threading.Event())
bundles = [{"sequence": sequence, "payload_bytes": 10} for sequence in range(3)]
mailbox.put(bundles[0])
assert compute_started.wait(1)
mailbox.put(bundles[1])
mailbox.put(bundles[2])
mailbox.finish()
received = []
consumer = threading.Thread(target=lambda: received.append(stage.take()))
consumer.start()
with stage.output_condition:
assert stage.output_condition.wait_for(lambda: stage.consumer_waiting, timeout=1)
release_compute.set()
consumer.join(timeout=1)
assert not consumer.is_alive()
assert received[0][0]["sequence"] == 0
assert not mailbox.dropped
assert stage.direct_handoffs == 1
for expected in (1, 2):
bundle, result = stage.take()
assert (bundle["sequence"], result) == (expected, expected * 10)
for bundle in bundles:
mailbox.release(bundle)
assert stage.take() is None
assert mailbox.bytes == 0 and stage.close()
def test_gpu_stage_propagates_failure_and_stops_waiting_for_input(pilot):
mailbox = pilot("pilot_queue").Mailbox(capacity=1)