outerbounds 0.3.171__py3-none-any.whl → 0.3.173rc0__py3-none-any.whl

outerbounds/apps/code_package/code_packager.py

@@ -0,0 +1,612 @@
+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.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"
+
+
+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, follow_symlinks=True
+    ) -> List[Tuple[str, str]]:
+        """
+        Walk a directory and yield tuples of (file_path, relative_arcname) for files
+        that match the given suffix filters.
+
+        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'])
+        follow_symlinks : bool, default True
+            Whether to follow symlinks (with cycle detection)
+
+        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
+        """
+        if suffixes is None:
+            suffixes = []
+
+        # Convert root to unicode to handle files/folders with non-ascii chars
+        root = str(root)
+
+        # Calculate the prefix length to strip from paths
+        prefixlen = len(os.path.dirname(root)) + 1  # +1 for the trailing slash
+
+        # Use a set to track visited symlinks to avoid cycles
+        seen = set()
+
+        def _walk_without_cycles(walk_root):
+            for parent, dirs, files in os.walk(walk_root, followlinks=follow_symlinks):
+                # If not following symlinks, we're done
+                if not follow_symlinks:
+                    yield parent, files
+                    continue
+
+                # When following symlinks, we need to check for cycles
+                for d_idx in range(
+                    len(dirs) - 1, -1, -1
+                ):  # Iterate backwards to safely remove
+                    d = dirs[d_idx]
+                    path = os.path.join(parent, d)
+                    if os.path.islink(path):
+                        # Break cycles by never following the same symlink twice
+                        reallink = os.path.realpath(path)
+                        if reallink in seen:
+                            # Remove from dirs to avoid following it
+                            dirs.pop(d_idx)
+                        else:
+                            seen.add(reallink)
+
+                yield parent, files
+
+        # Build the list of path tuples
+        result = []
+        for path, files in _walk_without_cycles(root):
+            # Skip hidden directories if requested
+            if exclude_hidden and "/." in path:
+                continue
+
+            for fname in files:
+                # Skip hidden files if requested, unless they have a specified suffix
+                if (
+                    (fname[0] == "." and fname in suffixes)
+                    or (fname[0] != "." or not exclude_hidden)
+                    and any(fname.endswith(suffix) for suffix in suffixes)
+                ):
+                    file_path = os.path.join(path, fname)
+                    rel_path = file_path[prefixlen:]
+                    result.append((file_path, rel_path))
+
+        return result
+
+    @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,
+                        follow_symlinks=True,
+                    ):
+                        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,
+        follow_symlinks: bool = True,
+    ) -> 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.
+
+        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
+        follow_symlinks : bool, default True
+            Whether to follow symlinks when walking the directory
+
+        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,
+                follow_symlinks=follow_symlinks,
+            ):
+                # 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