metaflow-prebuilt 0.1.0__tar.gz

metaflow_prebuilt-0.1.0/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: metaflow-prebuilt
+Version: 0.1.0
+Summary: Metaflow extension: pre-bake conda environments into Docker images for fast cold starts
+License: Apache-2.0
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: metaflow
+Provides-Extra: ecr
+Requires-Dist: boto3; extra == "ecr"
+Provides-Extra: codebuild
+Requires-Dist: boto3; extra == "codebuild"
+Provides-Extra: gcr
+Requires-Dist: google-cloud-storage; extra == "gcr"
+Provides-Extra: kaniko
+Requires-Dist: google-cloud-storage; extra == "kaniko"
+Requires-Dist: kubernetes; extra == "kaniko"
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: pytest-mock; extra == "dev"
+Requires-Dist: tox>=4; extra == "dev"

metaflow_prebuilt-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=45"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "metaflow-prebuilt"
+version = "0.1.0"
+description = "Metaflow extension: pre-bake conda environments into Docker images for fast cold starts"
+readme = "README.md"
+license = { text = "Apache-2.0" }
+requires-python = ">=3.8"
+dependencies = [
+    "metaflow",
+]
+
+[project.optional-dependencies]
+ecr = ["boto3"]
+codebuild = ["boto3"]
+gcr = ["google-cloud-storage"]
+kaniko = ["google-cloud-storage", "kubernetes"]
+dev = [
+    "pytest>=7",
+    "pytest-mock",
+    "tox>=4",
+]
+
+[project.entry-points."metaflow_prebuilt.build_services"]
+docker    = "metaflow_extensions.prebuilt.plugins.conda.services.docker_service:LocalDockerBuildService"
+kaniko    = "metaflow_extensions.prebuilt.plugins.conda.services.kaniko_service:KanikoBuildService"
+buildx    = "metaflow_extensions.prebuilt.plugins.conda.services.buildx_service:BuildxBuildService"
+codebuild = "metaflow_extensions.prebuilt.plugins.conda.services.codebuild_service:CodeBuildService"
+
+[project.entry-points."metaflow_prebuilt.image_registries"]
+ecr       = "metaflow_extensions.prebuilt.plugins.conda.registries.ecr_registry:ECRRegistry"
+gcr       = "metaflow_extensions.prebuilt.plugins.conda.registries.gcr_registry:GCRRegistry"
+dockerhub = "metaflow_extensions.prebuilt.plugins.conda.registries.dockerhub_registry:DockerHubRegistry"
+local     = "metaflow_extensions.prebuilt.plugins.conda.registries.local_registry:LocalRegistry"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

metaflow_prebuilt-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

metaflow_prebuilt-0.1.0/src/metaflow_extensions/prebuilt/mfextinit_prebuilt.py

@@ -0,0 +1,3 @@
+ENVIRONMENTS_DESC = [
+    ("prebuilt", ".plugins.conda.prebuilt_conda_environment.PrebuiltCondaEnvironment"),
+]

metaflow_prebuilt-0.1.0/src/metaflow_extensions/prebuilt/plugins/conda/build_service.py

@@ -0,0 +1,77 @@
+import importlib.metadata
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    pass
+
+
+def _resolve_entry_point(group: str, name: str) -> Any:
+    try:
+        eps = importlib.metadata.entry_points(group=group)
+    except TypeError:
+        # Python 3.8 compatibility
+        eps = importlib.metadata.entry_points().get(group, [])  # type: ignore[assignment]
+    mapping = {ep.name: ep for ep in eps}
+    if name not in mapping:
+        from metaflow.exception import MetaflowException  # noqa: PLC0415
+        raise MetaflowException(
+            "No entry point named %r in group %r. "
+            "Installed: %s. "
+            "Install the package that provides this backend or check "
+            "METAFLOW_PREBUILT_BUILD_SERVICE." % (name, group, sorted(mapping))
+        )
+    return mapping[name].load()
+
+
+class DockerBuildService(ABC):
+    """Abstract base class for Docker image build-and-push backends.
+
+    Implement this class and register it under the
+    ``metaflow_prebuilt.build_services`` entry point group to make it
+    selectable via ``METAFLOW_PREBUILT_BUILD_SERVICE=<name>``.
+    """
+
+    @abstractmethod
+    def build_and_push(
+        self,
+        dockerfile: str,
+        context_files: Dict[str, Any],
+        image_tag: str,
+        push_credentials: Dict[str, Any],
+        echo: Callable[..., None],
+    ) -> bool:
+        """Build a Docker image and push it to the registry.
+
+        Args:
+            dockerfile: Full Dockerfile content.
+            context_files: Files to include in the build context alongside the
+                Dockerfile. Keys are filenames; values are ``str`` (text) or
+                ``bytes`` (binary).
+            image_tag: Fully-qualified destination tag, e.g.
+                ``registry.example.com/ns/name:v28-abc123``.
+                Provided by ``ImageRegistry.push_tag()``.
+            push_credentials: Opaque dict from ``ImageRegistry.push_credentials()``.
+                Schema is defined by the paired registry implementation.
+            echo: Callable for user-visible progress output.
+
+        Returns:
+            ``True`` on success, ``False`` on recoverable failure.
+            Must NOT raise on recoverable failures — return ``False`` and
+            call ``echo()`` with a diagnostic. MAY raise on programmer errors.
+        """
+        ...
+
+    @classmethod
+    def from_config(cls) -> "DockerBuildService":
+        """Resolve and instantiate the build service selected by
+        ``METAFLOW_PREBUILT_BUILD_SERVICE`` (default: ``docker``).
+
+        Raises ``MetaflowException`` if the name is not registered.
+        """
+        name = os.environ.get("METAFLOW_PREBUILT_BUILD_SERVICE", "docker")
+        service_cls = _resolve_entry_point(
+            "metaflow_prebuilt.build_services", name
+        )
+        return service_cls()

