ob-metaflow 2.12.39.1__py2.py3-none-any.whl → 2.13.1.1__py2.py3-none-any.whl

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -19,6 +19,8 @@ from metaflow.metaflow_config import (
     KUBERNETES_GPU_VENDOR,
     KUBERNETES_IMAGE_PULL_POLICY,
     KUBERNETES_MEMORY,
+    KUBERNETES_LABELS,
+    KUBERNETES_ANNOTATIONS,
     KUBERNETES_NAMESPACE,
     KUBERNETES_NODE_SELECTOR,
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
@@ -34,7 +36,8 @@ from metaflow.sidecar import Sidecar
 from metaflow.unbounded_foreach import UBF_CONTROL
 
 from ..aws.aws_utils import get_docker_registry, get_ec2_instance_metadata
-from .kubernetes import KubernetesException, parse_kube_keyvalue_list
+from .kubernetes import KubernetesException
+from .kube_utils import validate_kube_labels, parse_kube_keyvalue_list
 
 from metaflow.metaflow_config import MAX_MEMORY_PER_TASK, MAX_CPU_PER_TASK
 
@@ -91,6 +94,10 @@ class KubernetesDecorator(StepDecorator):
     tolerations : List[str], default []
         The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
         Kubernetes tolerations to use when launching pod in Kubernetes.
+    labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
+        Kubernetes labels to use when launching pod in Kubernetes.
+    annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
+        Kubernetes annotations to use when launching pod in Kubernetes.
     use_tmpfs : bool, default False
         This enables an explicit tmpfs mount for this step.
     tmpfs_tempdir : bool, default True
@@ -133,6 +140,8 @@ class KubernetesDecorator(StepDecorator):
         "gpu_vendor": None,
         "tolerations": None,  # e.g., [{"key": "arch", "operator": "Equal", "value": "amd"},
         #                              {"key": "foo", "operator": "Equal", "value": "bar"}]
+        "labels": None,  # e.g. {"test-label": "value", "another-label":"value2"}
+        "annotations": None,  # e.g. {"note": "value", "another-note": "value2"}
         "use_tmpfs": None,
         "tmpfs_tempdir": True,
         "tmpfs_size": None,
@@ -219,6 +228,36 @@ class KubernetesDecorator(StepDecorator):
             self.attributes["memory"] = KUBERNETES_MEMORY
         if self.attributes["disk"] == self.defaults["disk"] and KUBERNETES_DISK:
             self.attributes["disk"] = KUBERNETES_DISK
+        # Label source precedence (decreasing):
+        # - System labels (set outside of decorator)
+        # - Decorator labels: @kubernetes(labels={})
+        # - Environment variable labels: METAFLOW_KUBERNETES_LABELS=
+        deco_labels = {}
+        if self.attributes["labels"] is not None:
+            deco_labels = self.attributes["labels"]
+
+        env_labels = {}
+        if KUBERNETES_LABELS:
+            env_labels = parse_kube_keyvalue_list(KUBERNETES_LABELS.split(","), False)
+
+        self.attributes["labels"] = {**env_labels, **deco_labels}
+
+        # Annotations
+        # annotation precedence (decreasing):
+        # - System annotations (set outside of decorator)
+        # - Decorator annotations: @kubernetes(annotations={})
+        # - Environment annotations: METAFLOW_KUBERNETES_ANNOTATIONS=
+        deco_annotations = {}
+        if self.attributes["annotations"] is not None:
+            deco_annotations = self.attributes["annotations"]
+
+        env_annotations = {}
+        if KUBERNETES_ANNOTATIONS:
+            env_annotations = parse_kube_keyvalue_list(
+                KUBERNETES_ANNOTATIONS.split(","), False
+            )
+
+        self.attributes["annotations"] = {**env_annotations, **deco_annotations}
 
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
@@ -386,6 +425,9 @@ class KubernetesDecorator(StepDecorator):
                     )
                 )
 
+        validate_kube_labels(self.attributes["labels"])
+        # TODO: add validation to annotations as well?
+
     def package_init(self, flow, step_name, environment):
         try:
             # Kubernetes is a soft dependency.
