mirror of
https://github.com/vale981/ray
synced 2025-03-11 13:46:40 -04:00
82 lines
3 KiB
Python
82 lines
3 KiB
Python
import sys
|
|
import pytest
|
|
|
|
from typing import List, Tuple
|
|
from ray.dashboard.modules.runtime_env.runtime_env_agent import UriType, ReferenceTable
|
|
from ray.runtime_env import RuntimeEnv
|
|
|
|
|
|
def test_reference_table():
|
|
expected_unused_uris = []
|
|
expected_unused_runtime_env = str()
|
|
|
|
def uris_parser(runtime_env) -> Tuple[str, UriType]:
|
|
result = list()
|
|
result.append((runtime_env.working_dir(), "working_dir"))
|
|
py_module_uris = runtime_env.py_modules()
|
|
for uri in py_module_uris:
|
|
result.append((uri, "py_modules"))
|
|
return result
|
|
|
|
def unused_uris_processor(unused_uris: List[Tuple[str, UriType]]) -> None:
|
|
nonlocal expected_unused_uris
|
|
assert expected_unused_uris
|
|
for unused in unused_uris:
|
|
assert unused in expected_unused_uris
|
|
expected_unused_uris.remove(unused)
|
|
assert not expected_unused_uris
|
|
|
|
def unused_runtime_env_processor(unused_runtime_env: str) -> None:
|
|
nonlocal expected_unused_runtime_env
|
|
assert expected_unused_runtime_env
|
|
assert expected_unused_runtime_env == unused_runtime_env
|
|
expected_unused_runtime_env = None
|
|
|
|
reference_table = ReferenceTable(
|
|
uris_parser, unused_uris_processor, unused_runtime_env_processor
|
|
)
|
|
runtime_env_1 = RuntimeEnv(
|
|
working_dir="s3://working_dir_1.zip",
|
|
py_modules=["s3://py_module_A.zip", "s3://py_module_B.zip"],
|
|
)
|
|
runtime_env_2 = RuntimeEnv(
|
|
working_dir="s3://working_dir_2.zip",
|
|
py_modules=["s3://py_module_A.zip", "s3://py_module_C.zip"],
|
|
)
|
|
# Add runtime env 1
|
|
reference_table.increase_reference(
|
|
runtime_env_1, runtime_env_1.serialize(), "raylet"
|
|
)
|
|
# Add runtime env 2
|
|
reference_table.increase_reference(
|
|
runtime_env_2, runtime_env_2.serialize(), "raylet"
|
|
)
|
|
# Add runtime env 1 by `client_server`, this will be skipped by reference table.
|
|
reference_table.increase_reference(
|
|
runtime_env_1, runtime_env_1.serialize(), "client_server"
|
|
)
|
|
|
|
# Remove runtime env 1
|
|
expected_unused_uris.append(("s3://working_dir_1.zip", "working_dir"))
|
|
expected_unused_uris.append(("s3://py_module_B.zip", "py_modules"))
|
|
expected_unused_runtime_env = runtime_env_1.serialize()
|
|
reference_table.decrease_reference(
|
|
runtime_env_1, runtime_env_1.serialize(), "raylet"
|
|
)
|
|
assert not expected_unused_uris
|
|
assert not expected_unused_runtime_env
|
|
|
|
# Remove runtime env 2
|
|
expected_unused_uris.append(("s3://working_dir_2.zip", "working_dir"))
|
|
expected_unused_uris.append(("s3://py_module_A.zip", "py_modules"))
|
|
expected_unused_uris.append(("s3://py_module_C.zip", "py_modules"))
|
|
expected_unused_runtime_env = runtime_env_2.serialize()
|
|
reference_table.decrease_reference(
|
|
runtime_env_2, runtime_env_2.serialize(), "raylet"
|
|
)
|
|
assert not expected_unused_uris
|
|
assert not expected_unused_runtime_env
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|