def test_create_callable_pointer(workers): """ Asserts that a callable pointer is correctly created. """ alice = workers["alice"] bob = workers["bob"] p = callable_pointer.create_callable_pointer( id=500, id_at_location=2, location=alice, owner=bob, tags="tags", description="description", register_pointer=True, ) assert len(alice.object_store._tensors) == 0 assert isinstance(bob.object_store.get_obj(500), callable_pointer.CallablePointer) p = callable_pointer.create_callable_pointer( id=501, id_at_location=2, location=alice, owner=bob, tags="tags", description="description", register_pointer=False, ) assert len(alice.object_store._tensors) == 0 assert isinstance(bob.object_store.get_obj(500), callable_pointer.CallablePointer) assert 501 not in bob.object_store._objects
def test_create_callable_pointer(workers): """ Asserts that a callable pointer is correctly created. """ alice = workers["alice"] bob = workers["bob"] callable_pointer.create_callable_pointer( id=500, id_at_location=2, location=alice, owner=bob, tags="tags", description="description", register_pointer=True, ) assert len(alice._objects) == 0 assert len(bob._objects) == 1 callable_pointer.create_callable_pointer( id=501, id_at_location=2, location=alice, owner=bob, tags="tags", description="description", register_pointer=False, ) assert len(alice._objects) == 0 assert len(bob._objects) == 1
def test_call_callable_pointer(workers): """ Tests that the correct result after an operation is returned when `callable_pointer` is called. """ def foo(x): """ Adds 2 to a given input `x`.""" return x + 2 alice = workers["alice"] bob = workers["bob"] id_alice = 100 id_bob = 200 foo_wrapper = ObjectWrapper(id=id_alice, obj=foo) alice.register_obj(foo_wrapper, id_alice) foo_ptr = callable_pointer.create_callable_pointer( id=id_bob, id_at_location=id_alice, location=alice, owner=bob, tags="tags", description="description", register_pointer=True, ) res = foo_ptr(4) assert res == 6
def test_get_obj_callable_pointer(workers): """ Asserts that correct object values are returned when `callable_pointer` is called. """ alice = workers["alice"] bob = workers["bob"] x = torch.tensor(5) x_ptr = x.send(alice) obj_ptr = callable_pointer.create_callable_pointer( id=1, id_at_location=x_ptr.id_at_location, location=alice, owner=bob, tags="tags", description="description", register_pointer=True, ) assert len(alice.object_store._tensors) == 1 assert isinstance(bob.object_store.get_obj(1), callable_pointer.CallablePointer) x_get = obj_ptr.get() assert len(alice.object_store._tensors) == 0 assert len(bob.object_store._tensors) == 0 assert 1 not in bob.object_store._objects assert x_get == x
def test_call_callable_pointer(workers): def foo(x): return x + 2 alice = workers["alice"] bob = workers["bob"] id_alice = 100 id_bob = 200 foo_wrapper = ObjectWrapper(id=id_alice, obj=foo) alice.register_obj(foo_wrapper, id_alice) foo_ptr = callable_pointer.create_callable_pointer( id=id_bob, id_at_location=id_alice, location=alice, owner=bob, tags="tags", description="description", register_pointer=True, ) res = foo_ptr(4) assert res == 6
def test_get_obj_callable_pointer(workers): alice = workers["alice"] bob = workers["bob"] x = torch.tensor(5) x_ptr = x.send(alice) obj_ptr = callable_pointer.create_callable_pointer( id=1, id_at_location=x_ptr.id_at_location, location=alice, owner=bob, tags="tags", description="description", register_pointer=True, ) assert len(alice._objects) == 1 assert len(bob._objects) == 1 x_get = obj_ptr.get() assert len(alice._objects) == 0 assert len(bob._objects) == 0 assert x_get == x
def create_pointer( object, owner: "BaseWorker", location: "BaseWorker", ptr_id: Union[int, str], id_at_location: Union[int, str] = None, garbage_collect_data=None, **kwargs, ): """Creates a callable pointer to the object wrapper instance Args: owner: A BaseWorker parameter to specify the worker on which the pointer is located. It is also where the pointer is registered if register is set to True. location: The BaseWorker object which points to the worker on which this pointer's object can be found. In nearly all cases, this is self.owner and so this attribute can usually be left blank. Very rarely you may know that you are about to move the Tensor to another worker so you can pre-initialize the location attribute of the pointer to some other worker, but this is a rare exception. ptr_id: A string or integer parameter to specify the id of the pointer. id_at_location: A string or integer id of the object being pointed to. Similar to location, this parameter is almost always self.id and so you can leave this parameter to None. The only exception is if you happen to know that the ID is going to be something different than self.id, but again this is very rare and most of the time, setting this means that you are probably doing something you shouldn't. garbage_collect_data: If True, delete the remote object when the pointer is deleted. Returns: A pointers.CallablePointer pointer to self. """ pointer = create_callable_pointer( owner=owner, location=location, id=ptr_id, id_at_location=id_at_location if id_at_location is not None else object.id, tags=object.tags, description=object.description, garbage_collect_data=False if garbage_collect_data is None else garbage_collect_data, ) return pointer