def __init__(self, dashboard_agent):
super().__init__(dashboard_agent)
self._runtime_env_dir = dashboard_agent.runtime_env_dir
self._logging_params = dashboard_agent.logging_params
self._per_job_logger_cache = dict()
# Cache the results of creating envs to avoid repeatedly calling into
# conda and other slow calls.
self._env_cache: Dict[str, CreatedEnvResult] = dict()
# Maps a serialized runtime env to a lock that is used
# to prevent multiple concurrent installs of the same env.
self._env_locks: Dict[str, asyncio.Lock] = dict()
# Keeps track of the URIs contained within each env so we can
# invalidate the env cache when a URI is deleted.
# This is a temporary mechanism until we have per-URI caching.
self._uris_to_envs: Dict[str, Set[str]] = defaultdict(set)
# Initialize internal KV to be used by the working_dir setup code.
_initialize_internal_kv(self._dashboard_agent.gcs_client)
assert _internal_kv_initialized()
self._pip_manager = PipManager(self._runtime_env_dir)
self._conda_manager = CondaManager(self._runtime_env_dir)
self._py_modules_manager = PyModulesManager(self._runtime_env_dir)
self._working_dir_manager = WorkingDirManager(self._runtime_env_dir)
self._container_manager = ContainerManager(dashboard_agent.temp_dir)
self._working_dir_uri_cache = URICache(
self._working_dir_manager.delete_uri, WORKING_DIR_CACHE_SIZE_BYTES)
self._py_modules_uri_cache = URICache(
self._py_modules_manager.delete_uri, PY_MODULES_CACHE_SIZE_BYTES)
self._conda_uri_cache = URICache(self._conda_manager.delete_uri,
CONDA_CACHE_SIZE_BYTES)
self._pip_uri_cache = URICache(self._pip_manager.delete_uri,
PIP_CACHE_SIZE_BYTES)
self._logger = default_logger
def test_delete_fn_called(self):
num_delete_fn_calls = 0
uris_to_sizes = {"a": 8, "b": 6, "c": 4, "d": 20}
def delete_fn(uri, logger):
nonlocal num_delete_fn_calls
num_delete_fn_calls += 1
return uris_to_sizes[uri]
cache = URICache(delete_fn, max_total_size_bytes=10, debug_mode=True)
cache.add("a", 8)
cache.add("b", 6)
cache.mark_unused("b")
# Total size is 14 > 10, so we need to delete "b".
assert num_delete_fn_calls == 1
cache.add("c", 4)
cache.mark_unused("c")
# Total size is 12 > 10, so we delete "c".
assert num_delete_fn_calls == 2
cache.mark_unused("a")
# Total size is 8 <= 10, so we shouldn't delete anything.
assert num_delete_fn_calls == 2
cache.add("d", 20)
# Total size is 28 > 10, so we delete "a".
assert num_delete_fn_calls == 3
cache.mark_unused("d")
# Total size is 20 > 10, so we delete "d".
assert num_delete_fn_calls == 4
def __init__(self, dashboard_agent):
super().__init__(dashboard_agent)
self._runtime_env_dir = dashboard_agent.runtime_env_dir
self._logging_params = dashboard_agent.logging_params
self._per_job_logger_cache = dict()
# Cache the results of creating envs to avoid repeatedly calling into
# conda and other slow calls.
self._env_cache: Dict[str, CreatedEnvResult] = dict()
# Maps a serialized runtime env to a lock that is used
# to prevent multiple concurrent installs of the same env.
self._env_locks: Dict[str, asyncio.Lock] = dict()
self._gcs_aio_client = self._dashboard_agent.gcs_aio_client
self._pip_plugin = PipPlugin(self._runtime_env_dir)
self._conda_plugin = CondaPlugin(self._runtime_env_dir)
self._py_modules_plugin = PyModulesPlugin(self._runtime_env_dir,
self._gcs_aio_client)
self._java_jars_plugin = JavaJarsPlugin(self._runtime_env_dir,
self._gcs_aio_client)
self._working_dir_plugin = WorkingDirPlugin(self._runtime_env_dir,
self._gcs_aio_client)
self._container_manager = ContainerManager(dashboard_agent.temp_dir)
# TODO(architkulkarni): "base plugins" and third-party plugins should all go
# through the same code path. We should never need to refer to
# self._xxx_plugin, we should just iterate through self._plugins.
self._base_plugins = [
self._working_dir_plugin,
self._pip_plugin,
self._conda_plugin,
self._py_modules_plugin,
self._java_jars_plugin,
]
self._uri_caches = {}
self._base_plugin_cache_managers = {}
for plugin in self._base_plugins:
# Set the max size for the cache. Defaults to 10 GB.
cache_size_env_var = f"RAY_RUNTIME_ENV_{plugin.name}_CACHE_SIZE_GB".upper(
)
cache_size_bytes = int(
(1024**3) * float(os.environ.get(cache_size_env_var, 10)))
self._uri_caches[plugin.name] = URICache(plugin.delete_uri,
cache_size_bytes)
self._base_plugin_cache_managers[plugin.name] = PluginCacheManager(
plugin, self._uri_caches[plugin.name])
self._reference_table = ReferenceTable(
self.uris_parser,
self.unused_uris_processor,
self.unused_runtime_env_processor,
)
self._runtime_env_plugin_manager = RuntimeEnvPluginManager()
self._logger = default_logger
def test_many_URIs(self):
uris_to_sizes = {str(i): i for i in range(1000)}
def delete_fn(uri, logger):
return uris_to_sizes[uri]
cache = URICache(delete_fn, debug_mode=True)
for i in range(1000):
cache.add(str(i), i)
for i in range(1000):
cache.mark_unused(str(i))
for i in range(1000):
assert str(i) in cache
def __init__(self, dashboard_agent):
super().__init__(dashboard_agent)
self._runtime_env_dir = dashboard_agent.runtime_env_dir
self._logging_params = dashboard_agent.logging_params
self._per_job_logger_cache = dict()
# Cache the results of creating envs to avoid repeatedly calling into
# conda and other slow calls.
self._env_cache: Dict[str, CreatedEnvResult] = dict()
# Maps a serialized runtime env to a lock that is used
# to prevent multiple concurrent installs of the same env.
self._env_locks: Dict[str, asyncio.Lock] = dict()
_initialize_internal_kv(self._dashboard_agent.gcs_client)
assert _internal_kv_initialized()
self._pip_manager = PipManager(self._runtime_env_dir)
self._conda_manager = CondaManager(self._runtime_env_dir)
self._py_modules_manager = PyModulesManager(self._runtime_env_dir)
self._java_jars_manager = JavaJarsManager(self._runtime_env_dir)
self._working_dir_manager = WorkingDirManager(self._runtime_env_dir)
self._container_manager = ContainerManager(dashboard_agent.temp_dir)
self._reference_table = ReferenceTable(
self.uris_parser,
self.unused_uris_processor,
self.unused_runtime_env_processor,
)
self._working_dir_uri_cache = URICache(
self._working_dir_manager.delete_uri, WORKING_DIR_CACHE_SIZE_BYTES)
self._py_modules_uri_cache = URICache(
self._py_modules_manager.delete_uri, PY_MODULES_CACHE_SIZE_BYTES)
self._java_jars_uri_cache = URICache(
self._java_jars_manager.delete_uri, JAVA_JARS_CACHE_SIZE_BYTES)
self._conda_uri_cache = URICache(self._conda_manager.delete_uri,
CONDA_CACHE_SIZE_BYTES)
self._pip_uri_cache = URICache(self._pip_manager.delete_uri,
PIP_CACHE_SIZE_BYTES)
self._logger = default_logger
def test_nonzero_cache_size(self):
uris_to_sizes = {"a": 4, "b": 4, "c": 4}
def delete_fn(uri, logger):
return uris_to_sizes[uri]
cache = URICache(delete_fn, max_total_size_bytes=10, debug_mode=True)
cache.add("a", 4)
cache.add("b", 4)
cache.mark_unused("a")
assert "a" in cache
cache.add("c", 4)
# Now we have total size 12, which exceeds the max size 10.
assert cache.get_total_size_bytes() == 8
# "a" was the only unused URI, so it must have been deleted.
assert "b" and "c" in cache and "a" not in cache
def test_zero_cache_size(self):
uris_to_sizes = {"5": 5, "3": 3}
def delete_fn(uri, logger):
return uris_to_sizes[uri]
cache = URICache(delete_fn, max_total_size_bytes=0, debug_mode=True)
cache.add("5", 5)
assert cache.get_total_size_bytes() == 5
cache.mark_unused("5")
assert cache.get_total_size_bytes() == 0
cache.add("3", 3)
cache.add("5", 5)
assert cache.get_total_size_bytes() == 8
cache.mark_unused("3")
cache.mark_unused("5")
assert cache.get_total_size_bytes() == 0
def test_mark_used(self):
uris_to_sizes = {"a": 3, "b": 3, "big": 300}
def delete_fn(uri, logger):
return uris_to_sizes[uri]
cache = URICache(delete_fn, max_total_size_bytes=10, debug_mode=True)
cache.add("a", 3)
cache.add("b", 3)
cache.mark_unused("a")
cache.mark_unused("b")
assert "a" in cache and "b" in cache
assert cache.get_total_size_bytes() == 6
cache.mark_used("a")
cache.add("big", 300)
# We are over capacity and the only unused URI is "b", so we delete it
assert "a" in cache and "big" in cache and "b" not in cache
assert cache.get_total_size_bytes() == 303
cache.mark_unused("big")
assert "big" not in cache
assert cache.get_total_size_bytes() == 3
def test_mark_used_nonadded_uri_error(self):
cache = URICache(debug_mode=True)
with pytest.raises(ValueError):
cache.mark_used("nonadded_uri")