@@ -441,7 +483,12 @@ class KubernetesDecorator(StepDecorator):
                         "=".join([key, str(val)]) if val else key
                         for key, val in v.items()
                     ]
-                elif k in ["tolerations", "persistent_volume_claims"]:
+                elif k in [
+                    "tolerations",
+                    "persistent_volume_claims",
+                    "labels",
+                    "annotations",
+                ]:
                     cli_args.command_options[k] = json.dumps(v)
                 else:
                     cli_args.command_options[k] = v

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -1,23 +1,20 @@
-import copy
 import json
 import math
 import random
-import sys
 import time
 
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_SECRETS
 from metaflow.tracing import inject_tracing_vars
-from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from metaflow.metaflow_config_funcs import init_config
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
+
+from .kube_utils import qos_requests_and_limits
 from .kubernetes_jobsets import (
     KubernetesJobSet,
 )  # We need this import for Kubernetes Client.
 
-from .kube_utils import qos_requests_and_limits
-
 
 class KubernetesJobException(MetaflowException):
     headline = "Kubernetes job error"
@@ -451,7 +448,7 @@ class RunningJob(object):
         def best_effort_kill():
             try:
                 self.kill()
-            except Exception as ex:
+            except Exception:
                 pass
 
         atexit.register(best_effort_kill)

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -1,4 +1,3 @@
-import copy
 import json
 import math
 import random
@@ -7,7 +6,6 @@ from collections import namedtuple
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_JOBSET_GROUP, KUBERNETES_JOBSET_VERSION
 from metaflow.tracing import inject_tracing_vars
-from metaflow.metaflow_config import KUBERNETES_SECRETS
 
 from .kube_utils import qos_requests_and_limits
 
@@ -257,7 +255,7 @@ class RunningJobSet(object):
         def best_effort_kill():
             try:
                 self.kill()
-            except Exception as ex:
+            except Exception:
                 pass
 
         atexit.register(best_effort_kill)
@@ -342,7 +340,7 @@ class RunningJobSet(object):
                 stdout=True,
                 tty=False,
             )
-        except Exception as e:
+        except Exception:
             with client.ApiClient() as api_client:
                 # If we are unable to kill the control pod then
                 # Delete the jobset to kill the subsequent pods.
@@ -862,6 +860,16 @@ class KubernetesJobSet(object):
         self._annotations = dict(self._annotations, **{name: value})
         return self
 
+    def labels(self, labels):
+        for k, v in labels.items():
+            self.label(k, v)
+        return self
+
+    def annotations(self, annotations):
+        for k, v in annotations.items():
+            self.annotation(k, v)
+        return self
+
     def secret(self, name):
         self.worker.secret(name)
         self.control.secret(name)
@@ -987,15 +995,24 @@ class KubernetesArgoJobSet(object):
         self._labels = dict(self._labels, **{name: value})
         return self
 
+    def labels(self, labels):
+        for k, v in labels.items():
+            self.label(k, v)
+        return self
+
     def annotation(self, name, value):
         self.worker.annotation(name, value)
         self.control.annotation(name, value)
         self._annotations = dict(self._annotations, **{name: value})
         return self
 
+    def annotations(self, annotations):
+        for k, v in annotations.items():
+            self.annotation(k, v)
+        return self
+
     def dump(self):
         client = self._kubernetes_sdk
