59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import mmap
|
|
import os
|
|
import pickle
|
|
import struct
|
|
|
|
from iqpilot.common.swaglog import cloudlog
|
|
|
|
SMALL_CHANNEL = "/dev/shm/iqpilot_smallmodel"
|
|
BIG_CHANNEL = "/dev/shm/iqpilot_bigmodel"
|
|
SHM_SIZE = 8 * 1024 * 1024
|
|
HEADER = struct.Struct("<QqQ")
|
|
|
|
|
|
class ModelChannel:
|
|
def __init__(self, path: str, create: bool):
|
|
if create:
|
|
fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o600)
|
|
os.ftruncate(fd, SHM_SIZE)
|
|
else:
|
|
fd = os.open(path, os.O_RDWR)
|
|
self.mm = mmap.mmap(fd, SHM_SIZE)
|
|
os.close(fd)
|
|
if create:
|
|
self.mm[:HEADER.size] = HEADER.pack(0, -1, 0)
|
|
|
|
def write(self, frame_id: int, payload: dict) -> None:
|
|
data = pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
|
|
if HEADER.size + len(data) > SHM_SIZE:
|
|
cloudlog.error(f"model payload {len(data)} bytes exceeds shm {SHM_SIZE}, dropping frame {frame_id}")
|
|
return
|
|
seq = HEADER.unpack(self.mm[:HEADER.size])[0]
|
|
HEADER.pack_into(self.mm, 0, seq + 1, frame_id, len(data))
|
|
self.mm[HEADER.size:HEADER.size + len(data)] = data
|
|
HEADER.pack_into(self.mm, 0, seq + 2, frame_id, len(data))
|
|
|
|
def peek_frame_id(self) -> int | None:
|
|
seq, frame_id, length = HEADER.unpack(self.mm[:HEADER.size])
|
|
if seq == 0 or seq % 2 != 0 or length == 0:
|
|
return None
|
|
return frame_id
|
|
|
|
def read(self) -> tuple[int, dict] | None:
|
|
seq1, frame_id, length = HEADER.unpack(self.mm[:HEADER.size])
|
|
if seq1 == 0 or seq1 % 2 != 0 or length == 0:
|
|
return None
|
|
data = bytes(self.mm[HEADER.size:HEADER.size + length])
|
|
seq2 = HEADER.unpack(self.mm[:HEADER.size])[0]
|
|
if seq1 != seq2:
|
|
return None
|
|
try:
|
|
return frame_id, pickle.loads(data)
|
|
except Exception:
|
|
return None
|