mirror of
https://github.com/vale981/ray
synced 2025-03-08 19:41:38 -05:00
69af9764b2
- Move the URI reference logic from raylet to agent. - Redefine the runtime env agent RPC to `CreateRuntimeEnvOrGet` and `DeleteRuntimeEnvIfPossible` - More details https://github.com/ray-project/ray/issues/21695#issuecomment-1032161528 Future works - We don't remove the `RuntimeEnvUris` from `RuntimeEnv` protobuf in current PR because gcs also uses those URIs to do GC by runtime_env_manager. We should also clear this. - Ray client server shouldn't interact with agent directly. Or Ray client server should also decrease the reference count. - Currently, `WorkerPool::HandleJobStarted` will be called multiple times for one job. So we should make sure this function is idempotent. Can we change this logic and make this function be called only once?
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(), 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__]))
|