metaflow_prebuilt-0.1.0/src/metaflow_extensions/prebuilt/plugins/conda/image_registry.py

@@ -0,0 +1,99 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Dict, TYPE_CHECKING
+
+from .build_service import _resolve_entry_point
+
+if TYPE_CHECKING:
+    from .env_descr import EnvID
+
+
+class ImageRegistry(ABC):
+    """Abstract base class for Docker image registries.
+
+    Implement this class and register it under the
+    ``metaflow_prebuilt.image_registries`` entry point group to make it
+    selectable via ``METAFLOW_PREBUILT_IMAGE_REGISTRY=<name>``.
+
+    For most registries ``push_tag()`` and ``pull_tag()`` both delegate to
+    ``image_tag()`` — subclasses only need to implement ``image_tag()``.
+    Override ``push_tag()`` and/or ``pull_tag()`` independently only when the
+    push address differs from the pull address (e.g. ``LocalRegistry``).
+    """
+
+    @abstractmethod
+    def image_tag(self, env_id: "EnvID") -> str:
+        """Canonical fully-qualified image tag for the given env_id.
+
+        Used as the default return value for both ``push_tag()`` and
+        ``pull_tag()``. Format: ``<registry>/<namespace>/<name>:<version>``.
+        """
+        ...
+
+    def push_tag(self, env_id: "EnvID") -> str:
+        """Tag the build service uses to push. Defaults to ``image_tag()``."""
+        return self.image_tag(env_id)
+
+    def pull_tag(self, env_id: "EnvID") -> str:
+        """Tag baked into the remote runner spec. Defaults to ``image_tag()``."""
+        return self.image_tag(env_id)
+
+    def image_tag_for_named(self, name: str) -> str:
+        """Full mutable tag for a ``@named_env(fetch_at_exec=True)`` env.
+
+        Tagged by alias name (not env_id hash) so subsequent deploys under the
+        same name overwrite the manifest — matching fetch_at_exec semantics.
+
+        Raises ``NotImplementedError`` by default; registries that support named
+        envs must override this method.
+        """
+        raise NotImplementedError(
+            "Registry %r does not implement named env tags. "
+            "Override image_tag_for_named() to use @named_env(fetch_at_exec=True)."
+            % type(self).__name__
+        )
+
+    def push_tag_for_named(self, name: str) -> str:
+        """Push-side tag for a named env. Defaults to ``image_tag_for_named()``."""
+        return self.image_tag_for_named(name)
+
+    def pull_tag_for_named(self, name: str) -> str:
+        """Pull-side tag for a named env. Defaults to ``push_tag_for_named()``."""
+        return self.push_tag_for_named(name)
+
+    @abstractmethod
+    def push_credentials(self) -> Dict[str, Any]:
+        """Credentials/config passed to ``DockerBuildService.build_and_push``.
+
+        Return ``{}`` when the build service uses ambient auth (e.g. docker login).
+        The schema is build-service-specific; see the ``DockerBuildService``
+        contract for which keys each service consumes.
+        """
+        ...
+
+    @abstractmethod
+    def pull_config(self, pull_tag: str) -> Dict[str, Any]:
+        """Attributes to inject into the remote runner decorator so it can pull
+        the image at ``pull_tag``.
+
+        Return ``{}`` when the runner has ambient pull credentials (e.g. an IAM
+        role for ECR + Batch, or a service account for GKE).
+
+        Non-empty example (private registry with K8s imagePullSecrets)::
+
+            {"image_pull_policy": "Always", "image_pull_secrets": "my-secret"}
+        """
+        ...
+
+    @classmethod
+    def from_config(cls) -> "ImageRegistry":
+        """Resolve and instantiate the registry selected by
+        ``METAFLOW_PREBUILT_IMAGE_REGISTRY`` (default: ``dockerhub``).
+
+        Raises ``MetaflowException`` if the name is not registered.
+        """
+        name = os.environ.get("METAFLOW_PREBUILT_IMAGE_REGISTRY", "dockerhub")
+        registry_cls = _resolve_entry_point(
+            "metaflow_prebuilt.image_registries", name
+        )
+        return registry_cls()

