metaflow 2.13__py2.py3-none-any.whl → 2.13.2__py2.py3-none-any.whl

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
@@ -498,6 +493,7 @@ class LazyOpen(BufferedIOBase):
         self._file = None
         self._buffer = None
         self._position = 0
+        self.requests = None
 
     def _ensure_file(self):
         if not self._file:
@@ -514,8 +510,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

@@ -1,6 +1,7 @@
 import functools
 import json
 import os
+import re
 import subprocess
 import tempfile
 import time
@@ -8,7 +9,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):
@@ -23,6 +24,8 @@ class MicromambaException(MetaflowException):
 
 GLIBC_VERSION = os.environ.get("CONDA_OVERRIDE_GLIBC", "2.38")
 
+_double_equal_match = re.compile("==(?=[<=>!~])")
+
 
 class Micromamba(object):
     def __init__(self, logger=None):
@@ -101,7 +104,8 @@ class Micromamba(object):
                 cmd.append("--channel=%s" % channel)
 
             for package, version in packages.items():
-                cmd.append("%s==%s" % (package, version))
+                version_string = "%s==%s" % (package, version)
+                cmd.append(_double_equal_match.sub("", version_string))
             if python:
                 cmd.append("python==%s" % python)
             # TODO: Ensure a human readable message is returned when the environment
@@ -323,7 +327,7 @@ class Micromamba(object):
                         stderr="\n".join(err),
                     )
                 )
