add test_reference_table

dashboard/modules/runtime_env/runtime_env_agent.py

@@ -77,8 +77,8 @@ class ReferenceTable:
 
  def __init__(
  self,
- uris_parser: Callable[[RuntimeEnv], None],
- unused_uris_callback: Callable[[List[Tuple[UriType, str]]], None],
+ uris_parser: Callable[[RuntimeEnv], Tuple[str, UriType]],
+ unused_uris_callback: Callable[[List[Tuple[str, UriType]]], None],
  unused_runtime_env_callback: Callable[[str], None],
  ):
  # Runtime Environment reference table. The key is serialized runtime env and
@@ -222,7 +222,7 @@ def uris_parser(self, runtime_env):
  result.append((uri, UriType.CONDA))
  return result
 
- def unused_uris_processor(self, unused_uris: List[Tuple[UriType, str]]) -> None:
+ def unused_uris_processor(self, unused_uris: List[Tuple[str, UriType]]) -> None:
  for uri, uri_type in unused_uris:
  if uri_type == UriType.WORKING_DIR:
  self._working_dir_uri_cache.mark_unused(uri)

dashboard/modules/runtime_env/tests/test_runtime_env_agent.py

@@ -0,0 +1,82 @@
+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:https://working_dir_1.zip",
+ py_modules=["s3:https://py_module_A.zip", "s3:https://py_module_B.zip"],
+ )
+ runtime_env_2 = RuntimeEnv(
+ working_dir="s3:https://working_dir_2.zip",
+ py_modules=["s3:https://py_module_A.zip", "s3:https://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:https://working_dir_1.zip", UriType.WORKING_DIR))
+ expected_unused_uris.append(("s3:https://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:https://working_dir_2.zip", UriType.WORKING_DIR))
+ expected_unused_uris.append(("s3:https://py_module_A.zip", UriType.PY_MODULES))
+ expected_unused_uris.append(("s3:https://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__]))