chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,242 @@
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
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_file,
|
||||
get_uri_for_package,
|
||||
install_wheel_package,
|
||||
is_whl_uri,
|
||||
package_exists,
|
||||
parse_uri,
|
||||
upload_package_if_needed,
|
||||
upload_package_to_gcs,
|
||||
)
|
||||
from ray._private.runtime_env.plugin import RuntimeEnvPlugin
|
||||
from ray._private.runtime_env.working_dir import set_pythonpath_in_context
|
||||
from ray._private.utils import get_directory_size_bytes
|
||||
from ray._raylet import GcsClient
|
||||
from ray.exceptions import RuntimeEnvSetupError
|
||||
|
||||
default_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _check_is_uri(s: str) -> bool:
|
||||
try:
|
||||
protocol, path = parse_uri(s)
|
||||
except ValueError:
|
||||
protocol, path = None, None
|
||||
|
||||
supported_extensions = (".zip", ".whl", ".tar.gz", ".tgz")
|
||||
if protocol in Protocol.remote_protocols() and not any(
|
||||
path.endswith(ext) for ext in supported_extensions
|
||||
):
|
||||
raise ValueError(
|
||||
"Only .zip, .whl, .tar.gz, and .tgz files supported for remote URIs."
|
||||
)
|
||||
|
||||
return protocol is not None
|
||||
|
||||
|
||||
def upload_py_modules_if_needed(
|
||||
runtime_env: Dict[str, Any],
|
||||
include_gitignore: bool,
|
||||
scratch_dir: Optional[str] = None,
|
||||
logger: Optional[logging.Logger] = default_logger,
|
||||
upload_fn=None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Uploads the entries in py_modules and replaces them with a list of URIs.
|
||||
|
||||
For each entry that is already a URI, this is a no-op.
|
||||
"""
|
||||
py_modules = runtime_env.get("py_modules")
|
||||
if py_modules is None:
|
||||
return runtime_env
|
||||
|
||||
if not isinstance(py_modules, list):
|
||||
raise TypeError(
|
||||
"py_modules must be a List of local paths, imported modules, or "
|
||||
f"URIs, got {type(py_modules)}."
|
||||
)
|
||||
|
||||
py_modules_uris = []
|
||||
for module in py_modules:
|
||||
if isinstance(module, str):
|
||||
# module_path is a local path or a URI.
|
||||
module_path = module
|
||||
elif isinstance(module, Path):
|
||||
module_path = str(module)
|
||||
elif isinstance(module, ModuleType):
|
||||
if not hasattr(module, "__path__"):
|
||||
# This is a single-file module.
|
||||
module_path = module.__file__
|
||||
else:
|
||||
# NOTE(edoakes): Python allows some installed Python packages to
|
||||
# be split into multiple directories. We could probably handle
|
||||
# this, but it seems tricky & uncommon. If it's a problem for
|
||||
# users, we can add this support on demand.
|
||||
if len(module.__path__) > 1:
|
||||
raise ValueError(
|
||||
"py_modules only supports modules whose __path__"
|
||||
" has length 1 or those who are single-file."
|
||||
)
|
||||
[module_path] = module.__path__
|
||||
else:
|
||||
raise TypeError(
|
||||
"py_modules must be a list of file paths, URIs, "
|
||||
f"or imported modules, got {type(module)}."
|
||||
)
|
||||
|
||||
if _check_is_uri(module_path):
|
||||
module_uri = module_path
|
||||
else:
|
||||
# module_path is a local path.
|
||||
if Path(module_path).is_dir() or Path(module_path).suffix == ".py":
|
||||
is_dir = Path(module_path).is_dir()
|
||||
excludes = runtime_env.get("excludes", None)
|
||||
if is_dir:
|
||||
module_uri = get_uri_for_directory(
|
||||
module_path,
|
||||
include_gitignore=include_gitignore,
|
||||
excludes=excludes,
|
||||
)
|
||||
else:
|
||||
module_uri = get_uri_for_file(module_path)
|
||||
if upload_fn is None:
|
||||
if scratch_dir is None:
|
||||
scratch_dir = os.getcwd()
|
||||
try:
|
||||
upload_package_if_needed(
|
||||
module_uri,
|
||||
scratch_dir,
|
||||
module_path,
|
||||
include_gitignore=include_gitignore,
|
||||
include_parent_dir=is_dir,
|
||||
excludes=excludes,
|
||||
logger=logger,
|
||||
)
|
||||
except Exception as e:
|
||||
from ray.util.spark.utils import is_in_databricks_runtime
|
||||
|
||||
if is_in_databricks_runtime():
|
||||
raise RuntimeEnvSetupError(
|
||||
f"Failed to upload module {module_path} to the Ray "
|
||||
f"cluster, please ensure there are only files under "
|
||||
f"the module path, notebooks under the path are "
|
||||
f"not allowed, original exception: {e}"
|
||||
) from e
|
||||
raise RuntimeEnvSetupError(
|
||||
f"Failed to upload module {module_path} to the Ray "
|
||||
f"cluster: {e}"
|
||||
) from e
|
||||
else:
|
||||
upload_fn(module_path, excludes=excludes)
|
||||
elif Path(module_path).suffix == ".whl":
|
||||
module_uri = get_uri_for_package(Path(module_path))
|
||||
if upload_fn is None:
|
||||
if not package_exists(module_uri):
|
||||
try:
|
||||
upload_package_to_gcs(
|
||||
module_uri, Path(module_path).read_bytes()
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeEnvSetupError(
|
||||
f"Failed to upload {module_path} to the Ray "
|
||||
f"cluster: {e}"
|
||||
) from e
|
||||
else:
|
||||
upload_fn(module_path, excludes=None, is_file=True)
|
||||
else:
|
||||
raise ValueError(
|
||||
"py_modules entry must be a .py file, "
|
||||
"a directory, or a .whl file; "
|
||||
f"got {module_path}"
|
||||
)
|
||||
|
||||
py_modules_uris.append(module_uri)
|
||||
|
||||
# TODO(architkulkarni): Expose a single URI for py_modules. This plugin
|
||||
# should internally handle the "sub-URIs", the individual modules.
|
||||
|
||||
runtime_env["py_modules"] = py_modules_uris
|
||||
return runtime_env
|
||||
|
||||
|
||||
class PyModulesPlugin(RuntimeEnvPlugin):
|
||||
|
||||
name = "py_modules"
|
||||
|
||||
def __init__(self, resources_dir: str, gcs_client: GcsClient):
|
||||
self._resources_dir = os.path.join(resources_dir, "py_modules_files")
|
||||
self._gcs_client = gcs_client
|
||||
try_to_create_directory(self._resources_dir)
|
||||
|
||||
def _get_local_dir_from_uri(self, uri: str):
|
||||
return get_local_dir_from_uri(uri, 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 pymodule 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) -> List[str]:
|
||||
return runtime_env.py_modules()
|
||||
|
||||
async def create(
|
||||
self,
|
||||
uri: str,
|
||||
runtime_env: "RuntimeEnv", # noqa: F821
|
||||
context: RuntimeEnvContext,
|
||||
logger: Optional[logging.Logger] = default_logger,
|
||||
) -> int:
|
||||
|
||||
module_dir = await download_and_unpack_package(
|
||||
uri, self._resources_dir, self._gcs_client, logger=logger
|
||||
)
|
||||
|
||||
if is_whl_uri(uri):
|
||||
wheel_uri = module_dir
|
||||
module_dir = self._get_local_dir_from_uri(uri)
|
||||
await install_wheel_package(
|
||||
wheel_uri=wheel_uri, target_dir=module_dir, logger=logger
|
||||
)
|
||||
|
||||
return get_directory_size_bytes(module_dir)
|
||||
|
||||
def modify_context(
|
||||
self,
|
||||
uris: List[str],
|
||||
runtime_env_dict: Dict,
|
||||
context: RuntimeEnvContext,
|
||||
logger: Optional[logging.Logger] = default_logger,
|
||||
):
|
||||
module_dirs = []
|
||||
for uri in uris:
|
||||
module_dir = self._get_local_dir_from_uri(uri)
|
||||
if not module_dir.exists():
|
||||
raise ValueError(
|
||||
f"Local directory {module_dir} for URI {uri} does "
|
||||
"not exist on the cluster. Something may have gone wrong while "
|
||||
"downloading, unpacking or installing the py_modules files."
|
||||
)
|
||||
module_dirs.append(str(module_dir))
|
||||
set_pythonpath_in_context(os.pathsep.join(module_dirs), context)
|
||||
Reference in New Issue
Block a user