def test_future_on_chained_exception(): f = RPCFuture() new = f.map(lambda v: v + 2) f.set_exception(TExc()) with pytest.raises(TExc): assert new.result()
def test_future_map_raise_exception(): f = RPCFuture() new = f.map(process_err) f.set_result(40) with pytest.raises(TExc): assert new.result()
def test_rpcfuture_attach(): rpc_fut = RPCFuture() fut = Future() rpc_fut.attach(fut) fut.set_result(42) assert rpc_fut.result(timeout=1) == 42
def test_propagates_only_once_2(): rpc_fut = RPCFuture() fut = Future() rpc_fut.attach(fut) fut.set_result(42) fut.cancel() assert rpc_fut.result(timeout=1) == 42
def load_model(self, model: Model, state: ModelState, devices: list) -> RPCFuture[SetDeviceReturnType]: log_dir = model.config.get(LOGGING, {}).get(DIRECTORY, "") if log_dir: os.makedirs(log_dir, exist_ok=True) self.logger.info("log dir: %s", os.path.abspath(log_dir)) self._start_logging_handler() incomplete_msg = get_error_msg_for_incomplete_config(model.config) if incomplete_msg: raise ValueError(incomplete_msg) # todo: move test_transforms elsewhere self.test_transforms = model.config.get(TESTING, {}).get( TRANSFORMS, {"Normalize": {}}) if not devices: devices = ["cpu"] cuda_visible_devices, handler_devices = self.get_cuda_and_handler_device_names( devices) os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(cuda_visible_devices) self.logger.info("Set CUDA_VISIBLE_DEVICES to '%s'", os.environ["CUDA_VISIBLE_DEVICES"]) server_conn, handler_conn = mp.Pipe() p = mp.Process( target=run_handler, name="Handler", kwargs={ "conn": handler_conn, "config": model.config, "model_file": model.code, "model_state": state.model_state, "optimizer_state": state.optimizer_state, "log_queue": self.log_queue, }, ) try: p.start() except Exception as e: self.logger.error(e) err_fut = RPCFuture() err_fut.set_exception(e) return err_fut else: self.handler = create_client(IHandler, server_conn) try: tik_fut = self.handler.set_devices(handler_devices) except Exception as e: self.logger.exception("set_devices failed") err_fut = RPCFuture() err_fut.set_exception(e) return err_fut else: self.logger.info("got tik_fut") fut = tik_fut.map(convert_to_SetDeviceReturnType) self.logger.info("converted tik_fut") return fut
def test_propagates_only_once_1(): rpc_fut = RPCFuture() fut = Future() rpc_fut.attach(fut) rpc_fut.cancel() rpc_fut.set_result(42) with pytest.raises(CancelledError): assert fut.result(timeout=1)
def load_model(self, config: dict, model_file: bytes, model_state: bytes, optimizer_state: bytes, devices: list) -> RPCFuture[SetDeviceReturnType]: logger.info("load model") logger.debug( "config %s\n model_file %s\n model_state %s\n optimizer_state %s\n devices %s\n", config, model_file, model_state, optimizer_state, devices, ) fut = RPCFuture() ret = SetDeviceReturnType((1, 1, 1, 10, 10), [(1, 1, 1, 10, 10), (1, 1, 1, 20, 20)], (0, 0, 0, 4, 4)) fut.set_result(ret) return fut
def dry_run( self, devices: List[torch.device], training_shape: Optional[Point] = None, valid_shapes: Optional[List[Point]] = None, shrinkage: Optional[Point] = None, ) -> RPCFuture[Tuple[List[torch.device], Point, List[Point], Point]]: fut = RPCFuture() self.dry_run_queue.put((devices, training_shape, valid_shapes, shrinkage, fut)) return fut
def test_rpcfuture_cancellation(): rpc_fut = RPCFuture() fut = Future() rpc_fut.attach(fut) rpc_fut.cancel() assert fut.cancelled()
def case_rpc_future() -> RPCFuture: return RPCFuture()
def forward(self, data: TikTensor) -> RPCFuture[TikTensor]: fut = RPCFuture() self.forward_queue.put((data, fut)) return fut
def submit(self, func) -> RPCFuture: f = RPCFuture() self._q.put((func, f)) return f
def test_future_map_returns_new_future(): f = RPCFuture() new = f.map(lambda v: v + 2) f.set_result(40) assert new.result() == 42
def case_rpc_future_generic() -> RPCFuture[int]: return RPCFuture()
def set_devices( self, devices: Collection[torch.device]) -> RPCFuture[Set[torch.device]]: fut = RPCFuture() self.device_change_queue.put((set(devices), fut)) return fut
def set_devices( self, device_names: Sequence[str] ) -> RPCFuture[Tuple[Point, List[Point], Point]]: fut = RPCFuture() self.new_device_names.put((device_names, fut)) return fut
def forward(self, batch: NDArray) -> RPCFuture[NDArray]: logger.info("forward for array of shape %s", batch.shape) fut = RPCFuture() fut.set_result(batch) return fut