wehub-resource-sync
541 lines
19 KiB
Python
541 lines
19 KiB
Python
import dataclasses
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
import ray._private.ray_constants as ray_constants
|
|
from ray._private.runtime_env.uri_cache import URICache
|
|
from ray._private.runtime_env.utils import (
|
|
SubprocessCalledProcessError,
|
|
check_output_cmd,
|
|
)
|
|
from ray._private.test_utils import (
|
|
chdir,
|
|
)
|
|
from ray.runtime_env import RuntimeEnv
|
|
from ray.runtime_env.runtime_env import (
|
|
RuntimeEnvConfig,
|
|
_merge_runtime_env,
|
|
)
|
|
|
|
|
|
def test_runtime_env_merge():
|
|
# Both are None.
|
|
parent = None
|
|
child = None
|
|
assert _merge_runtime_env(parent, child) == {}
|
|
|
|
parent = {}
|
|
child = None
|
|
assert _merge_runtime_env(parent, child) == {}
|
|
|
|
parent = None
|
|
child = {}
|
|
assert _merge_runtime_env(parent, child) == {}
|
|
|
|
parent = {}
|
|
child = {}
|
|
assert _merge_runtime_env(parent, child) == {}
|
|
|
|
# Only parent is given.
|
|
parent = {"conda": ["requests"], "env_vars": {"A": "1"}}
|
|
child = None
|
|
assert _merge_runtime_env(parent, child) == parent
|
|
|
|
# Only child is given.
|
|
parent = None
|
|
child = {"conda": ["requests"], "env_vars": {"A": "1"}}
|
|
assert _merge_runtime_env(parent, child) == child
|
|
|
|
# Successful case.
|
|
parent = {"conda": ["requests"], "env_vars": {"A": "1"}}
|
|
child = {"pip": ["requests"], "env_vars": {"B": "2"}}
|
|
assert _merge_runtime_env(parent, child) == {
|
|
"conda": ["requests"],
|
|
"pip": ["requests"],
|
|
"env_vars": {"A": "1", "B": "2"},
|
|
}
|
|
|
|
# Failure case
|
|
parent = {"pip": ["requests"], "env_vars": {"A": "1"}}
|
|
child = {"pip": ["colors"], "env_vars": {"B": "2"}}
|
|
assert _merge_runtime_env(parent, child) is None
|
|
|
|
# Failure case (env_vars)
|
|
parent = {"pip": ["requests"], "env_vars": {"A": "1"}}
|
|
child = {"conda": ["requests"], "env_vars": {"A": "2"}}
|
|
assert _merge_runtime_env(parent, child) is None
|
|
|
|
# override = True
|
|
parent = {"pip": ["requests"], "env_vars": {"A": "1"}}
|
|
child = {"pip": ["colors"], "env_vars": {"B": "2"}}
|
|
assert _merge_runtime_env(parent, child, override=True) == {
|
|
"pip": ["colors"],
|
|
"env_vars": {"A": "1", "B": "2"},
|
|
}
|
|
|
|
# override = True + env vars
|
|
parent = {"pip": ["requests"], "env_vars": {"A": "1"}}
|
|
child = {"pip": ["colors"], "conda": ["requests"], "env_vars": {"A": "2"}}
|
|
assert _merge_runtime_env(parent, child, override=True) == {
|
|
"pip": ["colors"],
|
|
"env_vars": {"A": "2"},
|
|
"conda": ["requests"],
|
|
}
|
|
|
|
|
|
def test_current_py_version_supported():
|
|
"""Test that the running python version is supported.
|
|
|
|
This is run as a check in the Ray `runtime_env` `conda` code
|
|
before downloading the Ray wheel into the conda environment.
|
|
If Ray wheels are not available for this python version, then
|
|
the `conda` environment installation will fail.
|
|
|
|
When a new python version is added to the Ray wheels, please update
|
|
`ray_constants.RUNTIME_ENV_CONDA_PY_VERSIONS`. In a subsequent commit,
|
|
once wheels have been built for the new python version, please update
|
|
the tests test_get_wheel_filename, test_get_master_wheel_url, and
|
|
(after the first Ray release with the new python version)
|
|
test_get_release_wheel_url.
|
|
"""
|
|
py_version = sys.version_info[:2]
|
|
assert py_version in ray_constants.RUNTIME_ENV_CONDA_PY_VERSIONS
|
|
|
|
|
|
def test_compatible_with_dataclasses():
|
|
"""Test that the output of RuntimeEnv.to_dict() can be used as a dataclass field."""
|
|
config = RuntimeEnvConfig(setup_timeout_seconds=1)
|
|
runtime_env = RuntimeEnv(
|
|
pip={
|
|
"packages": ["tensorflow", "requests"],
|
|
"pip_check": False,
|
|
"pip_version": "==23.3.2;python_version=='3.9.16'",
|
|
},
|
|
env_vars={"FOO": "BAR"},
|
|
config=config,
|
|
)
|
|
|
|
@dataclass
|
|
class RuntimeEnvDataClass:
|
|
runtime_env: Dict[str, Any]
|
|
|
|
dataclasses.asdict(RuntimeEnvDataClass(runtime_env.to_dict()))
|
|
|
|
@dataclass
|
|
class RuntimeEnvConfigDataClass:
|
|
config: Dict[str, Any]
|
|
|
|
dataclasses.asdict(RuntimeEnvConfigDataClass(config.to_dict()))
|
|
|
|
|
|
@pytest.mark.parametrize("runtime_env_class", [dict, RuntimeEnv])
|
|
def test_container_option_serialize(runtime_env_class):
|
|
runtime_env = runtime_env_class(
|
|
container={"image": "ray:latest", "run_options": ["--name=test"]}
|
|
)
|
|
job_config = ray.job_config.JobConfig(runtime_env=runtime_env)
|
|
job_config_serialized = job_config._serialize()
|
|
# job_config_serialized is JobConfig protobuf serialized string,
|
|
# job_config.runtime_env_info.serialized_runtime_env
|
|
# has container_option info
|
|
assert job_config_serialized.count(b"ray:latest") == 1
|
|
assert job_config_serialized.count(b"--name=test") == 1
|
|
|
|
|
|
class TestURICache:
|
|
def test_zero_cache_size(self):
|
|
uris_to_sizes = {"5": 5, "3": 3}
|
|
|
|
def delete_fn(uri, logger):
|
|
return uris_to_sizes[uri]
|
|
|
|
cache = URICache(delete_fn, max_total_size_bytes=0, debug_mode=True)
|
|
cache.add("5", 5)
|
|
assert cache.get_total_size_bytes() == 5
|
|
cache.mark_unused("5")
|
|
assert cache.get_total_size_bytes() == 0
|
|
cache.add("3", 3)
|
|
cache.add("5", 5)
|
|
assert cache.get_total_size_bytes() == 8
|
|
cache.mark_unused("3")
|
|
cache.mark_unused("5")
|
|
assert cache.get_total_size_bytes() == 0
|
|
|
|
def test_nonzero_cache_size(self):
|
|
uris_to_sizes = {"a": 4, "b": 4, "c": 4}
|
|
|
|
def delete_fn(uri, logger):
|
|
return uris_to_sizes[uri]
|
|
|
|
cache = URICache(delete_fn, max_total_size_bytes=10, debug_mode=True)
|
|
cache.add("a", 4)
|
|
cache.add("b", 4)
|
|
cache.mark_unused("a")
|
|
assert "a" in cache
|
|
cache.add("c", 4)
|
|
# Now we have total size 12, which exceeds the max size 10.
|
|
assert cache.get_total_size_bytes() == 8
|
|
# "a" was the only unused URI, so it must have been deleted.
|
|
assert "b" and "c" in cache and "a" not in cache
|
|
|
|
def test_mark_used_nonadded_uri_error(self):
|
|
cache = URICache(debug_mode=True)
|
|
with pytest.raises(ValueError):
|
|
cache.mark_used("nonadded_uri")
|
|
|
|
def test_mark_used(self):
|
|
uris_to_sizes = {"a": 3, "b": 3, "big": 300}
|
|
|
|
def delete_fn(uri, logger):
|
|
return uris_to_sizes[uri]
|
|
|
|
cache = URICache(delete_fn, max_total_size_bytes=10, debug_mode=True)
|
|
cache.add("a", 3)
|
|
cache.add("b", 3)
|
|
cache.mark_unused("a")
|
|
cache.mark_unused("b")
|
|
assert "a" in cache and "b" in cache
|
|
assert cache.get_total_size_bytes() == 6
|
|
|
|
cache.mark_used("a")
|
|
cache.add("big", 300)
|
|
# We are over capacity and the only unused URI is "b", so we delete it
|
|
assert "a" in cache and "big" in cache and "b" not in cache
|
|
assert cache.get_total_size_bytes() == 303
|
|
|
|
cache.mark_unused("big")
|
|
assert "big" not in cache
|
|
assert cache.get_total_size_bytes() == 3
|
|
|
|
def test_many_URIs(self):
|
|
uris_to_sizes = {str(i): i for i in range(1000)}
|
|
|
|
def delete_fn(uri, logger):
|
|
return uris_to_sizes[uri]
|
|
|
|
cache = URICache(delete_fn, debug_mode=True)
|
|
for i in range(1000):
|
|
cache.add(str(i), i)
|
|
for i in range(1000):
|
|
cache.mark_unused(str(i))
|
|
for i in range(1000):
|
|
assert str(i) in cache
|
|
|
|
def test_delete_fn_called(self):
|
|
num_delete_fn_calls = 0
|
|
uris_to_sizes = {"a": 8, "b": 6, "c": 4, "d": 20}
|
|
|
|
def delete_fn(uri, logger):
|
|
nonlocal num_delete_fn_calls
|
|
num_delete_fn_calls += 1
|
|
return uris_to_sizes[uri]
|
|
|
|
cache = URICache(delete_fn, max_total_size_bytes=10, debug_mode=True)
|
|
cache.add("a", 8)
|
|
cache.add("b", 6)
|
|
cache.mark_unused("b")
|
|
# Total size is 14 > 10, so we need to delete "b".
|
|
assert num_delete_fn_calls == 1
|
|
|
|
cache.add("c", 4)
|
|
cache.mark_unused("c")
|
|
# Total size is 12 > 10, so we delete "c".
|
|
assert num_delete_fn_calls == 2
|
|
|
|
cache.mark_unused("a")
|
|
# Total size is 8 <= 10, so we shouldn't delete anything.
|
|
assert num_delete_fn_calls == 2
|
|
|
|
cache.add("d", 20)
|
|
# Total size is 28 > 10, so we delete "a".
|
|
assert num_delete_fn_calls == 3
|
|
|
|
cache.mark_unused("d")
|
|
# Total size is 20 > 10, so we delete "d".
|
|
assert num_delete_fn_calls == 4
|
|
|
|
|
|
@pytest.fixture
|
|
def enable_dev_mode(local_env_var_enabled, monkeypatch):
|
|
enabled = "1" if local_env_var_enabled else "0"
|
|
monkeypatch.setenv("RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED", enabled)
|
|
yield
|
|
|
|
|
|
def test_subprocess_error():
|
|
ex = SubprocessCalledProcessError
|
|
with pytest.raises(subprocess.SubprocessError) as e:
|
|
raise ex(123, "abc")
|
|
assert "test_out" not in str(e.value)
|
|
assert "test_err" not in str(e.value)
|
|
with pytest.raises(subprocess.SubprocessError) as e:
|
|
raise ex(123, "abc", stderr="test_err")
|
|
assert "test_out" not in str(e.value)
|
|
assert "test_err" in str(e.value)
|
|
with pytest.raises(subprocess.SubprocessError) as e:
|
|
raise ex(123, "abc", output="test_out")
|
|
assert "test_out" in str(e.value)
|
|
assert "test_err" not in str(e.value)
|
|
with pytest.raises(subprocess.SubprocessError) as e:
|
|
raise ex(123, "abc", output="test_out", stderr="test_err")
|
|
assert "test_out" in str(e.value)
|
|
assert "test_err" in str(e.value)
|
|
|
|
|
|
def test_subprocess_error_with_last_n_lines():
|
|
stdout = "1\n2\n3\n4\n5\n"
|
|
stderr = "5\n4\n3\n2\n1\n"
|
|
exception = SubprocessCalledProcessError(888, "abc", output=stdout, stderr=stderr)
|
|
exception.LAST_N_LINES = 3
|
|
exception_str = str(exception)
|
|
assert "cmd" not in exception_str
|
|
assert "Last 3 lines" in exception_str
|
|
s = "".join([s.strip() for s in exception_str.splitlines()])
|
|
assert "345" in s
|
|
assert "321" in s
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_output_cmd():
|
|
cmd = "dir" if sys.platform.startswith("win") else "pwd"
|
|
logs = []
|
|
|
|
class _FakeLogger:
|
|
def __getattr__(self, item):
|
|
def _log(formatter, *args):
|
|
logs.append(formatter % args)
|
|
|
|
return _log
|
|
|
|
for _ in range(2):
|
|
output = await check_output_cmd([cmd], logger=_FakeLogger())
|
|
assert len(output) > 0
|
|
|
|
all_log_string = "\n".join(logs)
|
|
|
|
# Check the cmd index generator works.
|
|
assert "cmd[1]" in all_log_string
|
|
assert "cmd[2]" in all_log_string
|
|
|
|
# Test communicate fails.
|
|
with mock.patch(
|
|
"asyncio.subprocess.Process.communicate",
|
|
side_effect=Exception("fake exception"),
|
|
):
|
|
with pytest.raises(RuntimeError) as e:
|
|
await check_output_cmd([cmd], logger=_FakeLogger())
|
|
# Make sure the exception has cmd trace info.
|
|
assert "cmd[3]" in str(e.value)
|
|
|
|
# Test asyncio.create_subprocess_exec fails.
|
|
with pytest.raises(RuntimeError) as e:
|
|
await check_output_cmd(["not_exist_cmd"], logger=_FakeLogger())
|
|
# Make sure the exception has cmd trace info.
|
|
assert "cmd[4]" in str(e.value)
|
|
|
|
# Test returncode != 0.
|
|
with pytest.raises(SubprocessCalledProcessError) as e:
|
|
await check_output_cmd([cmd, "--abc"], logger=_FakeLogger())
|
|
# Make sure the exception has cmd trace info.
|
|
assert "cmd[5]" in str(e.value)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"option",
|
|
["pip_list", "pip_dict", "conda_name", "conda_dict", "container"],
|
|
)
|
|
def test_serialize_deserialize(option):
|
|
runtime_env = dict()
|
|
if option == "pip_list":
|
|
runtime_env["pip"] = ["pkg1", "pkg2"]
|
|
elif option == "pip_dict":
|
|
runtime_env["pip"] = {
|
|
"packages": ["pkg1", "pkg2"],
|
|
"pip_check": False,
|
|
"pip_version": "<22,>20",
|
|
}
|
|
elif option == "conda_name":
|
|
runtime_env["conda"] = "env_name"
|
|
elif option == "conda_dict":
|
|
runtime_env["conda"] = {"dependencies": ["dep1", "dep2"]}
|
|
elif option == "container":
|
|
runtime_env["container"] = {
|
|
"image": "anyscale/ray-ml:nightly-py38-cpu",
|
|
"worker_path": "/root/python/ray/_private/workers/default_worker.py",
|
|
"run_options": ["--cap-drop SYS_ADMIN", "--log-level=debug"],
|
|
}
|
|
else:
|
|
raise ValueError("unexpected option " + str(option))
|
|
|
|
typed_runtime_env = RuntimeEnv(**runtime_env)
|
|
serialized_runtime_env = typed_runtime_env.serialize()
|
|
cls_runtime_env = RuntimeEnv.deserialize(serialized_runtime_env)
|
|
cls_runtime_env_dict = cls_runtime_env.to_dict()
|
|
|
|
if "pip" in typed_runtime_env and isinstance(typed_runtime_env["pip"], list):
|
|
pip_config_in_cls_runtime_env = cls_runtime_env_dict.pop("pip")
|
|
pip_config_in_runtime_env = typed_runtime_env.pop("pip")
|
|
assert {
|
|
"packages": pip_config_in_runtime_env,
|
|
"pip_check": False,
|
|
} == pip_config_in_cls_runtime_env
|
|
|
|
assert cls_runtime_env_dict == typed_runtime_env
|
|
|
|
|
|
def test_runtime_env_interface():
|
|
# Test the interface related to working_dir
|
|
default_working_dir = "s3://bucket/key.zip"
|
|
modify_working_dir = "s3://bucket/key_A.zip"
|
|
runtime_env = RuntimeEnv(working_dir=default_working_dir)
|
|
runtime_env_dict = runtime_env.to_dict()
|
|
assert runtime_env.working_dir_uri() == default_working_dir
|
|
runtime_env["working_dir"] = modify_working_dir
|
|
runtime_env_dict["working_dir"] = modify_working_dir
|
|
assert runtime_env.working_dir_uri() == modify_working_dir
|
|
assert runtime_env.to_dict() == runtime_env_dict
|
|
|
|
runtime_env.pop("working_dir")
|
|
assert runtime_env.to_dict() == {}
|
|
|
|
# Test the interface related to py_modules
|
|
init_py_modules = ["s3://bucket/key_1.zip", "s3://bucket/key_2.zip"]
|
|
addition_py_modules = ["s3://bucket/key_3.zip", "s3://bucket/key_4.zip"]
|
|
runtime_env = RuntimeEnv(py_modules=init_py_modules)
|
|
runtime_env_dict = runtime_env.to_dict()
|
|
assert set(runtime_env.py_modules_uris()) == set(init_py_modules)
|
|
runtime_env["py_modules"].extend(addition_py_modules)
|
|
runtime_env_dict["py_modules"].extend(addition_py_modules)
|
|
assert set(runtime_env.py_modules_uris()) == set(
|
|
init_py_modules + addition_py_modules
|
|
)
|
|
assert runtime_env.to_dict() == runtime_env_dict
|
|
|
|
runtime_env.pop("py_modules")
|
|
assert runtime_env.to_dict() == {}
|
|
|
|
# Test the interface related to env_vars
|
|
init_env_vars = {"A": "a", "B": "b"}
|
|
update_env_vars = {"C": "c"}
|
|
runtime_env = RuntimeEnv(env_vars=init_env_vars)
|
|
runtime_env_dict = runtime_env.to_dict()
|
|
runtime_env["env_vars"].update(update_env_vars)
|
|
runtime_env_dict["env_vars"].update(update_env_vars)
|
|
init_env_vars_copy = init_env_vars.copy()
|
|
init_env_vars_copy.update(update_env_vars)
|
|
assert runtime_env["env_vars"] == init_env_vars_copy
|
|
assert runtime_env_dict == runtime_env.to_dict()
|
|
|
|
runtime_env.pop("env_vars")
|
|
assert runtime_env.to_dict() == {}
|
|
|
|
# Test the interface related to conda
|
|
conda_name = "conda"
|
|
modify_conda_name = "conda_A"
|
|
conda_config = {"dependencies": ["dep1", "dep2"]}
|
|
runtime_env = RuntimeEnv(conda=conda_name)
|
|
runtime_env_dict = runtime_env.to_dict()
|
|
assert runtime_env.has_conda()
|
|
assert runtime_env.conda_env_name() == conda_name
|
|
assert runtime_env.conda_config() is None
|
|
runtime_env["conda"] = modify_conda_name
|
|
runtime_env_dict["conda"] = modify_conda_name
|
|
assert runtime_env_dict == runtime_env.to_dict()
|
|
assert runtime_env.has_conda()
|
|
assert runtime_env.conda_env_name() == modify_conda_name
|
|
assert runtime_env.conda_config() is None
|
|
runtime_env["conda"] = conda_config
|
|
runtime_env_dict["conda"] = conda_config
|
|
assert runtime_env_dict == runtime_env.to_dict()
|
|
assert runtime_env.has_conda()
|
|
assert runtime_env.conda_env_name() is None
|
|
assert runtime_env.conda_config() == json.dumps(conda_config, sort_keys=True)
|
|
|
|
runtime_env.pop("conda")
|
|
assert runtime_env.to_dict() == {"_ray_commit": "{{RAY_COMMIT_SHA}}"}
|
|
|
|
# Test the interface related to pip
|
|
with tempfile.TemporaryDirectory() as tmpdir, chdir(tmpdir):
|
|
requirement_file = os.path.join(tmpdir, "requirements.txt")
|
|
requirement_packages = ["dep5", "dep6"]
|
|
with open(requirement_file, "wt") as f:
|
|
for package in requirement_packages:
|
|
f.write(package)
|
|
f.write("\n")
|
|
|
|
pip_packages = ["dep1", "dep2"]
|
|
addition_pip_packages = ["dep3", "dep4"]
|
|
runtime_env = RuntimeEnv(pip=pip_packages)
|
|
runtime_env_dict = runtime_env.to_dict()
|
|
assert runtime_env.has_pip()
|
|
assert set(runtime_env.pip_config()["packages"]) == set(pip_packages)
|
|
assert runtime_env.virtualenv_name() is None
|
|
runtime_env["pip"]["packages"].extend(addition_pip_packages)
|
|
runtime_env_dict["pip"]["packages"].extend(addition_pip_packages)
|
|
# The default value of pip_check is False
|
|
runtime_env_dict["pip"]["pip_check"] = False
|
|
assert runtime_env_dict == runtime_env.to_dict()
|
|
assert runtime_env.has_pip()
|
|
assert set(runtime_env.pip_config()["packages"]) == set(
|
|
pip_packages + addition_pip_packages
|
|
)
|
|
assert runtime_env.virtualenv_name() is None
|
|
runtime_env["pip"] = requirement_file
|
|
runtime_env_dict["pip"] = requirement_packages
|
|
assert runtime_env.has_pip()
|
|
assert set(runtime_env.pip_config()["packages"]) == set(requirement_packages)
|
|
assert runtime_env.virtualenv_name() is None
|
|
# The default value of pip_check is False
|
|
runtime_env_dict["pip"] = dict(
|
|
packages=runtime_env_dict["pip"], pip_check=False
|
|
)
|
|
assert runtime_env_dict == runtime_env.to_dict()
|
|
|
|
runtime_env.pop("pip")
|
|
assert runtime_env.to_dict() == {"_ray_commit": "{{RAY_COMMIT_SHA}}"}
|
|
|
|
# Test conflict
|
|
with pytest.raises(ValueError):
|
|
RuntimeEnv(pip=pip_packages, conda=conda_name)
|
|
|
|
runtime_env = RuntimeEnv(pip=pip_packages)
|
|
runtime_env["conda"] = conda_name
|
|
with pytest.raises(ValueError):
|
|
runtime_env.serialize()
|
|
|
|
# Test the interface related to container
|
|
container_init = {
|
|
"image": "anyscale/ray-ml:nightly-py38-cpu",
|
|
"run_options": ["--cap-drop SYS_ADMIN", "--log-level=debug"],
|
|
}
|
|
update_container = {"image": "test_modify"}
|
|
runtime_env = RuntimeEnv(container=container_init)
|
|
runtime_env_dict = runtime_env.to_dict()
|
|
assert runtime_env.has_py_container()
|
|
assert runtime_env.py_container_image() == container_init["image"]
|
|
assert runtime_env.py_container_run_options() == container_init["run_options"]
|
|
runtime_env["container"].update(update_container)
|
|
runtime_env_dict["container"].update(update_container)
|
|
container_copy = container_init
|
|
container_copy.update(update_container)
|
|
assert runtime_env_dict == runtime_env.to_dict()
|
|
assert runtime_env.has_py_container()
|
|
assert runtime_env.py_container_image() == container_copy["image"]
|
|
assert runtime_env.py_container_run_options() == container_copy["run_options"]
|
|
|
|
runtime_env.pop("container")
|
|
assert runtime_env.to_dict() == {}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|