From 1f11c0b832b2bfc25255fa9d85509a4b3a4bde91 Mon Sep 17 00:00:00 2001 From: vanminhpe Date: Fri, 7 Aug 2026 10:50:04 +0800 Subject: [PATCH] feat: add Modbus drivers for Jig, TPU, and RPU --- drivers/jig.py | 40 +++++++++ drivers/modbus_device.py | 186 +++++++++++++++++++++++++++++++++++++++ drivers/rpu.py | 43 +++++++++ drivers/tpu.py | 44 +++++++++ 4 files changed, 313 insertions(+) create mode 100644 drivers/jig.py create mode 100644 drivers/modbus_device.py create mode 100644 drivers/rpu.py create mode 100644 drivers/tpu.py diff --git a/drivers/jig.py b/drivers/jig.py new file mode 100644 index 0000000..e81aaed --- /dev/null +++ b/drivers/jig.py @@ -0,0 +1,40 @@ +import time +from dataclasses import dataclass + +from drivers.modbus_device import ModbusDevice, ModbusDeviceError + + +@dataclass +class Position: + x: int + y: int + z: int + + def __str__(self): + return f"(x={self.x}, y={self.y}, z={self.z})mm" + + +class JigMovementError(ModbusDeviceError): + pass + + +class Jig(ModbusDevice): + def __init__(self, port, registers, slave_id=0x10): + super().__init__(port, "Jig", registers, slave_id=slave_id) + + def move_to_position(self, position, wait_time=5): + try: + self.write_holding_register(self._addr("X POSITION"), position.x) + self.write_holding_register(self._addr("Y POSITION"), position.y) + self.write_holding_register(self._addr("Z POSITION"), position.z) + self.write_coil(self._addr("MOVE TO POSITION"), True, verify=False) + if wait_time > 0: + time.sleep(wait_time) + except ModbusDeviceError as exc: + raise JigMovementError(f"failed to move to {position}") from exc + + def get_current_position(self): + x = self.read_holding_register(self._addr("X POSITION")) + y = self.read_holding_register(self._addr("Y POSITION")) + z = self.read_holding_register(self._addr("Z POSITION")) + return Position(x, y, z) diff --git a/drivers/modbus_device.py b/drivers/modbus_device.py new file mode 100644 index 0000000..67df95c --- /dev/null +++ b/drivers/modbus_device.py @@ -0,0 +1,186 @@ +import struct +import time + +from pymodbus.client import ModbusSerialClient +from pymodbus.exceptions import ModbusException + + +class ModbusDeviceError(Exception): + pass + + +class ModbusConnectionError(ModbusDeviceError): + pass + + +class ModbusReadError(ModbusDeviceError): + pass + + +class ModbusWriteError(ModbusDeviceError): + pass + + +class ModbusDevice: + FLOAT_VERIFY_TOLERANCE = 0.01 + + def __init__(self, port, device_name, registers=None, slave_id=0x10, baudrate=9600, timeout=3, + retries=3, retry_delay_s=0.2): + self.port = port + self.device_name = device_name + self.registers = registers or {} + self.slave_id = slave_id + self.retries = retries + self.retry_delay_s = retry_delay_s + try: + self.client = ModbusSerialClient(port=port, baudrate=baudrate, parity="N", stopbits=1, bytesize=8, timeout=timeout) + if not self.client.connect(): + raise ModbusConnectionError(f"failed to connect to {device_name} on {port}") + except ModbusException as exc: + raise ModbusConnectionError(f"error connecting to {device_name} on {port}") from exc + + def _addr(self, label): + try: + return self.registers[label].address + except KeyError: + raise ModbusDeviceError(f"unknown register label for {self.device_name}: {label}") + + def _retry(self, action): + last_exc = None + for attempt in range(self.retries): + try: + return action() + except (ModbusDeviceError, ModbusException) as exc: + last_exc = exc + if attempt < self.retries - 1: + time.sleep(self.retry_delay_s) + raise last_exc + + # ---- raw reads (retried on failure) ---- + + def read_coil(self, address): + def _do(): + res = self.client.read_coils(address, count=1, device_id=self.slave_id) + if res.isError(): + raise ModbusReadError(f"{self.device_name}: failed to read coil {address}") + return res.bits[0] + return self._retry(_do) + + def read_holding_register(self, address): + def _do(): + res = self.client.read_holding_registers(address, count=1, device_id=self.slave_id) + if res.isError(): + raise ModbusReadError(f"{self.device_name}: failed to read holding {address}") + return res.registers[0] + return self._retry(_do) + + def read_holding_register_float(self, address): + def _do(): + res = self.client.read_holding_registers(address, count=2, device_id=self.slave_id) + if res.isError(): + raise ModbusReadError(f"{self.device_name}: failed to read holding float {address}") + return struct.unpack(">f", struct.pack(">HH", res.registers[0], res.registers[1]))[0] + return self._retry(_do) + + def read_input_register(self, address): + def _do(): + res = self.client.read_input_registers(address, count=1, device_id=self.slave_id) + if res.isError(): + raise ModbusReadError(f"{self.device_name}: failed to read input {address}") + return res.registers[0] + return self._retry(_do) + + def read_signed_input_register(self, address): + val = self.read_input_register(address) + return val - 65536 if val > 32767 else val + + def read_input_register_float(self, address): + def _do(): + res = self.client.read_input_registers(address, count=2, device_id=self.slave_id) + if res.isError(): + raise ModbusReadError(f"{self.device_name}: failed to read input float {address}") + return struct.unpack(">f", struct.pack(">HH", res.registers[0], res.registers[1]))[0] + return self._retry(_do) + + # ---- writes, verified by reading the value back (whole write+verify retried on mismatch) ---- + + def write_coil(self, address, value, verify=True): + value = bool(value) + + def _do(): + res = self.client.write_coil(address, value, device_id=self.slave_id) + if res.isError(): + raise ModbusWriteError(f"{self.device_name}: failed to write coil {address}") + if verify and self.read_coil(address) != value: + raise ModbusWriteError(f"{self.device_name}: write coil {address}={value} not verified") + self._retry(_do) + + def write_holding_register(self, address, value): + value = int(value) & 0xFFFF # two's-complement wrap so negative (signed) values fit the uint16 wire format + + def _do(): + res = self.client.write_register(address, value, device_id=self.slave_id) + if res.isError(): + raise ModbusWriteError(f"{self.device_name}: failed to write holding {address}") + if self.read_holding_register(address) != value: + raise ModbusWriteError(f"{self.device_name}: write holding {address}={value} not verified") + self._retry(_do) + + def read_signed_holding_register(self, address): + val = self.read_holding_register(address) + return val - 65536 if val > 32767 else val + + def write_holding_register_float(self, address, value): + def _do(): + reg1, reg2 = struct.unpack(">HH", struct.pack(">f", value)) + res = self.client.write_registers(address, [reg1, reg2], device_id=self.slave_id) + if res.isError(): + raise ModbusWriteError(f"{self.device_name}: failed to write holding float {address}") + readback = self.read_holding_register_float(address) + if abs(readback - value) > self.FLOAT_VERIFY_TOLERANCE: + raise ModbusWriteError( + f"{self.device_name}: write holding float {address}={value} not verified (read back {readback})" + ) + self._retry(_do) + + # ---- generic label-based access (covers every register in self.registers) ---- + + def _register(self, label): + try: + return self.registers[label] + except KeyError: + raise ModbusDeviceError(f"unknown register label for {self.device_name}: {label}") + + def read(self, label): + reg = self._register(label) + if reg.object_type == "coil": + return self.read_coil(reg.address) + if reg.object_type == "input": + if reg.data_type == "float": + return self.read_input_register_float(reg.address) + if reg.data_type == "int": + return self.read_signed_input_register(reg.address) + return self.read_input_register(reg.address) + if reg.object_type == "holding": + if reg.data_type == "float": + return self.read_holding_register_float(reg.address) + if reg.data_type == "int": + return self.read_signed_holding_register(reg.address) + return self.read_holding_register(reg.address) + raise ModbusDeviceError(f"{self.device_name}: unsupported object_type {reg.object_type!r} for {label!r}") + + def write(self, label, value): + reg = self._register(label) + if reg.object_type == "coil": + self.write_coil(reg.address, value) + return + if reg.object_type == "holding": + if reg.data_type == "float": + self.write_holding_register_float(reg.address, value) + else: + self.write_holding_register(reg.address, value) + return + raise ModbusDeviceError(f"{self.device_name}: cannot write to object_type {reg.object_type!r} for {label!r}") + + def close(self): + self.client.close() diff --git a/drivers/rpu.py b/drivers/rpu.py new file mode 100644 index 0000000..0ef8f6d --- /dev/null +++ b/drivers/rpu.py @@ -0,0 +1,43 @@ +import time + +from common.config_loader import load_garuda_registers +from drivers.modbus_device import ModbusDevice + +DEFAULT_DATASETS = ["Sec_normal", "sec_flash"] + + +class Rpu(ModbusDevice): + """Covers every register in Sec_normal + sec_flash by label via read()/write() + (inherited from ModbusDevice). Extra datasets can be added per-test via + extra_datasets. Only registers that need real logic beyond a 1:1 read/write + get a dedicated method below. + """ + + Q128_DIVISOR = 128.0 + REWRITE_PASSWORD_LO = 0xFEED + REWRITE_PASSWORD_HI = 0xCEE5 + REWRITE_WAIT_S = 5 + + def __init__(self, port, extra_datasets=None, slave_id=0x10): + registers = load_garuda_registers(DEFAULT_DATASETS + list(extra_datasets or [])) + super().__init__(port, "RPU", registers, slave_id=slave_id) + + def set_charging(self, enable): + self.write_coil(self._addr("Enable_charger"), enable) + + def set_runtime_voltage(self, volts): + self.write_holding_register(self._addr("Runtime_voltage_setting_Q128"), round(volts * self.Q128_DIVISOR)) + + def set_runtime_current(self, amps): + self.write_holding_register(self._addr("Runtime_current_setting_Q128"), round(amps * self.Q128_DIVISOR)) + + def read_error_code(self): + return (self.read("Error_code_hi") << 16) | self.read("Error_code_lo") + + def read_shadow_error_code(self): + return (self.read("Shadow_error_code_hi") << 16) | self.read("Shadow_error_code_lo") + + def rewrite_config(self): + self.write_holding_register(self._addr("PASSWORD_LO"), self.REWRITE_PASSWORD_LO) + self.write_holding_register(self._addr("PASSWORD_HI"), self.REWRITE_PASSWORD_HI) + time.sleep(self.REWRITE_WAIT_S) diff --git a/drivers/tpu.py b/drivers/tpu.py new file mode 100644 index 0000000..f816e6a --- /dev/null +++ b/drivers/tpu.py @@ -0,0 +1,44 @@ +import time + +from common.config_loader import load_fenghuang_registers +from drivers.modbus_device import ModbusDevice, ModbusDeviceError + +DEFAULT_DATASETS = ["Pri_normal", "Debug", "event_log"] + +# TPU_STATE values that allow modbus_to_config()/write_all_flash() to run +# (fenghuang-dsp/app/app_modbus/app_modbus.c update_modbus_input()). +# TPU_RUN_INV (charging) is NOT in this set - config writes are silently +# ignored while charging is active. +CONFIG_WRITABLE_STATES = {2, 10, 15, 16, 238} # READY, INITIALIZATION, DEBUG_CONFIG, TEST_MODE, ERROR + + +class Tpu(ModbusDevice): + """Covers every register in Pri_normal + Debug + event_log by label via + read()/write() (inherited from ModbusDevice). More datasets can be added + per-test via extra_datasets. Only registers that need real logic beyond a + 1:1 read/write get a dedicated method below. + """ + + REWRITE_WAIT_S = 5 + + def __init__(self, port, extra_datasets=None, slave_id=0x10): + registers = load_fenghuang_registers(DEFAULT_DATASETS + list(extra_datasets or [])) + super().__init__(port, "TPU", registers, slave_id=slave_id) + + def wait_until_config_writable(self, timeout_s=10, poll_interval_s=0.5): + deadline = time.time() + timeout_s + while time.time() < deadline: + if self.read("Main_state") in CONFIG_WRITABLE_STATES: + return + time.sleep(poll_interval_s) + raise ModbusDeviceError(f"TPU did not reach a config-writable state within {timeout_s}s") + + def read_error_code(self): + return (self.read("Error_code_HI") << 16) | self.read("Error_code_LO") + + def read_shadow_error_code(self): + return (self.read("Shadow_Error_code_HI") << 16) | self.read("Shadow_Error_code_LO") + + def rewrite_config(self): + self.write_coil(self._addr("rewrite_cfg_flash"), True, verify=False) + time.sleep(self.REWRITE_WAIT_S)