NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-25 17:10:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(shm): bugfix and code cleanup

Browse source
This commit is contained in:
RUFFY-369 2026-04-06 22:37:04 +05:30
parent 43a5cdcdfc
commit 48dcd64299
5 changed files with 38 additions and 83 deletions
Download patch file Download diff file Expand all files Collapse all files

40
atroposlib/api/shm_buffer.py
View file

@ -27,8 +27,7 @@ class SHMBufferConfig:
class ZeroCopySHMBuffer:
"""
High-performance circular buffer using multiprocessing.shared_memory.
Eliminates JSON serialization and HTTP overhead for trajectory transport.
Now expanded with TrajectoryID and metadata slots for universal Atropos use.
Eliminates serialization and HTTP overhead for trajectory transport.
"""
def __init__(
@ -96,9 +95,11 @@ class ZeroCopySHMBuffer:
raise ValueError("Invalid SHM Magic")
return read_idx, write_idx, max_size, entry_size
def _set_indices(self, read_idx: int, write_idx: int):
# We only update these two fields (Offsets: ReadIdx=8, WriteIdx=12)
struct.pack_into("II", self.buf, 8, read_idx, write_idx)
def _set_read_idx(self, idx: int):
struct.pack_into("I", self.buf, 8, idx)
def _set_write_idx(self, idx: int):
struct.pack_into("I", self.buf, 12, idx)
def write_trajectory(
self,
@ -122,36 +123,29 @@ class ZeroCopySHMBuffer:
# Calculate offset in data segment
offset = SHMBufferConfig.SIZE + (write_idx * self.slot_size)
# write Score (8)
# Pack Metadata and Rich attributes
struct.pack_into("d", self.buf, offset, float(score))
# write Token Length (4)
token_len = min(len(tokens), entry_size)
struct.pack_into("i", self.buf, offset + 8, token_len)
# write Instance ID (fixed len)
id_bytes = instance_id.encode('utf-8')[:self.instance_id_len]
struct.pack_into(f"{self.instance_id_len}s", self.buf, offset + 12, id_bytes)
# write Repetition ID (4)
struct.pack_into("i", self.buf, offset + 12 + self.instance_id_len, int(repetition_id))
# write Metadata (fixed len JSON)
meta_json = json.dumps(metadata or {}).encode('utf-8')[:self.metadata_len]
struct.pack_into(f"{self.metadata_len}s", self.buf, offset + 12 + self.instance_id_len + 4, meta_json)
# write Tokens (Numpy View)
# Copy tokens via Numpy View directly into SHM slot
token_offset = offset + 12 + self.instance_id_len + 4 + self.metadata_len
token_arr = np.array(tokens, dtype=np.int32)
# View the SHM as a numpy array for the specific token slot
shm_slot = np.ndarray((entry_size,), dtype=np.int32, buffer=self.buf, offset=token_offset)
shm_slot[:token_len] = token_arr[:token_len]
if token_len < entry_size:
shm_slot[token_len:] = 0 # Padding
shm_slot[token_len:] = 0
# Update write index
self._set_indices(read_idx, next_write)
self._set_write_idx(next_write)
return True
def read_next(self) -> Optional[Dict[str, Any]]:
@ -165,31 +159,25 @@ class ZeroCopySHMBuffer:
offset = SHMBufferConfig.SIZE + (read_idx * self.slot_size)
# Read Score and Token Length
# Unpack Metadata and Rich attributes
score = struct.unpack_from("d", self.buf, offset)[0]
token_len = struct.unpack_from("i", self.buf, offset + 8)[0]
token_len = min(token_len, entry_size)
token_len = min(struct.unpack_from("i", self.buf, offset + 8)[0], entry_size)
# Read Instance ID
id_bytes = struct.unpack_from(f"{self.instance_id_len}s", self.buf, offset + 12)[0]
instance_id = id_bytes.decode('utf-8', errors='ignore').strip('\x00')
# Read Repetition ID
repetition_id = struct.unpack_from("i", self.buf, offset + 12 + self.instance_id_len)[0]
# Read Metadata
meta_bytes = struct.unpack_from(f"{self.metadata_len}s", self.buf, offset + 12 + self.instance_id_len + 4)[0]
try:
metadata = json.loads(meta_bytes.decode('utf-8', errors='ignore').strip('\x00'))
except:
except (json.JSONDecodeError, UnicodeDecodeError):
metadata = {}
# Read Tokens (Numpy View)
token_offset = offset + 12 + self.instance_id_len + 4 + self.metadata_len
tokens_view = np.ndarray((token_len,), dtype=np.int32, buffer=self.buf, offset=token_offset)
# Advance read index
self._set_indices((read_idx + 1) % max_size, write_idx)
self._set_read_idx((read_idx + 1) % max_size)
return {
"tokens": tokens_view.tolist(),