Files
ray-project--ray/python/ray/tests/test_advanced_8.py
T
2026-07-13 13:17:40 +08:00

652 lines
23 KiB
Python

# coding: utf-8
import glob
import logging
import multiprocessing
import os
import sys
import tempfile
import time
from unittest import mock
import numpy as np
import pytest
import ray
import ray._private.gcs_utils as gcs_utils
import ray._private.ray_constants as ray_constants
import ray._private.utils
import ray.cluster_utils
import ray.util.accelerators
from ray._common.test_utils import wait_for_condition
from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX
from ray._private.test_utils import rocksdb_gcs_test_enabled
from ray.dashboard import k8s_utils
from ray.runtime_env import RuntimeEnv
import psutil
logger = logging.getLogger(__name__)
def test_export_after_shutdown(ray_start_regular):
# This test checks that we can use actor and remote function definitions
# across multiple Ray sessions.
@ray.remote
def f():
pass
@ray.remote
class Actor:
def method(self):
pass
ray.get(f.remote())
a = Actor.remote()
ray.get(a.method.remote())
ray.shutdown()
# Start Ray and use the remote function and actor again.
ray.init(num_cpus=1)
ray.get(f.remote())
a = Actor.remote()
ray.get(a.method.remote())
ray.shutdown()
# Start Ray again and make sure that these definitions can be exported from
# workers.
ray.init(num_cpus=2)
@ray.remote
def export_definitions_from_worker(remote_function, actor_class):
ray.get(remote_function.remote())
actor_handle = actor_class.remote()
ray.get(actor_handle.method.remote())
ray.get(export_definitions_from_worker.remote(f, Actor))
def test_invalid_unicode_in_worker_log(shutdown_only):
info = ray.init(num_cpus=1)
logs_dir = os.path.join(info["session_dir"], "logs")
# Wait till first worker log file is created.
while True:
log_file_paths = glob.glob(f"{logs_dir}/worker*.out")
if len(log_file_paths) == 0:
time.sleep(0.2)
else:
break
with open(log_file_paths[0], "wb") as f:
f.write(b"\xe5abc\nline2\nline3\n")
f.write(b"\xe5abc\nline2\nline3\n")
f.write(b"\xe5abc\nline2\nline3\n")
f.flush()
# Wait till the log monitor reads the file.
time.sleep(1.0)
# Make sure that nothing has died.
assert ray._private.services.remaining_processes_alive()
@pytest.mark.skipif(
rocksdb_gcs_test_enabled(),
reason=(
"Starts a second local Ray cluster while the fixture's head is alive; "
"both inherit RAY_gcs_storage_path and collide on the same RocksDB "
"directory (single-writer LOCK). Redis tolerates this via server-side "
"multiplexing; RocksDB cannot."
),
)
@pytest.mark.parametrize(
"ray_start_cluster",
[
{
"num_cpus": 0,
"num_nodes": 1,
"do_init": False,
}
],
indirect=True,
)
def test_ray_address_environment_variable(ray_start_cluster):
address = ray_start_cluster.address
# In this test we use zero CPUs to distinguish between starting a local
# ray cluster and connecting to an existing one.
# Make sure we connect to an existing cluster if
# RAY_ADDRESS is set to the cluster address.
os.environ["RAY_ADDRESS"] = address
ray.init()
assert "CPU" not in ray._private.state.cluster_resources()
ray.shutdown()
del os.environ["RAY_ADDRESS"]
# Make sure we connect to an existing cluster if
# RAY_ADDRESS is set to "auto".
os.environ["RAY_ADDRESS"] = "auto"
ray.init()
assert "CPU" not in ray._private.state.cluster_resources()
ray.shutdown()
del os.environ["RAY_ADDRESS"]
# Prefer `address` parameter to the `RAY_ADDRESS` environment variable,
# when `address` is not `auto`.
os.environ["RAY_ADDRESS"] = "test"
ray.init(address=address)
assert "CPU" not in ray._private.state.cluster_resources()
ray.shutdown()
del os.environ["RAY_ADDRESS"]
# Make sure we connect to the existing cluster with on args and RAY_ADDRESS
# is not set.
ray.init()
assert "CPU" not in ray._private.state.cluster_resources()
ray.shutdown()
# Make sure we start a new cluster if "local" is explicitly passed.
# is not set.
ray.init(address="local")
assert "CPU" in ray._private.state.cluster_resources()
ray.shutdown()
def test_ray_resources_environment_variable(shutdown_only):
os.environ[
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE
] = '{"custom1":1, "custom2":2, "CPU":3}'
ray.init(resources={"custom1": 3, "custom3": 3})
cluster_resources = ray.cluster_resources()
print(cluster_resources)
assert cluster_resources["custom1"] == 1
assert cluster_resources["custom2"] == 2
assert cluster_resources["custom3"] == 3
assert cluster_resources["CPU"] == 3
def test_ray_labels_environment_variables(shutdown_only):
os.environ[
ray_constants.LABELS_ENVIRONMENT_VARIABLE
] = '{"custom1":"1", "custom2":"2"}'
ray.init(labels={"custom1": "3", "custom3": "3"})
node_info = ray.nodes()[0]
assert node_info["Labels"]["custom1"] == "1"
assert node_info["Labels"]["custom2"] == "2"
assert node_info["Labels"]["custom3"] == "3"
@pytest.mark.parametrize(
"accelerator_type",
[ray.util.accelerators.NVIDIA_TESLA_V100, ray.util.accelerators.AWS_NEURON_CORE],
)
def test_accelerator_type_api(accelerator_type, shutdown_only):
resource_name = f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"
ray.init(num_cpus=4, resources={resource_name: 1})
quantity = 1
@ray.remote(accelerator_type=accelerator_type)
def decorated_func(quantity):
wait_for_condition(lambda: ray.available_resources()[resource_name] < quantity)
return True
assert ray.get(decorated_func.remote(quantity))
def via_options_func(quantity):
wait_for_condition(lambda: ray.available_resources()[resource_name] < quantity)
return True
assert ray.get(
ray.remote(via_options_func)
.options(accelerator_type=accelerator_type)
.remote(quantity)
)
@ray.remote(accelerator_type=accelerator_type)
class DecoratedActor:
def __init__(self):
pass
def initialized(self):
pass
class ActorWithOptions:
def __init__(self):
pass
def initialized(self):
pass
decorated_actor = DecoratedActor.remote()
# Avoid a race condition where the actor hasn't been initialized and
# claimed the resources yet.
ray.get(decorated_actor.initialized.remote())
wait_for_condition(lambda: ray.available_resources()[resource_name] < quantity)
quantity = ray.available_resources()[resource_name]
with_options = (
ray.remote(ActorWithOptions).options(accelerator_type=accelerator_type).remote()
)
ray.get(with_options.initialized.remote())
wait_for_condition(lambda: ray.available_resources()[resource_name] < quantity)
@pytest.mark.skipif(sys.platform == "win32", reason="not relevant for windows")
def test_get_system_memory():
# cgroups v1, set
with tempfile.NamedTemporaryFile("w") as memory_limit_file:
memory_limit_file.write("100")
memory_limit_file.flush()
assert (
ray._common.utils.get_system_memory(
memory_limit_filename=memory_limit_file.name,
memory_limit_filename_v2="__does_not_exist__",
)
== 100
)
# cgroups v1, high
with tempfile.NamedTemporaryFile("w") as memory_limit_file:
memory_limit_file.write(str(2**64))
memory_limit_file.flush()
psutil_memory_in_bytes = psutil.virtual_memory().total
assert (
ray._common.utils.get_system_memory(
memory_limit_filename=memory_limit_file.name,
memory_limit_filename_v2="__does_not_exist__",
)
== psutil_memory_in_bytes
)
# cgroups v2, set
with tempfile.NamedTemporaryFile("w") as memory_max_file:
memory_max_file.write("100\n")
memory_max_file.flush()
assert (
ray._common.utils.get_system_memory(
memory_limit_filename="__does_not_exist__",
memory_limit_filename_v2=memory_max_file.name,
)
== 100
)
# cgroups v2, not set
with tempfile.NamedTemporaryFile("w") as memory_max_file:
memory_max_file.write("max")
memory_max_file.flush()
psutil_memory_in_bytes = psutil.virtual_memory().total
assert (
ray._common.utils.get_system_memory(
memory_limit_filename="__does_not_exist__",
memory_limit_filename_v2=memory_max_file.name,
)
== psutil_memory_in_bytes
)
@pytest.mark.parametrize("in_k8s", [True, False])
@pytest.mark.parametrize("env_disable", [True, False])
@pytest.mark.parametrize("override_disable", [True, False])
@pytest.mark.parametrize("got_docker_cpus", [True, False])
def test_get_num_cpus(
in_k8s: bool,
env_disable: bool,
override_disable: bool,
got_docker_cpus: bool,
monkeypatch,
):
"""Tests
- Conditions under which ray._private.utils.get_num_cpus logs a warning about
docker.
- Fallback to multiprocessing.cpu_count if there's no docker count available.
"""
# Shouldn't get the log warning if we're in K8s, the env variable is set,
# the flag arg to get_num_cpus is set, or getting docker cpus fails.
# Otherwise, should get the log message.
should_not_log = any([in_k8s, env_disable, override_disable, not got_docker_cpus])
expected_warning = (
"Detecting docker specified CPUs. In "
"previous versions of Ray, CPU detection in containers "
"was incorrect. Please ensure that Ray has enough CPUs "
"allocated. As a temporary workaround to revert to the "
"prior behavior, set "
"`RAY_USE_MULTIPROCESSING_CPU_COUNT=1` as an env var "
"before starting Ray. Set the env var: "
"`RAY_DISABLE_DOCKER_CPU_WARNING=1` to mute this warning."
)
if got_docker_cpus:
mock_get_docker_cpus = mock.Mock(return_value=128)
else:
mock_get_docker_cpus = mock.Mock(side_effect=Exception())
if in_k8s:
monkeypatch.setenv("KUBERNETES_SERVICE_HOST", 1)
else:
try:
monkeypatch.delenv("KUBERNETES_SERVICE_HOST")
except KeyError:
pass
with mock.patch.multiple(
"ray._private.utils",
_get_docker_cpus=mock_get_docker_cpus,
ENV_DISABLE_DOCKER_CPU_WARNING=env_disable,
logger=mock.DEFAULT,
) as mocks:
num_cpus = ray._private.utils.get_num_cpus(override_disable)
if got_docker_cpus:
# Got the docker count of 128 CPUs in the giant mock container.
assert num_cpus == 128
else:
# Failed to get docker count and fell back to multiprocessing count.
assert num_cpus == multiprocessing.cpu_count()
if should_not_log:
mocks["logger"].warning.assert_not_called()
else:
mocks["logger"].warning.assert_called_with(expected_warning)
@pytest.mark.skipif(sys.platform == "win32", reason="not relevant for windows")
def test_detect_docker_cpus():
# No limits set
with tempfile.NamedTemporaryFile("w") as quota_file, tempfile.NamedTemporaryFile(
"w"
) as period_file, tempfile.NamedTemporaryFile("w") as cpuset_file:
quota_file.write("-1")
period_file.write("100000")
cpuset_file.write("0-63")
quota_file.flush()
period_file.flush()
cpuset_file.flush()
assert (
ray._private.utils._get_docker_cpus(
cpu_quota_file_name=quota_file.name,
cpu_period_file_name=period_file.name,
cpuset_file_name=cpuset_file.name,
)
== 64
)
# No cpuset used
with tempfile.NamedTemporaryFile("w") as quota_file, tempfile.NamedTemporaryFile(
"w"
) as period_file, tempfile.NamedTemporaryFile("w") as cpuset_file:
quota_file.write("-1")
period_file.write("100000")
cpuset_file.write("0-10,20,50-63")
quota_file.flush()
period_file.flush()
cpuset_file.flush()
assert (
ray._private.utils._get_docker_cpus(
cpu_quota_file_name=quota_file.name,
cpu_period_file_name=period_file.name,
cpuset_file_name=cpuset_file.name,
)
== 26
)
# Quota set
with tempfile.NamedTemporaryFile("w") as quota_file, tempfile.NamedTemporaryFile(
"w"
) as period_file, tempfile.NamedTemporaryFile("w") as cpuset_file:
quota_file.write("42")
period_file.write("100")
cpuset_file.write("0-63")
quota_file.flush()
period_file.flush()
cpuset_file.flush()
assert (
ray._private.utils._get_docker_cpus(
cpu_quota_file_name=quota_file.name,
cpu_period_file_name=period_file.name,
cpuset_file_name=cpuset_file.name,
)
== 0.42
)
# cgroups v2, cpu_quota set
with tempfile.NamedTemporaryFile("w") as cpu_max_file:
cpu_max_file.write("200000 100000")
cpu_max_file.flush()
assert (
ray._private.utils._get_docker_cpus(
cpu_quota_file_name="nope",
cpu_period_file_name="give_up",
cpuset_file_name="lose_hope",
cpu_max_file_name=cpu_max_file.name,
)
== 2.0
)
# cgroups v2, cpu_quota unset
with tempfile.NamedTemporaryFile("w") as cpu_max_file:
cpu_max_file.write("max 100000")
cpu_max_file.flush()
assert (
ray._private.utils._get_docker_cpus(
cpu_quota_file_name="nope",
cpu_period_file_name="give_up",
cpuset_file_name="lose_hope",
cpu_max_file_name=cpu_max_file.name,
)
is None
)
@pytest.mark.skipif(
sys.platform.startswith("win"), reason="No need to test on Windows."
)
@pytest.mark.parametrize("use_cgroups_v2", [True, False])
def test_k8s_cpu(use_cgroups_v2: bool):
"""Test all the functions in dashboard/k8s_utils.py.
Also test ray._private.utils.get_num_cpus when running in a K8s pod.
Files were obtained from within a K8s pod with 2 CPU request, CPU limit
unset, with 1 CPU of stress applied.
"""
# Some experimentally-obtained K8S CPU usage files for use in test_k8s_cpu.
PROCSTAT1 = """cpu 2945022 98 3329420 148744854 39522 0 118587 0 0 0
cpu0 370299 14 413841 18589778 5304 0 15288 0 0 0
cpu1 378637 10 414414 18589275 5283 0 14731 0 0 0
cpu2 367328 8 420914 18590974 4844 0 14416 0 0 0
cpu3 368378 11 423720 18572899 4948 0 14394 0 0 0
cpu4 369051 13 414615 18607285 4736 0 14383 0 0 0
cpu5 362958 10 415984 18576655 4590 0 16614 0 0 0
cpu6 362536 13 414430 18605197 4785 0 14353 0 0 0
cpu7 365833 15 411499 18612787 5028 0 14405 0 0 0
intr 1000694027 125 0 0 39 154 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1028 0 2160913 0 2779605 8 0 3981333 3665198 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
ctxt 1574979439
btime 1615208601
processes 857411
procs_running 6
procs_blocked 0
softirq 524311775 0 230142964 27143 63542182 0 0 171 74042767 0 156556548
""" # noqa
PROCSTAT2 = """cpu 2945152 98 3329436 148745483 39522 0 118587 0 0 0
cpu0 370399 14 413841 18589778 5304 0 15288 0 0 0
cpu1 378647 10 414415 18589362 5283 0 14731 0 0 0
cpu2 367329 8 420916 18591067 4844 0 14416 0 0 0
cpu3 368381 11 423724 18572989 4948 0 14395 0 0 0
cpu4 369052 13 414618 18607374 4736 0 14383 0 0 0
cpu5 362968 10 415986 18576741 4590 0 16614 0 0 0
cpu6 362537 13 414432 18605290 4785 0 14353 0 0 0
cpu7 365836 15 411502 18612878 5028 0 14405 0 0 0
intr 1000700905 125 0 0 39 154 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1028 0 2160923 0 2779605 8 0 3981353 3665218 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
ctxt 1574988760
btime 1615208601
processes 857411
procs_running 4
procs_blocked 0
softirq 524317451 0 230145523 27143 63542930 0 0 171 74043232 0 156558452
""" # noqa
CPUACCTUSAGE1 = "2268980984000"
CPUACCTUSAGE2 = "2270120061999"
CPU_STAT_1 = """usage_usec 2268980984
user_usec 5673216
system_usec 794353
nr_periods 168
nr_throttled 6
throttled_usec 638117
"""
CPU_STAT_2 = """usage_usec 2270120061
user_usec 5673216
system_usec 794353
nr_periods 168
nr_throttled 6
throttled_usec 638117
"""
cpu_file, cpu_v2_file, proc_stat_file = [
tempfile.NamedTemporaryFile("w+") for _ in range(3)
]
cpu_file.write(CPUACCTUSAGE1)
cpu_v2_file.write(CPU_STAT_1)
proc_stat_file.write(PROCSTAT1)
for file in cpu_file, cpu_v2_file, proc_stat_file:
file.flush()
if use_cgroups_v2:
# Should get a file not found for cpuacctusage if on cgroups v2
cpu_usage_file = "NO_SUCH_FILE"
else:
# If using cgroups v1, use the temp file we've just made
cpu_usage_file = cpu_file.name
with mock.patch(
"ray._private.utils.os.environ", {"KUBERNETES_SERVICE_HOST": "host"}
), mock.patch("ray.dashboard.k8s_utils.CPU_USAGE_PATH", cpu_usage_file), mock.patch(
"ray.dashboard.k8s_utils.CPU_USAGE_PATH_V2", cpu_v2_file.name
), mock.patch(
"ray.dashboard.k8s_utils.PROC_STAT_PATH", proc_stat_file.name
), mock.patch(
# get_num_cpus is tested elsewhere
"ray.dashboard.k8s_utils.get_num_cpus",
mock.Mock(return_value=2),
), mock.patch(
# Reset this global variable between tests.
"ray.dashboard.k8s_utils.last_system_usage",
None,
):
# Validate mocks:
# Confirm CPU_USAGE_PATH is found with cgroups v2, but not with v2.
from ray.dashboard.k8s_utils import CPU_USAGE_PATH
if use_cgroups_v2:
with pytest.raises(FileNotFoundError):
print(open(CPU_USAGE_PATH).read())
else:
print(open(CPU_USAGE_PATH).read())
# Test helpers
assert k8s_utils._cpu_usage() == 2268980984000
assert k8s_utils._system_usage() == 1551775030000000
assert k8s_utils._host_num_cpus() == 8
# No delta for first computation, return 0.
assert k8s_utils.cpu_percent() == 0.0
# Write new usage info obtained after 1 sec wait.
for file in cpu_file, cpu_v2_file, proc_stat_file:
file.truncate(0)
file.seek(0)
cpu_file.write(CPUACCTUSAGE2)
cpu_v2_file.write(CPU_STAT_2)
proc_stat_file.write(PROCSTAT2)
for file in cpu_file, cpu_v2_file, proc_stat_file:
file.flush()
# Files were extracted under 1 CPU of load on a 2 CPU pod
assert 50 < k8s_utils.cpu_percent() < 60
def test_sync_job_config(shutdown_only):
runtime_env = {"env_vars": {"key": "value"}}
ray.init(
job_config=ray.job_config.JobConfig(
runtime_env=runtime_env,
)
)
# Check that the job config is synchronized at the driver side.
job_config = ray._private.worker.global_worker.core_worker.get_job_config()
job_runtime_env = RuntimeEnv.deserialize(
job_config.runtime_env_info.serialized_runtime_env
)
assert job_runtime_env.env_vars() == runtime_env["env_vars"]
@ray.remote
def get_job_config():
job_config = ray._private.worker.global_worker.core_worker.get_job_config()
return job_config.SerializeToString()
# Check that the job config is synchronized at the worker side.
job_config = gcs_utils.JobConfig()
job_config.ParseFromString(ray.get(get_job_config.remote()))
job_runtime_env = RuntimeEnv.deserialize(
job_config.runtime_env_info.serialized_runtime_env
)
assert job_runtime_env.env_vars() == runtime_env["env_vars"]
def test_duplicated_arg(ray_start_cluster):
cluster = ray_start_cluster
cluster.add_node(num_cpus=1)
ray.init(address=cluster.address)
@ray.remote
def task_with_dup_arg(*args):
return sum(args)
# Basic verification.
arr = np.ones(1 * 1024 * 1024, dtype=np.uint8) # 1MB
ref = ray.put(arr)
assert np.array_equal(
ray.get(task_with_dup_arg.remote(ref, ref, ref)), sum([arr, arr, arr])
)
# Make sure it works when it is mixed with other args.
ref2 = ray.put(arr)
assert np.array_equal(
ray.get(task_with_dup_arg.remote(ref, ref2, ref)), sum([arr, arr, arr])
)
# Test complicated scenario with multi nodes.
cluster.add_node(num_cpus=1, resources={"worker_1": 1})
cluster.add_node(num_cpus=1, resources={"worker_2": 1})
cluster.wait_for_nodes()
@ray.remote
def create_remote_ref(arr):
return ray.put(arr)
@ray.remote
def task_with_dup_arg_ref(*args):
args = ray.get(list(args))
return sum(args)
ref1 = create_remote_ref.options(resources={"worker_1": 1}).remote(arr)
ref2 = create_remote_ref.options(resources={"worker_2": 1}).remote(arr)
ref3 = create_remote_ref.remote(arr)
np.array_equal(
ray.get(task_with_dup_arg_ref.remote(ref1, ref2, ref3, ref1, ref2, ref3)),
sum([arr] * 6),
)
if __name__ == "__main__":
sys.exit(pytest.main(["-sv", __file__]))