wehub-resource-sync
279 lines
10 KiB
Python
279 lines
10 KiB
Python
import logging
|
|
import os
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
import ray._private.ray_constants as ray_constants
|
|
from ray._common.utils import try_to_create_directory
|
|
from ray._private.runtime_env.context import RuntimeEnvContext
|
|
from ray._private.runtime_env.packaging import (
|
|
Protocol,
|
|
delete_package,
|
|
download_and_unpack_package,
|
|
get_local_dir_from_uri,
|
|
get_uri_for_directory,
|
|
get_uri_for_package,
|
|
parse_uri,
|
|
upload_package_if_needed,
|
|
upload_package_to_gcs,
|
|
)
|
|
from ray._private.runtime_env.plugin import RuntimeEnvPlugin
|
|
from ray._private.utils import get_directory_size_bytes
|
|
from ray._raylet import GcsClient
|
|
from ray.exceptions import RuntimeEnvSetupError
|
|
from ray.util.debug import log_once
|
|
|
|
default_logger = logging.getLogger(__name__)
|
|
|
|
_WIN32 = os.name == "nt"
|
|
_LOG_ONCE_DEFAULT_EXCLUDE_PREFIX = "runtime_env_default_exclude:"
|
|
|
|
|
|
def upload_working_dir_if_needed(
|
|
runtime_env: Dict[str, Any],
|
|
include_gitignore: bool,
|
|
scratch_dir: Optional[str] = None,
|
|
logger: Optional[logging.Logger] = default_logger,
|
|
upload_fn: Optional[Callable[[str, Optional[List[str]], bool], None]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Uploads the working_dir and replaces it with a URI.
|
|
|
|
If the working_dir is already a URI, this is a no-op.
|
|
|
|
Excludes are combined from:
|
|
- .gitignore and .rayignore files in the working_dir
|
|
- runtime_env["excludes"] field
|
|
- RAY_RUNTIME_ENV_DEFAULT_EXCLUDES constant, overridable via
|
|
RAY_OVERRIDE_RUNTIME_ENV_DEFAULT_EXCLUDES environment variable
|
|
"""
|
|
working_dir = runtime_env.get("working_dir")
|
|
if working_dir is None:
|
|
return runtime_env
|
|
|
|
if not isinstance(working_dir, str) and not isinstance(working_dir, Path):
|
|
raise TypeError(
|
|
"working_dir must be a string or Path (either a local path "
|
|
f"or remote URI), got {type(working_dir)}."
|
|
)
|
|
|
|
if isinstance(working_dir, Path):
|
|
working_dir = str(working_dir)
|
|
|
|
# working_dir is already a URI -- just pass it through.
|
|
try:
|
|
protocol, path = parse_uri(working_dir)
|
|
except ValueError:
|
|
protocol, path = None, None
|
|
|
|
if protocol is not None:
|
|
supported_extensions = (".zip", ".tar.gz", ".tgz")
|
|
if protocol in Protocol.remote_protocols() and not any(
|
|
path.endswith(ext) for ext in supported_extensions
|
|
):
|
|
raise ValueError(
|
|
"Only .zip, .tar.gz, and .tgz files supported for remote URIs."
|
|
)
|
|
return runtime_env
|
|
|
|
default_excludes = ray_constants.get_runtime_env_default_excludes()
|
|
user_excludes = runtime_env.get("excludes") or []
|
|
excludes = default_excludes + list(user_excludes)
|
|
# TODO(ricardo): 2026-01-07 Remove these warnings in a few releases. Added in
|
|
# case users rely on these directories being uploaded with their working_dir
|
|
# since this change would be difficult to debug.
|
|
logger = logger or default_logger
|
|
working_dir_path = Path(working_dir)
|
|
for d in default_excludes:
|
|
if (working_dir_path / d).exists() and log_once(
|
|
f"{_LOG_ONCE_DEFAULT_EXCLUDE_PREFIX}{d}"
|
|
):
|
|
logger.warning(
|
|
"Directory %r is now ignored by default when packaging the working "
|
|
"directory. To disable this behavior, set "
|
|
"the `RAY_OVERRIDE_RUNTIME_ENV_DEFAULT_EXCLUDES=''` environment "
|
|
"variable.",
|
|
d,
|
|
)
|
|
try:
|
|
working_dir_uri = get_uri_for_directory(
|
|
working_dir,
|
|
include_gitignore=include_gitignore,
|
|
excludes=excludes,
|
|
)
|
|
except ValueError: # working_dir is not a directory
|
|
package_path = Path(working_dir)
|
|
supported_local = (
|
|
package_path.suffix == ".zip"
|
|
or package_path.suffix == ".tgz"
|
|
or package_path.name.endswith(".tar.gz")
|
|
)
|
|
if not package_path.exists() or not supported_local:
|
|
raise ValueError(
|
|
f"directory {package_path} must be an existing "
|
|
"directory or a supported archive (.zip, .tar.gz, .tgz)"
|
|
)
|
|
|
|
pkg_uri = get_uri_for_package(package_path)
|
|
if upload_fn is not None:
|
|
upload_fn(working_dir, excludes=excludes, is_file=True)
|
|
else:
|
|
try:
|
|
upload_package_to_gcs(pkg_uri, package_path.read_bytes())
|
|
except Exception as e:
|
|
raise RuntimeEnvSetupError(
|
|
f"Failed to upload package {package_path} to the Ray cluster: {e}"
|
|
) from e
|
|
runtime_env["working_dir"] = pkg_uri
|
|
return runtime_env
|
|
if upload_fn is None:
|
|
if scratch_dir is None:
|
|
scratch_dir = os.getcwd()
|
|
try:
|
|
upload_package_if_needed(
|
|
working_dir_uri,
|
|
scratch_dir,
|
|
working_dir,
|
|
include_parent_dir=False,
|
|
excludes=excludes,
|
|
include_gitignore=include_gitignore,
|
|
logger=logger,
|
|
)
|
|
except Exception as e:
|
|
raise RuntimeEnvSetupError(
|
|
f"Failed to upload working_dir {working_dir} to the Ray cluster: {e}"
|
|
) from e
|
|
else:
|
|
upload_fn(working_dir, excludes=excludes)
|
|
|
|
runtime_env["working_dir"] = working_dir_uri
|
|
return runtime_env
|
|
|
|
|
|
def set_pythonpath_in_context(python_path: str, context: RuntimeEnvContext):
|
|
"""Insert the path as the first entry in PYTHONPATH in the runtime env.
|
|
|
|
This is compatible with users providing their own PYTHONPATH in env_vars,
|
|
and is also compatible with the existing PYTHONPATH in the cluster.
|
|
|
|
The import priority is as follows:
|
|
this python_path arg > env_vars PYTHONPATH > existing cluster env PYTHONPATH.
|
|
"""
|
|
if "PYTHONPATH" in context.env_vars:
|
|
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
|
|
if "PYTHONPATH" in os.environ:
|
|
python_path += os.pathsep + os.environ["PYTHONPATH"]
|
|
context.env_vars["PYTHONPATH"] = python_path
|
|
|
|
|
|
class WorkingDirPlugin(RuntimeEnvPlugin):
|
|
name = "working_dir"
|
|
|
|
# Note working_dir is not following the priority order of other plugins. Instead
|
|
# it's specially treated to happen before all other plugins.
|
|
priority = 5
|
|
|
|
def __init__(self, resources_dir: str, gcs_client: GcsClient):
|
|
self._resources_dir = os.path.join(resources_dir, "working_dir_files")
|
|
self._gcs_client = gcs_client
|
|
try_to_create_directory(self._resources_dir)
|
|
|
|
def delete_uri(
|
|
self, uri: str, logger: Optional[logging.Logger] = default_logger
|
|
) -> int:
|
|
"""Delete URI and return the number of bytes deleted."""
|
|
logger.info("Got request to delete working dir URI %s", uri)
|
|
local_dir = get_local_dir_from_uri(uri, self._resources_dir)
|
|
local_dir_size = get_directory_size_bytes(local_dir)
|
|
|
|
deleted = delete_package(uri, self._resources_dir)
|
|
if not deleted:
|
|
logger.warning(f"Tried to delete nonexistent URI: {uri}.")
|
|
return 0
|
|
|
|
return local_dir_size
|
|
|
|
def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F821
|
|
working_dir_uri = runtime_env.working_dir()
|
|
if working_dir_uri != "":
|
|
return [working_dir_uri]
|
|
return []
|
|
|
|
async def create(
|
|
self,
|
|
uri: Optional[str],
|
|
runtime_env: dict,
|
|
context: RuntimeEnvContext,
|
|
logger: logging.Logger = default_logger,
|
|
) -> int:
|
|
local_dir = await download_and_unpack_package(
|
|
uri,
|
|
self._resources_dir,
|
|
self._gcs_client,
|
|
logger=logger,
|
|
overwrite=True,
|
|
)
|
|
return get_directory_size_bytes(local_dir)
|
|
|
|
def modify_context(
|
|
self,
|
|
uris: List[str],
|
|
runtime_env_dict: Dict,
|
|
context: RuntimeEnvContext,
|
|
logger: Optional[logging.Logger] = default_logger,
|
|
):
|
|
if not uris:
|
|
return
|
|
|
|
# WorkingDirPlugin uses a single URI.
|
|
uri = uris[0]
|
|
local_dir = get_local_dir_from_uri(uri, self._resources_dir)
|
|
if not local_dir.exists():
|
|
raise ValueError(
|
|
f"Local directory {local_dir} for URI {uri} does "
|
|
"not exist on the cluster. Something may have gone wrong while "
|
|
"downloading or unpacking the working_dir."
|
|
)
|
|
|
|
if not _WIN32:
|
|
context.command_prefix += ["cd", str(local_dir), "&&"]
|
|
else:
|
|
# Include '/d' incase temp folder is on different drive than Ray install.
|
|
context.command_prefix += ["cd", "/d", f"{local_dir}", "&&"]
|
|
set_pythonpath_in_context(python_path=str(local_dir), context=context)
|
|
|
|
@contextmanager
|
|
def with_working_dir_env(self, uri):
|
|
"""
|
|
If uri is not None, add the local working directory to the environment variable
|
|
as "RAY_RUNTIME_ENV_CREATE_WORKING_DIR". This is useful for other plugins to
|
|
create their environment with reference to the working directory. For example
|
|
`pip -r ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR}/requirements.txt`
|
|
|
|
The environment variable is removed after the context manager exits.
|
|
"""
|
|
if uri is None:
|
|
yield
|
|
else:
|
|
local_dir = get_local_dir_from_uri(uri, self._resources_dir)
|
|
if not local_dir.exists():
|
|
raise ValueError(
|
|
f"Local directory {local_dir} for URI {uri} does "
|
|
"not exist on the cluster. Something may have gone wrong while "
|
|
"downloading or unpacking the working_dir."
|
|
)
|
|
key = ray_constants.RAY_RUNTIME_ENV_CREATE_WORKING_DIR_ENV_VAR
|
|
prev = os.environ.get(key)
|
|
# Windows backslash paths are weird. When it's passed to the env var, and
|
|
# when Pip expands it, the backslashes are interpreted as escape characters
|
|
# and messes up the whole path. So we convert it to forward slashes.
|
|
# This works at least for all Python applications, including pip.
|
|
os.environ[key] = local_dir.as_posix()
|
|
try:
|
|
yield
|
|
finally:
|
|
if prev is None:
|
|
del os.environ[key]
|
|
else:
|
|
os.environ[key] = prev
|