ob-metaflow-extensions 1.1.79__tar.gz → 1.1.80__tar.gz

{ob-metaflow-extensions-1.1.79 → ob-metaflow-extensions-1.1.80}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.79
+Version: 1.1.80
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

ob-metaflow-extensions-1.1.80/metaflow_extensions/outerbounds/config/__init__.py

@@ -0,0 +1,33 @@
+from metaflow.metaflow_config import from_conf
+
+DEFAULT_AWS_CLIENT_PROVIDER = "obp"
+
+DEFAULT_AZURE_CLIENT_PROVIDER = "obp"
+
+DEFAULT_GCP_CLIENT_PROVIDER = "obp"
+
+
+###
+# On Demand Docker image build configuration
+###
+# Image builder service url
+FAST_BAKERY_URL = from_conf("FAST_BAKERY_URL", None)
+
+
+###
+# Snowpark configuration
+###
+# Snowflake account to use with the @snowpark decorator
+SNOWPARK_ACCOUNT = from_conf("SNOWPARK_ACCOUNT")
+# Snowflake user to use with the @snowpark decorator
+SNOWPARK_USER = from_conf("SNOWPARK_USER")
+# Snowflake password to use with the @snowpark decorator
+SNOWPARK_PASSWORD = from_conf("SNOWPARK_PASSWORD")
+# Snowflake role to use with the @snowpark decorator
+SNOWPARK_ROLE = from_conf("SNOWPARK_ROLE")
+# Snowflake database to use with the @snowpark decorator
+SNOWPARK_DATABASE = from_conf("SNOWPARK_DATABASE")
+# Snowflake warehouse to use with the @snowpark decorator
+SNOWPARK_WAREHOUSE = from_conf("SNOWPARK_WAREHOUSE")
+# Snowflake schema to use with the @snowpark decorator
+SNOWPARK_SCHEMA = from_conf("SNOWPARK_SCHEMA")

{ob-metaflow-extensions-1.1.79 → ob-metaflow-extensions-1.1.80}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -178,7 +178,6 @@ class ObpGcpAuthProvider(object):
 
     @staticmethod
     def get_gs_storage_client(*args, **kwargs):
-
         import sys
         from metaflow_extensions.outerbounds.plugins.auth_server import get_token
 
@@ -240,8 +239,19 @@ class ObpGcpAuthProvider(object):
 
 
 GCP_CLIENT_PROVIDERS_DESC = [("obp", ".ObpGcpAuthProvider")]
