outerbounds 0.3.183rc1__py3-none-any.whl → 0.3.185__py3-none-any.whl

outerbounds/apps/code_package/code_packager.py

@@ -1,610 +0,0 @@
-import os
-import sys
-import time
-import tarfile
-import json
-from io import BytesIO
-from typing import List, Tuple, Dict, Any, Optional, Callable, Union
-
-from metaflow.datastore.content_addressed_store import ContentAddressedStore
-from metaflow.util import to_unicode
-from metaflow.metaflow_config import (
-    DATASTORE_SYSROOT_S3,
-    DATASTORE_SYSROOT_AZURE,
-    DATASTORE_SYSROOT_GS,
-    DATASTORE_SYSROOT_LOCAL,
-)
-
-# Default prefix for code packages in content addressed store
-CODE_PACKAGE_PREFIX = "apps-code-packages"
-
-
-# this is os.walk(follow_symlinks=True) with cycle detection
-def walk_without_cycles(top_root):
-    seen = set()
-
-    def _recurse(root):
-        for parent, dirs, files in os.walk(root):
-            for d in dirs:
-                path = os.path.join(parent, d)
-                if os.path.islink(path):
-                    # Breaking loops: never follow the same symlink twice
-                    #
-                    # NOTE: this also means that links to sibling links are
-                    # not followed. In this case:
-                    #
-                    #   x -> y
-                    #   y -> oo
-                    #   oo/real_file
-                    #
-                    # real_file is only included twice, not three times
-                    reallink = os.path.realpath(path)
-                    if reallink not in seen:
-                        seen.add(reallink)
-                        for x in _recurse(path):
-                            yield x
-            yield parent, files
-
-    for x in _recurse(top_root):
-        yield x
-
-
-def symlink_friendly_walk(root, exclude_hidden=True, suffixes=None):
-    if suffixes is None:
-        suffixes = []
-    root = to_unicode(root)  # handle files/folder with non ascii chars
-    prefixlen = len("%s/" % os.path.dirname(root))
-    for (
-        path,
-        files,
-    ) in walk_without_cycles(root):
-        if exclude_hidden and "/." in path:
-            continue
-        # path = path[2:] # strip the ./ prefix
-        # if path and (path[0] == '.' or './' in path):
-        #    continue
-        for fname in files:
-            if (fname[0] == "." and fname in suffixes) or (
-                fname[0] != "." and any(fname.endswith(suffix) for suffix in suffixes)
-            ):
-                p = os.path.join(path, fname)
-                yield p, p[prefixlen:]
-
-
-class CodePackager:
-    """
-    A datastore-agnostic class for packaging code.
-
-    This class handles creating a code package (tarball) for deployment
-    and provides methods for storing and retrieving it using Metaflow's
-    ContentAddressedStore directly.
-
-    Usage examples:
-    ```python
-    packager = CodePackager(
-        datastore_type: str = "s3",
-        datastore_root = None,
-        code_package_prefix = None,
-    )
-
-    package_url, package_key = packager.store(
-        paths_to_include = ["./"],
-        file_suffixes = [".py", ".txt", ".yaml", ".yml", ".json"],
-    )
-
-    package_url, package_key = packager.store(
-        package_create_fn = lambda: my_custom_package_create_fn(),
-    )
-    ```
-    """
-
-    def __init__(
-        self,
-        datastore_type: str = "s3",
-        datastore_root: Optional[str] = None,
-        code_package_prefix: Optional[str] = None,
-    ):
-        """
-        Initialize the CodePackager with datastore configuration.
-
-        Parameters
-        ----------
-        datastore_type : str, default "s3"
-            The type of datastore to use: "s3", "azure", "gs", or "local"
-        datastore_root : str, optional
-            Root path for the datastore. If not provided, uses the default for the datastore type.
-        code_package_prefix : str, optional
-            The prefix to use for storing code packages in the content addressed store.
-            If not provided, uses the CODE_PACKAGE_PREFIX configuration value.
-        """
-        self._datastore_type = datastore_type
-        self._datastore_root = datastore_root
-        self._code_package_prefix = code_package_prefix
-
-    def store(
-        self,
-        package_create_fn: Optional[Callable[[], bytes]] = None,
-        paths_to_include: Optional[List[str]] = None,
-        file_suffixes: Optional[List[str]] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[str, str]:
-        """
-        Create and store a code package using Metaflow's ContentAddressedStore.
-
-        This method can be called in two ways:
-        1. With paths_to_include and file_suffixes to use the default packaging
-        2. With a custom package_create_fn for custom packaging logic
-
-        Parameters
-        ----------
-        package_create_fn : Callable[[], bytes], optional
-            A function that creates and returns a package as bytes.
-            This allows for custom packaging logic without dependency on specific objects.
-        paths_to_include : List[str], optional
-            List of paths to include in the package. Used by default_package_create.
-        file_suffixes : List[str], optional
-            List of file suffixes to include. Used by default_package_create.
-        metadata : Dict[str, Any], optional
-            Metadata to include in the package when using default_package_create.
-
-        Returns
-        -------
-        Tuple[str, str]
-            A tuple containing (package_url, package_key) that identifies the location
-            and content-addressed key of the stored package.
-        """
-        # Prepare default values
-        _paths_to_include = paths_to_include or []
-        _file_suffixes = file_suffixes or [
-            ".py",
-            ".txt",
-            ".yaml",
-            ".yml",
-            ".json",
-            ".html",
-            ".css",
-            ".js",
-            ".jsx",
-            ".ts",
-            ".tsx",
-            ".md",
-            ".rst",
-        ]
-        _metadata = metadata or {}
-
-        # If no package_create_fn provided, use default_package_create
-        if package_create_fn is None:
-            _package_create_fn = lambda: self.default_package_create(
-                _paths_to_include, _file_suffixes, _metadata
-            )
-        else:
-            _package_create_fn = package_create_fn
-
-        # Create the package
-        code_package = _package_create_fn()
-
-        # Get the ContentAddressedStore for the specified datastore
-        ca_store = self.get_content_addressed_store(
-            datastore_type=self._datastore_type,
-            datastore_root=self._datastore_root,
-            prefix=(
-                str(self._code_package_prefix)
-                if self._code_package_prefix is not None
-                else str(CODE_PACKAGE_PREFIX)
-            ),
-        )
-
-        # Store the package using raw=True to ensure we can access it directly via URL
-        results = ca_store.save_blobs([code_package], raw=True, len_hint=1)
-        package_url, package_key = results[0].uri, results[0].key
-
-        return package_url, package_key
-
-    @staticmethod
-    def get_content_addressed_store(
-        datastore_type: str = "s3",
-        datastore_root: Optional[str] = None,
-        prefix: Optional[str] = None,
-    ) -> ContentAddressedStore:
-        """
-        Get a ContentAddressedStore instance for the specified datastore.
-
-        Parameters
-        ----------
-        datastore_type : str, default "s3"
-            Type of datastore: "s3", "azure", "gs", or "local"
-        datastore_root : str, optional
-            Root path for the datastore. If not provided, uses the default for the datastore type.
-        prefix : str, optional
-            Prefix to use when storing objects in the datastore.
-            If not provided, uses the CODE_PACKAGE_PREFIX configuration value.
-
-        Returns
-        -------
-        ContentAddressedStore
-            A ContentAddressedStore instance configured for the specified datastore
-        """
-        from metaflow.plugins import DATASTORES
-
-        datastore_impls = [i for i in DATASTORES if i.TYPE == datastore_type]
-        if len(datastore_impls) == 0:
-            raise ValueError(f"Unsupported datastore type: {datastore_type}")
-        if len(datastore_impls) > 1:
-            raise ValueError(
-                f"Multiple datastore implementations found for type: {datastore_type}"
-            )
-        datastore_impl = datastore_impls[0]
-        root = None
-        # Import the storage implementation based on datastore_type
-        if datastore_type == "s3":
-            root = datastore_root or DATASTORE_SYSROOT_S3
-        elif datastore_type == "azure":
-            root = datastore_root or DATASTORE_SYSROOT_AZURE
-        elif datastore_type == "gs":
-            root = datastore_root or DATASTORE_SYSROOT_GS
-        elif datastore_type == "local":
-            root = datastore_root or DATASTORE_SYSROOT_LOCAL
-
-        # Ensure prefix is a string
-        store_prefix = str(prefix) if prefix is not None else str(CODE_PACKAGE_PREFIX)
-
-        storage_impl = datastore_impl(root=root)
-        # Create and return a ContentAddressedStore
-        return ContentAddressedStore(prefix=store_prefix, storage_impl=storage_impl)
-
-    @staticmethod
-    def get_download_cmd(
-        package_url: str,
-        datastore_type: str,
-        python_cmd: str = "python",
-        target_file: str = "job.tar",
-        escape_quotes: bool = True,
-    ) -> str:
-        """
-        Generate a command to download the code package.
-
-        Parameters
-        ----------
-        package_url : str
-            The URL of the package to download
-        datastore_type : str
-            The type of datastore (s3, azure, gs, local)
-        python_cmd : str, optional
-            The Python command to use
-        target_file : str, optional
-            The target file name to save the package as
-        escape_quotes : bool, optional
-            Whether to escape quotes in the command
-
-        Returns
-        -------
-        str
-            A shell command string to download the package
-        """
-        if datastore_type == "s3":
-            from metaflow.plugins.aws.aws_utils import parse_s3_full_path
-
-            bucket, s3_object = parse_s3_full_path(package_url)
-            # Simplify the script and use single quotes to avoid shell escaping issues
-            script = 'import boto3, os; ep=os.getenv({quote}METAFLOW_S3_ENDPOINT_URL{quote}); boto3.client("s3", **({{"endpoint_url":ep}} if ep else {{}})).download_file({quote}{bucket}{quote}, {quote}{s3_object}{quote}, {quote}{target_file}{quote})'.format(
-                quote='\\"' if escape_quotes else '"',
-                bucket=bucket,
-                s3_object=s3_object,
-                target_file=target_file,
-            )
-            # Format the command with proper quoting
-            return f"{python_cmd} -c '{script}'"
-        elif datastore_type == "azure":
-            from metaflow.plugins.azure.azure_utils import parse_azure_full_path
-
-            container_name, blob = parse_azure_full_path(package_url)
-            # remove a trailing slash, if present
-            blob_endpoint = "${METAFLOW_AZURE_STORAGE_BLOB_SERVICE_ENDPOINT%/}"
-            return "download-azure-blob --blob-endpoint={blob_endpoint} --container={container} --blob={blob} --output-file={target}".format(
-                blob_endpoint=blob_endpoint,
-                blob=blob,
-                container=container_name,
-                target=target_file,
-            )
-        elif datastore_type == "gs":
-            from metaflow.plugins.gcp.gs_utils import parse_gs_full_path
-
-            bucket_name, gs_object = parse_gs_full_path(package_url)
-            return "download-gcp-object --bucket=%s --object=%s --output-file=%s" % (
-                bucket_name,
-                gs_object,
-                target_file,
-            )
-        elif datastore_type == "local":
-            # For local storage, simply copy the file
-            return "cp %s %s" % (package_url, target_file)
-        else:
-            raise NotImplementedError(
-                f"Download command not implemented for datastore type: {datastore_type}"
-            )
-
-    def get_package_commands(
-        self,
-        code_package_url: str,
-        python_cmd: str = "python",
-        target_file: str = "job.tar",
-        working_dir: str = "metaflow",
-        retries: int = 5,
-        escape_quotes: bool = True,
-    ) -> List[str]:
-        """
-        Get a complete list of shell commands to download and extract a code package.
-
-        This method generates a comprehensive set of shell commands for downloading
-        and extracting a code package, similar to MetaflowEnvironment.get_package_commands.
-
-        Parameters
-        ----------
-        code_package_url : str
-            The URL of the code package to download
-        python_cmd : str, optional
-            The Python command to use
-        target_file : str, optional
-            The target file name to save the package as
-        working_dir : str, optional
-            The directory to create and extract the package into
-        retries : int, optional
-            Number of download retries to attempt
-        escape_quotes : bool, optional
-            Whether to escape quotes in the command
-
-        Returns
-        -------
-        List[str]
-            List of shell commands to execute
-        """
-        # Use the datastore_type from initialization if not provided
-        datastore_type = self._datastore_type
-
-        # Helper function to create dependency installation command
-        def _get_install_dependencies_cmd():
-            base_cmd = "{} -m pip install -qqq --no-compile --no-cache-dir --disable-pip-version-check".format(
-                python_cmd
-            )
-
-            datastore_packages = {
-                "s3": ["boto3"],
-                "azure": [
-                    "azure-identity",
-                    "azure-storage-blob",
-                    "azure-keyvault-secrets",
-                    "simple-azure-blob-downloader",
-                ],
-                "gs": [
-                    "google-cloud-storage",
-                    "google-auth",
-                    "simple-gcp-object-downloader",
-                    "google-cloud-secret-manager",
-                ],
-                "local": [],
-            }
-
-            if datastore_type not in datastore_packages:
-                raise NotImplementedError(
-                    "Unknown datastore type: {}".format(datastore_type)
-                )
-
-            if not datastore_packages[datastore_type]:
-                return "# No dependencies required for local datastore"
-
-            cmd = "{} {}".format(
-                base_cmd, " ".join(datastore_packages[datastore_type] + ["requests"])
-            )
-            # Skip pip installs if we know packages might already be available
-            return "if [ -z $METAFLOW_SKIP_INSTALL_DEPENDENCIES ]; then {}; fi".format(
-                cmd
-            )
-
-        download_cmd = self.get_download_cmd(
-            code_package_url, datastore_type, python_cmd, target_file, escape_quotes
-        )
-
-        # Define the log functions for bash
-        bash_mflog = (
-            'function mflog() { echo "[$(date -u +"%Y-%m-%dT%H:%M:%SZ")]" "$@"; }'
-        )
-        bash_flush_logs = 'function flush_mflogs() { echo "[$(date -u +"%Y-%m-%dT%H:%M:%SZ")] Flushing logs"; }'
-
-        cmds = [
-            bash_mflog,
-            bash_flush_logs,
-            "mflog 'Setting up task environment.'",
-            _get_install_dependencies_cmd(),
-            f"mkdir -p {working_dir}",
-            f"cd {working_dir}",
-            "mkdir -p .metaflow",  # mute local datastore creation log
-            f"i=0; while [ $i -le {retries} ]; do "
-            "mflog 'Downloading code package...'; "
-            + download_cmd
-            + " && mflog 'Code package downloaded.' && break; "
-            "sleep 10; i=$((i+1)); "
-            "done",
-            f"if [ $i -gt {retries} ]; then "
-            "mflog 'Failed to download code package from %s "
-            f"after {retries+1} tries. Exiting...' && exit 1; "
-            "fi" % code_package_url,
-            "TAR_OPTIONS='--warning=no-timestamp' tar xf %s" % target_file,
-            "mflog 'Task is starting.'",
-            "flush_mflogs",
-        ]
-
-        return cmds
-
-    @staticmethod
-    def directory_walker(
-        root,
-        exclude_hidden=True,
-        suffixes=None,
-    ) -> List[Tuple[str, str]]:
-        """
-        Walk a directory and yield tuples of (file_path, relative_arcname) for files
-        that match the given suffix filters. It will follow symlinks, but not cycles.
-
-        This function is similar to MetaflowPackage._walk and handles symlinks safely.
-
-        Parameters
-        ----------
-        root : str
-            The root directory to walk
-        exclude_hidden : bool, default True
-            Whether to exclude hidden files and directories (those starting with '.')
-        suffixes : List[str], optional
-            List of file suffixes to include (e.g. ['.py', '.txt'])
-
-        Returns
-        -------
-        List[Tuple[str, str]]
-            List of tuples (file_path, relative_arcname) where:
-            - file_path is the full path to the file
-            - relative_arcname is the path to use within the archive
-        """
-        files = []
-        for file_path, rel_path in symlink_friendly_walk(
-            root, exclude_hidden, suffixes
-        ):
-            files.append((file_path, rel_path))
-        return files
-
-    @staticmethod
-    def default_package_create(
-        paths: List[str], suffixes: List[str], metadata: Optional[Dict[str, Any]] = None
-    ) -> bytes:
-        """
-        Create a default tarball package from specified paths.
-
-        Parameters
-        ----------
-        paths : List[str]
-            List of paths to include in the package
-        suffixes : List[str]
-            List of file suffixes to include
-        metadata : Dict[str, Any], optional
-            Metadata to include in the package
-
-        Returns
-        -------
-        bytes
-            The binary content of the tarball
-        """
-        buf = BytesIO()
-
-        with tarfile.open(fileobj=buf, mode="w:gz", compresslevel=3) as tar:
-            # Add metadata if provided
-            if metadata:
-                metadata_buf = BytesIO()
-                metadata_buf.write(json.dumps(metadata).encode("utf-8"))
-                metadata_buf.seek(0)
-                info = tarfile.TarInfo("metadata.json")
-                info.size = len(metadata_buf.getvalue())
-                info.mtime = 1747158696  # 13 May 2025 10:31:36 (so that we dont have a changing hash everytime we run)
-                tar.addfile(info, metadata_buf)
-
-            def no_mtime(tarinfo):
-                # a modification time change should not change the hash of
-                # the package. Only content modifications will.
-                # Setting this default to Dec 3, 2019
-                tarinfo.mtime = 1747158696  # 13 May 2025 10:31:36 (so that we dont have a changing hash everytime we run)
-                return tarinfo
-
-            # Add files from specified paths
-            for path in paths:
-                if os.path.isdir(path):
-                    # Use directory_walker for directories to handle symlinks properly
-                    for file_path, rel_path in CodePackager.directory_walker(
-                        path,
-                        exclude_hidden=True,
-                        suffixes=suffixes,
-                    ):
-                        tar.add(
-                            file_path,
-                            arcname=rel_path,
-                            filter=no_mtime,
-                            recursive=False,
-                        )
-                elif os.path.isfile(path):
-                    if any(path.endswith(suffix) for suffix in suffixes):
-                        tar.add(path, arcname=os.path.basename(path))
-
-        tarball = bytearray(buf.getvalue())
-        tarball[4:8] = [0] * 4  # Reset 4 bytes from offset 4 to account for ts
-        return tarball
-
-    @staticmethod
-    def _add_tar_file(tar, filename, buf):
-        tarinfo = tarfile.TarInfo(name=filename)
-        tarinfo.size = len(buf.getvalue())
-        buf.seek(0)
-        tarinfo.mtime = 1747158696  # 13 May 2025 10:31:36 (so that we dont have a changing hash everytime we run)
-        tar.addfile(tarinfo, fileobj=buf)
-
-    @classmethod
-    def package_directory(
-        cls,
-        directory_path: str,
-        suffixes: Optional[List[str]] = None,
-        exclude_hidden: bool = True,
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> bytes:
-        """
-        Package a directory and all of its contents that match the given suffixes.
-
-        This is a convenience method that works similarly to MetaflowPackage._walk
-        to package a directory for deployment. Will default follow_symlinks.
-
-        Parameters
-        ----------
-        directory_path : str
-            The directory to package
-        suffixes : List[str], optional
-            List of file suffixes to include (defaults to standard code extensions)
-        exclude_hidden : bool, default True
-            Whether to exclude hidden files and directories
-        metadata : Dict[str, Any], optional
-            Metadata to include in the package
-        Returns
-        -------
-        bytes
-            The binary content of the tarball
-        """
-        if not os.path.isdir(directory_path):
-            raise ValueError(f"The path '{directory_path}' is not a directory")
-
-        # Use default suffixes if none provided
-        if suffixes is None:
-            suffixes = [".py", ".txt", ".yaml", ".yml", ".json"]
-
-        buf = BytesIO()
-
-        def no_mtime(tarinfo):
-            # a modification time change should not change the hash of
-            # the package. Only content modifications will.
-            # Setting this to a fixed date so that we don't have a changing hash everytime we run
-            tarinfo.mtime = 1747158696  # 13 May 2025 10:31:36
-            return tarinfo
-
-        with tarfile.open(
-            fileobj=buf, mode="w:gz", compresslevel=3, dereference=True
-        ) as tar:
-            # Add metadata if provided
-            if metadata:
-                cls._add_tar_file(
-                    tar, "metadata.json", BytesIO(json.dumps(metadata).encode("utf-8"))
-                )
-
-            # Walk the directory and add matching files
-            for file_path, rel_path in cls.directory_walker(
-                directory_path,
-                exclude_hidden=exclude_hidden,
-                suffixes=suffixes,
-            ):
-                # Remove debug print statement
-                tar.add(file_path, arcname=rel_path, recursive=False, filter=no_mtime)
-
-        tarball = bytearray(buf.getvalue())
-        tarball[4:8] = [0] * 4  # Reset 4 bytes from offset 4 to account for ts
-        return tarball