Passing the same large argument by value repeatedly harms performance

When passing a large argument (>100KB) by value to a task, Ray will implicitly store the argument in the object store and the worker process will fetch the argument to the local object store from the caller’s object store before running the task. If we pass the same large argument to multiple tasks, Ray will end up storing multiple copies of the argument in the object store since Ray doesn’t do deduplication.

Instead of passing the large argument by value to multiple tasks, we should use ray.put() to store the argument to the object store once and get an ObjectRef, then pass the argument reference to tasks. This way, we make sure all tasks use the same copy of the argument, which is faster and uses less object store memory.

Example

Say we had to do some work:

@ray.remote
def do_work(a):
    # do some work with the large object a
    return np.sum(a)

A non-optimized way to do this would be to pass the large object by value to each task:

random.seed(42)
 
start = time.time()
a = np.random.rand(5000, 5000)
 
# Sending the big array to each remote task, which will
# its copy of the same data into its object store
result_ids = [do_work.remote(a) for x in range(10)]
 
results = math.fsum(ray.get(result_ids))
print(f" results = {results:.2f} and duration = {time.time() - start:.3f} sec")

results = 125012473.08 and duration = 0.680 sec

A better approach is to put the value in the object store and only send the reference:

start = time.time()
# Adding the big array into the object store
a_id_ref = ray.put(a)
 
# Now send the objectID ref
result_ids = [do_work.remote(a_id_ref) for x in range(10)]
results = math.fsum(ray.get(result_ids))
print(f" results = {results:.2f} and duration = {time.time() - start:.3f} sec")

results = 125012473.08 and duration = 0.275 sec