metaflow_prebuilt-0.1.0/src/metaflow_extensions/prebuilt/plugins/conda/prebuilt_build_install.py

@@ -0,0 +1,99 @@
+# pyright: strict, reportTypeCommentUsage=false, reportMissingTypeStubs=false
+"""Build-time conda env installer for `--environment=prebuilt`.
+
+Runs INSIDE the Cloudbuild docker build, with the MetaflowPackage
+tarball already extracted at `cwd`. The runtime task entry_point does
+exactly the same setup before invoking `remote_bootstrap.bootstrap_environment`
+— we just call `Conda.create_for_step` directly with the slightly
+different settings the build container needs.
+
+What the Dockerfile sets up for us:
+
+- `METAFLOW_CONDA_REMOTE_INSTALLER=""` → `Conda._ensure_remote_conda`'s
+  truthy check drops into `_ensure_micromamba` (public download from
+  micro.mamba.pm). No AWS creds required.
+- `METAFLOW_DATASTORE_SYSROOT_LOCAL=<cwd>` → `LocalStorage` reads the
+  resolved-env manifest from `<cwd>/.metaflow/<CONDA_MAGIC_FILE_V2>`.
+  We call `setup_conda_manifest()` (reused from `remote_bootstrap`)
+  to move the manifest from its packed location to that path.
+- `METAFLOW_EXTRACTED_ROOT=<cwd>` → activates the `.mf_install` marker
+  bypass in `metaflow.extension_support`, so the flattened nflx-*
+  distributions in the code package load cleanly without tripping
+  the path walker's per-distribution checks.
+- `PYTHONPATH=<cwd>/.mf_code` → metaflow + all extensions resolve from
+  the extracted code package.
+- `MAMBA_ROOT_PREFIX=<conda-root>` → micromamba creates the env at the
+  deterministic runtime-shared path.
+
+What's left for the shim:
+
+1.  `setup_conda_manifest()` — move the conda manifest from the
+    packaging-internal location to the LocalStorage path. Identical
+    call to the one `remote_bootstrap.bootstrap_environment` makes.
+2.  `Conda(echo, "local", mode="remote")` — `datastore_type="local"`
+    makes `Conda._storage = None`, which (after the upstream storage
+    guards in `_create` + `lazy_fetch_packages`) disables cache_info
+    URL emission and forces web downloads from `pkg.url`. Those
+    URLs go to conda.netflix.net / pypi.netflix.net over HTTPS.
+3.  `create_for_step` — same call the runtime makes via
+    `remote_bootstrap.bootstrap_environment`.
+"""
+import os
+import sys
+import time
+
+from metaflow.cli import echo_always
+
+from .conda import Conda
+from .env_descr import EnvID
+from .remote_bootstrap import setup_conda_manifest
+from .utils import arch_id
+
+
+def _echo(*args, **kwargs):
+    kwargs["err"] = False
+    echo_always(*args, **kwargs)
+
+
+def install_env(req_id: str, full_id: str) -> str:
+    start = time.time()
+    _echo("    Setting up Conda (build-time) ...", nl=False)
+
+    setup_conda_manifest()
+
+    my_conda = Conda(_echo, "local", mode="remote")
+    my_conda.binary("micromamba")
+    _echo(" done in %ds." % int(time.time() - start))
+
+    env_id = EnvID(req_id=req_id, full_id=full_id, arch=arch_id())
+    resolved_env = my_conda.environment(env_id)
+    if resolved_env is None:
+        raise RuntimeError(
+            "Cannot find cached environment for hash %s:%s in build-time "
+            "manifest. Verify the MetaflowPackage tarball was extracted at "
+            "the WORKDIR and that setup_conda_manifest moved the manifest "
+            "into place." % (req_id, full_id)
+        )
+
+    install_start = time.time()
+    env_path = my_conda.create_for_step(
+        "prebuilt_build", resolved_env, do_symlink=False
+    )
+    _echo(
+        "    Env installed at %s (%ds)" % (env_path, int(time.time() - install_start))
+    )
+    return env_path
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 3:
+        print(
+            "Usage: python -m metaflow_extensions.prebuilt.plugins.conda."
+            "prebuilt_build_install <req_id> <full_id>",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+    path = install_env(sys.argv[1], sys.argv[2])
+    print(path)
+    sys.stdout.flush()
+    os._exit(0)