def test_objects_count_remote(hook, start_proc):
kwargs = {"id": "fed", "host": "localhost", "port": 8764, "hook": hook}
process_remote_worker = start_proc(WebsocketServerWorker, kwargs)
time.sleep(0.1)
kwargs = {"id": "fed", "host": "localhost", "port": 8764, "hook": hook}
local_worker = WebsocketClientWorker(**kwargs)
x = torch.tensor([1, 2, 3]).send(local_worker)
nr_objects = local_worker.objects_count_remote()
assert nr_objects == 1
y = torch.tensor([4, 5, 6]).send(local_worker)
nr_objects = local_worker.objects_count_remote()
assert nr_objects == 2
x.get()
nr_objects = local_worker.objects_count_remote()
assert nr_objects == 1
# delete remote object before terminating the websocket connection
del y
time.sleep(0.1)
local_worker.ws.shutdown()
time.sleep(0.1)
local_worker.remove_worker_from_local_worker_registry()
process_remote_worker.terminate()
def test_connect_close(hook, start_proc):
kwargs = {"id": "fed", "host": "localhost", "port": 8763, "hook": hook}
process_remote_worker = start_proc(WebsocketServerWorker, **kwargs)
time.sleep(0.1)
kwargs = {"id": "fed", "host": "localhost", "port": 8763, "hook": hook}
local_worker = WebsocketClientWorker(**kwargs)
x = torch.tensor([1, 2, 3])
x_ptr = x.send(local_worker)
assert local_worker.objects_count_remote() == 1
local_worker.close()
time.sleep(0.1)
local_worker.connect()
assert local_worker.objects_count_remote() == 1
x_val = x_ptr.get()
assert (x_val == x).all()
local_worker.ws.shutdown()
time.sleep(0.1)
process_remote_worker.terminate()