| 121 | |
| 122 | |
| 123 | class PolicyController: |
| 124 | def __init__(self, player_name: str, agent_name: str, path: str, *, timeout: float): |
| 125 | self.player_name = player_name |
| 126 | self.agent_name = agent_name |
| 127 | self.path = path |
| 128 | self.timeout = max(float(timeout), 0.01) |
| 129 | self.errors: list[dict] = [] |
| 130 | self.invalid_actions = 0 |
| 131 | self.decisions = 0 |
| 132 | self.startup_error: str | None = None |
| 133 | self._next_request_id = 0 |
| 134 | self._start_worker() |
| 135 | |
| 136 | def _start_worker(self) -> None: |
| 137 | self.startup_error = None |
| 138 | ctx = mp.get_context("spawn") |
| 139 | self.command_queue = ctx.Queue() |
| 140 | self.result_queue = ctx.Queue() |
| 141 | self.startup_queue = ctx.Queue() |
| 142 | self.process = ctx.Process( |
| 143 | target=policy_worker, |
| 144 | args=( |
| 145 | self.command_queue, |
| 146 | self.result_queue, |
| 147 | self.startup_queue, |
| 148 | self.player_name, |
| 149 | self.agent_name, |
| 150 | self.path, |
| 151 | ), |
| 152 | ) |
| 153 | self.process.start() |
| 154 | startup_timeout = max(self.timeout, 10.0) |
| 155 | try: |
| 156 | startup_message = self.startup_queue.get(timeout=startup_timeout) |
| 157 | except queue.Empty: |
| 158 | self.startup_error = f"policy import exceeded {startup_timeout}s timeout" |
| 159 | self.close() |
| 160 | return |
| 161 | if "error" in startup_message: |
| 162 | self.startup_error = startup_message["error"] |
| 163 | self.close() |
| 164 | |
| 165 | def close(self) -> None: |
| 166 | with contextlib.suppress(Exception): |
| 167 | self.command_queue.put_nowait(None) |
| 168 | self.process.join(0.1) |
| 169 | if self.process.is_alive(): |
| 170 | self.process.terminate() |
| 171 | self.process.join(0.1) |
| 172 | if self.process.is_alive(): |
| 173 | self.process.kill() |
| 174 | self.process.join() |
| 175 | |
| 176 | def restart(self) -> None: |
| 177 | self.close() |
| 178 | self._start_worker() |
| 179 | |
| 180 | def decide(self, observation, action_space) -> int: |