metaflow 2.11.2__py2.py3-none-any.whl → 2.11.3__py2.py3-none-any.whl

metaflow/flowspec.py

@@ -17,7 +17,7 @@ from .exception import (
 )
 from .graph import FlowGraph
 from .unbounded_foreach import UnboundedForeachInput
-from .metaflow_config import INCLUDE_FOREACH_STACK
+from .metaflow_config import INCLUDE_FOREACH_STACK, MAXIMUM_FOREACH_VALUE_CHARS
 
 # For Python 3 compatibility
 try:
@@ -28,6 +28,8 @@ except NameError:
 
 from .datastore.inputs import Inputs
 
+INTERNAL_ARTIFACTS_SET = set(["_foreach_values"])
+
 
 class InvalidNextException(MetaflowException):
     headline = "Invalid self.next() transition detected"
@@ -446,7 +448,9 @@ class FlowSpec(object):
                 available_vars = (
                     (var, sha)
                     for var, sha in inp._datastore.items()
-                    if (var not in exclude) and (not hasattr(self, var))
+                    if (var not in exclude)
+                    and (not hasattr(self, var))
+                    and (var not in INTERNAL_ARTIFACTS_SET)
                 )
             for var, sha in available_vars:
                 _, previous_sha = to_merge.setdefault(var, (inp, sha))
@@ -529,7 +533,7 @@ class FlowSpec(object):
             )
 
         value = item if _is_primitive_type(item) else reprlib.Repr().repr(item)
-        return basestring(value)
+        return basestring(value)[:MAXIMUM_FOREACH_VALUE_CHARS]
 
     def next(self, *dsts: Callable[..., None], **kwargs) -> None:
         """

metaflow/metaflow_config.py

@@ -205,6 +205,8 @@ DEFAULT_CONTAINER_REGISTRY = from_conf("DEFAULT_CONTAINER_REGISTRY")
 # Controls whether to include foreach stack information in metadata.
 # TODO(Darin, 05/01/24): Remove this flag once we are confident with this feature.
 INCLUDE_FOREACH_STACK = from_conf("INCLUDE_FOREACH_STACK", False)
+# Maximum length of the foreach value string to be stored in each ForeachFrame.
+MAXIMUM_FOREACH_VALUE_CHARS = from_conf("MAXIMUM_FOREACH_VALUE_CHARS", 30)
 
 ###
 # Organization customizations
@@ -268,7 +270,13 @@ SFN_STATE_MACHINE_PREFIX = from_conf("SFN_STATE_MACHINE_PREFIX")
 # machine execution logs. This needs to be available when using the
 # `step-functions create --log-execution-history` command.
 SFN_EXECUTION_LOG_GROUP_ARN = from_conf("SFN_EXECUTION_LOG_GROUP_ARN")
-
+# Amazon S3 path for storing the results of AWS Step Functions Distributed Map
+SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH = from_conf(
+    "SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH",
+    os.path.join(DATASTORE_SYSROOT_S3, "sfn_distributed_map_output")
+    if DATASTORE_SYSROOT_S3
+    else None,
+)
 ###
 # Kubernetes configuration
 ###

metaflow/plugins/argo/argo_workflows.py

@@ -1540,7 +1540,6 @@ class ArgoWorkflows(object):
 
     # Return exit hook templates for workflow execution notifications.
     def _exit_hook_templates(self):
-        # TODO: Add details to slack message
         templates = []
         if self.notify_on_error:
             templates.append(self._slack_error_template())
@@ -1649,36 +1648,100 @@ class ArgoWorkflows(object):
 
         return links
 
+    def _get_slack_blocks(self, message):
+        """
+        Use Slack's Block Kit to add general information about the environment and
+        execution metadata, including a link to the UI and an optional message.
+        """
+        ui_link = "%s%s/argo-{{workflow.name}}" % (UI_URL, self.flow.name)
+        # fmt: off
+        if getattr(current, "project_name", None):
+            # Add @project metadata when available.
+            environment_details_block = {
+                "type": "section",
+                "text": {
+                    "type": "mrkdwn",
+                    "text": ":metaflow: Environment details"
+                },
+                "fields": [
+                    {
+                        "type": "mrkdwn",
+                        "text": "*Project:* %s" % current.project_name 
+                    },
+                    {
+                        "type": "mrkdwn",
+                        "text": "*Project Branch:* %s" % current.branch_name
+                    }
+                ]
+            }
+        else:
+            environment_details_block = {
+                "type": "section",
+                "text": {
+                    "type": "mrkdwn",
+                    "text": ":metaflow: Environment details"
+                }
+            }
+
+        blocks = [
+            environment_details_block,
+            {
+                "type": "context",
+                "elements": [
+                    {
+                        "type": "mrkdwn",
+                        "text": " :information_source: *<%s>*" % ui_link,
+                    }
+                ],
+            },
+            {
+                "type": "divider"
+            },
+        ]
+
+        if message:
+            blocks += [
+                {
+                    "type": "section",
+                    "text": {
+                        "type": "mrkdwn",
+                        "text": message
+                    }
+                }
+            ]
+        # fmt: on
+        return blocks
+
     def _slack_error_template(self):
         if self.notify_slack_webhook_url is None:
             return None
+
+        message = (
+            ":rotating_light: _%s/argo-{{workflow.name}}_ failed!" % self.flow.name
+        )
+        payload = {"text": message}
+        if UI_URL:
+            blocks = self._get_slack_blocks(message)
+            payload = {"text": message, "blocks": blocks}
+
         return Template("notify-slack-on-error").http(
-            Http("POST")
-            .url(self.notify_slack_webhook_url)
-            .body(
-                json.dumps(
-                    {
-                        "text": ":rotating_light: _%s/argo-{{workflow.name}}_ failed!"
-                        % self.flow.name
-                    }
-                )
-            )
+            Http("POST").url(self.notify_slack_webhook_url).body(json.dumps(payload))
         )
 
     def _slack_success_template(self):
         if self.notify_slack_webhook_url is None:
             return None
+
+        message = (
+            ":white_check_mark: _%s/argo-{{workflow.name}}_ succeeded!" % self.flow.name
+        )
+        payload = {"text": message}
+        if UI_URL:
+            blocks = self._get_slack_blocks(message)
+            payload = {"text": message, "blocks": blocks}
+
         return Template("notify-slack-on-success").http(
-            Http("POST")
-            .url(self.notify_slack_webhook_url)
-            .body(
-                json.dumps(
-                    {
-                        "text": ":white_check_mark: _%s/argo-{{workflow.name}}_ succeeded!"
-                        % self.flow.name
-                    }
-                )
-            )
+            Http("POST").url(self.notify_slack_webhook_url).body(json.dumps(payload))
         )
 
     def _compile_sensor(self):

metaflow/plugins/aws/batch/batch.py

@@ -11,6 +11,7 @@ from metaflow.plugins.datatools.s3.s3tail import S3Tail
 from metaflow.plugins.aws.aws_utils import sanitize_batch_tag
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import (
+    OTEL_ENDPOINT,
     SERVICE_INTERNAL_URL,
     DATATOOLS_S3ROOT,
     DATASTORE_SYSROOT_S3,
@@ -255,6 +256,7 @@ class Batch(object):
             .environment_variable("METAFLOW_DEFAULT_DATASTORE", "s3")
             .environment_variable("METAFLOW_DEFAULT_METADATA", DEFAULT_METADATA)
             .environment_variable("METAFLOW_CARD_S3ROOT", CARD_S3ROOT)
+            .environment_variable("METAFLOW_OTEL_ENDPOINT", OTEL_ENDPOINT)
             .environment_variable("METAFLOW_RUNTIME_ENVIRONMENT", "aws-batch")
         )
 

metaflow/plugins/aws/step_functions/dynamo_db_client.py

@@ -1,5 +1,8 @@
 import os
+import time
+
 import requests
+
 from metaflow.metaflow_config import SFN_DYNAMO_DB_TABLE
 
 
@@ -25,12 +28,31 @@ class DynamoDbClient(object):
     def save_parent_task_id_for_foreach_join(
         self, foreach_split_task_id, foreach_join_parent_task_id
     ):
-        return self._client.update_item(
-            TableName=self.name,
-            Key={"pathspec": {"S": foreach_split_task_id}},
-            UpdateExpression="ADD parent_task_ids_for_foreach_join :val",
-            ExpressionAttributeValues={":val": {"SS": [foreach_join_parent_task_id]}},
-        )
+        ex = None
+        for attempt in range(10):
+            try:
+                return self._client.update_item(
+                    TableName=self.name,
+                    Key={"pathspec": {"S": foreach_split_task_id}},
+                    UpdateExpression="ADD parent_task_ids_for_foreach_join :val",
+                    ExpressionAttributeValues={
+                        ":val": {"SS": [foreach_join_parent_task_id]}
+                    },
+                )
+            except self._client.exceptions.ClientError as error:
+                ex = error
+                if (
+                    error.response["Error"]["Code"]
+                    == "ProvisionedThroughputExceededException"
+                ):
+                    # hopefully, enough time for AWS to scale up! otherwise
+                    # ensure sufficient on-demand throughput for dynamo db
+                    # is provisioned ahead of time
+                    sleep_time = min((2**attempt) * 10, 60)
+                    time.sleep(sleep_time)
+                else:
+                    raise
+        raise ex
 
     def get_parent_task_ids_for_foreach_join(self, foreach_split_task_id):
         response = self._client.get_item(

metaflow/plugins/aws/step_functions/production_token.py

@@ -1,5 +1,5 @@
-import os
 import json
+import os
 import random
 import string
 import zlib

metaflow/plugins/aws/step_functions/step_functions.py

@@ -15,6 +15,7 @@ from metaflow.metaflow_config import (
     SFN_DYNAMO_DB_TABLE,
     SFN_EXECUTION_LOG_GROUP_ARN,
     SFN_IAM_ROLE,
+    SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH,
 )
 from metaflow.parameters import deploy_time_eval
 from metaflow.util import dict_to_cli_options, to_pascalcase
@@ -52,6 +53,7 @@ class StepFunctions(object):
         max_workers=None,
         workflow_timeout=None,
         is_project=False,
+        use_distributed_map=False,
     ):
         self.name = name
         self.graph = graph
@@ -70,6 +72,9 @@ class StepFunctions(object):
         self.max_workers = max_workers
         self.workflow_timeout = workflow_timeout
 
+        # https://aws.amazon.com/blogs/aws/step-functions-distributed-map-a-serverless-solution-for-large-scale-parallel-data-processing/
+        self.use_distributed_map = use_distributed_map
+
         self._client = StepFunctionsClient()
         self._workflow = self._compile()
         self._cron = self._cron()
@@ -365,17 +370,80 @@ class StepFunctions(object):
                     .parameter("SplitParentTaskId.$", "$.JobId")
                     .parameter("Parameters.$", "$.Parameters")
                     .parameter("Index.$", "$$.Map.Item.Value")
-                    .next(node.matching_join)
+                    .next(
+                        "%s_*GetManifest" % iterator_name
+                        if self.use_distributed_map
+                        else node.matching_join
+                    )
                     .iterator(
                         _visit(
                             self.graph[node.out_funcs[0]],
-                            Workflow(node.out_funcs[0]).start_at(node.out_funcs[0]),
+                            Workflow(node.out_funcs[0])
+                            .start_at(node.out_funcs[0])
+                            .mode(
+                                "DISTRIBUTED" if self.use_distributed_map else "INLINE"
+                            ),
                             node.matching_join,
                         )
                     )
                     .max_concurrency(self.max_workers)
-                    .output_path("$.[0]")
+                    # AWS Step Functions has a short coming for DistributedMap at the
+                    # moment that does not allow us to subset the output of for-each
+                    # to just a single element. We have to rely on a rather terrible
+                    # hack and resort to using ResultWriter to write the state to
+                    # Amazon S3 and process it in another task. But, well what can we
+                    # do...
+                    .result_writer(
+                        *(
+                            (
+                                (
+                                    SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH[len("s3://") :]
+                                    if SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH.startswith(
+                                        "s3://"
+                                    )
+                                    else SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH
+                                ).split("/", 1)
+                                + [""]
+                            )[:2]
+                            if self.use_distributed_map
+                            else (None, None)
+                        )
+                    )
+                    .output_path("$" if self.use_distributed_map else "$.[0]")
                 )
+                if self.use_distributed_map:
+                    workflow.add_state(
+                        State("%s_*GetManifest" % iterator_name)
+                        .resource("arn:aws:states:::aws-sdk:s3:getObject")
+                        .parameter("Bucket.$", "$.ResultWriterDetails.Bucket")
+                        .parameter("Key.$", "$.ResultWriterDetails.Key")
+                        .next("%s_*Map" % iterator_name)
+                        .result_selector("Body.$", "States.StringToJson($.Body)")
+                    )
+                    workflow.add_state(
+                        Map("%s_*Map" % iterator_name)
+                        .iterator(
+                            Workflow("%s_*PassWorkflow" % iterator_name)
+                            .mode("DISTRIBUTED")
+                            .start_at("%s_*Pass" % iterator_name)
+                            .add_state(
+                                Pass("%s_*Pass" % iterator_name)
+                                .end()
+                                .parameter("Output.$", "States.StringToJson($.Output)")
+                                .output_path("$.Output")
+                            )
+                        )
+                        .next(node.matching_join)
+                        .max_concurrency(1000)
+                        .item_reader(
+                            JSONItemReader()
+                            .resource("arn:aws:states:::s3:getObject")
+                            .parameter("Bucket.$", "$.Body.DestinationBucket")
+                            .parameter("Key.$", "$.Body.ResultFiles.SUCCEEDED.[0].Key")
+                        )
+                        .output_path("$.[0]")
+                    )
+
                 # Continue the traversal from the matching_join.
                 _visit(self.graph[node.matching_join], workflow, exit_node)
             # We shouldn't ideally ever get here.
@@ -444,7 +512,6 @@ class StepFunctions(object):
             "metaflow.owner": self.username,
             "metaflow.flow_name": self.flow.name,
             "metaflow.step_name": node.name,
-            "metaflow.run_id.$": "$$.Execution.Name",
             # Unfortunately we can't set the task id here since AWS Step
             # Functions lacks any notion of run-scoped task identifiers. We
             # instead co-opt the AWS Batch job id as the task id. This also
@@ -456,6 +523,10 @@ class StepFunctions(object):
             # `$$.State.RetryCount` resolves to an int dynamically and
             # AWS Batch job specification only accepts strings. We handle
             # retries/catch within AWS Batch to get around this limitation.
+            # And, we also cannot set the run id here since the run id maps to
+            # the execution name of the AWS Step Functions State Machine, which
+            # is different when executing inside a distributed map. We set it once
+            # in the start step and move it along to be consumed by all the children.
             "metaflow.version": self.environment.get_environment_info()[
                 "metaflow_version"
             ],
@@ -492,6 +563,12 @@ class StepFunctions(object):
             env["METAFLOW_S3_ENDPOINT_URL"] = S3_ENDPOINT_URL
 
         if node.name == "start":
+            # metaflow.run_id maps to AWS Step Functions State Machine Execution in all
+            # cases except for when within a for-each construct that relies on
+            # Distributed Map. To work around this issue, we pass the run id from the
+            # start step to all subsequent tasks.
+            attrs["metaflow.run_id.$"] = "$$.Execution.Name"
+
             # Initialize parameters for the flow in the `start` step.
             parameters = self._process_parameters()
             if parameters:
@@ -550,6 +627,8 @@ class StepFunctions(object):
                 env["METAFLOW_SPLIT_PARENT_TASK_ID"] = (
                     "$.Parameters.split_parent_task_id_%s" % node.split_parents[-1]
                 )
+                # Inherit the run id from the parent and pass it along to children.
+                attrs["metaflow.run_id.$"] = "$.Parameters.['metaflow.run_id']"
             else:
                 # Set appropriate environment variables for runtime replacement.
                 if len(node.in_funcs) == 1:
@@ -558,6 +637,8 @@ class StepFunctions(object):
                         % node.in_funcs[0]
                     )
                     env["METAFLOW_PARENT_TASK_ID"] = "$.JobId"
+                    # Inherit the run id from the parent and pass it along to children.
+                    attrs["metaflow.run_id.$"] = "$.Parameters.['metaflow.run_id']"
                 else:
                     # Generate the input paths in a quasi-compressed format.
                     # See util.decompress_list for why this is written the way
@@ -567,6 +648,8 @@ class StepFunctions(object):
                         "${METAFLOW_PARENT_%s_TASK_ID}" % (idx, idx)
                         for idx, _ in enumerate(node.in_funcs)
                     )
+                    # Inherit the run id from the parent and pass it along to children.
+                    attrs["metaflow.run_id.$"] = "$.[0].Parameters.['metaflow.run_id']"
                     for idx, _ in enumerate(node.in_funcs):
                         env["METAFLOW_PARENT_%s_TASK_ID" % idx] = "$.[%s].JobId" % idx
                         env["METAFLOW_PARENT_%s_STEP" % idx] = (
@@ -893,6 +976,12 @@ class Workflow(object):
         tree = lambda: defaultdict(tree)
         self.payload = tree()
 
+    def mode(self, mode):
+        self.payload["ProcessorConfig"] = {"Mode": mode}
+        if mode == "DISTRIBUTED":
+            self.payload["ProcessorConfig"]["ExecutionType"] = "STANDARD"
+        return self
+
     def start_at(self, start_at):
         self.payload["StartAt"] = start_at
         return self
@@ -940,10 +1029,18 @@ class State(object):
         self.payload["ResultPath"] = result_path
         return self
 
+    def result_selector(self, name, value):
+        self.payload["ResultSelector"][name] = value
+        return self
+
     def _partition(self):
         # This is needed to support AWS Gov Cloud and AWS CN regions
         return SFN_IAM_ROLE.split(":")[1]
 
+    def retry_strategy(self, retry_strategy):
+        self.payload["Retry"] = [retry_strategy]
+        return self
+
     def batch(self, job):
         self.resource(
             "arn:%s:states:::batch:submitJob.sync" % self._partition()
@@ -963,6 +1060,19 @@ class State(object):
         # tags may not be present in all scenarios
         if "tags" in job.payload:
             self.parameter("Tags", job.payload["tags"])
+        # set retry strategy for AWS Batch job submission to account for the
+        # measily 50 jobs / second queue admission limit which people can
+        # run into very quickly.
+        self.retry_strategy(
+            {
+                "ErrorEquals": ["Batch.AWSBatchException"],
+                "BackoffRate": 2,
+                "IntervalSeconds": 2,
+                "MaxDelaySeconds": 60,
+                "MaxAttempts": 10,
+                "JitterStrategy": "FULL",
+            }
+        )
         return self
 
     def dynamo_db(self, table_name, primary_key, values):
@@ -976,6 +1086,26 @@ class State(object):
         return self
 
 
+class Pass(object):
+    def __init__(self, name):
+        self.name = name
+        tree = lambda: defaultdict(tree)
+        self.payload = tree()
+        self.payload["Type"] = "Pass"
+
+    def end(self):
+        self.payload["End"] = True
+        return self
+
+    def parameter(self, name, value):
+        self.payload["Parameters"][name] = value
+        return self
+
+    def output_path(self, output_path):
+        self.payload["OutputPath"] = output_path
+        return self
+
+
 class Parallel(object):
     def __init__(self, name):
         self.name = name
@@ -1037,3 +1167,37 @@ class Map(object):
     def result_path(self, result_path):
         self.payload["ResultPath"] = result_path
         return self
+
+    def item_reader(self, item_reader):
+        self.payload["ItemReader"] = item_reader.payload
+        return self
+
+    def result_writer(self, bucket, prefix):
+        if bucket is not None and prefix is not None:
+            self.payload["ResultWriter"] = {
+                "Resource": "arn:aws:states:::s3:putObject",
+                "Parameters": {
+                    "Bucket": bucket,
+                    "Prefix": prefix,
+                },
+            }
+        return self
+
+
+class JSONItemReader(object):
+    def __init__(self):
+        tree = lambda: defaultdict(tree)
+        self.payload = tree()
+        self.payload["ReaderConfig"] = {"InputType": "JSON", "MaxItems": 1}
+
+    def resource(self, resource):
+        self.payload["Resource"] = resource
+        return self
+
+    def parameter(self, name, value):
+        self.payload["Parameters"][name] = value
+        return self
+
+    def output_path(self, output_path):
+        self.payload["OutputPath"] = output_path
+        return self

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -1,23 +1,23 @@
 import base64
-from metaflow._vendor import click
-from hashlib import sha1
 import json
 import re
+from hashlib import sha1
 
-from metaflow import current, decorators, parameters, JSONType
+from metaflow import JSONType, current, decorators, parameters
+from metaflow._vendor import click
+from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.metaflow_config import (
     SERVICE_VERSION_CHECK,
     SFN_STATE_MACHINE_PREFIX,
     UI_URL,
 )
-from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.package import MetaflowPackage
 from metaflow.plugins.aws.batch.batch_decorator import BatchDecorator
 from metaflow.tagging_util import validate_tags
 from metaflow.util import get_username, to_bytes, to_unicode, version_parse
 
+from .production_token import load_token, new_token, store_token
 from .step_functions import StepFunctions
-from .production_token import load_token, store_token, new_token
 
 VALID_NAME = re.compile(r"[^a-zA-Z0-9_\-\.]")
 
@@ -124,6 +124,12 @@ def step_functions(obj, name=None):
     help="Log AWS Step Functions execution history to AWS CloudWatch "
     "Logs log group.",
 )
+@click.option(
+    "--use-distributed-map/--no-use-distributed-map",
+    is_flag=True,
+    help="Use AWS Step Functions Distributed Map instead of Inline Map for "
+    "defining foreach tasks in Amazon State Language.",
+)
 @click.pass_obj
 def create(
     obj,
@@ -136,6 +142,7 @@ def create(
     max_workers=None,
     workflow_timeout=None,
     log_execution_history=False,
+    use_distributed_map=False,
 ):
     validate_tags(tags)
 
@@ -165,6 +172,7 @@ def create(
         max_workers,
         workflow_timeout,
         obj.is_project,
+        use_distributed_map,
     )
 
     if only_json:
@@ -273,7 +281,15 @@ def resolve_state_machine_name(obj, name):
 
 
 def make_flow(
-    obj, token, name, tags, namespace, max_workers, workflow_timeout, is_project
+    obj,
+    token,
+    name,
+    tags,
+    namespace,
+    max_workers,
+    workflow_timeout,
+    is_project,
+    use_distributed_map,
 ):
     if obj.flow_datastore.TYPE != "s3":
         raise MetaflowException("AWS Step Functions requires --datastore=s3.")
@@ -309,6 +325,7 @@ def make_flow(
         username=get_username(),
         workflow_timeout=workflow_timeout,
         is_project=is_project,
+        use_distributed_map=use_distributed_map,
     )
 
 

metaflow/plugins/aws/step_functions/step_functions_decorator.py

@@ -1,5 +1,5 @@
-import os
 import json
+import os
 import time
 
 from metaflow.decorators import StepDecorator

metaflow/plugins/cards/card_cli.py

@@ -17,6 +17,7 @@ import random
 from contextlib import contextmanager
 from functools import wraps
 from metaflow.exception import MetaflowNamespaceMismatch
+
 from .card_datastore import CardDatastore, NUM_SHORT_HASH_CHARS
 from .exception import (
     CardClassFoundException,
@@ -736,8 +737,7 @@ def create(
 
     if error_stack_trace is not None and mode != "refresh":
         rendered_content = error_card().render(task, stack_trace=error_stack_trace)
-
-    if (
+    elif (
         rendered_info.is_implemented
         and rendered_info.timed_out
         and mode != "refresh"

metaflow/plugins/pypi/bootstrap.py

@@ -45,7 +45,7 @@ if __name__ == "__main__":
     #     fi
     # fi
 
-    prefix = os.path.join(os.getcwd(), id_)
+    prefix = os.path.join(os.getcwd(), architecture, id_)
     pkgs_dir = os.path.join(os.getcwd(), ".pkgs")
     manifest_dir = os.path.join(os.getcwd(), DATASTORE_LOCAL_DIR, flow_name)
 

metaflow/plugins/pypi/conda_decorator.py

@@ -2,11 +2,13 @@ import importlib
 import json
 import os
 import platform
+import re
 import sys
 import tempfile
 
 from metaflow.decorators import FlowDecorator, StepDecorator
 from metaflow.extension_support import EXT_PKG
+from metaflow.metadata import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
 from metaflow.util import get_metaflow_root
 
@@ -241,7 +243,25 @@ class CondaStepDecorator(StepDecorator):
                 ),
             )
         )
-        # TODO: Register metadata
+
+        # Infer environment prefix from Python interpreter
+        match = re.search(
+            r"(?:.*\/)(metaflow\/[^/]+\/[^/]+)(?=\/bin\/python)", sys.executable
+        )
+        if match:
+            meta.register_metadata(
+                run_id,
+                step_name,
+                task_id,
+                [
+                    MetaDatum(
+                        field="conda_env_prefix",
+                        value=match.group(1),
+                        type="conda_env_prefix",
+                        tags=["attempt_id:{0}".format(retry_count)],
+                    )
+                ],
+            )
 
     def runtime_step_cli(
         self, cli_args, retry_count, max_user_code_retries, ubf_context

metaflow/plugins/pypi/conda_environment.py

@@ -6,10 +6,11 @@ import io
 import json
 import os
 import sys
+import tarfile
 import time
 from concurrent.futures import ThreadPoolExecutor
 from hashlib import sha256
-from io import BufferedIOBase
+from io import BufferedIOBase, BytesIO
 from itertools import chain
 from urllib.parse import unquote, urlparse
 
@@ -33,6 +34,7 @@ class CondaEnvironmentException(MetaflowException):
 
 class CondaEnvironment(MetaflowEnvironment):
     TYPE = "conda"
+    _filecache = None
 
     def __init__(self, flow):
         self.flow = flow
@@ -188,7 +190,7 @@ class CondaEnvironment(MetaflowEnvironment):
         if id_:
             # bootstrap.py is responsible for ensuring the validity of this executable.
             # -s is important! Can otherwise leak packages to other environments.
-            return os.path.join(id_, "bin/python -s")
+            return os.path.join("linux-64", id_, "bin/python -s")
         else:
             # for @conda/@pypi(disabled=True).
             return super().executable(step_name, default)
@@ -320,8 +322,23 @@ class CondaEnvironment(MetaflowEnvironment):
 
     @classmethod
     def get_client_info(cls, flow_name, metadata):
-        # TODO: Decide this method's fate
-        return None
+        if cls._filecache is None:
+            from metaflow.client.filecache import FileCache
+
+            cls._filecache = FileCache()
+
+        info = metadata.get("code-package")
+        prefix = metadata.get("conda_env_prefix")
+        if info is None or prefix is None:
+            return {}
+        info = json.loads(info)
+        _, blobdata = cls._filecache.get_data(
+            info["ds_type"], flow_name, info["location"], info["sha"]
+        )
+        with tarfile.open(fileobj=BytesIO(blobdata), mode="r:gz") as tar:
+            manifest = tar.extractfile(MAGIC_FILE)
+            info = json.loads(manifest.read().decode("utf-8"))
+            return info[prefix.split("/")[2]][prefix.split("/")[1]]
 
     def add_to_package(self):
         # Add manifest file to job package at the top level.

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.11.2"
+metaflow_version = "2.11.3"

{metaflow-2.11.2.dist-info → metaflow-2.11.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.11.2
+Version: 2.11.3
 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.11.2 ; extra == 'stubs'
+Requires-Dist: metaflow-stubs ==2.11.3 ; extra == 'stubs'
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)
 

{metaflow-2.11.2.dist-info → metaflow-2.11.3.dist-info}/RECORD

@@ -9,12 +9,12 @@ metaflow/decorators.py,sha256=EGL1_nkdxoYG5AZiOQ8sLGA1bprGK8ENwlSIOYQmLhs,21357
 metaflow/event_logger.py,sha256=joTVRqZPL87nvah4ZOwtqWX8NeraM_CXKXXGVpKGD8o,780
 metaflow/events.py,sha256=ahjzkSbSnRCK9RZ-9vTfUviz_6gMvSO9DGkJ86X80-k,5300
 metaflow/exception.py,sha256=KC1LHJQzzYkWib0DeQ4l_A2r8VaudywsSqIQuq1RDZU,4954
-metaflow/flowspec.py,sha256=wA6Ci-vZ_CrfLEEMnCqTEkNv14O5Bq5AFvH-8Xt1sDw,26684
+metaflow/flowspec.py,sha256=YMo2zNEy3UcclYys5rPpMk7bk0JvtABo-EpmFTi-IPQ,26874
 metaflow/graph.py,sha256=ZPxyG8uwVMk5YYgX4pQEQaPZtZM5Wy-G4NtJK73IEuA,11818
 metaflow/includefile.py,sha256=yHczcZ_U0SrasxSNhZb3DIBzx8UZnrJCl3FzvpEQLOA,19753
 metaflow/integrations.py,sha256=LlsaoePRg03DjENnmLxZDYto3NwWc9z_PtU6nJxLldg,1480
 metaflow/lint.py,sha256=_kYAbAtsP7IG1Rd0FqNbo8I8Zs66_0WXbaZJFARO3dE,10394
-metaflow/metaflow_config.py,sha256=4F1HEvEHiGpEj-J5ccYw5b39myh7X2xkp5AyPM5rnAY,19225
+metaflow/metaflow_config.py,sha256=_--DonM1Sj6PbQH7Am_0r-9oZFhI1DLJmiJgJkpEkJY,19662
 metaflow/metaflow_config_funcs.py,sha256=pCaiQ2ez9wXixJI3ehmf3QiW9lUqFrZnBZx1my_0wIg,4874
 metaflow/metaflow_current.py,sha256=sCENPBiji3LcPbwgOG0ukGd_yEc5tST8EowES8DzRtA,7430
 metaflow/metaflow_environment.py,sha256=JdsmQsYp1SDQniQ0-q1mKRrmzSFfYuzrf6jLEHmyaiM,7352
@@ -33,7 +33,7 @@ metaflow/task.py,sha256=rGBlG18vnHKC65CYiFkja0GsXYWmRVrNOeCD4TY_nL4,27032
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=RrjsvADLKxSqjL76CxKh_J4OJl840B9Ak3V-vXleGas,13429
 metaflow/vendor.py,sha256=LZgXrh7ZSDmD32D1T5jj3OKKpXIqqxKzdMAOc5V0SD4,5162
-metaflow/version.py,sha256=TE2PAReb-M2lZz1GS382YFZ7Fg6NCktRPZ4qBMZ9wyg,28
+metaflow/version.py,sha256=gJROyhqslzcmekVHrVdDybqVCjTbJnSNU5-tig72mrA,28
 metaflow/_vendor/__init__.py,sha256=y_CiwUD3l4eAKvTVDZeqgVujMy31cAM1qjAB-HfI-9s,353
 metaflow/_vendor/click/__init__.py,sha256=FkyGDQ-cbiQxP_lxgUspyFYS48f2S_pTcfKPz-d_RMo,2463
 metaflow/_vendor/click/_bashcomplete.py,sha256=9J98IHQYmCAr2Jup6TDshUr5FJEen-AoQCZR0K5nKxQ,12309
@@ -148,7 +148,7 @@ metaflow/plugins/airflow/sensors/s3_sensor.py,sha256=JUKoGNoTCtrO9MNEneEC7ldRNwg
 metaflow/plugins/argo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/argo/argo_client.py,sha256=MKKhMCbWOPzf6z5zQQiyDRHHkAXcO7ipboDZDqAAvOk,15849
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=bFpTxpTztYNe-uG0WbOL7HRca8JY2UyxIbBwQcodKm8,119689
+metaflow/plugins/argo/argo_workflows.py,sha256=gKGmrcVwEktaqa7dOh9pWwAktGUQwwTACDT5-uUpJ4s,121723
 metaflow/plugins/argo/argo_workflows_cli.py,sha256=sZTpgfmc50eT3e0qIxpVqUgWhTcYlO1HM4gU6Oaya8g,33259
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=CfKVoHCOsCCQMghhPE30xw15gacwp3hR23HCo9ZZFVg,6580
 metaflow/plugins/argo/process_input_paths.py,sha256=4SiUoxbnTX4rCt0RSLcxG5jysbyd8oU-5JT0UOgy-vk,555
@@ -156,22 +156,22 @@ metaflow/plugins/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hS
 metaflow/plugins/aws/aws_client.py,sha256=mO8UD6pxFaOnxDb3hTP3HB7Gqb_ZxoR-76LT683WHvI,4036
 metaflow/plugins/aws/aws_utils.py,sha256=pkkH8Cy9sF5tp3HoZ84wkN-84NmksgCdNN4cMSdsLaA,6455
 metaflow/plugins/aws/batch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/aws/batch/batch.py,sha256=aVZOWdD9d7Zvhr1RYcJTKT9DDQmaPaU7Zu7ASZbnHgA,17273
+metaflow/plugins/aws/batch/batch.py,sha256=ziYpcolvbAZ4DoI63-A9mZcBM5oKUcueeGMjPVLdCaI,17367
 metaflow/plugins/aws/batch/batch_cli.py,sha256=W4DW1Ldx4wK6pIdDOsNusUphpD-2a93qTBMRTRh5a40,11048
 metaflow/plugins/aws/batch/batch_client.py,sha256=_etUOsqz5d9tSmO9HGsNV_zzJz2Z25CDyqzOqWcORnI,27066
 metaflow/plugins/aws/batch/batch_decorator.py,sha256=C_K476zuU6-FF7DImijgKeGLdNXbBIcp4WC7K1bMDsc,15851
 metaflow/plugins/aws/secrets_manager/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.py,sha256=JtFUVu00Cg0FzAizgrPLXmrMqsT7YeQMkQlgeivUxcE,7986
 metaflow/plugins/aws/step_functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/aws/step_functions/dynamo_db_client.py,sha256=KF8omWGe93_Qiec41LNdSuIHa0hmv1lbX7L8IwOR8Kw,1532
+metaflow/plugins/aws/step_functions/dynamo_db_client.py,sha256=c1nFZc9v_dLTD3iD-75I6wrlM6oM7YmxpFPTRo8Tra0,2327
 metaflow/plugins/aws/step_functions/event_bridge_client.py,sha256=U9-tqKdih4KR-ZDRhFc-jHmYIcHgpS4swfgtTxNMB94,2690
-metaflow/plugins/aws/step_functions/production_token.py,sha256=jEsDH0etkzORgeDTUx6UDRjRsy8QQHYjHWEDLE-0rVY,1898
+metaflow/plugins/aws/step_functions/production_token.py,sha256=_o4emv3rozYZoWpaj1Y6UfKhTMlYpQc7GDDDBfZ2G7s,1898
 metaflow/plugins/aws/step_functions/schedule_decorator.py,sha256=Ab1rW8O_no4HNZm4__iBmFDCDW0Z8-TgK4lnxHHA6HI,1940
 metaflow/plugins/aws/step_functions/set_batch_environment.py,sha256=ibiGWFHDjKcLfprH3OsX-g2M9lUsh6J-bp7v2cdLhD4,1294
-metaflow/plugins/aws/step_functions/step_functions.py,sha256=FjymSGA3-lTlaneZK5CKiC0yPBwVNblVWbmh1KZfvn0,44545
-metaflow/plugins/aws/step_functions/step_functions_cli.py,sha256=eHy9RGaSYOnN6VcCC_Nhvq5QBQPBlreX4Xem8Vu3R-w,24041
+metaflow/plugins/aws/step_functions/step_functions.py,sha256=V079UmCANB6clTzmOZe8Gq7hImb12hd_bw3DCLdMeZw,51673
+metaflow/plugins/aws/step_functions/step_functions_cli.py,sha256=KlH9jJL0VfsT0JqBhLwaWdYjaccU8UEArKAFnIJbSoU,24426
 metaflow/plugins/aws/step_functions/step_functions_client.py,sha256=DKpNwAIWElvWjFANs5Ku3rgzjxFoqAD6k-EF8Xhkg3Q,4754
-metaflow/plugins/aws/step_functions/step_functions_decorator.py,sha256=R1C1EYdoYJUILXdpV-vdXOlyEBNiyDWfNXsTcG8URww,3791
+metaflow/plugins/aws/step_functions/step_functions_decorator.py,sha256=9hw_MX36RyFp6IowuAYaJzJg9UC5KCe1FNt1PcG7_J0,3791
 metaflow/plugins/azure/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/azure/azure_exceptions.py,sha256=uvxE3E3nsbQq1dxCx1Yl9O54frbnMS5Elk8Z4qQ2Oh4,404
 metaflow/plugins/azure/azure_tail.py,sha256=JAqV4mC42bMpR0O7m6X4cpFuh0peV1ufs_jJXrmicTc,3362
@@ -179,7 +179,7 @@ metaflow/plugins/azure/azure_utils.py,sha256=efJnGl15s2HxxCtmLUsYDDsn1ek2tUl-5lr
 metaflow/plugins/azure/blob_service_client_factory.py,sha256=j1FJa06h6V-QcYIir9tcoEb23XcufcecyJQRsGhD2Qw,6489
 metaflow/plugins/azure/includefile_support.py,sha256=Wv3g3RlGtLbxAh3Reg0BDLWwqavYibQNCDWddlH7XCE,4706
 metaflow/plugins/cards/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/cards/card_cli.py,sha256=Nd-tEQJ2mnJGgDlFfDiZRrcvD_f2CKiMkK0H8bxRjDY,34625
+metaflow/plugins/cards/card_cli.py,sha256=wE5IbcS2za2Tv8SrLmTjPDub48879g_JkNMyLlnh_4w,34627
 metaflow/plugins/cards/card_client.py,sha256=LnRDqpyheDIehiHufpvMVTwE7cYEVIiMq8ggqZF4I2E,9368
 metaflow/plugins/cards/card_creator.py,sha256=E_NCmWPK6DzkqigtpUpeddCDbjnKF6dJcE6IvWzwiyA,7740
 metaflow/plugins/cards/card_datastore.py,sha256=3K19wE0CZVvOpuYUytftIYYnHHn3pMZJE87FMD6OYlM,14244
@@ -254,9 +254,9 @@ metaflow/plugins/metadata/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-
 metaflow/plugins/metadata/local.py,sha256=YhLJC5zjVJrvQFIyQ92ZBByiUmhCC762RUX7ITX12O8,22428
 metaflow/plugins/metadata/service.py,sha256=ihq5F7KQZlxvYwzH_-jyP2aWN_I96i2vp92j_d697s8,20204
 metaflow/plugins/pypi/__init__.py,sha256=0YFZpXvX7HCkyBFglatual7XGifdA1RwC3U4kcizyak,1037
-metaflow/plugins/pypi/bootstrap.py,sha256=PBIStzaEuhAl0xBd0Ap6oUaWMDwvz5bZEWo6h2SFru0,5073
-metaflow/plugins/pypi/conda_decorator.py,sha256=mkHzyWSEPSbIH6qk8UUvIjtDxzyl7tQxpLffO9ZBFLA,13547
-metaflow/plugins/pypi/conda_environment.py,sha256=FeHwE9xkI3nYv_OZ_iFzMI2Z2s5BG7llFSi-ndPeId8,17867
+metaflow/plugins/pypi/bootstrap.py,sha256=nCe8FadqfIM19yj64m4JWdv_QEnQEp01bzQZrxzo5bs,5087
+metaflow/plugins/pypi/conda_decorator.py,sha256=-bPxNtZKjxqOo4sj89uIp8ZVrCIontWhAp7wwRFjYpg,14189
+metaflow/plugins/pypi/conda_environment.py,sha256=gfSa50ukez72aB16RZmILPuc8GXCfw0W-sbmWzcmfsE,18575
 metaflow/plugins/pypi/micromamba.py,sha256=wlVN2fm4WXFh3jVNtpDfu4XEz6VJKbmFNp0QvqlMIuI,12179
 metaflow/plugins/pypi/pip.py,sha256=paL-hbj5j-vcYOeSm04OiBix-TpoKMF7bRfGVdXfGZs,11685
 metaflow/plugins/pypi/pypi_decorator.py,sha256=syWk_oSQhIK9Y7OeOINMG2XVyxh9sj5uJhapwAXRBDw,5583
@@ -297,9 +297,9 @@ metaflow/tutorials/07-worldview/README.md,sha256=5vQTrFqulJ7rWN6r20dhot9lI2sVj9W
 metaflow/tutorials/07-worldview/worldview.ipynb,sha256=ztPZPI9BXxvW1QdS2Tfe7LBuVzvFvv0AToDnsDJhLdE,2237
 metaflow/tutorials/08-autopilot/README.md,sha256=GnePFp_q76jPs991lMUqfIIh5zSorIeWznyiUxzeUVE,1039
 metaflow/tutorials/08-autopilot/autopilot.ipynb,sha256=DQoJlILV7Mq9vfPBGW-QV_kNhWPjS5n6SJLqePjFYLY,3191
-metaflow-2.11.2.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-metaflow-2.11.2.dist-info/METADATA,sha256=-g4FgX7It1THhPXL0GVND89oCcx21LYN2cj5cbcZ0_o,5906
-metaflow-2.11.2.dist-info/WHEEL,sha256=-G_t0oGuE7UD0DrSpVZnq1hHMBV9DD2XkS5v7XpmTnk,110
-metaflow-2.11.2.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-metaflow-2.11.2.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-metaflow-2.11.2.dist-info/RECORD,,
+metaflow-2.11.3.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+metaflow-2.11.3.dist-info/METADATA,sha256=h3w2Ot3jIufP4RiQdkcOOhicH_MoMEXgF_34jpbsIZQ,5906
+metaflow-2.11.3.dist-info/WHEEL,sha256=-G_t0oGuE7UD0DrSpVZnq1hHMBV9DD2XkS5v7XpmTnk,110
+metaflow-2.11.3.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+metaflow-2.11.3.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+metaflow-2.11.3.dist-info/RECORD,,