ob-metaflow 2.9.10.1__py2.py3-none-any.whl → 2.10.2.6__py2.py3-none-any.whl

metaflow/plugins/aws/batch/batch_decorator.py

@@ -91,8 +91,8 @@ class BatchDecorator(StepDecorator):
         "gpu": None,
         "memory": None,
         "image": None,
-        "queue": BATCH_JOB_QUEUE,  # Will be replaced with an available queue if not provided.
-        "iam_role": ECS_S3_ACCESS_IAM_ROLE,  # Required
+        "queue": BATCH_JOB_QUEUE,
+        "iam_role": ECS_S3_ACCESS_IAM_ROLE,
         "execution_role": ECS_FARGATE_EXECUTION_ROLE,
         "shared_memory": None,
         "max_swap": None,
@@ -149,13 +149,6 @@ class BatchDecorator(StepDecorator):
         if flow_datastore.TYPE != "s3":
             raise BatchException("The *@batch* decorator requires --datastore=s3.")
 
-        # Require iam_role
-        if self.attributes["iam_role"] is None:
-            # TODO: Unify messaging on various configuration options.
-            raise BatchException(
-                "The *@batch* decorator requires an IAM Role that allows AWS Batch job to communicate with Amazon S3 datastore.\n"
-                'You can specify it either with @batch(iam_role="role-name") or by setting METAFLOW_ECS_S3_ACCESS_IAM_ROLE in your configuration.'
-            )
         # Set internal state.
         self.logger = logger
         self.environment = environment

metaflow/plugins/datatools/s3/s3.py

@@ -513,6 +513,10 @@ class S3(object):
 
         if run:
             # 1. use a (current) run ID with optional customizations
+            if DATATOOLS_S3ROOT is None:
+                raise MetaflowS3URLException(
+                    "DATATOOLS_S3ROOT is not configured when trying to use S3 storage"
+                )
             parsed = urlparse(DATATOOLS_S3ROOT)
             if not bucket:
                 bucket = parsed.netloc

metaflow/plugins/env_escape/client.py

@@ -1,4 +1,5 @@
 import fcntl
+import gc
 import os
 import importlib
 import itertools
@@ -60,6 +61,7 @@ class Client(object):
         # have an exception
         self._poller = None
         self._poller_lock = threading.Lock()
+        self._active_pid = os.getpid()
         self._server_process = None
         self._socket_path = None
 
@@ -70,11 +72,20 @@ class Client(object):
         # The client launches the server when created; we use
         # Unix sockets for now
         server_module = ".".join([__package__, "server"])
-        self._socket_path = "/tmp/%s_%d" % (server_config, os.getpid())
+        self._socket_path = "/tmp/%s_%d" % (server_config, self._active_pid)
         if os.path.exists(self._socket_path):
             raise RuntimeError("Existing socket: %s" % self._socket_path)
         env = os.environ.copy()
         env["PYTHONPATH"] = pythonpath
