ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/secrets/secrets_spec.py

@@ -0,0 +1,101 @@
+from metaflow.exception import MetaflowException
+from metaflow.plugins.secrets.utils import get_default_secrets_backend_type
+
+
+class SecretSpec:
+    def __init__(self, secrets_backend_type, secret_id, options={}, role=None):
+        self._secrets_backend_type = secrets_backend_type
+        self._secret_id = secret_id
+        self._options = options
+        self._role = role
+
+    @property
+    def secrets_backend_type(self):
+        return self._secrets_backend_type
+
+    @property
+    def secret_id(self):
+        return self._secret_id
+
+    @property
+    def options(self):
+        return self._options
+
+    @property
+    def role(self):
+        return self._role
+
+    def to_json(self):
+        """Mainly used for testing... not the same as the input dict in secret_spec_from_dict()!"""
+        return {
+            "secrets_backend_type": self.secrets_backend_type,
+            "secret_id": self.secret_id,
+            "options": self.options,
+            "role": self.role,
+        }
+
+    def __str__(self):
+        return "%s (%s)" % (self._secret_id, self._secrets_backend_type)
+
+    @staticmethod
+    def secret_spec_from_str(secret_spec_str, role):
+        # "." may be used in secret_id one day (provider specific). HOWEVER, it provides the best UX for
+        # non-conflicting cases (i.e. for secret ids that don't contain "."). This is true for all AWS
+        # Secrets Manager secrets.
+        #
+        # So we skew heavily optimize for best upfront UX for the present (1/2023).
+        #
+        # If/when a certain secret backend supports "." secret names, we can figure out a solution at that time.
+        # At a minimum, dictionary style secret spec may be used with no code changes (see secret_spec_from_dict()).
+        # Other options could be:
+        #  - accept and document that "." secret_ids don't work in Metaflow (across all possible providers)
+        #  - add a Metaflow config variable that specifies the separator (default ".")
+        #  - smarter spec parsing, that errors on secrets that look ambiguous. "aws-secrets-manager.XYZ" could mean:
+        #    + secret_id "XYZ" in aws-secrets-manager backend, OR
+        #    + secret_id "aws-secrets-manager.XYZ" default backend (if it is defined).
+        #    + in this case, user can simply set "azure-key-vault.aws-secrets-manager.XYZ" instead!
+        parts = secret_spec_str.split(".", maxsplit=1)
+        if len(parts) == 1:
+            secrets_backend_type = get_default_secrets_backend_type()
+            secret_id = parts[0]
+        else:
+            secrets_backend_type = parts[0]
+            secret_id = parts[1]
+        return SecretSpec(
+            secrets_backend_type, secret_id=secret_id, options={}, role=role
+        )
+
+    @staticmethod
+    def secret_spec_from_dict(secret_spec_dict, role):
+        if "type" not in secret_spec_dict:
+            secrets_backend_type = get_default_secrets_backend_type()
+        else:
+            secrets_backend_type = secret_spec_dict["type"]
+            if not isinstance(secrets_backend_type, str):
+                raise MetaflowException(
+                    "Bad @secrets specification - 'type' must be a string - found %s"
+                    % type(secrets_backend_type)
+                )
+        secret_id = secret_spec_dict.get("id")
+        if not isinstance(secret_id, str):
+            raise MetaflowException(
+                "Bad @secrets specification - 'id' must be a string - found %s"
+                % type(secret_id)
+            )
+        options = secret_spec_dict.get("options", {})
+        if not isinstance(options, dict):
+            raise MetaflowException(
+                "Bad @secrets specification - 'option' must be a dict - found %s"
+                % type(options)
+            )
+        role_for_source = secret_spec_dict.get("role", None)
+        if role_for_source is not None:
+            if not isinstance(role_for_source, str):
+                raise MetaflowException(
+                    "Bad @secrets specification - 'role' must be a str - found %s"
+                    % type(role_for_source)
+                )
+            role = role_for_source
+        return SecretSpec(
+            secrets_backend_type, secret_id=secret_id, options=options, role=role
+        )