-            except (TypeError, ValueError) as ve:
+            except (TypeError, ValueError):
                 pass
             raise MicromambaException(
                 msg.format(
@@ -339,23 +343,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

metaflow/runner/metaflow_runner.py

@@ -7,6 +7,8 @@ from typing import Dict, Iterator, Optional, Tuple
 
 from metaflow import Run
 
+from metaflow.plugins import get_runner_cli
+
 from .utils import (
     temporary_fifo,
     handle_timeout,
@@ -187,7 +189,27 @@ class ExecutingRun(object):
             yield position, line
 
 
-class Runner(object):
+class RunnerMeta(type):
+    def __new__(mcs, name, bases, dct):
+        cls = super().__new__(mcs, name, bases, dct)
+
+        def _injected_method(subcommand_name, runner_subcommand):
+            def f(self, *args, **kwargs):
+                return runner_subcommand(self, *args, **kwargs)
+
+            f.__doc__ = runner_subcommand.__doc__ or ""
+            f.__name__ = subcommand_name
+
+            return f
+
+        for runner_subcommand in get_runner_cli():
+            method_name = runner_subcommand.name.replace("-", "_")
+            setattr(cls, method_name, _injected_method(method_name, runner_subcommand))
+
+        return cls
+
+
+class Runner(metaclass=RunnerMeta):
     """
     Metaflow's Runner API that presents a programmatic interface
     to run flows and perform other operations either synchronously or asynchronously.
@@ -337,7 +359,7 @@ class Runner(object):
 
             return self.__get_executing_run(attribute_file_fd, command_obj)
 
-    def resume(self, **kwargs):
+    def resume(self, **kwargs) -> ExecutingRun:
         """
         Blocking resume execution of the run.
         This method will wait until the resumed run has completed execution.
@@ -400,7 +422,7 @@ class Runner(object):
 
             return await self.__async_get_executing_run(attribute_file_fd, command_obj)
 
-    async def async_resume(self, **kwargs):
+    async def async_resume(self, **kwargs) -> ExecutingRun:
         """
         Non-blocking resume execution of the run.
         This method will return as soon as the resume has launched.

metaflow/runtime.py

@@ -125,7 +125,7 @@ class NativeRuntime(object):
         self._clone_run_id = clone_run_id
         self._clone_only = clone_only
         self._cloned_tasks = []
-        self._cloned_task_index = set()
+        self._ran_or_scheduled_task_index = set()
         self._reentrant = reentrant
         self._run_url = None
 
@@ -297,7 +297,7 @@ class NativeRuntime(object):
                     task.ubf_context = ubf_context
                 new_task_id = task.task_id
                 self._cloned_tasks.append(task)
-                self._cloned_task_index.add(cloned_task_pathspec_index)
+                self._ran_or_scheduled_task_index.add(cloned_task_pathspec_index)
                 task_pathspec = "{}/{}/{}".format(self._run_id, step_name, new_task_id)
             else:
                 task_pathspec = "{}/{}/{}".format(self._run_id, step_name, new_task_id)
@@ -384,8 +384,10 @@ class NativeRuntime(object):
                 and step_name != "_parameters"
                 and (step_name not in self._steps_to_rerun)
             ):
-                # "_unbounded_foreach" is a special flag to indicate that the transition is an unbounded foreach.
-                # Both parent and splitted children tasks will have this flag set. The splitted control/mapper tasks
+                # "_unbounded_foreach" is a special flag to indicate that the transition
+                # is an unbounded foreach.
+                # Both parent and splitted children tasks will have this flag set.
+                # The splitted control/mapper tasks
                 # are not foreach types because UBF is always followed by a join step.
                 is_ubf_task = (
                     "_unbounded_foreach" in task_ds and task_ds["_unbounded_foreach"]
@@ -647,10 +649,18 @@ class NativeRuntime(object):
     # Store the parameters needed for task creation, so that pushing on items
     # onto the run_queue is an inexpensive operation.
     def _queue_push(self, step, task_kwargs, index=None):
-        # If the to-be-pushed task is already cloned before, we don't need
-        # to re-run it.
-        if index and index in self._cloned_task_index:
-            return
+        # In the case of cloning, we set all the cloned tasks as the
+        # finished tasks when pushing tasks using _queue_tasks. This means that we
+        # could potentially try to push the same task multiple times (for example
+        # if multiple parents of a join are cloned). We therefore keep track of what
+        # has executed (been cloned) or what has been scheduled and avoid scheduling
+        # it again.
+        if index:
+            if index in self._ran_or_scheduled_task_index:
+                # It has already run or been scheduled
+                return
+            # Note that we are scheduling this to run
+            self._ran_or_scheduled_task_index.add(index)
         self._run_queue.insert(0, (step, task_kwargs))
         # For foreaches, this will happen multiple time but is ok, becomes a no-op
         self._unprocessed_steps.discard(step)

metaflow/user_configs/config_parameters.py

@@ -183,7 +183,7 @@ class DelayEvaluator(collections.abc.Mapping):
 
     def __getattr__(self, name):
         if self._access is None:
-            raise AttributeError()
+            raise AttributeError(name)
         self._access.append(name)
         return self
 
@@ -336,6 +336,8 @@ class Config(Parameter, collections.abc.Mapping):
         self.parser = parser
         self._computed_value = None
 
+        self._delayed_evaluator = None
+
     def load_parameter(self, v):
         if v is None:
             return None
@@ -344,22 +346,37 @@ class Config(Parameter, collections.abc.Mapping):
     def _store_value(self, v: Any) -> None:
         self._computed_value = v
 
+    def _init_delayed_evaluator(self) -> None:
+        if self._delayed_evaluator is None:
+            self._delayed_evaluator = DelayEvaluator(self.name.lower())
+
     # Support <config>.<var> syntax
     def __getattr__(self, name):
-        return DelayEvaluator(self.name.lower()).__getattr__(name)
+        # Need to return a new DelayEvaluator everytime because the evaluator will
+        # contain the "path" (ie: .name) and can be further accessed.
+        return getattr(DelayEvaluator(self.name.lower()), name)
 
-    # Next three methods are to implement mapping to support **<config> syntax
+    # Next three methods are to implement mapping to support **<config> syntax. We
+    # need to be careful, however, to also support a regular `config["key"]` syntax
+    # which calls into `__getitem__` and therefore behaves like __getattr__ above.
     def __iter__(self):
-        return iter(DelayEvaluator(self.name.lower()))
+        self._init_delayed_evaluator()
+        yield from self._delayed_evaluator
 
     def __len__(self):
-        return len(DelayEvaluator(self.name.lower()))
+        self._init_delayed_evaluator()
+        return len(self._delayed_evaluator)
 
     def __getitem__(self, key):
+        self._init_delayed_evaluator()
+        if key.startswith(UNPACK_KEY):
+            return self._delayed_evaluator[key]
         return DelayEvaluator(self.name.lower())[key]
 
 
 def resolve_delayed_evaluator(v: Any, ignore_errors: bool = False) -> Any:
+    # NOTE: We don't ignore errors in downstream calls because we want to have either
+    # all or nothing for the top-level call by the user.
     try:
         if isinstance(v, DelayEvaluator):
             return v()
@@ -397,7 +414,7 @@ def unpack_delayed_evaluator(
         else:
             # k.startswith(UNPACK_KEY)
             try:
-                result.update(resolve_delayed_evaluator(v[k]))
+                result.update(resolve_delayed_evaluator(v))
             except Exception as e:
                 if ignore_errors:
                     continue

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.13"
+metaflow_version = "2.13.2"

{metaflow-2.13.dist-info → metaflow-2.13.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.13
+Version: 2.13.2
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.13; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.13.2; extra == "stubs"
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)