+        # When coming from a conda environment, LD_LIBRARY_PATH may be set to
+        # first include the Conda environment's library. When breaking out to
+        # the underlying python, we need to reset it to the original LD_LIBRARY_PATH
+        ld_lib_path = env.get("LD_LIBRARY_PATH")
+        orig_ld_lib_path = env.get("MF_ORIG_LD_LIBRARY_PATH")
+        if ld_lib_path is not None and orig_ld_lib_path is not None:
+            env["LD_LIBRARY_PATH"] = orig_ld_lib_path
+            if orig_ld_lib_path is not None:
+                del env["MF_ORIG_LD_LIBRARY_PATH"]
         self._server_process = Popen(
             [
                 python_executable,
@@ -381,8 +392,18 @@ class Client(object):
         return name
 
     def _communicate(self, msg):
-        with self._poller_lock:
-            return self._locked_communicate(msg)
+        if os.getpid() != self._active_pid:
+            raise RuntimeError(
+                "You cannot use the environment escape across process boundaries."
+            )
+        # We also disable the GC because in some rare cases, it may try to delete
+        # a remote object while we are communicating which will cause a deadlock
+        try:
+            gc.disable()
+            with self._poller_lock:
+                return self._locked_communicate(msg)
+        finally:
+            gc.enable()
 
     def _locked_communicate(self, msg):
         self._channel.send(msg)

metaflow/plugins/env_escape/stub.py

@@ -26,7 +26,6 @@ LOCAL_ATTRS = (
             "___remote_class_name___",
             "___identifier___",
             "___connection___",
-            "___refcount___",
             "___local_overrides___" "__class__",
             "__init__",
             "__del__",
@@ -78,6 +77,7 @@ class StubMetaClass(type):
 
 def with_metaclass(meta, *bases):
     """Create a base class with a metaclass."""
+
     # Compatibility 2/3. Remove when only 3 support
     class metaclass(type):
         def __new__(cls, name, this_bases, d):
@@ -99,7 +99,6 @@ class Stub(with_metaclass(StubMetaClass, object)):
         "___identifier___",
         "___connection___",
         "__weakref__",
-        "___refcount___",
     ]
 
     # def __iter__(self):  # FIXME: Keep debugger QUIET!!
@@ -109,14 +108,10 @@ class Stub(with_metaclass(StubMetaClass, object)):
         self.___remote_class_name___ = remote_class_name
         self.___identifier___ = identifier
         self.___connection___ = connection
-        self.___refcount___ = 1
 
     def __del__(self):
         try:
-            pass
-            self.___refcount___ -= 1
-            if self.___refcount___ == 0:
-                fwd_request(self, OP_DEL)
+            fwd_request(self, OP_DEL)
         except Exception:
             # raised in a destructor, most likely on program termination,
             # when the connection might have already been closed.
@@ -262,7 +257,6 @@ def create_class(
     setattr_overrides,
     class_methods,
 ):
-
     class_dict = {"__slots__": ()}
     for name, doc in class_methods.items():
         method_type = NORMAL_METHOD

metaflow/plugins/kubernetes/kubernetes.py

@@ -15,6 +15,7 @@ from metaflow.metaflow_config import (
     ARGO_EVENTS_SERVICE_ACCOUNT,
     ARGO_EVENTS_INTERNAL_WEBHOOK_URL,
     AWS_SECRETS_MANAGER_DEFAULT_REGION,
+    ARGO_EVENTS_WEBHOOK_AUTH,
     AZURE_STORAGE_BLOB_SERVICE_ENDPOINT,
     CARD_AZUREROOT,
     CARD_GSROOT,
@@ -35,6 +36,8 @@ from metaflow.metaflow_config import (
     S3_SERVER_SIDE_ENCRYPTION,
     OTEL_ENDPOINT,
 )
+from metaflow.metaflow_config_funcs import config_values
+
 from metaflow.mflog import (
     BASH_SAVE_LOGS,
     bash_capture_logs,
@@ -261,6 +264,12 @@ class Kubernetes(object):
             # see get_datastore_root_from_config in datastore/local.py).
         )
 
+        # Temporary passing of *some* environment variables. Do not rely on this
+        # mechanism as it will be removed in the near future
+        for k, v in config_values():
+            if k.startswith("METAFLOW_CONDA_") or k.startswith("METAFLOW_DEBUG_"):
+                job.environment_variable(k, v)
+
         if S3_SERVER_SIDE_ENCRYPTION is not None:
             job.environment_variable(
                 "METAFLOW_S3_SERVER_SIDE_ENCRYPTION", S3_SERVER_SIDE_ENCRYPTION
@@ -280,6 +289,10 @@ class Kubernetes(object):
         job.environment_variable(
             "METAFLOW_ARGO_EVENTS_SERVICE_ACCOUNT", ARGO_EVENTS_SERVICE_ACCOUNT
         )
+        job.environment_variable(
+            "METAFLOW_ARGO_EVENTS_WEBHOOK_AUTH",
+            ARGO_EVENTS_WEBHOOK_AUTH,
+        )
 
         tmpfs_enabled = use_tmpfs or (tmpfs_size and not use_tmpfs)
         if tmpfs_enabled and tmpfs_tempdir:

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -143,8 +143,7 @@ def step(
     node = ctx.obj.graph[step_name]
 
     # Construct entrypoint CLI
-    if executable is None:
-        executable = ctx.obj.environment.executable(step_name)
+    executable = ctx.obj.environment.executable(step_name, executable)
 
     # Set environment
     env = {}

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -58,14 +58,21 @@ class KubernetesDecorator(StepDecorator):
         Docker image to use when launching on Kubernetes. If not specified, and
         METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
         not, a default Docker image mapping to the current version of Python is used.
+    image_pull_policy: str, default: KUBERNETES_IMAGE_PULL_POLICY
+        If given, the imagePullPolicy to be applied to the Docker image of the step.
     service_account : str, default: METAFLOW_KUBERNETES_SERVICE_ACCOUNT
         Kubernetes service account to use when launching pod in Kubernetes.
-    namespace : str, default: METAFLOW_KUBERNETES_NAMESPACE
-        Kubernetes namespace to use when launching pod in Kubernetes.
     secrets : List[str], optional
         Kubernetes secrets to use when launching pod in Kubernetes. These
         secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
         in Metaflow configuration.
+    namespace : str, default: METAFLOW_KUBERNETES_NAMESPACE
+        Kubernetes namespace to use when launching pod in Kubernetes.
+    gpu: int, optional
+        Number of GPUs required for this step. A value of zero implies that
+        the scheduled node should not have GPUs.
+    gpu_vendor: str, default: KUBERNETES_GPU_VENDOR
+        The vendor of the GPUs to be used for this step.
     tolerations : List[str], default: METAFLOW_KUBERNETES_TOLERATIONS
         Kubernetes tolerations to use when launching pod in Kubernetes.
     use_tmpfs: bool, default: False

metaflow/plugins/pypi/__init__.py

@@ -0,0 +1,29 @@
+from metaflow import metaflow_config
+from metaflow.exception import MetaflowException
+
+MAGIC_FILE = "conda.manifest"
+
+
+# TODO: This can be lifted all the way into metaflow config
+def _datastore_packageroot(datastore, echo):
+    datastore_type = datastore.TYPE
+    datastore_packageroot = getattr(
+        metaflow_config,
+        "CONDA_PACKAGE_{datastore_type}ROOT".format(
+            datastore_type=datastore_type.upper()
+        ),
+        None,
+    )
+    if datastore_packageroot is None:
+        datastore_sysroot = datastore.get_datastore_root_from_config(echo)
+        if datastore_sysroot is None:
+            # TODO: Throw a more evocative error message
+            raise MetaflowException(
+                msg="METAFLOW_DATASTORE_SYSROOT_{datastore_type} must be set!".format(
+                    datastore_type=datastore_type.upper()
+                )
+            )
+        datastore_packageroot = "{datastore_sysroot}/conda".format(
+            datastore_sysroot=datastore_sysroot
+        )
+    return datastore_packageroot

metaflow/plugins/pypi/bootstrap.py

@@ -0,0 +1,131 @@
+import bz2
+import io
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+
+from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
+from metaflow.plugins import DATASTORES
+from metaflow.util import which
+
+from . import MAGIC_FILE, _datastore_packageroot
+
+# Bootstraps a valid conda virtual environment composed of conda and pypi packages
+
+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)
+
+    # if [[ "$OS" == "Linux" ]]; then
+    #     PLATFORM="linux"
+    #     if [[ "$ARCH" == "aarch64" ]]; then
+    #         ARCH="aarch64";
+    #     elif [[ $ARCH == "ppc64le" ]]; then
+    #         ARCH="ppc64le";
+    #     else
+    #         ARCH="64";
+    #     fi
+    # fi
+
+    # if [[ "$OS" == "Darwin" ]]; then
+    #     PLATFORM="osx";
+    #     if [[ "$ARCH" == "arm64" ]]; then
+    #         ARCH="arm64";
+    #     else
+    #         ARCH="64"
+    #     fi
+    # fi
+
+    prefix = os.path.join(os.getcwd(), id_)
+    pkgs_dir = os.path.join(os.getcwd(), ".pkgs")
+    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]
+
+    # Download Conda packages.
+    conda_pkgs_dir = os.path.join(pkgs_dir, "conda")
+    with storage.load_bytes([package["path"] for package in env["conda"]]) as results:
+        for key, tmpfile, _ in results:
+            # Ensure that conda packages go into architecture specific folders.
+            # The path looks like REPO/CHANNEL/CONDA_SUBDIR/PACKAGE. We trick
+            # Micromamba into believing that all packages are coming from a local
+            # channel - the only hurdle is ensuring that packages are organised
+            # properly.
+
+            # TODO: consider RAM disk
+            dest = os.path.join(conda_pkgs_dir, "/".join(key.split("/")[-2:]))
+            os.makedirs(os.path.dirname(dest), exist_ok=True)
+            shutil.move(tmpfile, dest)
+
+    # Create Conda environment.
+    cmds = [
+        # TODO: check if mamba or conda are already available on the image
+        # TODO: search for micromamba everywhere
+        f"""if ! command -v ./micromamba >/dev/null 2>&1; then
+            wget -qO- https://micro.mamba.pm/api/micromamba/{architecture}/latest | python -c "import sys, bz2; sys.stdout.buffer.write(bz2.decompress(sys.stdin.buffer.read()))" | tar -xv bin/micromamba --strip-components=1 ;
+            if ! command -v ./micromamba >/dev/null 2>&1; then
+                echo "Failed to install Micromamba!";
+                exit 1;
+            fi;
+        fi""",
+        # Create a conda environment through Micromamba.
+        f'''tmpfile=$(mktemp);
+        echo "@EXPLICIT" > "$tmpfile";
+        ls -d {conda_pkgs_dir}/*/* >> "$tmpfile";
+        ./micromamba create --yes --offline --no-deps --safety-checks=disabled --no-extra-safety-checks --prefix {prefix} --file "$tmpfile";
+        rm "$tmpfile"''',
+    ]
+
+    # Download PyPI packages.
+    if "pypi" in env:
+        pypi_pkgs_dir = os.path.join(pkgs_dir, "pypi")
+        with storage.load_bytes(
+            [package["path"] for package in env["pypi"]]
+        ) as results:
+            for key, tmpfile, _ in results:
+                dest = os.path.join(pypi_pkgs_dir, os.path.basename(key))
+                os.makedirs(os.path.dirname(dest), exist_ok=True)
+                shutil.move(tmpfile, dest)
+
+        # Install PyPI packages.
+        cmds.extend(
+            [
+                f"""./micromamba run --prefix {prefix} pip --disable-pip-version-check install --root-user-action=ignore --no-compile {pypi_pkgs_dir}/*.whl"""
+            ]
+        )
+
+    for cmd in cmds:
+        result = subprocess.run(
+            cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+        if result.returncode != 0:
+            print(f"Bootstrap failed while executing: {cmd}")
+            print("Stdout:", result.stdout.decode())
+            print("Stderr:", result.stderr.decode())
+            sys.exit(1)