ob-metaflow-extensions 1.1.45rc3__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -0,0 +1,391 @@
+import hashlib
+import json
+import os
+import threading
+import time
+import uuid
+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.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.conda_environment import CondaEnvironment
+from metaflow.plugins.pypi.pypi_decorator import PyPIStepDecorator
+from metaflow import decorators
+
+from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
+
+BAKERY_METAFILE = ".imagebakery-cache"
+
+import fcntl
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor
+from functools import wraps
+
+
+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)
+            call_args.pop("ref", None)
+            # invalidate cache when moving from one deployment to another
+            call_args.update({"fast_bakery_url": FAST_BAKERY_URL})
+            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 baking image"
+
+    def __init__(self, msg):
+        super(DockerEnvironmentException, self).__init__(msg)
+
+
+class DockerEnvironment(MetaflowEnvironment):
+    TYPE = "fast-bakery"
+    _filecache = None
+    _force_rebuild = False
+
+    def __init__(self, flow):
+        self.skipped_steps = set()
+        self.flow = flow
+
+        self.results = {}
+        self.images_baked = 0
+
+    def set_local_root(self, local_root):
+        self.local_root = local_root
+
+    def decospecs(self):
+        # Due to conflicts with the CondaEnvironment fallback and bakery,
+        # we can not simply attach 'conda' or 'pypi' to all steps here.
+        # Instead we do this on a per-step basis in init_environment
+        return ("fast_bakery_internal",) + super().decospecs()
+
+    def validate_environment(self, logger, datastore_type):
+        self.datastore_type = datastore_type
+        self.logger = logger
+
+        # 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 _step_executes_remotely(step)
+        }
+        # Attach environment decorator as needed. This is done on a step-by-step basis
+        # as we require a conda decorator for fallback steps, but prefer pypi for the baked ones.
+        for step in self.flow:
+            # Mixing @pypi/@conda in a single step is not supported yet.
+            # We validate this before attaching any new ones as the OSS Conda environment requires an implicit conda decorator for pypi environments which would fail the validation.
+            if sum(1 for deco in step.decorators if _is_env_deco(deco)) > 1:
+                raise MetaflowException(
+                    "Mixing and matching PyPI packages and Conda packages within a\n"
+                    "step is not yet supported. Use one of @pypi or @conda only for the *%s* step."
+                    % step.name
+                )
+            if step.name in self.skipped_steps:
+                # Conda fallback requires a conda decorator as the default for a step
+                decorators._attach_decorators_to_step(step, ["conda"])
+            else:
+                if not _step_has_environment_deco(step):
+                    # We default to PyPI for steps that are going to be baked.
+                    decorators._attach_decorators_to_step(step, ["pypi"])
+                    # init the attached decorator
+            # Initialize the decorator we attached.
+            # This is crucial for the conda decorator to work properly in the fallback environment
+            decorators._init(self.flow)
+            for deco in step.decorators:
+                if _is_env_deco(deco):
+                    deco.step_init(
+                        self.flow,
+                        None,  # not passing graph as it is not available, and not required by conda/pypi decorators
+                        step.name,
+                        step.decorators,
+                        self,
+                        self.datastore,
+                        echo,
+                    )
+
+        steps_to_bake = [
+            step
+            for step in self.flow
+            if step.name not in self.skipped_steps and not self.is_disabled(step)
+        ]
+        if steps_to_bake:
+            self.logger("🚀 Baking container image(s) ...")
+            start_time = time.time()
+            self.results = self._bake(steps_to_bake)
+            for step in steps_to_bake:
+                for d in step.decorators:
+                    if _is_remote_deco(d):
+                        d.attributes["image"] = self.results[step.name].container_image
+                        d.attributes["executable"] = self.results[step.name].python_path
+            if self.images_baked > 0:
+                bake_time = time.time() - start_time
+                self.logger(
+                    f"🎉 All container image(s) baked in {bake_time:.2f} seconds!"
+                )
+            else:
+                self.logger("🎉 All container image(s) baked!")
+
+        if self.skipped_steps:
+            self.delegate = CondaEnvironment(self.flow)
+            self.delegate._force_rebuild = self._force_rebuild
+            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) -> Dict[str, FastBakeryApiResponse]:
+        metafile_path = get_fastbakery_metafile_path(self.local_root, self.flow.name)
+        if self._force_rebuild:
+            # clear the metafile if force rebuilding, effectively skipping the cache.
+            try:
+                os.remove(metafile_path)
+            except Exception:
+                pass
+
+        logger_lock = threading.Lock()
+
+        @cache_request(metafile_path)
+        def _cached_bake(
+            ref=None,
+            python=None,
+            pypi_packages=None,
+            conda_packages=None,
+            base_image=None,
+        ):
+            try:
+                bakery = FastBakery(url=FAST_BAKERY_URL)
+                bakery._reset_payload()
+                bakery.python_version(python)
+                bakery.pypi_packages(pypi_packages)
+                bakery.conda_packages(conda_packages)
+                bakery.base_image(base_image)
+                if self._force_rebuild:
+                    bakery.ignore_cache()
+
+                with logger_lock:
+                    self.logger(f"🍳 Baking [{ref}] ...")
+                    self.logger(f"     🐍 Python: {python}")
+
+                    if pypi_packages:
+                        self.logger(f"     📦 PyPI packages:")
+                        for package, version in pypi_packages.items():
+                            self.logger(f"        🔧 {package}: {version}")
+
+                    if conda_packages:
+                        self.logger(f"     📦 Conda packages:")
+                        for package, version in conda_packages.items():
+                            self.logger(f"        🔧 {package}: {version}")
+
+                    self.logger(f"     🏗️  Base image: {base_image}")
+
+                start_time = time.time()
+                res = bakery.bake()
+                # TODO: Get actual bake time from bakery
+                bake_time = time.time() - start_time
+
+                with logger_lock:
+                    self.logger(f"🏁 Baked [{ref}] in {bake_time:.2f} seconds!")
+                self.images_baked += 1
+                return res
+            except FastBakeryException as ex:
+                raise DockerEnvironmentException(f"Bake [{ref}] failed: {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 _is_env_deco(d)),
+                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 {})
+
+            requested = {
+                "python": python,
+                "pypi_packages": (
+                    packages if isinstance(dependencies, PyPIStepDecorator) else None
+                ),
+                "conda_packages": (
+                    packages if isinstance(dependencies, CondaStepDecorator) else None
+                ),
+                "base_image": base_image,
+            }
+            dedup_key = hashlib.sha256(
+                json.dumps(requested).encode("utf-8")
+            ).hexdigest()
+
+            return step.name, dedup_key, requested
+
+        with ThreadPoolExecutor() as executor:
+            prepared_args = list(executor.map(prepare_step, steps))
+            # Deduplicate the requests for baking images of steps.
+            # We do not want to bake the same image twice.
+            dedup_requests = {}
+            for step_name, key, args in prepared_args:
+                if key not in dedup_requests:
+                    dedup_requests[key] = {"step_names": set(), "args": args}
+                dedup_requests[key]["step_names"].add(step_name)
+
+            # unique futures
+            futures = []
+            for i, kv in enumerate(dedup_requests.items(), 1):
+                key, value = kv
+                future = executor.submit(
+                    _cached_bake, **{**value["args"], "ref": f"#{i:02d}"}
+                )
+                futures.append({"step_names": value["step_names"], "future": future})
+
+            results = {}
+            for item in futures:
+                for step_name in item["step_names"]:
+                    results[step_name] = item["future"].result()
+
+            return results
+
+    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 get_package_commands(
+        self, codepackage_url, datastore_type, code_package_metadata=None
+    ):
+        # we must set the skip install flag at this stage in order to skip package downloads,
+        # doing so in bootstrap_commands is too late in the lifecycle.
+        return [
+            "export METAFLOW_SKIP_INSTALL_DEPENDENCIES=$FASTBAKERY_IMAGE",
+        ] + super().get_package_commands(
+            codepackage_url, datastore_type, code_package_metadata=code_package_metadata
+        )
+
+    def bootstrap_commands(self, step_name, datastore_type):
+        if step_name in self.skipped_steps:
+            return self.delegate.bootstrap_commands(step_name, datastore_type)
+        return 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)
+
+
+def _is_remote_deco(deco):
+    return isinstance(deco, (BatchDecorator, KubernetesDecorator))
+
+
+def _step_executes_remotely(step):
+    "Check if a step is going to execute remotely or locally"
+    return any(_is_remote_deco(deco) for deco in step.decorators)
+
+
+def _is_env_deco(deco):
+    "Check if a decorator is a known environment decorator (conda/pypi)"
+    return isinstance(deco, (PyPIStepDecorator, CondaStepDecorator))
+
+
+def _step_has_environment_deco(step):
+    "Check if a step has a virtual environment decorator"
+    return any(_is_env_deco(deco) for deco in step.decorators)

metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -0,0 +1,188 @@
+from typing import Dict, Optional
+import requests
+import time
+
+
+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):
+        if not url:
+            raise FastBakeryException("Specifying a url is required.")
+        self.url = url
+        self.headers = {"Content-Type": "application/json", "Connection": "keep-alive"}
+        self._reset_payload()
+
+    def _reset_payload(self):
+        self._payload = {}
+        from metaflow_extensions.outerbounds.remote_config import init_config
+        from os import environ
+
+        conf = init_config()
+        if "OBP_PERIMETER" in conf:
+            perimeter = conf["OBP_PERIMETER"]
+        else:
+            # if the perimeter is not in metaflow config, try to get it from the environment
+            perimeter = environ.get("OBP_PERIMETER", "")
+
+        self._payload["perimeterName"] = perimeter
+
+    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,
+            "overwriteExistingLayers": True,  # Used primarily to rewrite possibly corrupted layers.
+        }
+        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}" if ver != "" else f"{pkg}")
+            )
+
+        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
+
+        retryable_status_codes = [409]
+
+        for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
+            response = requests.post(self.url, json=payload, headers=headers)
+
+            if response.status_code not in retryable_status_codes:
+                break
+
+            if attempt < 2:  # Don't sleep after the last attempt
+                sleep_time = 0.5 * (attempt + 1)
+                time.sleep(sleep_time)
+
+        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}")

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.")

metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py

@@ -0,0 +1,50 @@
+import os
+from metaflow.decorators import StepDecorator
+from metaflow.metadata_provider.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)