-        import json
 
         data = json.dumps(
             client.ApiClient().sanitize_for_serialization(

metaflow/plugins/pypi/bootstrap.py

@@ -8,12 +8,14 @@ import subprocess
 import sys
 import tarfile
 import time
-
-import requests
-
+from urllib.error import URLError
+from urllib.request import urlopen
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.plugins import DATASTORES
+from metaflow.plugins.pypi.utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL
 from metaflow.util import which
+from urllib.request import Request
+import warnings
 
 from . import MAGIC_FILE, _datastore_packageroot
 
@@ -32,11 +34,6 @@ def timer(func):
 
 
 if __name__ == "__main__":
-    if len(sys.argv) != 5:
-        print("Usage: bootstrap.py <flow_name> <id> <datastore_type> <architecture>")
-        sys.exit(1)
-    _, flow_name, id_, datastore_type, architecture = sys.argv
-
     # TODO: Detect architecture on the fly when dealing with arm architectures.
     # ARCH=$(uname -m)
     # OS=$(uname)
@@ -61,30 +58,6 @@ if __name__ == "__main__":
     #     fi
     # fi
 
-    prefix = os.path.join(os.getcwd(), architecture, id_)
-    pkgs_dir = os.path.join(os.getcwd(), ".pkgs")
-    conda_pkgs_dir = os.path.join(pkgs_dir, "conda")
-    pypi_pkgs_dir = os.path.join(pkgs_dir, "pypi")
-    manifest_dir = os.path.join(os.getcwd(), DATASTORE_LOCAL_DIR, flow_name)
-
-    datastores = [d for d in DATASTORES if d.TYPE == datastore_type]
-    if not datastores:
-        print(f"No datastore found for type: {datastore_type}")
-        sys.exit(1)
-
-    storage = datastores[0](
-        _datastore_packageroot(datastores[0], lambda *args, **kwargs: None)
-    )
-
-    # Move MAGIC_FILE inside local datastore.
-    os.makedirs(manifest_dir, exist_ok=True)
-    shutil.move(
-        os.path.join(os.getcwd(), MAGIC_FILE),
-        os.path.join(manifest_dir, MAGIC_FILE),
-    )
-    with open(os.path.join(manifest_dir, MAGIC_FILE)) as f:
-        env = json.load(f)[id_][architecture]
-
     def run_cmd(cmd):
         result = subprocess.run(
             cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
@@ -107,29 +80,55 @@ if __name__ == "__main__":
             return micromamba_path
 
         # Download and extract in one go
-        # TODO: Serve from cloudflare
-        url = f"https://micro.mamba.pm/api/micromamba/{architecture}/2.0.4"
+        url = MICROMAMBA_URL.format(platform=architecture, version="2.0.4")
+        mirror_url = MICROMAMBA_MIRROR_URL.format(
+            platform=architecture, version="2.0.4"
+        )
 
         # Prepare directory once
         os.makedirs(os.path.dirname(micromamba_path), exist_ok=True)
 
-        # Stream and process directly to file
-        with requests.get(url, stream=True, timeout=30) as response:
-            if response.status_code != 200:
-                raise Exception(
-                    f"Failed to download micromamba: HTTP {response.status_code}"
-                )
-
-            decompressor = bz2.BZ2Decompressor()
-
-            # Process in memory without temporary files
-            tar_content = decompressor.decompress(response.raw.read())
+        # Download and decompress in one go
+        def _download_and_extract(url):
+            headers = {
+                "Accept-Encoding": "gzip, deflate, br",
+                "Connection": "keep-alive",
+                "User-Agent": "python-urllib",
+            }
 
-            with tarfile.open(fileobj=io.BytesIO(tar_content), mode="r:") as tar:
-                member = tar.getmember("bin/micromamba")
-                # Extract directly to final location
-                with open(micromamba_path, "wb") as f:
-                    f.write(tar.extractfile(member).read())
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    req = Request(url, headers=headers)
+
+                    with urlopen(req) as response:
+                        decompressor = bz2.BZ2Decompressor()
+                        with warnings.catch_warnings():
+                            warnings.filterwarnings(
+                                "ignore", category=DeprecationWarning
+                            )
+                            with tarfile.open(
+                                fileobj=io.BytesIO(
+                                    decompressor.decompress(response.read())
+                                ),
+                                mode="r:",
+                            ) as tar:
+                                member = tar.getmember("bin/micromamba")
+                                tar.extract(member, micromamba_dir)
+                    break
+                except (URLError, IOError) as e:
+                    if attempt == max_retries - 1:
+                        raise Exception(
+                            f"Failed to download micromamba after {max_retries} attempts: {e}"
+                        )
+                    time.sleep(2**attempt)
+
+        try:
+            # first try from mirror
+            _download_and_extract(mirror_url)
+        except Exception:
+            # download from mirror failed, try official source before failing.
+            _download_and_extract(url)
 
         # Set executable permission
         os.chmod(micromamba_path, 0o755)
@@ -140,7 +139,6 @@ if __name__ == "__main__":
 
     @timer
     def download_conda_packages(storage, packages, dest_dir):
-
         def process_conda_package(args):
             # Ensure that conda packages go into architecture specific folders.
             # The path looks like REPO/CHANNEL/CONDA_SUBDIR/PACKAGE. We trick
@@ -169,7 +167,6 @@ if __name__ == "__main__":
 
     @timer
     def download_pypi_packages(storage, packages, dest_dir):
-
         def process_pypi_package(args):
             key, tmpfile, dest_dir = args
             dest = os.path.join(dest_dir, os.path.basename(key))
@@ -208,7 +205,6 @@ if __name__ == "__main__":
 
     @timer
     def install_pypi_packages(prefix, pypi_pkgs_dir):
-
         cmd = f"""set -e;
             export PATH=$PATH:$(pwd)/micromamba;
             export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;
@@ -272,4 +268,41 @@ if __name__ == "__main__":
                 # wait for conda environment to be created
                 futures["conda_env"].result()
 
-    setup_environment(architecture, storage, env, prefix, conda_pkgs_dir, pypi_pkgs_dir)
+    if len(sys.argv) != 5:
+        print("Usage: bootstrap.py <flow_name> <id> <datastore_type> <architecture>")
+        sys.exit(1)
+
+    try:
+        _, flow_name, id_, datastore_type, architecture = sys.argv
+
+        prefix = os.path.join(os.getcwd(), architecture, id_)
+        pkgs_dir = os.path.join(os.getcwd(), ".pkgs")
+        conda_pkgs_dir = os.path.join(pkgs_dir, "conda")
+        pypi_pkgs_dir = os.path.join(pkgs_dir, "pypi")
+        manifest_dir = os.path.join(os.getcwd(), DATASTORE_LOCAL_DIR, flow_name)
+
+        datastores = [d for d in DATASTORES if d.TYPE == datastore_type]
+        if not datastores:
+            print(f"No datastore found for type: {datastore_type}")
+            sys.exit(1)
+
+        storage = datastores[0](
+            _datastore_packageroot(datastores[0], lambda *args, **kwargs: None)
+        )
+
+        # Move MAGIC_FILE inside local datastore.
+        os.makedirs(manifest_dir, exist_ok=True)
+        shutil.move(
+            os.path.join(os.getcwd(), MAGIC_FILE),
+            os.path.join(manifest_dir, MAGIC_FILE),
+        )
+        with open(os.path.join(manifest_dir, MAGIC_FILE)) as f:
+            env = json.load(f)[id_][architecture]
+
+        setup_environment(
+            architecture, storage, env, prefix, conda_pkgs_dir, pypi_pkgs_dir
+        )
+
+    except Exception as e:
+        print(f"Error: {str(e)}", file=sys.stderr)
+        sys.exit(1)

metaflow/plugins/pypi/conda_environment.py

@@ -7,20 +7,15 @@ import json
 import os
 import tarfile
 import threading
-import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from functools import wraps
 from hashlib import sha256
 from io import BufferedIOBase, BytesIO
-from itertools import chain
 from urllib.parse import unquote, urlparse
 
-import requests
-
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import get_pinned_conda_libs
 from metaflow.metaflow_environment import MetaflowEnvironment
-from metaflow.metaflow_profile import profile
 
 from . import MAGIC_FILE, _datastore_packageroot
 from .utils import conda_platform
@@ -501,6 +496,7 @@ class LazyOpen(BufferedIOBase):
         self._file = None
         self._buffer = None
         self._position = 0
+        self.requests = None
 
     def _ensure_file(self):
         if not self._file:
@@ -517,8 +513,13 @@ class LazyOpen(BufferedIOBase):
                 raise ValueError("Both filename and url are missing")
 
     def _download_to_buffer(self):
+        if self.requests is None:
+            # TODO: Remove dependency on requests
+            import requests
+
+            self.requests = requests
         # TODO: Stream it in chunks?
-        response = requests.get(self.url, stream=True)
+        response = self.requests.get(self.url, stream=True)
         response.raise_for_status()
         return response.content
 

metaflow/plugins/pypi/micromamba.py

@@ -8,7 +8,7 @@ import time
 from metaflow.exception import MetaflowException
 from metaflow.util import which
 
-from .utils import conda_platform
+from .utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL, conda_platform
 
 
 class MicromambaException(MetaflowException):
@@ -323,7 +323,7 @@ class Micromamba(object):
                         stderr="\n".join(err),
                     )
                 )
-            except (TypeError, ValueError) as ve:
+            except (TypeError, ValueError):
                 pass
             raise MicromambaException(
                 msg.format(
@@ -339,23 +339,37 @@ def _install_micromamba(installation_location):
     # Unfortunately no 32bit binaries are available for micromamba, which ideally
     # shouldn't be much of a problem in today's world.
     platform = conda_platform()
-    try:
-        subprocess.Popen(f"mkdir -p {installation_location}", shell=True).wait()
-        # https://mamba.readthedocs.io/en/latest/micromamba-installation.html#manual-installation
-        # requires bzip2
-        result = subprocess.Popen(
-            f"curl -Ls https://micro.mamba.pm/api/micromamba/{platform}/1.5.7 | tar -xvj -C {installation_location} bin/micromamba",
-            shell=True,
-            stderr=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-        )
-        _, err = result.communicate()
-        if result.returncode != 0:
-            raise MicromambaException(
-                f"Micromamba installation '{result.args}' failed:\n{err.decode()}"
-            )
+    url = MICROMAMBA_URL.format(platform=platform, version="1.5.7")
+    mirror_url = MICROMAMBA_MIRROR_URL.format(platform=platform, version="1.5.7")
+    os.makedirs(installation_location, exist_ok=True)
 
-    except subprocess.CalledProcessError as e:
-        raise MicromambaException(
-            "Micromamba installation failed:\n{}".format(e.stderr.decode())
-        )
+    def _download_and_extract(url):
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                # https://mamba.readthedocs.io/en/latest/micromamba-installation.html#manual-installation
+                # requires bzip2
+                result = subprocess.Popen(
+                    f"curl -Ls {url} | tar -xvj -C {installation_location} bin/micromamba",
+                    shell=True,
+                    stderr=subprocess.PIPE,
+                    stdout=subprocess.PIPE,
+                )
+                _, err = result.communicate()
+                if result.returncode != 0:
+                    raise MicromambaException(
+                        f"Micromamba installation '{result.args}' failed:\n{err.decode()}"
+                    )
+            except subprocess.CalledProcessError as e:
+                if attempt == max_retries - 1:
+                    raise MicromambaException(
+                        "Micromamba installation failed:\n{}".format(e.stderr.decode())
+                    )
+                time.sleep(2**attempt)
+
+    try:
+        # prioritize downloading from mirror
+        _download_and_extract(mirror_url)
+    except Exception:
+        # download from official source as a fallback
+        _download_and_extract(url)

metaflow/plugins/pypi/pip.py

@@ -4,7 +4,6 @@ import re
 import shutil
 import subprocess
 import tempfile
-import time
 from concurrent.futures import ThreadPoolExecutor
 from itertools import chain, product
 from urllib.parse import unquote
@@ -107,9 +106,8 @@ class Pip(object):
             except PipPackageNotFound as ex:
                 # pretty print package errors
                 raise PipException(
-                    "Could not find a binary distribution for %s \n"
-                    "for the platform %s\n\n"
-                    "Note that ***@pypi*** does not currently support source distributions"
+                    "Unable to find a binary distribution compatible with %s for %s.\n\n"
+                    "Note: ***@pypi*** does not currently support source distributions"
                     % (ex.package_spec, platform)
                 )
 

metaflow/plugins/pypi/utils.py

@@ -1,4 +1,3 @@
-import os
 import platform
 import sys
 
@@ -17,10 +16,13 @@ else:
     from metaflow._vendor.packaging import tags
     from metaflow._vendor.packaging.utils import parse_wheel_filename
 
-from urllib.parse import unquote, urlparse
+from urllib.parse import unquote
 
 from metaflow.exception import MetaflowException
 
+MICROMAMBA_URL = "https://micro.mamba.pm/api/micromamba/{platform}/{version}"
+MICROMAMBA_MIRROR_URL = "https://micromamba.outerbounds.sh/{platform}/{version}.tar.bz2"
+
 
 def conda_platform():
     # Returns the conda platform for the Python interpreter