metaflow/plugins/secrets/utils.py

@@ -0,0 +1,74 @@
+import os
+import re
+from metaflow.exception import MetaflowException
+
+DISALLOWED_SECRETS_ENV_VAR_PREFIXES = ["METAFLOW_"]
+
+
+def get_default_secrets_backend_type():
+    from metaflow.metaflow_config import DEFAULT_SECRETS_BACKEND_TYPE
+
+    if DEFAULT_SECRETS_BACKEND_TYPE is None:
+        raise MetaflowException(
+            "No default secrets backend type configured, but needed by @secrets. "
+            "Set METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE."
+        )
+    return DEFAULT_SECRETS_BACKEND_TYPE
+
+
+def validate_env_vars_across_secrets(all_secrets_env_vars):
+    vars_injected_by = {}
+    for secret_spec, env_vars in all_secrets_env_vars:
+        for k in env_vars:
+            if k in vars_injected_by:
+                raise MetaflowException(
+                    "Secret '%s' will inject '%s' as env var, and it is also added by '%s'"
+                    % (secret_spec, k, vars_injected_by[k])
+                )
+            vars_injected_by[k] = secret_spec
+
+
+def validate_env_vars_vs_existing_env(all_secrets_env_vars):
+    for secret_spec, env_vars in all_secrets_env_vars:
+        for k in env_vars:
+            if k in os.environ:
+                raise MetaflowException(
+                    "Secret '%s' will inject '%s' as env var, but it already exists in env"
+                    % (secret_spec, k)
+                )
+
+
+def validate_env_vars(env_vars):
+    for k, v in env_vars.items():
+        if not isinstance(k, str):
+            raise MetaflowException("Found non string key %s (%s)" % (str(k), type(k)))
+        if not isinstance(v, str):
+            raise MetaflowException(
+                "Found non string value %s (%s)" % (str(v), type(v))
+            )
+        if not re.fullmatch("[a-zA-Z_][a-zA-Z0-9_]*", k):
+            raise MetaflowException("Found invalid env var name '%s'." % k)
+        for disallowed_prefix in DISALLOWED_SECRETS_ENV_VAR_PREFIXES:
+            if k.startswith(disallowed_prefix):
+                raise MetaflowException(
+                    "Found disallowed env var name '%s' (starts with '%s')."
+                    % (k, disallowed_prefix)
+                )
+
+
+def get_secrets_backend_provider(secrets_backend_type):
+    from metaflow.plugins import SECRETS_PROVIDERS
+
+    try:
+        provider_cls = [
+            pc for pc in SECRETS_PROVIDERS if pc.TYPE == secrets_backend_type
+        ][0]
+        return provider_cls()
+    except IndexError:
+        raise MetaflowException(
+            "Unknown secrets backend type %s (available types: %s)"
+            % (
+                secrets_backend_type,
+                ", ".join(pc.TYPE for pc in SECRETS_PROVIDERS if pc.TYPE != "inline"),
+            )
+        )

metaflow/plugins/tag_cli.py

@@ -225,10 +225,7 @@ def _get_client_run_obj(obj, run_id, user_namespace):
 
 
 def _set_current(obj):
-    current._set_env(
-        metadata_str="%s@%s"
-        % (obj.metadata.__class__.TYPE, obj.metadata.__class__.INFO)
-    )
+    current._set_env(metadata_str=obj.metadata.metadata_str())
 
 
 @click.group()
