wehub-resource-sync
370 lines
11 KiB
Python
370 lines
11 KiB
Python
# coding: utf-8
|
|
import gc
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
import ray.cluster_utils
|
|
from ray._common.constants import HEAD_NODE_RESOURCE_NAME
|
|
from ray._common.test_utils import run_string_as_driver
|
|
from ray._private.test_utils import (
|
|
client_test_enabled,
|
|
wait_for_pid_to_exit,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def test_background_tasks_with_max_calls(shutdown_only):
|
|
ray.init(
|
|
# TODO (Alex): We need to fix
|
|
# https://github.com/ray-project/ray/issues/20203 to remove this flag.
|
|
num_cpus=2,
|
|
_system_config={"worker_cap_initial_backoff_delay_ms": 0},
|
|
)
|
|
|
|
num_tasks = 3 if sys.platform == "win32" else 10
|
|
|
|
@ray.remote
|
|
def g():
|
|
time.sleep(0.1)
|
|
return 0
|
|
|
|
@ray.remote(max_calls=1, max_retries=0)
|
|
def f():
|
|
return [g.remote()]
|
|
|
|
nested = ray.get([f.remote() for _ in range(num_tasks)])
|
|
|
|
# Should still be able to retrieve these objects, since f's workers will
|
|
# wait for g to finish before exiting.
|
|
ray.get([x[0] for x in nested])
|
|
|
|
@ray.remote(max_calls=1, max_retries=0)
|
|
def f():
|
|
return os.getpid(), g.remote()
|
|
|
|
nested = ray.get([f.remote() for _ in range(num_tasks)])
|
|
while nested:
|
|
pid, g_id = nested.pop(0)
|
|
assert ray.get(g_id) == 0
|
|
del g_id
|
|
# Necessary to dereference the object via GC, so the worker can exit.
|
|
gc.collect()
|
|
wait_for_pid_to_exit(pid)
|
|
|
|
|
|
def test_actor_killing(shutdown_only):
|
|
# This is to test create and kill an actor immediately
|
|
ray.init(num_cpus=1)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def foo(self):
|
|
return None
|
|
|
|
worker_1 = Actor.remote()
|
|
ray.kill(worker_1)
|
|
worker_2 = Actor.remote()
|
|
assert ray.get(worker_2.foo.remote()) is None
|
|
ray.kill(worker_2)
|
|
|
|
worker_1 = Actor.options(max_restarts=1, max_task_retries=-1).remote()
|
|
ray.kill(worker_1, no_restart=False)
|
|
assert ray.get(worker_1.foo.remote()) is None
|
|
|
|
ray.kill(worker_1, no_restart=False)
|
|
worker_2 = Actor.remote()
|
|
assert ray.get(worker_2.foo.remote()) is None
|
|
|
|
|
|
def test_internal_kv(ray_start_regular):
|
|
import ray.experimental.internal_kv as kv
|
|
|
|
assert kv._internal_kv_get("k1") is None
|
|
assert kv._internal_kv_put("k1", "v1") is False
|
|
assert kv._internal_kv_put("k1", "v1") is True
|
|
assert kv._internal_kv_get("k1") == b"v1"
|
|
assert kv._internal_kv_exists(b"k1") is True
|
|
assert kv._internal_kv_exists(b"k2") is False
|
|
|
|
assert kv._internal_kv_get("k1", namespace="n") is None
|
|
assert kv._internal_kv_put("k1", "v1", namespace="n") is False
|
|
assert kv._internal_kv_put("k1", "v1", namespace="n") is True
|
|
assert kv._internal_kv_put("k1", "v2", True, namespace="n") is True
|
|
assert kv._internal_kv_get("k1", namespace="n") == b"v2"
|
|
|
|
assert kv._internal_kv_del("k1") == 1
|
|
assert kv._internal_kv_del("k1") == 0
|
|
assert kv._internal_kv_get("k1") is None
|
|
|
|
assert kv._internal_kv_put("k2", "v2", namespace="n") is False
|
|
assert kv._internal_kv_put("k3", "v3", namespace="n") is False
|
|
|
|
assert set(kv._internal_kv_list("k", namespace="n")) == {b"k1", b"k2", b"k3"}
|
|
assert kv._internal_kv_del("k", del_by_prefix=True, namespace="n") == 3
|
|
assert kv._internal_kv_del("x", del_by_prefix=True, namespace="n") == 0
|
|
assert kv._internal_kv_get("k1", namespace="n") is None
|
|
assert kv._internal_kv_get("k2", namespace="n") is None
|
|
assert kv._internal_kv_get("k3", namespace="n") is None
|
|
|
|
with pytest.raises(ray.exceptions.RaySystemError):
|
|
kv._internal_kv_put("@namespace_", "x", True)
|
|
with pytest.raises(ray.exceptions.RaySystemError):
|
|
kv._internal_kv_get("@namespace_", namespace="n")
|
|
with pytest.raises(ray.exceptions.RaySystemError):
|
|
kv._internal_kv_del("@namespace_def", namespace="n")
|
|
with pytest.raises(ray.exceptions.RaySystemError):
|
|
kv._internal_kv_list("@namespace_abc", namespace="n")
|
|
|
|
|
|
def test_exit_logging():
|
|
log = run_string_as_driver(
|
|
"""
|
|
import ray
|
|
|
|
@ray.remote
|
|
class A:
|
|
def pid(self):
|
|
import os
|
|
return os.getpid()
|
|
|
|
|
|
a = A.remote()
|
|
ray.get(a.pid.remote())
|
|
"""
|
|
)
|
|
assert "Traceback" not in log
|
|
|
|
|
|
def test_worker_sys_path_contains_driver_script_directory(tmp_path, monkeypatch):
|
|
package_folder = tmp_path / "package"
|
|
package_folder.mkdir()
|
|
init_file = package_folder / "__init__.py"
|
|
init_file.write_text("")
|
|
|
|
module1_file = package_folder / "module1.py"
|
|
module1_file.write_text(
|
|
f"""
|
|
import sys
|
|
import ray
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
def sys_path():
|
|
return sys.path
|
|
|
|
remote_sys_path = ray.get(sys_path.remote())
|
|
assert r'{str(package_folder)}' in remote_sys_path, remote_sys_path
|
|
"""
|
|
)
|
|
|
|
# Ray's handling of sys.path does not work with PYTHONSAFEPATH.
|
|
env = os.environ.copy()
|
|
if env.get("PYTHONSAFEPATH", "") != "":
|
|
env["PYTHONSAFEPATH"] = "" # Set to empty string to disable.
|
|
subprocess.check_call([sys.executable, str(module1_file)], env=env)
|
|
|
|
# If the driver script is run via `python -m`,
|
|
# the script directory is not included in sys.path.
|
|
module2_file = package_folder / "module2.py"
|
|
module2_file.write_text(
|
|
f"""
|
|
import sys
|
|
import ray
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
def sys_path():
|
|
return sys.path
|
|
|
|
remote_sys_path = ray.get(sys_path.remote())
|
|
assert r'{str(package_folder)}' not in remote_sys_path, remote_sys_path
|
|
"""
|
|
)
|
|
monkeypatch.chdir(str(tmp_path))
|
|
subprocess.check_call([sys.executable, "-m", "package.module2"], env=env)
|
|
|
|
|
|
def test_worker_kv_calls(monkeypatch, shutdown_only):
|
|
monkeypatch.setenv("TEST_RAY_COLLECT_KV_FREQUENCY", "1")
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
def get_kv_metrics():
|
|
from time import sleep
|
|
|
|
sleep(2)
|
|
return ray._private.utils._CALLED_FREQ
|
|
|
|
freqs = ray.get(get_kv_metrics.remote())
|
|
# So far we have the following gets
|
|
"""
|
|
b'cluster' b'CLUSTER_METADATA'
|
|
b'tracing' b'tracing_startup_hook'
|
|
b'fun' b'RemoteFunction:01000000:\x00\x00\x00\x00\x00\x00\x00\x01'
|
|
"""
|
|
# !!!If you want to increase this number, please let ray-core knows this!!!
|
|
assert freqs["internal_kv_get"] == 3
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Fails on Windows.")
|
|
@pytest.mark.skipif(sys.version_info >= (3, 14), reason="Flaky on Python 3.14")
|
|
@pytest.mark.parametrize("root_process_no_site", [0, 1])
|
|
@pytest.mark.parametrize("root_process_no_user_site", [0, 1])
|
|
def test_site_flag_inherited(
|
|
shutdown_only, monkeypatch, root_process_no_site, root_process_no_user_site
|
|
):
|
|
# The flags we're testing could prevent Python workers in the test environment
|
|
# from locating site packages; the workers would fail to find Ray and would thus
|
|
# fail to start. To prevent that, set the PYTHONPATH env. The env will be inherited
|
|
# by the Python workers, so that they are able to import Ray.
|
|
monkeypatch.setenv("PYTHONPATH", ":".join(sys.path))
|
|
|
|
@ray.remote
|
|
def get_flags():
|
|
return sys.flags.no_site, sys.flags.no_user_site
|
|
|
|
with patch.multiple(
|
|
"ray._private.services",
|
|
_no_site=Mock(return_value=root_process_no_site),
|
|
_no_user_site=Mock(return_value=root_process_no_user_site),
|
|
):
|
|
ray.init()
|
|
worker_process_no_site, worker_process_no_user_site = ray.get(
|
|
get_flags.remote()
|
|
)
|
|
assert worker_process_no_site == root_process_no_site
|
|
assert worker_process_no_user_site == root_process_no_user_site
|
|
|
|
|
|
@pytest.mark.parametrize("preload", [True, False])
|
|
def test_preload_workers(ray_start_cluster, preload):
|
|
"""
|
|
Verify preload_python_modules actually preloads modules in the Ray workers.
|
|
Also verify that it does not crash if a non-existent module is provided.
|
|
"""
|
|
cluster = ray_start_cluster
|
|
|
|
# Specifying imports not currently imported by default_worker.py
|
|
expect_succeed_imports = ["html.parser", "webbrowser"]
|
|
expect_fail_imports = ["fake_module_expect_ModuleNotFoundError"]
|
|
|
|
if preload:
|
|
cluster.add_node(
|
|
_system_config={
|
|
"preload_python_modules": [
|
|
*expect_succeed_imports,
|
|
*expect_fail_imports,
|
|
]
|
|
}
|
|
)
|
|
else:
|
|
cluster.add_node()
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class Latch:
|
|
"""
|
|
Used to ensure two separate worker processes.
|
|
"""
|
|
|
|
def __init__(self, count):
|
|
self.count = count
|
|
|
|
def decr(self):
|
|
self.count -= 1
|
|
|
|
def is_ready(self):
|
|
return self.count <= 0
|
|
|
|
def wait_latch(latch):
|
|
latch.decr.remote()
|
|
while not ray.get(latch.is_ready.remote()):
|
|
time.sleep(0.01)
|
|
|
|
def assert_correct_imports():
|
|
import sys
|
|
|
|
imported_modules = set(sys.modules.keys())
|
|
|
|
if preload:
|
|
for expected_import in expect_succeed_imports:
|
|
assert (
|
|
expected_import in imported_modules
|
|
), f"Expected {expected_import} to be in {imported_modules}"
|
|
for unexpected_import in expect_fail_imports:
|
|
assert (
|
|
unexpected_import not in imported_modules
|
|
), f"Expected {unexpected_import} to not be in {imported_modules}"
|
|
else:
|
|
for unexpected_import in expect_succeed_imports:
|
|
assert (
|
|
unexpected_import not in imported_modules
|
|
), f"Expected {unexpected_import} to not be in {imported_modules}"
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class Actor:
|
|
def verify_imports(self, latch):
|
|
wait_latch(latch)
|
|
assert_correct_imports()
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def verify_imports(latch):
|
|
wait_latch(latch)
|
|
assert_correct_imports()
|
|
|
|
latch = Latch.remote(2)
|
|
actor = Actor.remote()
|
|
futures = [verify_imports.remote(latch), actor.verify_imports.remote(latch)]
|
|
ray.get(futures)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="only server mode")
|
|
def test_gcs_port_env(shutdown_only):
|
|
try:
|
|
with unittest.mock.patch.dict(os.environ):
|
|
os.environ["RAY_GCS_SERVER_PORT"] = "12345"
|
|
ray.init()
|
|
except RuntimeError:
|
|
pass
|
|
# it's ok to throw runtime error for port conflicts
|
|
|
|
|
|
def test_head_node_resource(ray_start_cluster):
|
|
"""Test that the special head node resource is set."""
|
|
cluster = ray_start_cluster
|
|
# head node
|
|
cluster.add_node(num_cpus=1)
|
|
ray.init(address=cluster.address)
|
|
|
|
assert ray.cluster_resources()[HEAD_NODE_RESOURCE_NAME] == 1
|
|
|
|
# worker node
|
|
cluster.add_node(num_cpus=1)
|
|
|
|
assert ray.cluster_resources()[HEAD_NODE_RESOURCE_NAME] == 1
|
|
|
|
|
|
def test_head_node_resource_ray_init(shutdown_only):
|
|
ray.init()
|
|
|
|
assert ray.cluster_resources()[HEAD_NODE_RESOURCE_NAME] == 1
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="grpc deadlock with ray client")
|
|
def test_head_node_resource_ray_start(call_ray_start):
|
|
ray.init(address=call_ray_start)
|
|
|
|
assert ray.cluster_resources()[HEAD_NODE_RESOURCE_NAME] == 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|