yyc12345/BasaltMeter
1
0
Fork 0
Code Activity

fix while stopper

Browse Source
This commit is contained in:
yyc12345
2026-01-08 20:04:24 +08:00
parent bfaee2028e
commit 886e0caab2
6 changed files with 65 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
BasaltTrainer/.style.yapf
View File

@@ -1,2 +0,0 @@
# The column limit.
column_limit=79

26
BasaltTrainer/cmd_server.py
View File

@@ -68,15 +68,15 @@ class CmdServer:
# Send handshake request to Presenter # Send handshake request to Presenter
self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.HANDSHAKE_REQUEST)) self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.HANDSHAKE_REQUEST))
# And the payload data # And the payload data
self.__pipe_operator.write(HANDSHAKE_REQUEST_PACKER.pack( self.__pipe_operator.write(
payload.pixel_kind, HANDSHAKE_REQUEST_PACKER.pack(
payload.width, payload.pixel_kind, payload.width, payload.height
payload.height )
)) )
# Wait for handshake response from Presenter (code 0x62) # Wait for handshake response from Presenter (code 0x62)
code_bytes = self.__pipe_operator.read(CODE_PACKER.size) code_bytes = self.__pipe_operator.read(CODE_PACKER.size)
(code, ) = CODE_PACKER.unpack(code_bytes) (code,) = CODE_PACKER.unpack(code_bytes)
if ProtocolCode(code) != ProtocolCode.HANDSHAKE_RESPONSE: if ProtocolCode(code) != ProtocolCode.HANDSHAKE_RESPONSE:
raise RuntimeError("expect handshake response code, but got another") raise RuntimeError("expect handshake response code, but got another")
@@ -84,8 +84,7 @@ class CmdServer:
self.__status = ServerStatus.Running self.__status = ServerStatus.Running
return return
def tick(self, data_receiver: Callable[[], None], def tick(self, data_receiver: Callable[[], None], request_stop: bool) -> bool:
request_stop: bool) -> bool:
""" """
Tick function called every frame to wait for data ready from Presenter and send response. 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. Returns True if a stop code was received (meaning the process should stop), False otherwise.
@@ -101,7 +100,7 @@ class CmdServer:
while True: while True:
# Wait for code from Presenter # Wait for code from Presenter
code_bytes = self.__pipe_operator.read(CODE_PACKER.size) code_bytes = self.__pipe_operator.read(CODE_PACKER.size)
(code, ) = CODE_PACKER.unpack(code_bytes) (code,) = CODE_PACKER.unpack(code_bytes)
# Analyse code # Analyse code
match ProtocolCode(code): match ProtocolCode(code):
@@ -109,7 +108,9 @@ class CmdServer:
# Receive data # Receive data
data_receiver() data_receiver()
# Send data received symbol # Send data received symbol
self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.DATA_RECEIVED)) self.__pipe_operator.write(
CODE_PACKER.pack(ProtocolCode.DATA_RECEIVED)
)
return False return False
case ProtocolCode.ACTIVELY_STOP: case ProtocolCode.ACTIVELY_STOP:
# Presenter requested stop. # Presenter requested stop.
@@ -119,7 +120,6 @@ class CmdServer:
case _: case _:
raise RuntimeError("unexpected protocol code when running") raise RuntimeError("unexpected protocol code when running")
def __wait_stop(self) -> None: def __wait_stop(self) -> None:
# Send stop request code # Send stop request code
self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP_REQUEST)) self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP_REQUEST))
@@ -128,7 +128,7 @@ class CmdServer:
while True: while True:
# Accept code # Accept code
code_bytes = self.__pipe_operator.read(CODE_PACKER.size) code_bytes = self.__pipe_operator.read(CODE_PACKER.size)
(code, ) = CODE_PACKER.unpack(code_bytes) (code,) = CODE_PACKER.unpack(code_bytes)
# Check whether it is stop response # Check whether it is stop response
match ProtocolCode(code): match ProtocolCode(code):
@@ -138,5 +138,3 @@ class CmdServer:
return return
case _: case _:
raise RuntimeError("unexpected protocol code when waiting quit") raise RuntimeError("unexpected protocol code when waiting quit")

33
BasaltTrainer/main.py
View File

@@ -1,39 +1,36 @@
from cmd_server import CmdServer, HandshakePayload, PixelKind from cmd_server import CmdServer, HandshakePayload, PixelKind
from while_stopper import WhileStopper from while_stopper import INSTANCE as STOPPER
import logging import logging
import signal
def receive_data() -> None: def receive_data() -> None:
logging.info('Data received') logging.info("Data received")
def main(): def main():
server = CmdServer() server = CmdServer()
print('Please launch BasaltPresenter now.') print("Please launch BasaltPresenter now.")
print('Then press Enter to continue...') print("Then press Enter to continue...")
input() input()
logging.info('Waiting BasaltPresenter...') logging.info("Waiting BasaltPresenter...")
server.wait_handshake(HandshakePayload(PixelKind.GRAY_U8, 600, 600)) server.wait_handshake(HandshakePayload(PixelKind.GRAY_U8, 600, 600))
logging.info('Start to running.') logging.info("Start to running.")
stopper = WhileStopper() STOPPER.register()
stopper.register()
while True: while True:
if server.tick(receive_data, stopper.is_stop_requested()): if server.tick(receive_data, STOPPER.is_stop_requested()):
break break
stopper.unregister() STOPPER.unregister()
logging.info('Program stop.') logging.info("Program stop.")
print('Press Enter to exit...')
input()
if __name__ == "__main__": if __name__ == "__main__":
logging.basicConfig(level=logging.DEBUG, logging.basicConfig(
format='[%(asctime)s] [%(levelname)s] %(message)s', level=logging.DEBUG,
datefmt='%Y-%m-%d %H:%M:%S') format="[%(asctime)s] [%(levelname)s] %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
main() main()

11
BasaltTrainer/pipe_operator.py
View File

@@ -35,13 +35,14 @@ class PipeOperator:
self.pipe_handle = win32pipe.CreateNamedPipe( self.pipe_handle = win32pipe.CreateNamedPipe(
self.pipe_name, self.pipe_name,
win32pipe.PIPE_ACCESS_DUPLEX, win32pipe.PIPE_ACCESS_DUPLEX,
win32pipe.PIPE_TYPE_BYTE | win32pipe.PIPE_READMODE_BYTE win32pipe.PIPE_TYPE_BYTE
| win32pipe.PIPE_READMODE_BYTE
| win32pipe.PIPE_WAIT, | win32pipe.PIPE_WAIT,
1, # Number of pipe instances 1, # Number of pipe instances
65536, # Output buffer size 65536, # Output buffer size
65536, # Input buffer size 65536, # Input buffer size
0, # Default timeout 0, # Default timeout
None # Security attributes None, # Security attributes
) )
if self.pipe_handle == win32file.INVALID_HANDLE_VALUE: if self.pipe_handle == win32file.INVALID_HANDLE_VALUE:
raise RuntimeError("Failed to create named pipe.") raise RuntimeError("Failed to create named pipe.")
@@ -96,8 +97,7 @@ class PipeOperator:
data: Any data: Any
(hr, data) = win32file.ReadFile(self.pipe_handle, size) (hr, data) = win32file.ReadFile(self.pipe_handle, size)
if hr != 0: # Error occurred if hr != 0: # Error occurred
raise RuntimeError( raise RuntimeError(f"Failed to read from pipe, error code: {hr}")
f"Failed to read from pipe, error code: {hr}")
data = bytes(data) data = bytes(data)
if len(data) != size: if len(data) != size:
raise RuntimeError( raise RuntimeError(
@@ -147,8 +147,7 @@ class PipeOperator:
total_written = 0 total_written = 0
while total_written < len(data): while total_written < len(data):
try: try:
bytes_written = os.write(self.pipe_handle, bytes_written = os.write(self.pipe_handle, data[total_written:])
data[total_written:])
total_written += bytes_written total_written += bytes_written
except OSError as e: except OSError as e:
raise RuntimeError(f"Failed to write to named pipe: {e}") raise RuntimeError(f"Failed to write to named pipe: {e}")

2
BasaltTrainer/pyproject.toml
View File

@@ -26,3 +26,5 @@ name = "pytorch-cu126"
url = "https://download.pytorch.org/whl/cu126" url = "https://download.pytorch.org/whl/cu126"
explicit = true explicit = true
[too.ruff]
line-length = 88

51
BasaltTrainer/while_stopper.py
View File

@@ -9,44 +9,49 @@ else:
class WhileStopper: class WhileStopper:
_is_registered: bool
is_registered: bool _stop_requested: bool
stop_requested: bool
def __init__(self) -> None: def __init__(self) -> None:
self.is_registered = False self._is_registered = False
self.stop_requested = False self._stop_requested = False
def is_stop_requested(self) -> bool: def is_stop_requested(self) -> bool:
if not self.is_registered: if not self._is_registered:
raise RuntimeError('unexpected stopper status') raise RuntimeError("unexpected stopper status")
return self.stop_requested return self._stop_requested
def register(self): def register(self):
if self.is_registered: if self._is_registered:
raise RuntimeError('unexpected stopper status') raise RuntimeError("unexpected stopper status")
self.stop_requested = False self._stop_requested = False
if IS_WINDOWS: if IS_WINDOWS:
win32api.SetConsoleCtrlHandler(self.__win_handler, True) win32api.SetConsoleCtrlHandler(_win_handler, True)
else: else:
signal.signal(signal.SIGINT, self.__posix_handler) signal.signal(signal.SIGINT, _posix_handler)
self.is_registered = True self._is_registered = True
def unregister(self): def unregister(self):
if not self.is_registered: if not self._is_registered:
raise RuntimeError('unexpected stopper status') raise RuntimeError("unexpected stopper status")
if IS_WINDOWS: if IS_WINDOWS:
win32api.SetConsoleCtrlHandler(None, False) win32api.SetConsoleCtrlHandler(_win_handler, False)
else: else:
signal.signal(signal.SIGINT, signal.SIG_DFL) signal.signal(signal.SIGINT, signal.SIG_DFL)
self.is_registered = False self._is_registered = False
def __win_handler(self, ctrl_type: int) -> bool:
self.stop_requested = True
return True
def __posix_handler(self, signal, frame) -> None: INSTANCE: WhileStopper = WhileStopper()
self.stop_requested = True """The singleton of WhileStopper"""
def _win_handler(ctrl_type: int) -> bool:
INSTANCE._stop_requested = True
return True
def _posix_handler(signal, frame) -> None:
INSTANCE._stop_requested = True