wehub-resource-sync
775 lines
22 KiB
Python
775 lines
22 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import List
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private import ray_constants
|
|
from ray._private.runtime_env.context import RuntimeEnvContext
|
|
from ray._private.runtime_env.plugin import RuntimeEnvPlugin
|
|
from ray._private.test_utils import external_redis_test_enabled
|
|
from ray.exceptions import RuntimeEnvSetupError
|
|
from ray.runtime_env.runtime_env import RuntimeEnv
|
|
|
|
MY_PLUGIN_CLASS_PATH = "ray.tests.test_runtime_env_plugin.MyPlugin"
|
|
MY_PLUGIN_NAME = "MyPlugin"
|
|
|
|
|
|
class MyPlugin(RuntimeEnvPlugin):
|
|
name = MY_PLUGIN_NAME
|
|
env_key = "MY_PLUGIN_TEST_ENVIRONMENT_KEY"
|
|
|
|
@staticmethod
|
|
def validate(runtime_env: RuntimeEnv) -> str:
|
|
value = runtime_env[MY_PLUGIN_NAME]
|
|
if value == "fail":
|
|
raise ValueError("not allowed")
|
|
return value
|
|
|
|
def modify_context(
|
|
self,
|
|
uris: List[str],
|
|
runtime_env: RuntimeEnv,
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> None:
|
|
plugin_config_dict = runtime_env[MY_PLUGIN_NAME]
|
|
ctx.env_vars[MyPlugin.env_key] = str(plugin_config_dict["env_value"])
|
|
ctx.command_prefix += [
|
|
"echo",
|
|
plugin_config_dict["tmp_content"],
|
|
">",
|
|
plugin_config_dict["tmp_file"],
|
|
"&&",
|
|
]
|
|
ctx.py_executable = (
|
|
plugin_config_dict["prefix_command"] + " " + ctx.py_executable
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
'[{"class":"' + MY_PLUGIN_CLASS_PATH + '"}]',
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_simple_env_modification_plugin(set_runtime_env_plugins, ray_start_regular):
|
|
_, tmp_file_path = tempfile.mkstemp()
|
|
|
|
@ray.remote
|
|
def f():
|
|
import psutil
|
|
|
|
with open(tmp_file_path, "r") as f:
|
|
content = f.read().strip()
|
|
return {
|
|
"env_value": os.environ[MyPlugin.env_key],
|
|
"tmp_content": content,
|
|
"nice": psutil.Process().nice(),
|
|
}
|
|
|
|
with pytest.raises(RuntimeEnvSetupError, match="not allowed"):
|
|
ray.get(f.options(runtime_env={MY_PLUGIN_NAME: "fail"}).remote())
|
|
|
|
if os.name != "nt":
|
|
output = ray.get(
|
|
f.options(
|
|
runtime_env={
|
|
MY_PLUGIN_NAME: {
|
|
"env_value": 42,
|
|
"tmp_file": tmp_file_path,
|
|
"tmp_content": "hello",
|
|
# See https://en.wikipedia.org/wiki/Nice_(Unix)
|
|
"prefix_command": "nice -n 19",
|
|
}
|
|
}
|
|
).remote()
|
|
)
|
|
|
|
assert output == {"env_value": "42", "tmp_content": "hello", "nice": 19}
|
|
|
|
|
|
MY_PLUGIN_FOR_HANG_CLASS_PATH = "ray.tests.test_runtime_env_plugin.MyPluginForHang"
|
|
MY_PLUGIN_FOR_HANG_NAME = "MyPluginForHang"
|
|
my_plugin_setup_times = 0
|
|
|
|
|
|
# This plugin will hang when first setup, second setup will ok
|
|
class MyPluginForHang(RuntimeEnvPlugin):
|
|
name = MY_PLUGIN_FOR_HANG_NAME
|
|
env_key = "MY_PLUGIN_FOR_HANG_TEST_ENVIRONMENT_KEY"
|
|
|
|
@staticmethod
|
|
def validate(runtime_env_dict: dict) -> str:
|
|
return "True"
|
|
|
|
async def create(
|
|
self,
|
|
uri: str,
|
|
runtime_env: dict,
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> float:
|
|
global my_plugin_setup_times
|
|
my_plugin_setup_times += 1
|
|
|
|
# first setup
|
|
if my_plugin_setup_times == 1:
|
|
# sleep forever
|
|
await asyncio.sleep(3600)
|
|
|
|
def modify_context(
|
|
self,
|
|
uris: List[str],
|
|
plugin_config_dict: dict,
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> None:
|
|
global my_plugin_setup_times
|
|
ctx.env_vars[MyPluginForHang.env_key] = str(my_plugin_setup_times)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
'[{"class":"' + MY_PLUGIN_FOR_HANG_CLASS_PATH + '"}]',
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_plugin_hang(set_runtime_env_plugins, ray_start_regular):
|
|
env_key = MyPluginForHang.env_key
|
|
|
|
@ray.remote(num_cpus=0.1)
|
|
def f():
|
|
return os.environ[env_key]
|
|
|
|
refs = [
|
|
f.options(
|
|
# Avoid hitting the cache of runtime_env
|
|
runtime_env={MY_PLUGIN_FOR_HANG_NAME: {"name": "f1"}}
|
|
).remote(),
|
|
f.options(runtime_env={MY_PLUGIN_FOR_HANG_NAME: {"name": "f2"}}).remote(),
|
|
]
|
|
|
|
def condition():
|
|
for ref in refs:
|
|
try:
|
|
res = ray.get(ref, timeout=1)
|
|
print("result:", res)
|
|
assert int(res) == 2
|
|
return True
|
|
except Exception as error:
|
|
print(f"Got error: {error}")
|
|
pass
|
|
return False
|
|
|
|
wait_for_condition(condition, timeout=60)
|
|
|
|
|
|
DUMMY_PLUGIN_CLASS_PATH = "ray.tests.test_runtime_env_plugin.DummyPlugin"
|
|
DUMMY_PLUGIN_NAME = "DummyPlugin"
|
|
HANG_PLUGIN_CLASS_PATH = "ray.tests.test_runtime_env_plugin.HangPlugin"
|
|
HANG_PLUGIN_NAME = "HangPlugin"
|
|
|
|
|
|
class DummyPlugin(RuntimeEnvPlugin):
|
|
name = DUMMY_PLUGIN_NAME
|
|
|
|
@staticmethod
|
|
def validate(runtime_env_dict: dict) -> str:
|
|
return 1
|
|
|
|
|
|
class HangPlugin(DummyPlugin):
|
|
name = HANG_PLUGIN_NAME
|
|
|
|
async def create(
|
|
self,
|
|
uri: str,
|
|
runtime_env: "RuntimeEnv",
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger, # noqa: F821
|
|
) -> float:
|
|
await asyncio.sleep(3600)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
'[{"class":"' + DUMMY_PLUGIN_CLASS_PATH + '"},'
|
|
'{"class":"' + HANG_PLUGIN_CLASS_PATH + '"}]',
|
|
],
|
|
indirect=True,
|
|
)
|
|
@pytest.mark.skipif(external_redis_test_enabled(), reason="Failing in redis mode.")
|
|
def test_plugin_timeout(set_runtime_env_plugins, start_cluster):
|
|
@ray.remote(num_cpus=0.1)
|
|
def f():
|
|
return True
|
|
|
|
refs = [
|
|
f.options(
|
|
runtime_env={
|
|
HANG_PLUGIN_NAME: {"name": "f1"},
|
|
"config": {"setup_timeout_seconds": 1},
|
|
}
|
|
).remote(),
|
|
f.options(runtime_env={DUMMY_PLUGIN_NAME: {"name": "f2"}}).remote(),
|
|
f.options(
|
|
runtime_env={
|
|
HANG_PLUGIN_NAME: {"name": "f3"},
|
|
"config": {"setup_timeout_seconds": -1},
|
|
}
|
|
).remote(),
|
|
]
|
|
|
|
def condition():
|
|
good_fun_num = 0
|
|
bad_fun_num = 0
|
|
for ref in refs:
|
|
try:
|
|
res = ray.get(ref, timeout=1)
|
|
print("result:", res)
|
|
if res:
|
|
good_fun_num += 1
|
|
return True
|
|
except RuntimeEnvSetupError:
|
|
bad_fun_num += 1
|
|
return bad_fun_num == 1 and good_fun_num == 2
|
|
|
|
wait_for_condition(condition, timeout=60)
|
|
|
|
|
|
FAULT_PLUGIN_CLASS_PATH = "ray.tests.test_runtime_env_plugin.FaultPlugin"
|
|
FAULT_PLUGIN_NAME = "FaultPlugin"
|
|
FAULT_PLUGIN_KEY = "FAULT_PLUGIN_KEY"
|
|
|
|
|
|
class FaultPlugin(DummyPlugin):
|
|
name = FAULT_PLUGIN_NAME
|
|
|
|
async def create(
|
|
self,
|
|
uri: str,
|
|
runtime_env: "RuntimeEnv",
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger, # noqa: F821
|
|
) -> float:
|
|
action = os.environ.get(FAULT_PLUGIN_KEY, "raise")
|
|
if action == "raise":
|
|
raise RuntimeError(
|
|
"Ever tried. Ever failed. No matter. Try again. Fail again. Fail "
|
|
"better. -- Waiting for Godot, Samuel Beckett"
|
|
)
|
|
elif action == "sleep":
|
|
await asyncio.sleep(3600)
|
|
elif action == "ok":
|
|
return
|
|
else:
|
|
raise ValueError(f"unknown action {action}")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
'[{"class":"' + FAULT_PLUGIN_CLASS_PATH + '"}]',
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_task_fails_on_rt_env_failure(
|
|
set_runtime_env_plugins,
|
|
monkeypatch,
|
|
ray_start_cluster,
|
|
):
|
|
"""
|
|
Simulate runtime env failure on a node. The task should fail but the raylet should
|
|
not die. See https://github.com/ray-project/ray/pull/46991
|
|
"""
|
|
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
|
|
with monkeypatch.context() as m:
|
|
m.setenv(FAULT_PLUGIN_KEY, "raise")
|
|
fault_node = cluster.add_node(num_cpus=1)
|
|
|
|
@ray.remote(num_cpus=0.1, runtime_env={FAULT_PLUGIN_NAME: {}})
|
|
def s(a, b):
|
|
return a + b
|
|
|
|
two = ray.put(2)
|
|
three = ray.put(3)
|
|
|
|
mortal = s.remote(two, three)
|
|
immortal = s.options(max_retries=-1).remote(two, three)
|
|
|
|
with pytest.raises(RuntimeEnvSetupError) as e:
|
|
ray.get(mortal)
|
|
assert "Samuel Beckett" in str(e.value)
|
|
# Note: even immortal task will fail because a rt env failure cancels tasks.
|
|
with pytest.raises(RuntimeEnvSetupError) as e:
|
|
ray.get(immortal)
|
|
assert "Samuel Beckett" in str(e.value)
|
|
|
|
# Assert that the raylet is still alive.
|
|
for _ in range(5):
|
|
time.sleep(1)
|
|
assert (
|
|
fault_node.all_processes[ray_constants.PROCESS_TYPE_RAYLET][
|
|
0
|
|
].process.poll()
|
|
is None
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
'[{"class":"' + FAULT_PLUGIN_CLASS_PATH + '"}]',
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_actor_fails_on_rt_env_failure(
|
|
set_runtime_env_plugins,
|
|
monkeypatch,
|
|
ray_start_cluster,
|
|
):
|
|
"""
|
|
Simulate runtime env failure on a node. The actor should be dead, but the raylet
|
|
should not die. See https://github.com/ray-project/ray/pull/46991
|
|
"""
|
|
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
|
|
with monkeypatch.context() as m:
|
|
m.setenv(FAULT_PLUGIN_KEY, "raise")
|
|
fault_node = cluster.add_node(num_cpus=1)
|
|
|
|
@ray.remote(num_cpus=0.1, runtime_env={FAULT_PLUGIN_NAME: {}})
|
|
class Actor:
|
|
def __init__(self, a, b):
|
|
self.a = a
|
|
self.b = b
|
|
|
|
def ping(self):
|
|
return self.a + self.b
|
|
|
|
two = ray.put(2)
|
|
three = ray.put(3)
|
|
|
|
mortal = Actor.remote(two, three)
|
|
immortal = Actor.options(max_restarts=-1).remote(two, three)
|
|
|
|
with pytest.raises(RuntimeEnvSetupError) as e:
|
|
ray.get(mortal.ping.remote())
|
|
assert "Samuel Beckett" in str(e.value)
|
|
|
|
# Note: even immortal actor will die because a rt env failure cancels tasks.
|
|
with pytest.raises(RuntimeEnvSetupError) as e:
|
|
ray.get(immortal.ping.remote())
|
|
assert "Samuel Beckett" in str(e.value)
|
|
|
|
# Assert that the raylet is still alive.
|
|
for _ in range(5):
|
|
time.sleep(1)
|
|
assert (
|
|
fault_node.all_processes[ray_constants.PROCESS_TYPE_RAYLET][
|
|
0
|
|
].process.poll()
|
|
is None
|
|
)
|
|
|
|
|
|
PRIORITY_TEST_PLUGIN1_CLASS_PATH = (
|
|
"ray.tests.test_runtime_env_plugin.PriorityTestPlugin1"
|
|
)
|
|
PRIORITY_TEST_PLUGIN1_NAME = "PriorityTestPlugin1"
|
|
PRIORITY_TEST_PLUGIN2_CLASS_PATH = (
|
|
"ray.tests.test_runtime_env_plugin.PriorityTestPlugin2"
|
|
)
|
|
PRIORITY_TEST_PLUGIN2_NAME = "PriorityTestPlugin2"
|
|
PRIORITY_TEST_ENV_VAR_NAME = "PriorityTestEnv"
|
|
|
|
|
|
class PriorityTestPlugin1(RuntimeEnvPlugin):
|
|
name = PRIORITY_TEST_PLUGIN1_NAME
|
|
priority = 11
|
|
env_value = " world"
|
|
|
|
@staticmethod
|
|
def validate(runtime_env_dict: dict) -> str:
|
|
return None
|
|
|
|
def modify_context(
|
|
self,
|
|
uris: List[str],
|
|
plugin_config_dict: dict,
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> None:
|
|
if PRIORITY_TEST_ENV_VAR_NAME in ctx.env_vars:
|
|
ctx.env_vars[PRIORITY_TEST_ENV_VAR_NAME] += PriorityTestPlugin1.env_value
|
|
else:
|
|
ctx.env_vars[PRIORITY_TEST_ENV_VAR_NAME] = PriorityTestPlugin1.env_value
|
|
|
|
|
|
class PriorityTestPlugin2(RuntimeEnvPlugin):
|
|
name = PRIORITY_TEST_PLUGIN2_NAME
|
|
priority = 10
|
|
env_value = "hello"
|
|
|
|
@staticmethod
|
|
def validate(runtime_env_dict: dict) -> str:
|
|
return None
|
|
|
|
def modify_context(
|
|
self,
|
|
uris: List[str],
|
|
plugin_config_dict: dict,
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> None:
|
|
if PRIORITY_TEST_ENV_VAR_NAME in ctx.env_vars:
|
|
raise RuntimeError(
|
|
f"Env var {PRIORITY_TEST_ENV_VAR_NAME} has been set to "
|
|
f"{ctx.env_vars[PRIORITY_TEST_ENV_VAR_NAME]}."
|
|
)
|
|
ctx.env_vars[PRIORITY_TEST_ENV_VAR_NAME] = PriorityTestPlugin2.env_value
|
|
|
|
|
|
priority_test_plugin_config_without_priority = [
|
|
{
|
|
"class": PRIORITY_TEST_PLUGIN1_CLASS_PATH,
|
|
},
|
|
{
|
|
"class": PRIORITY_TEST_PLUGIN2_CLASS_PATH,
|
|
},
|
|
]
|
|
|
|
|
|
priority_test_plugin_config = [
|
|
{
|
|
"class": PRIORITY_TEST_PLUGIN1_CLASS_PATH,
|
|
"priority": 1,
|
|
},
|
|
{
|
|
"class": PRIORITY_TEST_PLUGIN2_CLASS_PATH,
|
|
"priority": 0,
|
|
},
|
|
]
|
|
|
|
priority_test_plugin_bad_config = [
|
|
{
|
|
"class": PRIORITY_TEST_PLUGIN1_CLASS_PATH,
|
|
"priority": 0,
|
|
# Only used to distinguish the bad config in test body.
|
|
"tag": "bad",
|
|
},
|
|
{
|
|
"class": PRIORITY_TEST_PLUGIN2_CLASS_PATH,
|
|
"priority": 1,
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
json.dumps(priority_test_plugin_config_without_priority),
|
|
json.dumps(priority_test_plugin_config),
|
|
json.dumps(priority_test_plugin_bad_config),
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_plugin_priority(set_runtime_env_plugins, ray_start_regular):
|
|
config = set_runtime_env_plugins
|
|
_, tmp_file_path = tempfile.mkstemp()
|
|
|
|
@ray.remote
|
|
def f():
|
|
import os
|
|
|
|
return os.environ.get(PRIORITY_TEST_ENV_VAR_NAME)
|
|
|
|
if "bad" in config:
|
|
with pytest.raises(RuntimeEnvSetupError, match="has been set"):
|
|
value = ray.get(
|
|
f.options(
|
|
runtime_env={
|
|
PRIORITY_TEST_PLUGIN1_NAME: {},
|
|
PRIORITY_TEST_PLUGIN2_NAME: {},
|
|
}
|
|
).remote()
|
|
)
|
|
else:
|
|
value = ray.get(
|
|
f.options(
|
|
runtime_env={
|
|
PRIORITY_TEST_PLUGIN1_NAME: {},
|
|
PRIORITY_TEST_PLUGIN2_NAME: {},
|
|
}
|
|
).remote()
|
|
)
|
|
assert value is not None
|
|
assert value == "hello world"
|
|
|
|
|
|
def test_unexpected_field_warning(shutdown_only):
|
|
"""Test that an unexpected runtime_env field doesn't error."""
|
|
ray.init(runtime_env={"unexpected_field": "value"})
|
|
|
|
@ray.remote
|
|
def f():
|
|
return True
|
|
|
|
# Run a task to trigger runtime_env creation.
|
|
assert ray.get(f.remote())
|
|
|
|
# Check that the warning is logged.
|
|
session_dir = ray._private.worker.global_worker.node.address_info["session_dir"]
|
|
log_path = Path(session_dir) / "logs"
|
|
|
|
# Check that a warning appears in some "runtime_env_setup*.log"
|
|
wait_for_condition(
|
|
lambda: any(
|
|
"unexpected_field is not recognized" in open(f).read()
|
|
for f in log_path.glob("runtime_env_setup*.log")
|
|
)
|
|
)
|
|
|
|
|
|
URI_CACHING_TEST_PLUGIN_CLASS_PATH = (
|
|
"ray.tests.test_runtime_env_plugin.UriCachingTestPlugin"
|
|
)
|
|
URI_CACHING_TEST_PLUGIN_NAME = "UriCachingTestPlugin"
|
|
URI_CACHING_TEST_DIR = Path(tempfile.gettempdir()) / "runtime_env_uri_caching_test"
|
|
uri_caching_test_file_path = URI_CACHING_TEST_DIR / "uri_caching_test_file.json"
|
|
URI_CACHING_TEST_DIR.mkdir(parents=True, exist_ok=True)
|
|
uri_caching_test_file_path.write_text("{}")
|
|
|
|
|
|
def get_plugin_usage_data():
|
|
with open(uri_caching_test_file_path, "r") as f:
|
|
data = json.loads(f.read())
|
|
return data
|
|
|
|
|
|
class UriCachingTestPlugin(RuntimeEnvPlugin):
|
|
"""A plugin that fakes taking up local disk space when creating its environment.
|
|
|
|
This plugin is used to test that the URI caching is working correctly.
|
|
Example:
|
|
runtime_env = {"UriCachingTestPlugin": {"uri": "file:///a", "size_bytes": 10}}
|
|
"""
|
|
|
|
name = URI_CACHING_TEST_PLUGIN_NAME
|
|
|
|
def __init__(self):
|
|
# Keeps track of the "disk space" each URI takes up for the
|
|
# UriCachingTestPlugin.
|
|
self.uris_to_sizes = {}
|
|
self.modify_context_call_count = 0
|
|
self.create_call_count = 0
|
|
|
|
def write_plugin_usage_data(self) -> None:
|
|
with open(uri_caching_test_file_path, "w") as f:
|
|
data = {
|
|
"uris_to_sizes": self.uris_to_sizes,
|
|
"modify_context_call_count": self.modify_context_call_count,
|
|
"create_call_count": self.create_call_count,
|
|
}
|
|
f.write(json.dumps(data))
|
|
|
|
def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F811
|
|
return [runtime_env[self.name]["uri"]]
|
|
|
|
async def create(
|
|
self,
|
|
uri,
|
|
runtime_env: "RuntimeEnv", # noqa: F821
|
|
context: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> float:
|
|
self.create_call_count += 1
|
|
created_size_bytes = runtime_env[self.name]["size_bytes"]
|
|
self.uris_to_sizes[uri] = created_size_bytes
|
|
self.write_plugin_usage_data()
|
|
return created_size_bytes
|
|
|
|
def modify_context(
|
|
self,
|
|
uris: List[str],
|
|
runtime_env: "RuntimeEnv",
|
|
context: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> None:
|
|
self.modify_context_call_count += 1
|
|
self.write_plugin_usage_data()
|
|
|
|
def delete_uri(self, uri: str, logger: logging.Logger) -> int:
|
|
size = self.uris_to_sizes.pop(uri)
|
|
self.write_plugin_usage_data()
|
|
return size
|
|
|
|
|
|
# Set scope to "class" to force this to run before start_cluster, whose scope
|
|
# is "function". We need these env vars to be set before Ray is started.
|
|
@pytest.fixture(scope="class")
|
|
def uri_cache_size_100_gb():
|
|
var = f"RAY_RUNTIME_ENV_{URI_CACHING_TEST_PLUGIN_NAME}_CACHE_SIZE_GB".upper()
|
|
with mock.patch.dict(
|
|
os.environ,
|
|
{
|
|
var: "100",
|
|
},
|
|
):
|
|
print("Set URI cache size for UriCachingTestPlugin to 100 GB")
|
|
yield
|
|
|
|
|
|
def gb_to_bytes(size_gb: int) -> int:
|
|
return size_gb * 1024 * 1024 * 1024
|
|
|
|
|
|
class TestGC:
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
json.dumps([{"class": URI_CACHING_TEST_PLUGIN_CLASS_PATH}]),
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_uri_caching(
|
|
self, set_runtime_env_plugins, start_cluster, uri_cache_size_100_gb
|
|
):
|
|
cluster, address = start_cluster
|
|
|
|
ray.init(address=address)
|
|
|
|
def reinit():
|
|
ray.shutdown()
|
|
# TODO(architkulkarni): Currently, reinit the driver will generate the same
|
|
# job id. And if we reinit immediately after shutdown, raylet may
|
|
# process new job started before old job finished in some cases. This
|
|
# inconsistency could disorder the URI reference and delete a valid
|
|
# runtime env. We sleep here to walk around this issue.
|
|
time.sleep(5)
|
|
ray.init(address=address)
|
|
|
|
@ray.remote
|
|
def f():
|
|
return True
|
|
|
|
# Run a task to trigger runtime_env creation.
|
|
ref1 = f.options(
|
|
runtime_env={
|
|
URI_CACHING_TEST_PLUGIN_NAME: {
|
|
"uri": "file:///tmp/test_uri_1",
|
|
"size_bytes": gb_to_bytes(50),
|
|
}
|
|
}
|
|
).remote()
|
|
ray.get(ref1)
|
|
# Check that the URI was "created on disk".
|
|
print(get_plugin_usage_data())
|
|
wait_for_condition(
|
|
lambda: get_plugin_usage_data()
|
|
== {
|
|
"uris_to_sizes": {"file:///tmp/test_uri_1": gb_to_bytes(50)},
|
|
"modify_context_call_count": 1,
|
|
"create_call_count": 1,
|
|
}
|
|
)
|
|
|
|
# Shutdown ray to stop the worker and remove the runtime_env reference.
|
|
reinit()
|
|
|
|
# Run a task with a different runtime env.
|
|
ref2 = f.options(
|
|
runtime_env={
|
|
URI_CACHING_TEST_PLUGIN_NAME: {
|
|
"uri": "file:///tmp/test_uri_2",
|
|
"size_bytes": gb_to_bytes(51),
|
|
}
|
|
}
|
|
).remote()
|
|
ray.get(ref2)
|
|
# This should delete the old URI and create a new one, because 50 + 51 > 100
|
|
# and the cache size limit is 100.
|
|
wait_for_condition(
|
|
lambda: get_plugin_usage_data()
|
|
== {
|
|
"uris_to_sizes": {"file:///tmp/test_uri_2": gb_to_bytes(51)},
|
|
"modify_context_call_count": 2,
|
|
"create_call_count": 2,
|
|
}
|
|
)
|
|
|
|
reinit()
|
|
|
|
# Run a task with the cached runtime env, to check that the runtime env is not
|
|
# created anew.
|
|
ref3 = f.options(
|
|
runtime_env={
|
|
URI_CACHING_TEST_PLUGIN_NAME: {
|
|
"uri": "file:///tmp/test_uri_2",
|
|
"size_bytes": gb_to_bytes(51),
|
|
}
|
|
}
|
|
).remote()
|
|
ray.get(ref3)
|
|
# modify_context should still be called even if create() is not called.
|
|
# Example: for a "conda" plugin, even if the conda env is already created
|
|
# and cached, we still need to call modify_context to add "conda activate" to
|
|
# the RuntimeEnvContext.command_prefix for the worker.
|
|
wait_for_condition(
|
|
lambda: get_plugin_usage_data()
|
|
== {
|
|
"uris_to_sizes": {"file:///tmp/test_uri_2": gb_to_bytes(51)},
|
|
"modify_context_call_count": 3,
|
|
"create_call_count": 2,
|
|
}
|
|
)
|
|
|
|
reinit()
|
|
|
|
# Run a task with a new runtime env
|
|
ref4 = f.options(
|
|
runtime_env={
|
|
URI_CACHING_TEST_PLUGIN_NAME: {
|
|
"uri": "file:///tmp/test_uri_3",
|
|
"size_bytes": gb_to_bytes(10),
|
|
}
|
|
}
|
|
).remote()
|
|
ray.get(ref4)
|
|
# The last two URIs should still be present in the cache, because 51 + 10 < 100.
|
|
wait_for_condition(
|
|
lambda: get_plugin_usage_data()
|
|
== {
|
|
"uris_to_sizes": {
|
|
"file:///tmp/test_uri_2": gb_to_bytes(51),
|
|
"file:///tmp/test_uri_3": gb_to_bytes(10),
|
|
},
|
|
"modify_context_call_count": 4,
|
|
"create_call_count": 3,
|
|
}
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|