refactor(perception): share bounded streaming scheduler
This commit is contained in:
DCCONSTRUCTIONS
@@ -1,74 +1,5 @@
|
||||
"""One producer, one graph, bounded pending queue and active payload accounting."""
|
||||
"""Compatibility import; the bounded mailbox belongs to the shared runtime."""
|
||||
|
||||
from collections import deque
|
||||
from threading import Condition
|
||||
from k1link.perception.streaming_queue import StreamMailbox as Mailbox
|
||||
|
||||
|
||||
class Mailbox:
|
||||
def __init__(self, capacity=2, byte_limit=16 * 1024 * 1024):
|
||||
self.capacity = capacity
|
||||
self.byte_limit = byte_limit
|
||||
self.condition = Condition()
|
||||
self.pending = deque()
|
||||
self.bytes = 0
|
||||
self.peak_bytes = 0
|
||||
self.peak_pending = 0
|
||||
self.external_pending = 0
|
||||
self.dropped = []
|
||||
self.done = False
|
||||
self.error = None
|
||||
|
||||
def put(self, bundle):
|
||||
size = bundle["payload_bytes"]
|
||||
with self.condition:
|
||||
while self.pending and (
|
||||
len(self.pending) + self.external_pending >= self.capacity
|
||||
or self.bytes + size > self.byte_limit
|
||||
):
|
||||
old = self.pending.popleft()
|
||||
self.bytes -= old["payload_bytes"]
|
||||
self.dropped.append({"sequence": old["sequence"], "reason": "pending-overflow"})
|
||||
if len(self.pending) + self.external_pending >= self.capacity:
|
||||
self.dropped.append({"sequence": bundle["sequence"], "reason": "pending-overflow"})
|
||||
return
|
||||
if self.bytes + size > self.byte_limit:
|
||||
self.dropped.append({"sequence": bundle["sequence"], "reason": "byte-budget"})
|
||||
return
|
||||
self.pending.append(bundle)
|
||||
self.bytes += size
|
||||
self.peak_bytes = max(self.peak_bytes, self.bytes)
|
||||
self.peak_pending = max(self.peak_pending, len(self.pending) + self.external_pending)
|
||||
self.condition.notify_all()
|
||||
|
||||
def reserve_completed(self):
|
||||
"""A finished GPU result shares the SAME pending budget as ingress."""
|
||||
with self.condition:
|
||||
if self.external_pending:
|
||||
raise ValueError("only one completed GPU slot is permitted")
|
||||
while len(self.pending) + 1 > self.capacity:
|
||||
old = self.pending.popleft()
|
||||
self.bytes -= old["payload_bytes"]
|
||||
self.dropped.append({"sequence": old["sequence"], "reason": "handoff-overflow"})
|
||||
self.external_pending = 1
|
||||
self.peak_pending = max(self.peak_pending, len(self.pending) + 1)
|
||||
|
||||
def take_completed(self):
|
||||
with self.condition:
|
||||
if self.external_pending != 1:
|
||||
raise ValueError("completed GPU slot accounting mismatch")
|
||||
self.external_pending = 0
|
||||
|
||||
def take(self):
|
||||
with self.condition:
|
||||
self.condition.wait_for(lambda: self.pending or self.done)
|
||||
return self.pending.popleft() if self.pending else None
|
||||
|
||||
def release(self, bundle):
|
||||
with self.condition:
|
||||
self.bytes -= bundle["payload_bytes"]
|
||||
|
||||
def finish(self, error=None):
|
||||
with self.condition:
|
||||
self.done = True
|
||||
self.error = error
|
||||
self.condition.notify_all()
|
||||
__all__ = ["Mailbox"]
|
||||
|
||||
@@ -1,91 +1,5 @@
|
||||
"""One serial GPU stage can overlap one chronological CPU fusion stage.
|
||||
"""Compatibility import; no separate per-LAB GPU scheduling implementation."""
|
||||
|
||||
No GPU model concurrency. Ingress and completed-GPU results share two pending
|
||||
slots. Active payloads remain in the shared mailbox byte budget.
|
||||
"""
|
||||
from k1link.perception.streaming_scheduler import SerialGpuStage as GpuStage
|
||||
|
||||
import threading
|
||||
import traceback
|
||||
|
||||
|
||||
class GpuStage:
|
||||
def __init__(self, mailbox, compute, stop):
|
||||
self.mailbox = mailbox
|
||||
self.compute = compute
|
||||
self.stop = stop
|
||||
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()
|
||||
|
||||
def _run(self):
|
||||
try:
|
||||
while not self.stop.is_set():
|
||||
if not self.output_slot.acquire(timeout=0.05):
|
||||
continue
|
||||
bundle = self.mailbox.take()
|
||||
if bundle is None:
|
||||
break
|
||||
result = self.compute(bundle)
|
||||
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:
|
||||
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()
|
||||
self.mailbox.finish(self.mailbox.error)
|
||||
self.thread.join(timeout=2)
|
||||
return not self.thread.is_alive()
|
||||
__all__ = ["GpuStage"]
|
||||
|
||||
Reference in New Issue
Block a user