mirror of
https://github.com/vale981/ray
synced 2025-03-09 04:46:38 -04:00
83 lines
3 KiB
Python
83 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(), UriType.WORKING_DIR))
|
||
|
|
py_module_uris = runtime_env.py_modules()
|
||
|
|
for uri in py_module_uris:
|
||
|
|
result.append((uri, UriType.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", UriType.WORKING_DIR))
|
||
|
|
expected_unused_uris.append(("s3://py_module_B.zip", UriType.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", UriType.WORKING_DIR))
|
||
|
|
expected_unused_uris.append(("s3://py_module_A.zip", UriType.PY_MODULES))
|
||
|
|
expected_unused_uris.append(("s3://py_module_C.zip", UriType.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__]))
|