-CLIS_DESC = [("nvcf", ".nvcf.nvcf_cli.cli")]
-STEP_DECORATORS_DESC = [("nvidia", ".nvcf.nvcf_decorator.NvcfDecorator")]
+CLIS_DESC = [
+    ("nvcf", ".nvcf.nvcf_cli.cli"),
+    ("fast-bakery", ".fast_bakery.fast_bakery_cli.cli"),
+    ("snowpark", ".snowpark.snowpark_cli.cli"),
+]
+STEP_DECORATORS_DESC = [
+    ("nvidia", ".nvcf.nvcf_decorator.NvcfDecorator"),
+    (
+        "fast_bakery_internal",
+        ".fast_bakery.fast_bakery_decorator.InternalFastBakeryDecorator",
+    ),
+    ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
+]
 FLOW_DECORATORS_DESC = [("nim", ".nim.NimDecorator")]
 TOGGLE_STEP_DECORATOR = [
     "-batch",
@@ -250,3 +260,7 @@ TOGGLE_STEP_DECORATOR = [
 ]
 
 TOGGLE_CLI = ["-batch", "-step-functions", "-airflow"]
+
+ENVIRONMENTS_DESC = [
+    ("fast-bakery", ".fast_bakery.docker_environment.DockerEnvironment")
+]

ob-metaflow-extensions-1.1.80/metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -0,0 +1,268 @@
+import hashlib
+import json
+import os
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import (
+    FAST_BAKERY_URL,
+    get_pinned_conda_libs,
+)
+from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.plugins.pypi.conda_environment import CondaEnvironment
+from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
+from metaflow.plugins.aws.batch.batch_decorator import BatchDecorator
+from metaflow.plugins.kubernetes.kubernetes_decorator import KubernetesDecorator
+from metaflow.plugins.pypi.conda_decorator import CondaStepDecorator
+from metaflow.plugins.pypi.pypi_decorator import PyPIStepDecorator
+
+BAKERY_METAFILE = ".imagebakery-cache"
+
+import json
+import os
+import fcntl
+from functools import wraps
+from concurrent.futures import ThreadPoolExecutor
+
+
+# TODO - ensure that both @conda/@pypi are not assigned to the same step
+
+
+def cache_request(cache_file):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            call_args = kwargs.copy()
+            call_args.update(zip(func.__code__.co_varnames, args))
+            call_args.pop("self", None)
+            cache_key = hashlib.md5(
+                json.dumps(call_args, sort_keys=True).encode("utf-8")
+            ).hexdigest()
+
+            try:
+                with open(cache_file, "r") as f:
+                    cache = json.load(f)
+                    if cache_key in cache:
+                        return FastBakeryApiResponse(cache[cache_key])
+            except (FileNotFoundError, json.JSONDecodeError):
+                cache = {}
+
+            result = func(*args, **kwargs)
+
+            try:
+                with open(cache_file, "r+") as f:
+                    fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+                    try:
+                        f.seek(0)
+                        cache = json.load(f)
+                    except json.JSONDecodeError:
+                        cache = {}
+
+                    cache[cache_key] = result.response
+
+                    f.seek(0)
+                    f.truncate()
+                    json.dump(cache, f)
+            except FileNotFoundError:
+                # path to cachefile might not exist.
+                os.makedirs(os.path.dirname(cache_file), exist_ok=True)
+                with open(cache_file, "w") as f:
+                    fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+                    json.dump({cache_key: result.response}, f)
+
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+class DockerEnvironmentException(MetaflowException):
+    headline = "Ran into an error while setting up the environment"
+
+    def __init__(self, msg):
+        super(DockerEnvironmentException, self).__init__(msg)
+
+
+class DockerEnvironment(MetaflowEnvironment):
+    TYPE = "fast-bakery"
+    _filecache = None
+
+    def __init__(self, flow):
+        self.skipped_steps = set()
+        self.flow = flow
+
+        self.bakery = FastBakery(url=FAST_BAKERY_URL)
+        self.results = {}
+
+    def set_local_root(self, local_root):
+        self.local_root = local_root
+
+    def decospecs(self):
+        return ("conda", "fast_bakery_internal") + super().decospecs()
+
+    def validate_environment(self, echo, datastore_type):
+        self.datastore_type = datastore_type
+        self.echo = echo
+
+        # Avoiding circular imports.
+        from metaflow.plugins import DATASTORES
+
+        self.datastore = [d for d in DATASTORES if d.TYPE == self.datastore_type][0]
+
+    def init_environment(self, echo):
+        self.skipped_steps = {
+            step.name
+            for step in self.flow
+            if not any(
+                isinstance(deco, (BatchDecorator, KubernetesDecorator))
+                for deco in step.decorators
+            )
+        }
+
+        steps_to_bake = [
+            step for step in self.flow if step.name not in self.skipped_steps
+        ]
+        if steps_to_bake:
+            echo("Baking container image(s) ...")
+            self.results = self._bake(steps_to_bake, echo)
+            for step in self.flow:
+                for d in step.decorators:
+                    if isinstance(d, (BatchDecorator, KubernetesDecorator)):
+                        d.attributes["image"] = self.results[step.name].container_image
+                        d.attributes["executable"] = self.results[step.name].python_path
+            echo("Container image(s) baked!")
+
+        if self.skipped_steps:
+            self.delegate = CondaEnvironment(self.flow)
+            self.delegate.set_local_root(self.local_root)
+            self.delegate.validate_environment(echo, self.datastore_type)
+            self.delegate.init_environment(echo, self.skipped_steps)
+
+    def _bake(self, steps, echo) -> Dict[str, FastBakeryApiResponse]:
+        metafile_path = get_fastbakery_metafile_path(self.local_root, self.flow.name)
+
+        @cache_request(metafile_path)
+        def _cached_bake(
+            python=None, pypi_packages=None, conda_packages=None, base_image=None
+        ):
+            self.bakery._reset_payload()
+            self.bakery.python_version(python)
+            self.bakery.pypi_packages(pypi_packages)
+            self.bakery.conda_packages(conda_packages)
+            self.bakery.base_image(base_image)
+            # self.bakery.ignore_cache()
+            try:
+                res = self.bakery.bake()
+                if res.baking_stats:
+                    echo(
+                        "baked image in: %s milliseconds"
+                        % res.baking_stats.solver_stats.duration_ms
+                    )
+                return res
+            except FastBakeryException as ex:
+                raise DockerEnvironmentException(str(ex))
+
+        def prepare_step(step):
+            base_image = next(
+                (
+                    d.attributes.get("image")
+                    for d in step.decorators
+                    if isinstance(d, (KubernetesDecorator))
+                ),
+                None,
+            )
+            dependencies = next(
+                (
+                    d
+                    for d in step.decorators
+                    if isinstance(d, (CondaStepDecorator, PyPIStepDecorator))
+                ),
+                None,
+            )
+            python = next(
+                (
+                    d.attributes["python"]
+                    for d in step.decorators
+                    if isinstance(d, CondaStepDecorator)
+                ),
+                None,
+            )
+            pypi_deco = next(
+                (d for d in step.decorators if isinstance(d, PyPIStepDecorator)), None
+            )
+            # if pypi decorator is set and user has specified a python version, we must create a new environment.
+            # otherwise rely on the base environment
+            if pypi_deco is not None:
+                python = (
+                    pypi_deco.attributes["python"]
+                    if pypi_deco.is_attribute_user_defined("python")
+                    else None
+                )
+
+            packages = get_pinned_conda_libs(python, self.datastore_type)
+            packages.update(dependencies.attributes["packages"] if dependencies else {})
+
+            return {
+                "python": python,
+                "pypi_packages": (
+                    packages if isinstance(dependencies, PyPIStepDecorator) else None
+                ),
+                "conda_packages": (
+                    packages if isinstance(dependencies, CondaStepDecorator) else None
+                ),
+                "base_image": base_image,
+            }
+
+        with ThreadPoolExecutor() as executor:
+            return {
+                step.name: _cached_bake(**args)
+                for step, args in zip(steps, executor.map(prepare_step, steps))
+            }
+
+    def executable(self, step_name, default=None):
+        if step_name in self.skipped_steps:
+            return self.delegate.executable(step_name, default)
+        # default is set to the right executable
+        if default is not None:
+            return default
+        if default is None and step_name in self.results:
+            # try to read pythonpath from results. This can happen immediately after baking.
+            return self.results[step_name].python_path
+        # we lack a default and baking results. fallback to parent executable.
+        return super().executable(step_name, default)
+
+    def interpreter(self, step_name):
+        if step_name in self.skipped_steps:
+            return self.delegate.interpreter(step_name)
+        return None
+
+    def is_disabled(self, step):
+        for decorator in step.decorators:
+            # @conda decorator is guaranteed to exist thanks to self.decospecs
+            if decorator.name in ["conda", "pypi"]:
+                # handle @conda/@pypi(disabled=True)
+                disabled = decorator.attributes["disabled"]
+                return str(disabled).lower() == "true"
+        return False
+
+    def pylint_config(self):
+        config = super().pylint_config()
+        # Disable (import-error) in pylint
+        config.append("--disable=F0401")
+        return config
+
+    def bootstrap_commands(self, step_name, datastore_type):
+        if step_name in self.skipped_steps:
+            return self.delegate.bootstrap_commands(step_name, datastore_type)
+        # Bootstrap conda and execution environment for step
+        # we set the environment flag for skipping bootstrap dependencies, as these are
+        # provided in all baked images.
+        return [
+            "export METAFLOW_SKIP_INSTALL_DEPENDENCIES=$FASTBAKERY_IMAGE",
+        ] + super().bootstrap_commands(step_name, datastore_type)
+
+
+def get_fastbakery_metafile_path(local_root, flow_name):
+    return os.path.join(local_root, flow_name, BAKERY_METAFILE)

ob-metaflow-extensions-1.1.80/metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -0,0 +1,160 @@
+from typing import Dict, Optional
+import requests
+
+
+class FastBakeryException(Exception):
+    pass
+
+
+class SolverStats:
+    def __init__(self, stats) -> None:
+        self.stats = stats
+
+    @property
+    def duration_ms(self):
+        return self.stats["durationMs"]
+
+    @property
+    def packages_in_solved_environment(self):
+        return self.stats["packagesInSolvedEnvironment"]
+
+
+class BakingStats:
+    def __init__(self, stats) -> None:
+        self.stats = stats
+
+    @property
+    def solver_stats(self) -> Optional[SolverStats]:
+        if "solverStats" not in self.stats:
+            return None
+        return SolverStats(self.stats["solverStats"])
+
+
+class FastBakeryApiResponse:
+    def __init__(self, response) -> None:
+        self.response = response
+
+    @property
+    def python_path(self) -> Optional[str]:
+        if not self.success:
+            return None
+
+        return self.response["success"]["pythonPath"]
+
+    @property
+    def container_image(self) -> Optional[str]:
+        if not self.success:
+            return None
+
+        return self.response["success"]["containerImage"]
+
+    @property
+    def success(self) -> bool:
+        return "success" in self.response
+
+    @property
+    def baking_stats(self) -> Optional[BakingStats]:
+        if not self.success:
+            return None
+
+        if "bakingStats" not in self.response["success"]:
+            return None
+
+        if self.response["success"]["bakingStats"] is None:
+            return None
+
+        return BakingStats(self.response["success"]["bakingStats"])
+
+    @property
+    def failure(self) -> bool:
+        return "failure" in self.response
+
+
+class FastBakery:
+    def __init__(self, url: str):
+        self.url = url
+        self.headers = {"Content-Type": "application/json", "Connection": "keep-alive"}
+        self._reset_payload()
+
+    def _reset_payload(self):
+        self._payload = {}
+
+    def python_version(self, version: str):
+        self._payload["pythonVersion"] = version
+        return self
+
+    def pypi_packages(self, packages: Dict[str, str]):
+        self._payload.setdefault("pipRequirements", []).extend(
+            self._format_packages(packages)
+        )
+        return self
+
+    def conda_packages(self, packages: Dict[str, str]):
+        self._payload.setdefault("condaMatchspecs", []).extend(
+            self._format_packages(packages)
+        )
+        return self
+
+    def base_image(self, image: str):
+        self._payload["baseImage"] = {"imageReference": image}
+        return self
+
+    def image_kind(self, kind: str):
+        self._payload["imageKind"] = kind
+        return self
+
+    def ignore_cache(self):
+        self._payload["cacheBehavior"] = {
+            "responseMaxAgeSeconds": 0,
+            "layerMaxAgeSeconds": 0,
+            "baseImageMaxAgeSeconds": 0,
+        }
+        return self
+
+    @staticmethod
+    def _format_packages(packages: Dict[str, str]) -> list:
+        if not packages:
+            return []
+
+        def format_package(pkg: str, ver: str) -> str:
+            return (
+                f"{pkg}{ver}"
+                if any(ver.startswith(c) for c in [">", "<", "~", "@", "="])
+                else f"{pkg}=={ver}"
+            )
+
+        return [format_package(pkg, ver) for pkg, ver in packages.items()]
+
+    def bake(self) -> FastBakeryApiResponse:
+        if "imageKind" not in self._payload:
+            self._payload["imageKind"] = "oci-zstd"  # Set default if not specified
+
+        res = self._make_request(self._payload)
+        self._reset_payload()
+        return res
+
+    def _make_request(self, payload: Dict) -> FastBakeryApiResponse:
+        try:
+            from metaflow.metaflow_config import SERVICE_HEADERS
+
+            headers = {**self.headers, **(SERVICE_HEADERS or {})}
+        except ImportError:
+            headers = self.headers
+        response = requests.post(self.url, json=payload, headers=headers)
+        self._handle_error_response(response)
+        return FastBakeryApiResponse(response.json())
+
+    @staticmethod
+    def _handle_error_response(response: requests.Response):
+        if response.status_code >= 500:
+            raise FastBakeryException(f"Server error: {response.text}")
+
+        body = response.json()
+        status_code = body.get("error", {}).get("statusCode", response.status_code)
+        if status_code >= 400:
+            try:
+                raise FastBakeryException(
+                    f"*{body['error']['details']['kind']}*\n{body['error']['details']['message']}"
+                )
+            except KeyError:
+                raise FastBakeryException(f"Unexpected error: {body}")

ob-metaflow-extensions-1.1.80/metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py

@@ -0,0 +1,54 @@
+import json
+import os
+from metaflow._vendor import click
+from metaflow.cli import echo_always as echo
+from metaflow.plugins.datastores.local_storage import LocalStorage
+
+from .docker_environment import get_fastbakery_metafile_path
+from .fast_bakery import FastBakeryApiResponse
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Commands related to Fast Bakery support.")
+@click.pass_context
+def fast_bakery(ctx):
+    path = LocalStorage.get_datastore_root_from_config(echo, create_on_absent=False)
+    ctx.obj.metafile_path = get_fastbakery_metafile_path(path, ctx.obj.flow.name)
+
+
+@fast_bakery.command(help="Purge local Fast Bakery cache.")
+@click.pass_obj
+def purge(obj):
+    try:
+        os.remove(obj.metafile_path)
+        echo("Local Fast Bakery cache purged.")
+    except FileNotFoundError:
+        echo("No local Fast Bakery cache found.")
+
+
+@fast_bakery.command(help="List the cached images")
+@click.pass_obj
+def images(obj):
+    current_cache = None
+    try:
+        with open(obj.metafile_path, "r") as f:
+            current_cache = json.load(f)
+    except FileNotFoundError:
+        pass
+
+    if current_cache:
+        echo("List of locally cached image tags:\n")
+
+        for val in current_cache.values():
+            response = FastBakeryApiResponse(val)
+            echo(response.container_image)
+
+        echo(
+            "In order to clear the cached images, you can use the command\n *fast-bakery purge*"
+        )
+    else:
+        echo("No locally cached images.")

ob-metaflow-extensions-1.1.80/metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py

@@ -0,0 +1,50 @@
+import os
+from metaflow.decorators import StepDecorator
+from metaflow.metadata.metadata import MetaDatum
+
+
+class InternalFastBakeryDecorator(StepDecorator):
+    """
+    Internal decorator to support Fast bakery
+    """
+
+    name = "fast_bakery_internal"
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_retries,
+        ubf_context,
+        inputs,
+    ):
+        # task_pre_step may run locally if fallback is activated for @catch
+        # decorator. In that scenario, we skip collecting Kubernetes execution
+        # metadata. A rudimentary way to detect non-local execution is to
+        # check for the existence of METAFLOW_KUBERNETES_WORKLOAD environment
+        # variable.
+        meta = {}
+        if "METAFLOW_KUBERNETES_WORKLOAD" in os.environ:
+            image = os.environ.get("FASTBAKERY_IMAGE")
+            if image:
+                meta["fast-bakery-image-name"] = image
+
+        if len(meta) > 0:
+            entries = [
+                MetaDatum(
+                    field=k,
+                    value=v,
+                    type=k,
+                    tags=["attempt_id:{0}".format(retry_count)],
+                )
+                for k, v in meta.items()
+                if v is not None
+            ]
+            # Register book-keeping metadata for debugging.
+            metadata.register_metadata(run_id, step_name, task_id, entries)