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

metaflow/plugins/aws/batch/batch.py

@@ -59,14 +59,24 @@ class Batch(object):
         self._client = BatchClient()
         atexit.register(lambda: self.job.kill() if hasattr(self, "job") else None)
 
-    def _command(self, environment, code_package_url, step_name, step_cmds, task_spec):
+    def _command(
+        self,
+        environment,
+        code_package_metadata,
+        code_package_url,
+        step_name,
+        step_cmds,
+        task_spec,
+    ):
         mflog_expr = export_mflog_env_vars(
             datastore_type="s3",
             stdout_path=STDOUT_PATH,
             stderr_path=STDERR_PATH,
             **task_spec
         )
-        init_cmds = environment.get_package_commands(code_package_url, "s3")
+        init_cmds = environment.get_package_commands(
+            code_package_url, "s3", code_package_metadata
+        )
         init_expr = " && ".join(init_cmds)
         step_expr = bash_capture_logs(
             " && ".join(environment.bootstrap_commands(step_name, "s3") + step_cmds)
@@ -167,6 +177,7 @@ class Batch(object):
         step_name,
         step_cli,
         task_spec,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -210,7 +221,12 @@ class Batch(object):
             .job_queue(queue)
             .command(
                 self._command(
-                    self.environment, code_package_url, step_name, [step_cli], task_spec
+                    self.environment,
+                    code_package_metadata,
+                    code_package_url,
+                    step_name,
+                    [step_cli],
+                    task_spec,
                 )
             )
             .image(image)
@@ -249,6 +265,7 @@ class Batch(object):
             )
             .task_id(attrs.get("metaflow.task_id"))
             .environment_variable("AWS_DEFAULT_REGION", self._client.region())
+            .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
             .environment_variable("METAFLOW_CODE_DS", code_package_ds)
@@ -334,6 +351,7 @@ class Batch(object):
         step_name,
         step_cli,
         task_spec,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -374,6 +392,7 @@ class Batch(object):
             step_name,
             step_cli,
             task_spec,
+            code_package_metadata,
             code_package_sha,
             code_package_url,
             code_package_ds,

metaflow/plugins/aws/batch/batch_cli.py

@@ -100,6 +100,7 @@ def kill(ctx, run_id, user, my_runs):
     "Metaflow."
 )
 @click.argument("step-name")
+@click.argument("code-package-metadata")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
 @click.option("--executable", help="Executable requirement for AWS Batch.")
@@ -185,6 +186,7 @@ def kill(ctx, run_id, user, my_runs):
 def step(
     ctx,
     step_name,
+    code_package_metadata,
     code_package_sha,
     code_package_url,
     executable=None,
@@ -317,6 +319,7 @@ def step(
                 step_name,
                 step_cli,
                 task_spec,
+                code_package_metadata,
                 code_package_sha,
                 code_package_url,
                 ctx.obj.flow_datastore.TYPE,

metaflow/plugins/aws/batch/batch_decorator.py

@@ -14,6 +14,7 @@ from metaflow.metaflow_config import (
     DATASTORE_LOCAL_DIR,
     ECS_FARGATE_EXECUTION_ROLE,
     ECS_S3_ACCESS_IAM_ROLE,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
 )
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.sidecar import Sidecar
@@ -126,6 +127,7 @@ class BatchDecorator(StepDecorator):
         "gpu": "0",
         "memory": "4096",
     }
+    package_metadata = None
     package_url = None
     package_sha = None
     run_time_limit = None
@@ -135,8 +137,6 @@ class BatchDecorator(StepDecorator):
     target_platform = "linux-64"
 
     def init(self):
-        super(BatchDecorator, self).init()
-
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
             # If metaflow-config specifies a docker image, just use that.
@@ -228,6 +228,7 @@ class BatchDecorator(StepDecorator):
             # to execute on AWS Batch anymore. We can execute possible fallback
             # code locally.
             cli_args.commands = ["batch", "step"]
+            cli_args.command_args.append(self.package_metadata)
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
             cli_args.command_options.update(self.attributes)
@@ -403,9 +404,16 @@ class BatchDecorator(StepDecorator):
     @classmethod
     def _save_package_once(cls, flow_datastore, package):
         if cls.package_url is None:
-            cls.package_url, cls.package_sha = flow_datastore.save_data(
-                [package.blob], len_hint=1
-            )[0]
+            if not FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+                cls.package_url, cls.package_sha = flow_datastore.save_data(
+                    [package.blob], len_hint=1
+                )[0]
+                cls.package_metadata = package.package_metadata
+            else:
+                # Blocks until the package is uploaded
+                cls.package_url = package.package_url()
+                cls.package_sha = package.package_sha()
+                cls.package_metadata = package.package_metadata
 
 
 def _setup_multinode_environment():

metaflow/plugins/aws/step_functions/step_functions.py

@@ -40,6 +40,7 @@ class StepFunctions(object):
         name,
         graph,
         flow,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         production_token,
@@ -59,6 +60,7 @@ class StepFunctions(object):
         self.name = name
         self.graph = graph
         self.flow = flow
+        self.code_package_metadata = code_package_metadata
         self.code_package_sha = code_package_sha
         self.code_package_url = code_package_url
         self.production_token = production_token
@@ -853,6 +855,7 @@ class StepFunctions(object):
                     node, input_paths, self.code_package_url, user_code_retries
                 ),
                 task_spec=task_spec,
+                code_package_metadata=self.code_package_metadata,
                 code_package_sha=self.code_package_sha,
                 code_package_url=self.code_package_url,
                 code_package_ds=self.flow_datastore.TYPE,
@@ -913,7 +916,7 @@ class StepFunctions(object):
             "with": [
                 decorator.make_decorator_spec()
                 for decorator in node.decorators
-                if not decorator.statically_defined
+                if not decorator.statically_defined and decorator.inserted_by is None
             ]
         }
         # FlowDecorators can define their own top-level options. They are

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -7,6 +7,7 @@ from metaflow import JSONType, current, decorators, parameters
 from metaflow._vendor import click
 from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.metaflow_config import (
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     SERVICE_VERSION_CHECK,
     SFN_STATE_MACHINE_PREFIX,
     UI_URL,
@@ -331,16 +332,26 @@ def make_flow(
     )
 
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=obj.flow_datastore if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE else None,
     )
-    package_url, package_sha = obj.flow_datastore.save_data(
-        [obj.package.blob], len_hint=1
-    )[0]
+    # This blocks until the package is created
+    if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+        package_url = obj.package.package_url()
+        package_sha = obj.package.package_sha()
+    else:
+        package_url, package_sha = obj.flow_datastore.save_data(
+            [obj.package.blob], len_hint=1
+        )[0]
 
     return StepFunctions(
         name,
         obj.graph,
         obj.flow,
+        obj.package.package_metadata,
         package_sha,
         package_url,
         token,

metaflow/plugins/cards/card_decorator.py

@@ -361,11 +361,6 @@ class CardDecorator(StepDecorator):
 
         return list(self._options(top_level_options))
 
-    def task_exception(
-        self, exception, step_name, flow, graph, retry_count, max_user_code_retries
-    ):
-        self._cleanup(step_name)
-
     def _cleanup(self, step_name):
         self._increment_completed_counter()
         if self.task_finished_decos == self.total_decos_on_step[step_name]:

metaflow/plugins/kubernetes/kubernetes.py

@@ -90,6 +90,7 @@ class Kubernetes(object):
         step_name,
         task_id,
         attempt,
+        code_package_metadata,
         code_package_url,
         step_cmds,
     ):
@@ -104,7 +105,7 @@ class Kubernetes(object):
             stderr_path=STDERR_PATH,
         )
         init_cmds = self._environment.get_package_commands(
-            code_package_url, self._datastore.TYPE
+            code_package_url, self._datastore.TYPE, code_package_metadata
         )
         init_expr = " && ".join(init_cmds)
         step_expr = bash_capture_logs(
@@ -165,6 +166,7 @@ class Kubernetes(object):
         task_id,
         attempt,
         user,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -232,6 +234,7 @@ class Kubernetes(object):
                 qos=qos,
                 security_context=security_context,
             )
+            .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
             .environment_variable("METAFLOW_CODE_DS", code_package_ds)
@@ -415,6 +418,7 @@ class Kubernetes(object):
             step_name=step_name,
             task_id=_tskid,
             attempt=attempt,
+            code_package_metadata=code_package_metadata,
             code_package_url=code_package_url,
             step_cmds=[
                 step_cli.replace(
@@ -463,6 +467,7 @@ class Kubernetes(object):
         task_id,
         attempt,
         user,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -511,6 +516,7 @@ class Kubernetes(object):
                     step_name=step_name,
                     task_id=task_id,
                     attempt=attempt,
+                    code_package_metadata=code_package_metadata,
                     code_package_url=code_package_url,
                     step_cmds=[step_cli],
                 ),
@@ -539,6 +545,7 @@ class Kubernetes(object):
                 qos=qos,
                 security_context=security_context,
             )
+            .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
             .environment_variable("METAFLOW_CODE_DS", code_package_ds)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -41,6 +41,7 @@ def kubernetes():
 )
 @tracing.cli("kubernetes/step")
 @click.argument("step-name")
+@click.argument("code-package-metadata")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
 @click.option(
@@ -161,6 +162,7 @@ def kubernetes():
 def step(
     ctx,
     step_name,
+    code_package_metadata,
     code_package_sha,
     code_package_url,
     executable=None,
@@ -304,6 +306,7 @@ def step(
                 task_id=task_id,
                 attempt=str(retry_count),
                 user=util.get_username(),
+                code_package_metadata=code_package_metadata,
                 code_package_sha=code_package_sha,
                 code_package_url=code_package_url,
                 code_package_ds=ctx.obj.flow_datastore.TYPE,

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -11,6 +11,7 @@ from metaflow.metadata_provider import MetaDatum
 from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
     DATASTORE_LOCAL_DIR,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     KUBERNETES_CONTAINER_IMAGE,
     KUBERNETES_CONTAINER_REGISTRY,
     KUBERNETES_CPU,
@@ -168,6 +169,7 @@ class KubernetesDecorator(StepDecorator):
         "qos": KUBERNETES_QOS,
         "security_context": None,
     }
+    package_metadata = None
     package_url = None
     package_sha = None
     run_time_limit = None
@@ -177,8 +179,6 @@ class KubernetesDecorator(StepDecorator):
     target_platform = KUBERNETES_CONDA_ARCH or "linux-64"
 
     def init(self):
-        super(KubernetesDecorator, self).init()
-
         if not self.attributes["namespace"]:
             self.attributes["namespace"] = KUBERNETES_NAMESPACE
         if not self.attributes["service_account"]:
@@ -471,6 +471,7 @@ class KubernetesDecorator(StepDecorator):
             # to execute on Kubernetes anymore. We can execute possible fallback
             # code locally.
             cli_args.commands = ["kubernetes", "step"]
+            cli_args.command_args.append(self.package_metadata)
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
 
@@ -642,9 +643,16 @@ class KubernetesDecorator(StepDecorator):
     @classmethod
     def _save_package_once(cls, flow_datastore, package):
         if cls.package_url is None:
-            cls.package_url, cls.package_sha = flow_datastore.save_data(
-                [package.blob], len_hint=1
-            )[0]
+            if not FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+                cls.package_url, cls.package_sha = flow_datastore.save_data(
+                    [package.blob], len_hint=1
+                )[0]
+                cls.package_metadata = package.package_metadata
+            else:
+                # Blocks until the package is uploaded
+                cls.package_url = package.package_url()
+                cls.package_sha = package.package_sha()
+                cls.package_metadata = package.package_metadata
 
 
 # TODO: Unify this method with the multi-node setup in @batch

metaflow/plugins/package_cli.py

@@ -9,35 +9,30 @@ def cli():
 
 
 @cli.group(help="Commands related to code packages.")
+@click.option(
+    "--timeout", default=60, help="Timeout for package operations in seconds."
+)
 @click.pass_obj
-def package(obj):
+def package(obj, timeout):
     # Prepare the package before any of the sub-commands are invoked.
+    # We explicitly will *not* upload it to the datastore.
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=None,
     )
+    obj.package_op_timeout = timeout
 
 
-@package.command(help="Output information about the current code package.")
+@package.command(help="Output information about the code package.")
 @click.pass_obj
 def info(obj):
-    obj.echo("Status of the current working directory:", fg="magenta", bold=False)
-    obj.echo_always(
-        "Hash: *%s*" % sha1(obj.package.blob).hexdigest(),
-        highlight="green",
-        highlight_bold=False,
-    )
-    obj.echo_always(
-        "Package size: *%d* KB" % (len(obj.package.blob) / 1024),
-        highlight="green",
-        highlight_bold=False,
-    )
-    num = sum(1 for _ in obj.package.path_tuples())
-    obj.echo_always(
-        "Number of files: *%d*" % num, highlight="green", highlight_bold=False
-    )
+    obj.echo_always(obj.package.show())
 
 
-@package.command(help="List files included in the code package.")
+@package.command(help="List all files included in the code package.")
 @click.option(
     "--archive/--no-archive",
     default=False,
@@ -47,8 +42,10 @@ def info(obj):
 )
 @click.pass_obj
 def list(obj, archive=False):
+    _ = obj.package.blob_with_timeout(timeout=obj.package_op_timeout)
+    # We now have all the information about the blob
     obj.echo(
-        "Files included in the code package " "(change with --package-suffixes):",
+        "Files included in the code package (change with --package-suffixes):",
         fg="magenta",
         bold=False,
     )
@@ -58,10 +55,15 @@ def list(obj, archive=False):
         obj.echo_always("\n".join(path for path, _ in obj.package.path_tuples()))
 
 
-@package.command(help="Save the current code package in a tar file")
+@package.command(help="Save the current code package to a file.")
 @click.argument("path")
 @click.pass_obj
 def save(obj, path):
     with open(path, "wb") as f:
-        f.write(obj.package.blob)
-    obj.echo("Code package saved in *%s*." % path, fg="magenta", bold=False)
+        f.write(obj.package.blob())
+    obj.echo(
+        "Code package saved in *%s* with metadata: %s"
+        % (path, obj.package.package_metadata),
+        fg="magenta",
+        bold=False,
+    )

metaflow/plugins/parallel_decorator.py

@@ -36,8 +36,10 @@ class ParallelDecorator(StepDecorator):
     defaults = {}
     IS_PARALLEL = True
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(ParallelDecorator, self).__init__(attributes, statically_defined)
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
+        super(ParallelDecorator, self).__init__(
+            attributes, statically_defined, inserted_by
+        )
 
     def runtime_step_cli(
         self, cli_args, retry_count, max_user_code_retries, ubf_context

metaflow/plugins/pypi/bootstrap.py

@@ -12,9 +12,10 @@ import platform
 from urllib.error import URLError
 from urllib.request import urlopen
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, CONDA_USE_FAST_INIT
+from metaflow.packaging_sys import MetaflowCodeContent, ContentType
 from metaflow.plugins import DATASTORES
 from metaflow.plugins.pypi.utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL
-from metaflow.util import which, get_metaflow_root
+from metaflow.util import which
 from urllib.request import Request
 import warnings
 
@@ -365,8 +366,13 @@ if __name__ == "__main__":
 
         # Move MAGIC_FILE inside local datastore.
         os.makedirs(manifest_dir, exist_ok=True)
+        path_to_manifest = MetaflowCodeContent.get_filename(
+            MAGIC_FILE, ContentType.OTHER_CONTENT
+        )
+        if path_to_manifest is None:
+            raise RuntimeError(f"Cannot find {MAGIC_FILE} in the package")
         shutil.move(
-            os.path.join(get_metaflow_root(), MAGIC_FILE),
+            path_to_manifest,
             os.path.join(manifest_dir, MAGIC_FILE),
         )
         with open(os.path.join(manifest_dir, MAGIC_FILE)) as f: