Python ray.put() Examples

def test_numpy_subclass_serialization(ray_start_regular):
    class MyNumpyConstant(np.ndarray):
        def __init__(self, value):
            super().__init__()
            self.constant = value

        def __str__(self):
            print(self.constant)

    constant = MyNumpyConstant(123)

    def explode(x):
        raise RuntimeError("Expected error.")

    ray.register_custom_serializer(
        type(constant), serializer=explode, deserializer=explode)

    try:
        ray.put(constant)
        assert False, "Should never get here!"
    except (RuntimeError, IndexError):
        print("Correct behavior, proof that customer serializer was used.") 

def test_asyncio_actor_async_get(ray_start_regular_shared):
    @ray.remote
    def remote_task():
        return 1

    @ray.remote
    class AsyncGetter:
        async def get(self):
            return await remote_task.remote()

        async def plasma_get(self, plasma_object):
            return await plasma_object[0]

    plasma_object = ray.put(2)
    getter = AsyncGetter.remote()
    assert ray.get(getter.get.remote()) == 1
    assert ray.get(getter.plasma_get.remote([plasma_object])) == 2 

def testPythonMode(self):
    reload(test_functions)
    ray.init(start_ray_local=True, driver_mode=ray.PYTHON_MODE)

    @ray.remote
    def f():
      return np.ones([3, 4, 5])
    xref = f.remote()
    assert_equal(xref, np.ones([3, 4, 5])) # remote functions should return by value
    assert_equal(xref, ray.get(xref)) # ray.get should be the identity
    y = np.random.normal(size=[11, 12])
    assert_equal(y, ray.put(y)) # ray.put should be the identity

    # make sure objects are immutable, this example is why we need to copy
    # arguments before passing them into remote functions in python mode
    aref = test_functions.python_mode_f.remote()
    assert_equal(aref, np.array([0, 0]))
    bref = test_functions.python_mode_g.remote(aref)
    assert_equal(aref, np.array([0, 0])) # python_mode_g should not mutate aref
    assert_equal(bref, np.array([1, 0]))

    ray.worker.cleanup() 

def testRegisterClass(self):
    ray.init(start_ray_local=True, num_workers=0)

    # Check that putting an object of a class that has not been registered
    # throws an exception.
    class TempClass(object):
      pass
    self.assertRaises(Exception, lambda : ray.put(Foo))
    # Check that registering a class that Ray cannot serialize efficiently
    # raises an exception.
    self.assertRaises(Exception, lambda : ray.register_class(type(True)))
    # Check that registering the same class with pickle works.
    ray.register_class(type(float), pickle=True)
    self.assertEqual(ray.get(ray.put(float)), float)

    ray.worker.cleanup() 

def test_multi_node_stats(shutdown_only):
    cluster = Cluster()
    for _ in range(2):
        cluster.add_node(num_cpus=1)

    ray.init(address=cluster.address)

    @ray.remote(num_cpus=1)
    class Actor:
        def __init__(self):
            self.ref = ray.put(np.zeros(100000))

        def ping(self):
            pass

    # Each actor will be on a different node.
    a = Actor.remote()
    b = Actor.remote()
    ray.get(a.ping.remote())
    ray.get(b.ping.remote())

    # Verify we have collected stats across the nodes.
    info = memory_summary()
    print(info)
    assert count(info, PUT_OBJ) == 2, info 

def test_raylet_crash_when_get(ray_start_regular):
    def sleep_to_kill_raylet():
        # Don't kill raylet before default workers get connected.
        time.sleep(2)
        ray.worker._global_node.kill_raylet()

    object_id = ray.put(None)
    ray.internal.free(object_id)
    while ray.worker.global_worker.core_worker.object_exists(object_id):
        time.sleep(1)

    thread = threading.Thread(target=sleep_to_kill_raylet)
    thread.start()
    with pytest.raises(ray.exceptions.UnreconstructableError):
        ray.get(object_id)
    thread.join() 

def full_grad(theta):
    theta_id = ray.put(theta)
    grad_ids = [grad.remote(theta_id, xs_id, ys_id) for (xs_id, ys_id) in batch_ids]
    return sum(ray.get(grad_ids)).astype("float64") # This conversion is necessary for use with fmin_l_bfgs_b.

  # From the perspective of scipy.optimize.fmin_l_bfgs_b, full_loss is simply a
  # function which takes some parameters theta, and computes a loss. Similarly,
  # full_grad is a function which takes some parameters theta, and computes the
  # gradient of the loss. Internally, these functions use Ray to distribute the
  # computation of the loss and the gradient over the data that is represented
  # by the remote object IDs x_batches and y_batches and which is potentially
  # distributed over a cluster. However, these details are hidden from
  # scipy.optimize.fmin_l_bfgs_b, which simply uses it to run the L-BFGS
  # algorithm.

  # Load the mnist data and turn the data into remote objects. 

