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
