| 141 | |
| 142 | |
| 143 | class MySimulatorMaster(SimulatorMaster, Callback): |
| 144 | def __init__(self, pipe_c2s, pipe_s2c, gpus): |
| 145 | """ |
| 146 | Args: |
| 147 | gpus (list[int]): the gpus used to run inference |
| 148 | """ |
| 149 | super(MySimulatorMaster, self).__init__(pipe_c2s, pipe_s2c) |
| 150 | self.queue = queue.Queue(maxsize=BATCH_SIZE * 8 * 2) |
| 151 | self._gpus = gpus |
| 152 | |
| 153 | def _setup_graph(self): |
| 154 | # Create predictors on the available predictor GPUs. |
| 155 | num_gpu = len(self._gpus) |
| 156 | predictors = [self.trainer.get_predictor( |
| 157 | ['state'], ['policy', 'pred_value'], |
| 158 | self._gpus[k % num_gpu]) |
| 159 | for k in range(PREDICTOR_THREAD)] |
| 160 | self.async_predictor = MultiThreadAsyncPredictor( |
| 161 | predictors, batch_size=PREDICT_BATCH_SIZE) |
| 162 | |
| 163 | def _before_train(self): |
| 164 | self.async_predictor.start() |
| 165 | logger.info("Starting MySimulatorMaster ...") |
| 166 | start_proc_mask_signal(self) |
| 167 | |
| 168 | def _on_state(self, state, client): |
| 169 | """ |
| 170 | Launch forward prediction for the new state given by some client. |
| 171 | """ |
| 172 | def cb(outputs): |
| 173 | try: |
| 174 | distrib, value = outputs.result() |
| 175 | except CancelledError: |
| 176 | logger.info("Client {} cancelled.".format(client.ident.decode('utf-8'))) |
| 177 | return |
| 178 | assert np.all(np.isfinite(distrib)), distrib |
| 179 | action = np.random.choice(len(distrib), p=distrib) |
| 180 | client.memory.append(TransitionExperience( |
| 181 | state, action, reward=None, value=value, prob=distrib[action])) |
| 182 | self.send_queue.put([client.ident, dumps(action)]) |
| 183 | self.async_predictor.put_task([state], cb) |
| 184 | |
| 185 | def _process_msg(self, client, state, reward, isOver): |
| 186 | """ |
| 187 | Process a message sent from some client. |
| 188 | """ |
| 189 | # in the first message, only state is valid, |
| 190 | # reward&isOver should be discarded |
| 191 | if len(client.memory) > 0: |
| 192 | client.memory[-1].reward = reward |
| 193 | if isOver: |
| 194 | # should clear client's memory and put to queue |
| 195 | self._parse_memory(0, client, True) |
| 196 | else: |
| 197 | if len(client.memory) == LOCAL_TIME_MAX + 1: |
| 198 | R = client.memory[-1].value |
| 199 | self._parse_memory(R, client, False) |
| 200 | # feed state and return action |
no outgoing calls
no test coverage detected
searching dependent graphs…