write shit
This commit is contained in:
113
BasaltTrainer/command_server.py
Normal file
113
BasaltTrainer/command_server.py
Normal file
@@ -0,0 +1,113 @@
|
||||
import struct
|
||||
from enum import IntEnum, Enum, auto
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple, Optional
|
||||
from pipe_operator import PipeOperator
|
||||
|
||||
PIPE_NAME: str = "ed0e3f1f-d214-4880-9562-640bce15e72e"
|
||||
|
||||
class ProtocolCode(IntEnum):
|
||||
HANDSHAKE_REQUEST = 0x61 # Trainer -> Presenter
|
||||
HANDSHAKE_RESPONSE = 0x62 # Presenter -> Trainer
|
||||
DATA_READY = 0x01 # Presenter -> Trainer
|
||||
DATA_RECEIVED = 0x02 # Trainer -> Presenter
|
||||
REQUEST_STOP = 0x71 # Presenter -> Trainer (request stop)
|
||||
STOP = 0x71 # Trainer -> Presenter (confirm stop)
|
||||
|
||||
class PixelKind(IntEnum):
|
||||
GRAY_FLOAT32 = 0x01 # Grayscale represented by one float32
|
||||
GRAY_U8 = 0x02 # Grayscale represented by one u8
|
||||
RGB_FLOAT32 = 0x03 # RGB represented by three float32
|
||||
RGB_U8 = 0x04 # RGB represented by three u8
|
||||
|
||||
@dataclass
|
||||
class ImageProperties:
|
||||
pixel_kind: PixelKind
|
||||
width: int
|
||||
height: int
|
||||
|
||||
class ServerStatus(Enum):
|
||||
Ready = auto()
|
||||
Running = auto()
|
||||
Stop = auto()
|
||||
|
||||
|
||||
CODE_PACKER: struct.Struct = struct.Struct("=B")
|
||||
HANDSHAKE_RESPONSE_PACKER: struct.Struct = struct.Struct("=BII")
|
||||
|
||||
class CommandServer:
|
||||
"""
|
||||
Command server implementation for the Trainer side according to the protocol.
|
||||
"""
|
||||
|
||||
pipe_operator: PipeOperator
|
||||
statue: ServerStatus
|
||||
|
||||
def __init__(self):
|
||||
self.pipe_operator = PipeOperator(PIPE_NAME)
|
||||
self.handshaked = ServerStatus.Ready
|
||||
|
||||
def __del__(self):
|
||||
"""Cleanup resources when object is destroyed."""
|
||||
self.pipe_operator.close()
|
||||
|
||||
def wait_handshake(self) -> ImageProperties:
|
||||
"""
|
||||
Wait for handshake from Presenter, send request first and wait for response with data properties.
|
||||
Returns a tuple of (pixel_kind, width, height) from the Presenter.
|
||||
"""
|
||||
if self.statue != ServerStatus.Ready:
|
||||
raise RuntimeError("unexpected server status")
|
||||
|
||||
# Send handshake request to Presenter (code 0x61)
|
||||
self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.HANDSHAKE_REQUEST))
|
||||
|
||||
# Wait for handshake response from Presenter (code 0x62)
|
||||
code_bytes = self.pipe_operator.read(CODE_PACKER.size)
|
||||
(code, ) = CODE_PACKER.unpack(code_bytes)
|
||||
if ProtocolCode(code) != ProtocolCode.HANDSHAKE_RESPONSE:
|
||||
raise RuntimeError("expect handshake response code, but got another")
|
||||
|
||||
# Read data properties from Presenter (pixel_kind, width, height)
|
||||
handshake_response_payload = self.pipe_operator.read(HANDSHAKE_RESPONSE_PACKER.size)
|
||||
(raw_pixel_kind, width, height) = HANDSHAKE_RESPONSE_PACKER.unpack(handshake_response_payload)
|
||||
|
||||
# Set status and return
|
||||
self.statue = ServerStatus.Running
|
||||
return ImageProperties(PixelKind(raw_pixel_kind), width, height)
|
||||
|
||||
def tick(self, request_stop: bool) -> bool:
|
||||
"""
|
||||
Tick function called every frame to wait for data ready from Presenter and send response.
|
||||
Returns True if a stop code was received (meaning the process should stop), False otherwise.
|
||||
"""
|
||||
if self.statue != ServerStatus.Running:
|
||||
raise RuntimeError("unexpected server status")
|
||||
|
||||
# If there is stop requested, we post it first and return
|
||||
if request_stop:
|
||||
self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP))
|
||||
self.statue = ServerStatus.Stop
|
||||
return True
|
||||
|
||||
# Wait for code from Presenter
|
||||
code_bytes = self.pipe_operator.read(CODE_PACKER.size)
|
||||
(code, ) = CODE_PACKER.unpack(code_bytes)
|
||||
|
||||
# Analyse code
|
||||
match ProtocolCode(code):
|
||||
case ProtocolCode.DATA_READY:
|
||||
# Receive data
|
||||
print('Data received')
|
||||
# Send data received symbol
|
||||
self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.DATA_RECEIVED))
|
||||
case ProtocolCode.REQUEST_STOP:
|
||||
# Presenter requested stop.
|
||||
# Agree with it, send code and return.
|
||||
self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP))
|
||||
self.statue = ServerStatus.Stop
|
||||
return True
|
||||
case _:
|
||||
raise RuntimeError("unexpected protocol code when running")
|
||||
|
||||
return False
|
||||
@@ -51,6 +51,10 @@ class PipeOperator:
|
||||
else:
|
||||
# POSIX implementation
|
||||
self.pipe_name = f"/tmp/{name}"
|
||||
# If there is an existing FIFO file, remove it
|
||||
if os.path.exists(self.pipe_name):
|
||||
os.unlink(self.pipe_name)
|
||||
# Create pipe
|
||||
try:
|
||||
os.mkfifo(self.pipe_name)
|
||||
except OSError as e:
|
||||
@@ -155,49 +159,3 @@ class PipeOperator:
|
||||
except OSError as e:
|
||||
raise RuntimeError(f"Failed to write to named pipe: {e}")
|
||||
|
||||
|
||||
# class PipeServer:
|
||||
# """
|
||||
# A convenience class for creating a pipe server that can handle connections in a separate thread.
|
||||
# """
|
||||
|
||||
# def __init__(self, name: str):
|
||||
# self._name = name
|
||||
# self._pipe_operator: Optional[PipeOperator] = None
|
||||
# self._server_thread: Optional[threading.Thread] = None
|
||||
# self._stop_event = threading.Event()
|
||||
|
||||
# def start_server(self, handler_func):
|
||||
# """
|
||||
# Start the pipe server in a separate thread.
|
||||
|
||||
# :param handler_func: Function to handle the connected pipe (takes PipeOperator as parameter)
|
||||
# """
|
||||
# self._server_thread = threading.Thread(target=self._server_worker,
|
||||
# args=(handler_func, ))
|
||||
# self._server_thread.start()
|
||||
|
||||
# def _server_worker(self, handler_func):
|
||||
# """Internal method to handle server operations."""
|
||||
# try:
|
||||
# # Create pipe server
|
||||
# pipe_op = PipeOperator(self._name, is_server=True)
|
||||
# self._pipe_operator = pipe_op
|
||||
|
||||
# # Call the handler function with the connected pipe
|
||||
# handler_func(pipe_op)
|
||||
# except Exception as e:
|
||||
# print(f"Error in pipe server: {e}")
|
||||
# finally:
|
||||
# if self._pipe_operator:
|
||||
# self._pipe_operator.close()
|
||||
|
||||
# def stop_server(self):
|
||||
# """Stop the pipe server."""
|
||||
# self._stop_event.set()
|
||||
# if self._server_thread:
|
||||
# self._server_thread.join()
|
||||
|
||||
# def get_pipe(self) -> Optional[PipeOperator]:
|
||||
# """Get the connected pipe operator (only valid after client connects)."""
|
||||
# return self._pipe_operator
|
||||
|
||||
Reference in New Issue
Block a user