expose resolved dependencies (#1740)

metaflow/plugins/pypi/bootstrap.py

@@ -45,7 +45,7 @@
     #     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 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 @@ def task_pre_step(
                 ),
             )
         )
-        # 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 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 @@ def __init__(self, msg):
 
 class CondaEnvironment(MetaflowEnvironment):
     TYPE = "conda"
+    _filecache = None
 
     def __init__(self, flow):
         self.flow = flow
@@ -188,7 +190,7 @@ def executable(self, step_name, default=None):
         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 @@ def pylint_config(self):
 
     @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.