metaflow 2.15.20__py2.py3-none-any.whl → 2.16.0__py2.py3-none-any.whl

metaflow/plugins/pypi/conda_decorator.py

@@ -1,5 +1,3 @@
-import importlib
-import json
 import os
 import platform
 import re
@@ -7,12 +5,9 @@ import sys
 import tempfile
 
 from metaflow.decorators import FlowDecorator, StepDecorator
-from metaflow.extension_support import EXT_PKG
 from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
-from metaflow.util import get_metaflow_root
-
-from ...info_file import INFO_FILE
+from metaflow.packaging_sys import ContentType
 
 
 class CondaStepDecorator(StepDecorator):
@@ -45,26 +40,31 @@ class CondaStepDecorator(StepDecorator):
         "python": None,
         "disabled": None,
     }
+
+    _metaflow_home = None
+    _addl_env_vars = None
+
     # To define conda channels for the whole solve, users can specify
     # CONDA_CHANNELS in their environment. For pinning specific packages to specific
     # conda channels, users can specify channel::package as the package name.
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super(CondaStepDecorator, self).__init__(attributes, statically_defined)
+        super(CondaStepDecorator, self).__init__(
+            attributes, statically_defined, inserted_by
+        )
 
     def init(self):
-        super(CondaStepDecorator, self).init()
-
         # Support legacy 'libraries=' attribute for the decorator.
         self.attributes["packages"] = {
             **self.attributes["libraries"],
             **self.attributes["packages"],
         }
-        del self.attributes["libraries"]
+        # Keep because otherwise make_decorator_spec will fail
+        self.attributes["libraries"] = {}
         if self.attributes["packages"]:
             self._attributes_with_user_values.add("packages")
 
@@ -152,67 +152,17 @@ class CondaStepDecorator(StepDecorator):
     def runtime_init(self, flow, graph, package, run_id):
         if self.disabled:
             return
-        # Create a symlink to metaflow installed outside the virtual environment.
-        self.metaflow_dir = tempfile.TemporaryDirectory(dir="/tmp")
-        os.symlink(
-            os.path.join(get_metaflow_root(), "metaflow"),
-            os.path.join(self.metaflow_dir.name, "metaflow"),
-        )
-
-        info = os.path.join(get_metaflow_root(), os.path.basename(INFO_FILE))
-        # Symlink the INFO file as well to properly propagate down the Metaflow version
-        if os.path.isfile(info):
-            os.symlink(
-                info, os.path.join(self.metaflow_dir.name, os.path.basename(INFO_FILE))
+        # We need to make all the code package available to the user code in
+        # a temporary directory which will be added to the PYTHONPATH.
+        if self.__class__._metaflow_home is None:
+            # Do this ONCE per flow
+            self.__class__._metaflow_home = tempfile.TemporaryDirectory(dir="/tmp")
+            package.extract_into(
+                self.__class__._metaflow_home.name, ContentType.ALL_CONTENT
+            )
+            self.__class__._addl_env_vars = package.get_post_extract_env_vars(
+                package.package_metadata, self.__class__._metaflow_home.name
             )
-        else:
-            # If there is no info file, we will actually create one in this new
-            # place because we won't be able to properly resolve the EXT_PKG extensions
-            # the same way as outside conda (looking at distributions, etc.). In a
-            # Conda environment, as shown below (where we set self.addl_paths), all
-            # EXT_PKG extensions are PYTHONPATH extensions. Instead of re-resolving,
-            # we use the resolved information that is written out to the INFO file.
-            with open(
-                os.path.join(self.metaflow_dir.name, os.path.basename(INFO_FILE)),
-                mode="wt",
-                encoding="utf-8",
-            ) as f:
-                f.write(
-                    json.dumps(
-                        self.environment.get_environment_info(include_ext_info=True)
-                    )
-                )
-
-        # Support metaflow extensions.
-        self.addl_paths = None
-        try:
-            m = importlib.import_module(EXT_PKG)
-        except ImportError:
-            # No additional check needed because if we are here, we already checked
-            # for other issues when loading at the toplevel.
-            pass
-        else:
-            custom_paths = list(set(m.__path__))
-            # For some reason, at times, unique paths appear multiple times. We
-            # simplify to avoid un-necessary links.
-
-            if len(custom_paths) == 1:
-                # Regular package; we take a quick shortcut here.
-                os.symlink(
-                    custom_paths[0],
-                    os.path.join(self.metaflow_dir.name, EXT_PKG),
-                )
-            else:
-                # This is a namespace package, we therefore create a bunch of
-                # directories so that we can symlink in those separately, and we will
-                # add those paths to the PYTHONPATH for the interpreter. Note that we
-                # don't symlink to the parent of the package because that could end up
-                # including more stuff we don't want
-                self.addl_paths = []
-                for p in custom_paths:
-                    temp_dir = tempfile.mkdtemp(dir=self.metaflow_dir.name)
-                    os.symlink(p, os.path.join(temp_dir, EXT_PKG))
-                    self.addl_paths.append(temp_dir)
 
         # # Also install any environment escape overrides directly here to enable
         # # the escape to work even in non metaflow-created subprocesses
@@ -291,11 +241,15 @@ class CondaStepDecorator(StepDecorator):
         if self.disabled:
             return
         # Ensure local installation of Metaflow is visible to user code
-        python_path = self.metaflow_dir.name
-        if self.addl_paths is not None:
-            addl_paths = os.pathsep.join(self.addl_paths)
-            python_path = os.pathsep.join([addl_paths, python_path])
-        cli_args.env["PYTHONPATH"] = python_path
+        python_path = self.__class__._metaflow_home.name
+        addl_env_vars = {}
+        if self.__class__._addl_env_vars is not None:
+            for key, value in self.__class__._addl_env_vars.items():
+                if key == "PYTHONPATH":
+                    addl_env_vars[key] = os.pathsep.join([value, python_path])
+                else:
+                    addl_env_vars[key] = value
+        cli_args.env.update(addl_env_vars)
         if self.interpreter:
             # https://github.com/conda/conda/issues/7707
             # Also ref - https://github.com/Netflix/metaflow/pull/178
@@ -306,7 +260,9 @@ class CondaStepDecorator(StepDecorator):
     def runtime_finished(self, exception):
         if self.disabled:
             return
-        self.metaflow_dir.cleanup()
+        if self.__class__._metaflow_home is not None:
+            self.__class__._metaflow_home.cleanup()
+            self.__class__._metaflow_home = None
 
 
 class CondaFlowDecorator(FlowDecorator):
@@ -339,22 +295,23 @@ class CondaFlowDecorator(FlowDecorator):
         "disabled": None,
     }
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super(CondaFlowDecorator, self).__init__(attributes, statically_defined)
+        super(CondaFlowDecorator, self).__init__(
+            attributes, statically_defined, inserted_by
+        )
 
     def init(self):
-        super(CondaFlowDecorator, self).init()
-
         # Support legacy 'libraries=' attribute for the decorator.
         self.attributes["packages"] = {
             **self.attributes["libraries"],
             **self.attributes["packages"],
         }
-        del self.attributes["libraries"]
+        # Keep because otherwise make_decorator_spec will fail
+        self.attributes["libraries"] = {}
         if self.attributes["python"]:
             self.attributes["python"] = str(self.attributes["python"])
 

metaflow/plugins/pypi/conda_environment.py

@@ -17,6 +17,7 @@ from metaflow.debug import debug
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import get_pinned_conda_libs
 from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.packaging_sys import ContentType
 
 from . import MAGIC_FILE, _datastore_packageroot
 from .utils import conda_platform
@@ -35,6 +36,7 @@ class CondaEnvironment(MetaflowEnvironment):
     _force_rebuild = False
 
     def __init__(self, flow):
+        super().__init__(flow)
         self.flow = flow
 
     def set_local_root(self, local_root):
@@ -335,7 +337,7 @@ class CondaEnvironment(MetaflowEnvironment):
                     environment[decorator.name] = {
                         k: copy.deepcopy(decorator.attributes[k])
                         for k in decorator.attributes
-                        if k != "disabled"
+                        if k not in ("disabled", "libraries")
                     }
                 else:
                     return {}
@@ -474,7 +476,9 @@ class CondaEnvironment(MetaflowEnvironment):
         files = []
         manifest = self.get_environment_manifest_path()
         if os.path.exists(manifest):
-            files.append((manifest, os.path.basename(manifest)))
+            files.append(
+                (manifest, os.path.basename(manifest), ContentType.OTHER_CONTENT)
+            )
         return files
 
     def bootstrap_commands(self, step_name, datastore_type):

metaflow/plugins/pypi/pypi_decorator.py

@@ -24,12 +24,12 @@ class PyPIStepDecorator(StepDecorator):
     name = "pypi"
     defaults = {"packages": {}, "python": None, "disabled": None}  # wheels
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super().__init__(attributes, statically_defined)
+        super().__init__(attributes, statically_defined, inserted_by)
 
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         # The init_environment hook for Environment creates the relevant virtual
@@ -128,12 +128,12 @@ class PyPIFlowDecorator(FlowDecorator):
     name = "pypi_base"
     defaults = {"packages": {}, "python": None, "disabled": None}
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super().__init__(attributes, statically_defined)
+        super().__init__(attributes, statically_defined, inserted_by)
 
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options

metaflow/plugins/test_unbounded_foreach_decorator.py

@@ -56,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(

metaflow/plugins/timeout_decorator.py

@@ -38,7 +38,6 @@ class TimeoutDecorator(StepDecorator):
     defaults = {"seconds": 0, "minutes": 0, "hours": 0}
 
     def init(self):
-        super().init()
         # 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

metaflow/plugins/uv/bootstrap.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import subprocess
 import sys
 import time
@@ -6,6 +7,7 @@ import time
 from metaflow.util import which
 from metaflow.info_file 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
 
@@ -93,6 +95,15 @@ if __name__ == "__main__":
         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(

metaflow/plugins/uv/uv_environment.py

@@ -2,6 +2,7 @@ import os
 
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.packaging_sys import ContentType
 
 
 class UVException(MetaflowException):
@@ -12,6 +13,7 @@ class UVEnvironment(MetaflowEnvironment):
     TYPE = "uv"
 
     def __init__(self, flow):
+        super().__init__(flow)
         self.flow = flow
 
     def validate_environment(self, logger, datastore_type):
@@ -43,8 +45,8 @@ class UVEnvironment(MetaflowEnvironment):
         pyproject_path = _find("pyproject.toml")
         uv_lock_path = _find("uv.lock")
         files = [
-            (uv_lock_path, "uv.lock"),
-            (pyproject_path, "pyproject.toml"),
+            (uv_lock_path, "uv.lock", ContentType.OTHER_CONTENT),
+            (pyproject_path, "pyproject.toml", ContentType.OTHER_CONTENT),
         ]
         return files
 

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:

metaflow/runner/click_api.py

@@ -46,7 +46,6 @@ from metaflow.exception import MetaflowException
 from metaflow.includefile import FilePathClass
 from metaflow.metaflow_config import CLICK_API_PROCESS_CONFIG
 from metaflow.parameters import JSONTypeClass, flow_context
-from metaflow.user_configs.config_decorators import CustomFlowDecorator
 from metaflow.user_configs.config_options import (
     ConfigValue,
     ConvertDictOrStr,
@@ -55,6 +54,7 @@ from metaflow.user_configs.config_options import (
     MultipleTuple,
     config_options_with_config_input,
 )
+from metaflow.user_decorators.user_flow_decorator import FlowMutator
 
 # Define a recursive type alias for JSON
 JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None]
@@ -264,12 +264,12 @@ def extract_flow_class_from_file(flow_file: str) -> FlowSpec:
             loaded_modules[flow_file] = module
 
         classes = inspect.getmembers(
-            module, lambda x: inspect.isclass(x) or isinstance(x, CustomFlowDecorator)
+            module, lambda x: inspect.isclass(x) or isinstance(x, FlowMutator)
         )
         flow_cls = None
 
         for _, kls in classes:
-            if isinstance(kls, CustomFlowDecorator):
+            if isinstance(kls, FlowMutator):
                 kls = kls._flow_cls
             if (
                 kls is not FlowSpec
@@ -512,7 +512,7 @@ class MetaflowAPI(object):
 
         # At this point, we are like in start() in cli.py -- we obtained the
         # properly processed config_options which we can now use to process
-        # the config decorators (including CustomStep/FlowDecorators)
+        # the config decorators (including StepMutator/FlowMutator)
         # Note that if CLICK_API_PROCESS_CONFIG is False, we still do this because
         # it will init all parameters (config_options will be None)
         # We ignore any errors if we don't check the configs in the click API.
@@ -658,6 +658,7 @@ if __name__ == "__main__":
         .kubernetes()
         .step(
             step_name="process",
+            code_package_metadata="some_version",
             code_package_sha="some_sha",
             code_package_url="some_url",
         )

metaflow/runner/subprocess_manager.py

@@ -9,6 +9,8 @@ import tempfile
 import threading
 from typing import Callable, Dict, Iterator, List, Optional, Tuple
 
+from metaflow.packaging_sys import MetaflowCodeContent
+from metaflow.util import get_metaflow_root
 from .utils import check_process_exited
 
 
@@ -150,7 +152,12 @@ class SubprocessManager(object):
         int
             The process ID of the subprocess.
         """
-
+        updated_env = MetaflowCodeContent.get_env_vars_for_packaged_metaflow(
+            get_metaflow_root()
+        )
+        if updated_env:
+            env = env or {}
+            env.update(updated_env)
         command_obj = CommandManager(command, env, cwd)
         pid = command_obj.run(show_output=show_output)
         self.commands[pid] = command_obj
@@ -181,7 +188,12 @@ class SubprocessManager(object):
         int
             The process ID of the subprocess.
         """
-
+        updated_env = MetaflowCodeContent.get_env_vars_for_packaged_metaflow(
+            get_metaflow_root()
+        )
+        if updated_env:
+            env = env or {}
+            env.update(updated_env)
         command_obj = CommandManager(command, env, cwd)
         pid = await command_obj.async_run()
         self.commands[pid] = command_obj

metaflow/runtime.py

@@ -16,6 +16,7 @@ import time
 import subprocess
 from datetime import datetime
 from io import BytesIO
+from itertools import chain
 from functools import partial
 from concurrent import futures
 
@@ -24,7 +25,7 @@ from contextlib import contextmanager
 
 from . import get_namespace
 from .metadata_provider import MetaDatum
-from .metaflow_config import MAX_ATTEMPTS, UI_URL
+from .metaflow_config import FEAT_ALWAYS_UPLOAD_CODE_PACKAGE, MAX_ATTEMPTS, UI_URL
 from .exception import (
     MetaflowException,
     MetaflowInternalError,
@@ -95,6 +96,7 @@ class NativeRuntime(object):
         max_num_splits=MAX_NUM_SPLITS,
         max_log_size=MAX_LOG_SIZE,
         resume_identifier=None,
+        skip_decorator_hooks=False,
     ):
         if run_id is None:
             self._run_id = metadata.new_run_id()
@@ -107,6 +109,7 @@ class NativeRuntime(object):
         self._flow_datastore = flow_datastore
         self._metadata = metadata
         self._environment = environment
+        self._package = package
         self._logger = logger
         self._max_workers = max_workers
         self._active_tasks = dict()  # Key: step name;
@@ -128,6 +131,7 @@ class NativeRuntime(object):
         self._ran_or_scheduled_task_index = set()
         self._reentrant = reentrant
         self._run_url = None
+        self._skip_decorator_hooks = skip_decorator_hooks
 
         # If steps_to_rerun is specified, we will not clone them in resume mode.
         self._steps_to_rerun = steps_to_rerun or {}
@@ -179,9 +183,10 @@ class NativeRuntime(object):
         # finished.
         self._control_num_splits = {}  # control_task -> num_splits mapping
 
-        for step in flow:
-            for deco in step.decorators:
-                deco.runtime_init(flow, graph, package, self._run_id)
+        if not self._skip_decorator_hooks:
+            for step in flow:
+                for deco in step.decorators:
+                    deco.runtime_init(flow, graph, package, self._run_id)
 
     def _new_task(self, step, input_paths=None, **kwargs):
         if input_paths is None:
@@ -192,7 +197,7 @@ class NativeRuntime(object):
         if step in self._steps_to_rerun:
             may_clone = False
 
-        if step == "_parameters":
+        if step == "_parameters" or self._skip_decorator_hooks:
             decos = []
         else:
             decos = getattr(self._flow, step).decorators
@@ -566,9 +571,10 @@ class NativeRuntime(object):
                 raise
             finally:
                 # on finish clean tasks
-                for step in self._flow:
-                    for deco in step.decorators:
-                        deco.runtime_finished(exception)
+                if not self._skip_decorator_hooks:
+                    for step in self._flow:
+                        for deco in step.decorators:
+                            deco.runtime_finished(exception)
                 self._run_exit_hooks()
 
         # assert that end was executed and it was successful
@@ -667,7 +673,6 @@ class NativeRuntime(object):
     # Given the current task information (task_index), the type of transition,
     # and the split index, return the new task index.
     def _translate_index(self, task, next_step, type, split_index=None):
-
         match = re.match(r"^(.+)\[(.*)\]$", task.task_index)
         if match:
             _, foreach_index = match.groups()
@@ -1005,6 +1010,22 @@ class NativeRuntime(object):
             # Initialize the task (which can be expensive using remote datastores)
             # before launching the worker so that cost is amortized over time, instead
             # of doing it during _queue_push.
+            if (
+                FEAT_ALWAYS_UPLOAD_CODE_PACKAGE
+                and "METAFLOW_CODE_SHA" not in os.environ
+            ):
+                # We check if the code package is uploaded and, if so, we set the
+                # environment variables that will cause the metadata service to
+                # register the code package with the task created in _new_task below
+                code_sha = self._package.package_sha(timeout=0.01)
+                if code_sha:
+                    os.environ["METAFLOW_CODE_SHA"] = code_sha
+                    os.environ["METAFLOW_CODE_URL"] = self._package.package_url()
+                    os.environ["METAFLOW_CODE_DS"] = self._flow_datastore.TYPE
+                    os.environ["METAFLOW_CODE_METADATA"] = (
+                        self._package.package_metadata
+                    )
+
             task = self._new_task(step, **task_kwargs)
             self._launch_worker(task)
 
@@ -1556,6 +1577,7 @@ class CLIArgs(object):
     def __init__(self, task):
         self.task = task
         self.entrypoint = list(task.entrypoint)
+        step_obj = getattr(self.task.flow, self.task.step)
         self.top_level_options = {
             "quiet": True,
             "metadata": self.task.metadata_type,
@@ -1567,8 +1589,12 @@ class CLIArgs(object):
             "datastore-root": self.task.datastore_sysroot,
             "with": [
                 deco.make_decorator_spec()
-                for deco in self.task.decos
-                if not deco.statically_defined
+                for deco in chain(
+                    self.task.decos,
+                    step_obj.wrappers,
+                    step_obj.config_decorators,
+                )
+                if not deco.statically_defined and deco.inserted_by is None
             ],
         }