def _run(self, driver_quota, a_quota, b_quota):
        print("*** Testing ***", driver_quota, a_quota, b_quota)
        try:
            ray.init(
                num_cpus=1,
                object_store_memory=300 * MB,
                driver_object_store_memory=driver_quota)
            z = ray.put("hi", weakref=True)
            a = LightActor._remote(object_store_memory=a_quota)
            b = GreedyActor._remote(object_store_memory=b_quota)
            for _ in range(5):
                r_a = a.sample.remote()
                for _ in range(20):
                    new_oid = b.sample.remote()
                    ray.get(new_oid)
                ray.get(r_a)
            ray.get(z)
        except Exception as e:
            print("Raised exception", type(e), e)
            raise e
        finally:
            print(ray.worker.global_worker.core_worker.
                  dump_object_store_memory_usage())
            ray.shutdown() 

def test_internal_config_when_connecting(ray_start_cluster):
    config = json.dumps({
        "object_pinning_enabled": 0,
        "initial_reconstruction_timeout_milliseconds": 200
    })
    cluster = ray.cluster_utils.Cluster()
    cluster.add_node(
        _internal_config=config, object_store_memory=100 * 1024 * 1024)
    cluster.wait_for_nodes()

    # Specifying _internal_config when connecting to a cluster is disallowed.
    with pytest.raises(ValueError):
        ray.init(address=cluster.address, _internal_config=config)

    # Check that the config was picked up (object pinning is disabled).
    ray.init(address=cluster.address)
    oid = ray.put(np.zeros(40 * 1024 * 1024, dtype=np.uint8))

    for _ in range(5):
        ray.put(np.zeros(40 * 1024 * 1024, dtype=np.uint8))

    # This would not raise an exception if object pinning was enabled.
    with pytest.raises(ray.exceptions.UnreconstructableError):
        ray.get(oid) 

def test_cache(ray_start_regular):
    A = np.random.rand(1, 1000000)
    v = np.random.rand(1000000)
    A_id = ray.put(A)
    v_id = ray.put(v)
    a = time.time()
    for i in range(100):
        A.dot(v)
    b = time.time() - a
    c = time.time()
    for i in range(100):
        ray.get(A_id).dot(ray.get(v_id))
    d = time.time() - c

    if d > 1.5 * b:
        print("WARNING: The caching test was too slow. "
              "d = {}, b = {}".format(d, b)) 

def test_put_pins_object(ray_start_object_store_memory):
    x_id = ray.put("HI")
    x_binary = x_id.binary()
    assert ray.get(ray.ObjectID(x_binary)) == "HI"

    # x cannot be evicted since x_id pins it
    for _ in range(10):
        ray.put(np.zeros(10 * 1024 * 1024))
    assert ray.get(x_id) == "HI"
    assert ray.get(ray.ObjectID(x_binary)) == "HI"

    # now it can be evicted since x_id pins it but x_binary does not
    del x_id
    for _ in range(10):
        ray.put(np.zeros(10 * 1024 * 1024))
    assert not ray.worker.global_worker.core_worker.object_exists(
        ray.ObjectID(x_binary))

    # weakref put
    y_id = ray.put("HI", weakref=True)
    for _ in range(10):
        ray.put(np.zeros(10 * 1024 * 1024))
    with pytest.raises(ray.exceptions.UnreconstructableError):
        ray.get(y_id) 

def test_load_balancing_with_dependencies(ray_start_cluster):
    # This test ensures that tasks are being assigned to all raylets in a
    # roughly equal manner even when the tasks have dependencies.
    cluster = ray_start_cluster
    num_nodes = 3
    for _ in range(num_nodes):
        cluster.add_node(num_cpus=1)
    ray.init(address=cluster.address)

    @ray.remote
    def f(x):
        time.sleep(0.010)
        return ray.worker.global_worker.node.unique_id

    # This object will be local to one of the raylets. Make sure
    # this doesn't prevent tasks from being scheduled on other raylets.
    x = ray.put(np.zeros(1000000))

    attempt_to_load_balance(f, [x], 100, num_nodes, 25) 

def test_transfer_performance(benchmark, ray_start_cluster_head, object_number,
                              data_size):
    cluster = ray_start_cluster_head
    cluster.add_node(resources={"my_resource": 1}, object_store_memory=10**9)

    @ray.remote(resources={"my_resource": 1})
    class ObjectActor:
        def f(self, object_ids):
            ray.get(object_ids)

    # setup remote actor
    actor = ObjectActor.remote()
    actor.f.remote([])

    data = bytes(1) * data_size
    object_ids = [ray.put(data) for _ in range(object_number)]

    benchmark(benchmark_transfer_object, actor, object_ids) 

def test_local_refcounts(ray_start_regular):
    oid1 = ray.put(None)
    check_refcounts({oid1: (1, 0)})
    oid1_copy = copy.copy(oid1)
    check_refcounts({oid1: (2, 0)})
    del oid1
    check_refcounts({oid1_copy: (1, 0)})
    del oid1_copy
    check_refcounts({}) 

def test_basic_pinning(one_worker_100MiB):
    @ray.remote
    def f(array):
        return np.sum(array)

    @ray.remote
    class Actor(object):
        def __init__(self):
            # Hold a long-lived reference to a ray.put object's ID. The object
            # should not be garbage collected while the actor is alive because
            # the object is pinned by the raylet.
            self.large_object = ray.put(
                np.zeros(25 * 1024 * 1024, dtype=np.uint8))

        def get_large_object(self):
            return ray.get(self.large_object)

    actor = Actor.remote()

    # Fill up the object store with short-lived objects. These should be
    # evicted before the long-lived object whose reference is held by
    # the actor.
    for batch in range(10):
        intermediate_result = f.remote(
            np.zeros(10 * 1024 * 1024, dtype=np.uint8))
        ray.get(intermediate_result)

    # The ray.get below would fail with only LRU eviction, as the object
    # that was ray.put by the actor would have been evicted.
    ray.get(actor.get_large_object.remote()) 

def test_serialization_final_fallback(ray_start_regular):
    pytest.importorskip("catboost")
    # This test will only run when "catboost" is installed.
    from catboost import CatBoostClassifier

    model = CatBoostClassifier(
        iterations=2,
        depth=2,
        learning_rate=1,
        loss_function="Logloss",
        logging_level="Verbose")

    reconstructed_model = ray.get(ray.put(model))
    assert set(model.get_params().items()) == set(
        reconstructed_model.get_params().items()) 

def test_duplicate_args(ray_start_regular):
    @ray.remote
    def f(arg1,
          arg2,
          arg1_duplicate,
          kwarg1=None,
          kwarg2=None,
          kwarg1_duplicate=None):
        assert arg1 == kwarg1
        assert arg1 != arg2
        assert arg1 == arg1_duplicate
        assert kwarg1 != kwarg2
        assert kwarg1 == kwarg1_duplicate

    # Test by-value arguments.
    arg1 = [1]
    arg2 = [2]
    ray.get(
        f.remote(
            arg1, arg2, arg1, kwarg1=arg1, kwarg2=arg2, kwarg1_duplicate=arg1))

    # Test by-reference arguments.
    arg1 = ray.put([1])
    arg2 = ray.put([2])
    ray.get(
        f.remote(
            arg1, arg2, arg1, kwarg1=arg1, kwarg2=arg2, kwarg1_duplicate=arg1)) 

def test_get_dict(ray_start_regular):
    d = {str(i): ray.put(i) for i in range(5)}
    for i in range(5, 10):
        d[str(i)] = i
    result = ray.experimental.get(d)
    expected = {str(i): i for i in range(10)}
    assert result == expected 

def test_get_multiple(ray_start_regular):
    object_ids = [ray.put(i) for i in range(10)]
    assert ray.get(object_ids) == list(range(10))

    # Get a random choice of object IDs with duplicates.
    indices = list(np.random.choice(range(10), 5))
    indices += indices
    results = ray.get([object_ids[i] for i in indices])
    assert results == indices 

def test_put_api(ray_start_regular):

    for obj in test_values:
        assert ray.get(ray.put(obj)) == obj

    # Test putting object IDs.
    x_id = ray.put(0)
    for obj in [[x_id], (x_id, ), {x_id: x_id}]:
        assert ray.get(ray.put(obj)) == obj 

def warmup():
    x = np.zeros(10**6, dtype=np.uint8)
    for _ in range(5):
        for _ in range(5):
            ray.put(x)
        for _ in range(5):
            ray.get([dummy_task.remote(0) for _ in range(1000)]) 

def test_passing_arguments_by_value_out_of_the_box(ray_start_regular):
    @ray.remote
    def f(x):
        return x

    # Test passing lambdas.

    def temp():
        return 1

    assert ray.get(f.remote(temp))() == 1
    assert ray.get(f.remote(lambda x: x + 1))(3) == 4

    # Test sets.
    assert ray.get(f.remote(set())) == set()
    s = {1, (1, 2, "hi")}
    assert ray.get(f.remote(s)) == s

    # Test types.
    assert ray.get(f.remote(int)) == int
    assert ray.get(f.remote(float)) == float
    assert ray.get(f.remote(str)) == str

    class Foo:
        def __init__(self):
            pass

    # Make sure that we can put and get a custom type. Note that the result
    # won't be "equal" to Foo.
    ray.get(ray.put(Foo)) 

def test_deserialized_from_buffer_immutable(ray_start_regular):
    x = np.full((2, 2), 1.)
    o = ray.put(x)
    y = ray.get(o)
    with pytest.raises(
            ValueError, match="assignment destination is read-only"):
        y[0, 0] = 9. 

def test_pinned_object_call_site(ray_start_regular):
    # Local ref only.
    x_id = ray.put(np.zeros(100000))
    info = memory_summary()
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, LOCAL_REF) == 1, info
    assert count(info, PINNED_IN_MEMORY) == 0, info
    assert count(info, "test_memstat.py") == 1, info

    # Local ref + pinned buffer.
    buf = ray.get(x_id)
    info = memory_summary()
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, LOCAL_REF) == 0, info
    assert count(info, PINNED_IN_MEMORY) == 1, info
    assert count(info, "test_memstat.py") == 1, info

    # Just pinned buffer.
    del x_id
    info = memory_summary()
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, LOCAL_REF) == 0, info
    assert count(info, PINNED_IN_MEMORY) == 1, info
    assert count(info, "test_memstat.py") == 1, info

    # Nothing.
    del buf
    info = memory_summary()
    print(info)
    assert num_objects(info) == 0, info 

def test_nested_object_refs(ray_start_regular):
    x_id = ray.put(np.zeros(100000))
    y_id = ray.put([x_id])
    z_id = ray.put([y_id])
    del x_id, y_id
    info = memory_summary()
    print(info)
    assert num_objects(info) == 3, info
    assert count(info, LOCAL_REF) == 1, info
    assert count(info, CAPTURED_IN_OBJECT) == 2, info
    del z_id 

def test_worker_task_refs(ray_start_regular):
    @ray.remote
    def f(y):
        x_id = ray.put("HI")
        info = memory_summary()
        del x_id
        return info

    x_id = f.remote(np.zeros(100000))
    info = ray.get(x_id)
    print(info)
    assert num_objects(info) == 4, info
    # Task argument plus task return ids.
    assert count(info, TASK_CALL_OBJ) == 2, info
    assert count(info, DRIVER_PID) == 1, info
    assert count(info, WORKER_PID) == 1, info
    assert count(info, LOCAL_REF) == 2, info
    assert count(info, PINNED_IN_MEMORY) == 1, info
    assert count(info, PUT_OBJ) == 1, info
    assert count(info, DESER_TASK_ARG) == 1, info
    assert count(info, UNKNOWN_SIZE) == 1, info
    assert count(info, "test_memstat.py:f") == 1, info
    assert count(info, "test_memstat.py:test_worker_task_refs") == 2, info

    info = memory_summary()
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, DRIVER_PID) == 1, info
    assert count(info, TASK_CALL_OBJ) == 1, info
    assert count(info, UNKNOWN_SIZE) == 0, info
    assert count(info, x_id.hex()) == 1, info

    del x_id
    info = memory_summary()
    assert num_objects(info) == 0, info 

def test_getting_and_putting(ray_start_sharded):
    for n in range(8):
        x = np.zeros(10**n)

        for _ in range(100):
            ray.put(x)

        x_id = ray.put(x)
        for _ in range(1000):
            ray.get(x_id)

    assert ray.services.remaining_processes_alive() 

def testTooLargeAllocation(self):
        try:
            ray.init(num_cpus=1, driver_object_store_memory=100 * MB)
            ray.put(np.zeros(50 * MB, dtype=np.uint8), weakref=True)
            self.assertRaises(
                OBJECT_TOO_LARGE,
                lambda: ray.put(np.zeros(200 * MB, dtype=np.uint8)))
        finally:
            ray.shutdown() 

def test_lease_request_leak(shutdown_only):
    ray.init(
        num_cpus=1,
        _internal_config=json.dumps({
            "initial_reconstruction_timeout_milliseconds": 200
        }))
    assert len(ray.objects()) == 0

    @ray.remote
    def f(x):
        time.sleep(0.1)
        return

    # Submit pairs of tasks. Tasks in a pair can reuse the same worker leased
    # from the raylet.
    tasks = []
    for _ in range(10):
        oid = ray.put(1)
        for _ in range(2):
            tasks.append(f.remote(oid))
        del oid
    ray.get(tasks)

    time.sleep(
        1)  # Sleep for an amount longer than the reconstruction timeout.
    assert len(ray.objects()) == 0, ray.objects() 

def testDriverPutEvictedCannotReconstruct(self):
        x_id = ray.put(np.zeros(1 * 1024 * 1024), weakref=True)
        ray.get(x_id)
        for _ in range(20):
            ray.put(np.zeros(10 * 1024 * 1024))
        self.assertRaises(ray.exceptions.UnreconstructableError,
                          lambda: ray.get(x_id))