fix(perception): rendezvous completed gpu handoff
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user