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)