@@ -507,9 +504,9 @@ def tag_list(
 
     if not group_by_run and not group_by_tag:
         # We list all the runs that match to print them out if needed.
-        system_tags_by_some_grouping[
-            ",".join(pathspecs)
-        ] = system_tags_by_some_grouping.get("_", set())
+        system_tags_by_some_grouping[",".join(pathspecs)] = (
+            system_tags_by_some_grouping.get("_", set())
+        )
         all_tags_by_some_grouping[",".join(pathspecs)] = all_tags_by_some_grouping.get(
             "_", set()
         )

metaflow/plugins/test_unbounded_foreach_decorator.py

@@ -8,8 +8,14 @@ import sys
 from metaflow.cli_args import cli_args
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
-from metaflow.unbounded_foreach import UnboundedForeachInput, UBF_CONTROL, UBF_TASK
+from metaflow.unbounded_foreach import (
+    UnboundedForeachInput,
+    UBF_CONTROL,
+    UBF_TASK,
+    CONTROL_TASK_TAG,
+)
 from metaflow.util import to_unicode
+from metaflow.metadata_provider import MetaDatum
 
 
 class InternalTestUnboundedForeachInput(UnboundedForeachInput):
@@ -50,9 +56,9 @@ class InternalTestUnboundedForeachDecorator(StepDecorator):
     name = "unbounded_test_foreach_internal"
     results_dict = {}
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         super(InternalTestUnboundedForeachDecorator, self).__init__(
-            attributes, statically_defined
+            attributes, statically_defined, inserted_by
         )
 
     def step_init(
@@ -60,13 +66,42 @@ class InternalTestUnboundedForeachDecorator(StepDecorator):
     ):
         self.environment = environment
 
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        if ubf_context == UBF_CONTROL:
+            metadata.register_metadata(
+                run_id,
+                step_name,
+                task_id,
+                [
+                    MetaDatum(
+                        field="internal_task_type",
+                        value=CONTROL_TASK_TAG,
+                        type="internal_task_type",
+                        tags=["attempt_id:{0}".format(0)],
+                    )
+                ],
+            )
+            self.input_paths = [obj.pathspec for obj in inputs]
+
     def control_task_step_func(self, flow, graph, retry_count):
         from metaflow import current
 
         run_id = current.run_id
         step_name = current.step_name
         control_task_id = current.task_id
-        (_, split_step_name, split_task_id) = control_task_id.split("-")[1:]
         # If we are running inside Conda, we use the base executable FIRST;
         # the conda environment will then be used when runtime_step_cli is
         # called. This is so that it can properly set up all the metaflow
@@ -94,10 +129,10 @@ class InternalTestUnboundedForeachDecorator(StepDecorator):
         mapper_tasks = []
 
         for i in range(foreach_num_splits):
-            task_id = "%s-%d" % (control_task_id.replace("control-", "test-ubf-"), i)
+            task_id = "%s-%d" % (control_task_id, i)
             pathspec = "%s/%s/%s" % (run_id, step_name, task_id)
             mapper_tasks.append(to_unicode(pathspec))
-            input_paths = "%s/%s/%s" % (run_id, split_step_name, split_task_id)
+            input_paths = ",".join(self.input_paths)
 
             # Override specific `step` kwargs.
             kwargs = cli_args.step_kwargs

metaflow/plugins/timeout_decorator.py

@@ -4,6 +4,7 @@ import traceback
 from metaflow.exception import MetaflowException
 from metaflow.decorators import StepDecorator
 from metaflow.unbounded_foreach import UBF_CONTROL
+from metaflow.metaflow_config import DEFAULT_RUNTIME_LIMIT
 
 
 class TimeoutException(MetaflowException):
@@ -36,8 +37,7 @@ class TimeoutDecorator(StepDecorator):
     name = "timeout"
     defaults = {"seconds": 0, "minutes": 0, "hours": 0}
 
-    def __init__(self, *args, **kwargs):
-        super(TimeoutDecorator, self).__init__(*args, **kwargs)
+    def init(self):
         # Initialize secs in __init__ so other decorators could safely use this
         # value without worrying about decorator order.
         # Convert values in attributes to type:int since they can be type:str
@@ -99,7 +99,7 @@ class TimeoutDecorator(StepDecorator):
 
 
 def get_run_time_limit_for_task(step_decos):
-    run_time_limit = 5 * 24 * 60 * 60  # 5 days.
+    run_time_limit = DEFAULT_RUNTIME_LIMIT
     for deco in step_decos:
         if isinstance(deco, TimeoutDecorator):
             run_time_limit = deco.secs

metaflow/plugins/uv/bootstrap.py

@@ -0,0 +1,128 @@
+import os
+import shutil
+import subprocess
+import sys
+import time
+
+from metaflow.util import which
+from metaflow.meta_files import read_info_file
+from metaflow.metaflow_config import get_pinned_conda_libs
+from metaflow.packaging_sys import MetaflowCodeContent, ContentType
+from urllib.request import Request, urlopen
+from urllib.error import URLError
+
+# TODO: support version/platform/architecture selection.
+UV_URL = "https://github.com/astral-sh/uv/releases/download/0.6.11/uv-x86_64-unknown-linux-gnu.tar.gz"
+
+if __name__ == "__main__":
+
+    def run_cmd(cmd, stdin_str=None):
+        result = subprocess.run(
+            cmd,
+            shell=True,
+            input=stdin_str,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        if result.returncode != 0:
+            print(f"Bootstrap failed while executing: {cmd}")
+            print("Stdout:", result.stdout)
+            print("Stderr:", result.stderr)
+            sys.exit(1)
+
+    def install_uv():
+        import tarfile
+
+        uv_install_path = os.path.join(os.getcwd(), "uv_install")
+        if which("uv"):
+            return
+
+        print("Installing uv...")
+
+        # Prepare directory once
+        os.makedirs(uv_install_path, exist_ok=True)
+
+        # Download and decompress in one go
+        headers = {
+            "Accept-Encoding": "gzip, deflate, br",
+            "Connection": "keep-alive",
+            "User-Agent": "python-urllib",
+        }
+
+        def _tar_filter(member: tarfile.TarInfo, path):
+            if os.path.basename(member.name) != "uv":
+                return None  # skip
+            member.path = os.path.basename(member.path)
+            return member
+
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                req = Request(UV_URL, headers=headers)
+                with urlopen(req) as response:
+                    with tarfile.open(fileobj=response, mode="r:gz") as tar:
+                        tar.extractall(uv_install_path, filter=_tar_filter)
+                break
+            except (URLError, IOError) as e:
+                if attempt == max_retries - 1:
+                    raise Exception(
+                        f"Failed to download UV after {max_retries} attempts: {e}"
+                    )
+                time.sleep(2**attempt)
+
+        # Update PATH only once at the end
+        os.environ["PATH"] += os.pathsep + uv_install_path
+
+    def get_dependencies(datastore_type):
+        # return required dependencies for Metaflow that must be added to the UV environment.
+        pinned = get_pinned_conda_libs(None, datastore_type)
+
+        # return only dependency names instead of pinned versions
+        return pinned.keys()
+
+    def skip_metaflow_dependencies():
+        skip_pkgs = ["metaflow", "ob-metaflow"]
+        info = read_info_file()
+        if info is not None:
+            try:
+                skip_pkgs.extend([ext_name for ext_name in info["ext_info"][0].keys()])
+            except Exception:
+                print(
+                    "Failed to read INFO. Metaflow-related packages might get installed during runtime."
+                )
+
+        return skip_pkgs
+
+    def sync_uv_project(datastore_type):
+        # Move the files to the current directory so uv can find them.
+        for filename in ["uv.lock", "pyproject.toml"]:
+            path_to_file = MetaflowCodeContent.get_filename(
+                filename, ContentType.OTHER_CONTENT
+            )
+            if path_to_file is None:
+                raise RuntimeError(f"Could not find {filename} in the package.")
+            shutil.move(path_to_file, os.path.join(os.getcwd(), filename))
+
+        print("Syncing uv project...")
+        dependencies = " ".join(get_dependencies(datastore_type))
+        skip_pkgs = " ".join(
+            [f"--no-install-package {dep}" for dep in skip_metaflow_dependencies()]
+        )
+        cmd = f"""set -e;
+            uv sync --frozen {skip_pkgs};
+            uv pip install {dependencies} --strict
+            """
+        run_cmd(cmd)
+
+    if len(sys.argv) != 2:
+        print("Usage: bootstrap.py <datastore_type>")
+        sys.exit(1)
+
+    try:
+        datastore_type = sys.argv[1]
+        install_uv()
+        sync_uv_project(datastore_type)
+    except Exception as e:
+        print(f"Error: {str(e)}", file=sys.stderr)
+        sys.exit(1)

metaflow/plugins/uv/uv_environment.py

@@ -0,0 +1,72 @@
+import os
+
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.packaging_sys import ContentType
+
+
+class UVException(MetaflowException):
+    headline = "uv error"
+
+
+class UVEnvironment(MetaflowEnvironment):
+    TYPE = "uv"
+
+    def __init__(self, flow):
+        super().__init__(flow)
+        self.flow = flow
+
+    def validate_environment(self, logger, datastore_type):
+        self.datastore_type = datastore_type
+        self.logger = logger
+
+    def init_environment(self, echo, only_steps=None):
+        self.logger("Bootstrapping uv...")
+
+    def executable(self, step_name, default=None):
+        return "uv run --no-sync python"
+
+    def add_to_package(self):
+        # NOTE: We treat uv.lock and pyproject.toml as regular project assets and ship these along user code as part of the code package
+        # These are the minimal required files to reproduce the UV environment on the remote platform.
+        def _find(filename):
+            current_dir = os.getcwd()
+            while True:
+                file_path = os.path.join(current_dir, filename)
+                if os.path.isfile(file_path):
+                    return file_path
+                parent_dir = os.path.dirname(current_dir)
+                if parent_dir == current_dir:  # Reached root
+                    raise UVException(
+                        f"Could not find {filename} in current directory or any parent directory"
+                    )
+                current_dir = parent_dir
+
+        pyproject_path = _find("pyproject.toml")
+        uv_lock_path = _find("uv.lock")
+        files = [
+            (uv_lock_path, "uv.lock", ContentType.OTHER_CONTENT),
+            (pyproject_path, "pyproject.toml", ContentType.OTHER_CONTENT),
+        ]
+        return files
+
+    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):
+        return [
+            "echo 'Bootstrapping uv project...'",
+            "flush_mflogs",
+            # We have to prevent the tracing module from loading, as the bootstrapping process
+            # uses the internal S3 client which would fail to import tracing due to the required
+            # dependencies being bundled into the conda environment, which is yet to be
+            # initialized at this point.
+            'DISABLE_TRACING=True python -m metaflow.plugins.uv.bootstrap "%s"'
+            % datastore_type,
+            "echo 'uv project bootstrapped.'",
+            "flush_mflogs",
+            "export PATH=$PATH:$(pwd)/uv_install",
+        ]

metaflow/procpoll.py

@@ -31,7 +31,7 @@ class LinuxProcPoll(ProcPoll):
         self._poll.unregister(fd)
 
     def poll(self, timeout):
-        for (fd, event) in self._poll.poll(timeout):
+        for fd, event in self._poll.poll(timeout):
             yield ProcPollEvent(
                 fd=fd,
                 can_read=bool(event & select.POLLIN),

metaflow/pylint_wrapper.py

@@ -28,7 +28,11 @@ class PyLint(object):
         return self._run is not None
 
     def run(self, logger=None, warnings=False, pylint_config=[]):
-        args = [self._fname]
+        args = [
+            self._fname,
+            "--signature-mutators",
+            "metaflow.user_decorators.user_step_decorator.user_step_decorator",
+        ]
         if not warnings:
             args.append("--errors-only")
         if pylint_config: