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

metaflow/client/filecache.py

@@ -1,20 +1,24 @@
 from __future__ import print_function
 from collections import OrderedDict
+import json
 import os
 import sys
 import time
 from tempfile import NamedTemporaryFile
 from hashlib import sha1
 
+from urllib.parse import urlparse
+
 from metaflow.datastore import FlowDataStore
 from metaflow.datastore.content_addressed_store import BlobCache
+from metaflow.datastore.flow_datastore import MetadataCache
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import (
     CLIENT_CACHE_PATH,
     CLIENT_CACHE_MAX_SIZE,
     CLIENT_CACHE_MAX_FLOWDATASTORE_COUNT,
-    CLIENT_CACHE_MAX_TASKDATASTORE_COUNT,
 )
+from metaflow.metaflow_profile import from_start
 
 from metaflow.plugins import DATASTORES
 
@@ -61,8 +65,8 @@ class FileCache(object):
         # when querying for sizes of artifacts. Once we have queried for the size
         # of one artifact in a TaskDatastore, caching this means that any
         # queries on that same TaskDatastore will be quick (since we already
-        # have all the metadata)
-        self._task_metadata_caches = OrderedDict()
+        # have all the metadata). We keep track of this in a file so it persists
+        # across processes.
 
     @property
     def cache_dir(self):
@@ -83,10 +87,9 @@ class FileCache(object):
     def get_log_legacy(
         self, ds_type, location, logtype, attempt, flow_name, run_id, step_name, task_id
     ):
-
         ds_cls = self._get_datastore_storage_impl(ds_type)
         ds_root = ds_cls.path_join(*ds_cls.path_split(location)[:-5])
-        cache_id = self._flow_ds_id(ds_type, ds_root, flow_name)
+        cache_id = self.flow_ds_id(ds_type, ds_root, flow_name)
 
         token = (
             "%s.cached"
@@ -310,13 +313,25 @@ class FileCache(object):
         self._objects = sorted(objects, reverse=False)
 
     @staticmethod
-    def _flow_ds_id(ds_type, ds_root, flow_name):
-        return ".".join([ds_type, ds_root, flow_name])
+    def flow_ds_id(ds_type, ds_root, flow_name):
+        p = urlparse(ds_root)
+        sanitized_root = (p.netloc + p.path).replace("/", "_")
+        return ".".join([ds_type, sanitized_root, flow_name])
 
     @staticmethod
-    def _task_ds_id(ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt):
+    def task_ds_id(ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt):
+        p = urlparse(ds_root)
+        sanitized_root = (p.netloc + p.path).replace("/", "_")
         return ".".join(
-            [ds_type, ds_root, flow_name, run_id, step_name, task_id, str(attempt)]
+            [
+                ds_type,
+                sanitized_root,
+                flow_name,
+                run_id,
+                step_name,
+                task_id,
+                str(attempt),
+            ]
         )
 
     def _garbage_collect(self):
@@ -352,7 +367,7 @@ class FileCache(object):
         return storage_impl[0]
 
     def _get_flow_datastore(self, ds_type, ds_root, flow_name):
-        cache_id = self._flow_ds_id(ds_type, ds_root, flow_name)
+        cache_id = self.flow_ds_id(ds_type, ds_root, flow_name)
         cached_flow_datastore = self._store_caches.get(cache_id)
 
         if cached_flow_datastore:
@@ -367,9 +382,14 @@ class FileCache(object):
                 ds_root=ds_root,
             )
             blob_cache = self._blob_caches.setdefault(
-                cache_id, FileBlobCache(self, cache_id)
+                cache_id,
+                (
+                    FileBlobCache(self, cache_id),
+                    TaskMetadataCache(self, ds_type, ds_root, flow_name),
+                ),
             )
-            cached_flow_datastore.ca_store.set_blob_cache(blob_cache)
+            cached_flow_datastore.ca_store.set_blob_cache(blob_cache[0])
+            cached_flow_datastore.set_metadata_cache(blob_cache[1])
             self._store_caches[cache_id] = cached_flow_datastore
             if len(self._store_caches) > CLIENT_CACHE_MAX_FLOWDATASTORE_COUNT:
                 cache_id_to_remove, _ = self._store_caches.popitem(last=False)
@@ -380,32 +400,52 @@ class FileCache(object):
         self, ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt
     ):
         flow_ds = self._get_flow_datastore(ds_type, ds_root, flow_name)
-        cached_metadata = None
-        if attempt is not None:
-            cache_id = self._task_ds_id(
-                ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt
+
+        return flow_ds.get_task_datastore(run_id, step_name, task_id, attempt=attempt)
+
+
+class TaskMetadataCache(MetadataCache):
+    def __init__(self, filecache, ds_type, ds_root, flow_name):
+        self._filecache = filecache
+        self._ds_type = ds_type
+        self._ds_root = ds_root
+        self._flow_name = flow_name
+
+    def _path(self, run_id, step_name, task_id, attempt):
+        if attempt is None:
+            raise MetaflowException(
+                "Attempt number must be specified to use task metadata cache. Raise an issue "
+                "on Metaflow GitHub if you see this message.",
             )
-            cached_metadata = self._task_metadata_caches.get(cache_id)
-            if cached_metadata:
-                od_move_to_end(self._task_metadata_caches, cache_id)
-                return flow_ds.get_task_datastore(
-                    run_id,
-                    step_name,
-                    task_id,
-                    attempt=attempt,
-                    data_metadata=cached_metadata,
-                )
-        # If we are here, we either have attempt=None or nothing in the cache
-        task_ds = flow_ds.get_task_datastore(
-            run_id, step_name, task_id, attempt=attempt
+        cache_id = self._filecache.task_ds_id(
+            self._ds_type,
+            self._ds_root,
+            self._flow_name,
+            run_id,
+            step_name,
+            task_id,
+            attempt,
         )
-        cache_id = self._task_ds_id(
-            ds_type, ds_root, flow_name, run_id, step_name, task_id, task_ds.attempt
+        token = (
+            "%s.cached"
+            % sha1(
+                os.path.join(
+                    run_id, step_name, task_id, str(attempt), "metadata"
+                ).encode("utf-8")
+            ).hexdigest()
+        )
+        return os.path.join(self._filecache.cache_dir, cache_id, token[:2], token)
+
+    def load_metadata(self, run_id, step_name, task_id, attempt):
+        d = self._filecache.read_file(self._path(run_id, step_name, task_id, attempt))
+        if d:
+            return json.loads(d)
+
+    def store_metadata(self, run_id, step_name, task_id, attempt, metadata_dict):
+        self._filecache.create_file(
+            self._path(run_id, step_name, task_id, attempt),
+            json.dumps(metadata_dict).encode("utf-8"),
         )
-        self._task_metadata_caches[cache_id] = task_ds.ds_metadata
-        if len(self._task_metadata_caches) > CLIENT_CACHE_MAX_TASKDATASTORE_COUNT:
-            self._task_metadata_caches.popitem(last=False)
-        return task_ds
 
 
 class FileBlobCache(BlobCache):

metaflow/clone_util.py

@@ -1,5 +1,5 @@
 import time
-from .metadata import MetaDatum
+from .metadata_provider import MetaDatum
 
 
 def clone_task_helper(
@@ -66,6 +66,12 @@ def clone_task_helper(
                 type="attempt",
                 tags=metadata_tags,
             ),
+            MetaDatum(
+                field="attempt_ok",
+                value="True",  # During clone, the task is always considered successful.
+                type="internal_attempt_status",
+                tags=metadata_tags,
+            ),
         ],
     )
     output.done()

metaflow/cmd/code/__init__.py

@@ -0,0 +1,231 @@
+import os
+import shutil
+import sys
+from subprocess import PIPE, CompletedProcess, run
+from tempfile import TemporaryDirectory
+from typing import Any, Callable, List, Mapping, Optional, cast
+
+from metaflow import Run
+from metaflow.util import walk_without_cycles
+from metaflow._vendor import click
+from metaflow.cli import echo_always
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Access, compare, and manage code associated with Metaflow runs.")
+def code():
+    pass
+
+
+def echo(line: str) -> None:
+    echo_always(line, err=True, fg="magenta")
+
+
+def extract_code_package(runspec: str) -> TemporaryDirectory:
+    try:
+        mf_run = Run(runspec, _namespace_check=False)
+        echo(f"✅  Run *{runspec}* found, downloading code..")
+    except Exception as e:
+        echo(f"❌  Run **{runspec}** not found")
+        raise e
+
+    if mf_run.code is None:
+        echo(
+            f"❌  Run **{runspec}** doesn't have a code package. Maybe it's a local run?"
+        )
+        raise RuntimeError("no code package found")
+
+    return mf_run.code.extract()
+
+
+def perform_diff(
+    source_dir: str,
+    target_dir: Optional[str] = None,
+    output: bool = False,
+    **kwargs: Mapping[str, Any],
+) -> Optional[List[str]]:
+    if target_dir is None:
+        target_dir = os.getcwd()
+
+    diffs = []
+    for dirpath, _, filenames in walk_without_cycles(source_dir):
+        for fname in filenames:
+            # NOTE: the paths below need to be set up carefully
+            # for the `patch` command to work. Better not to touch
+            # the directories below. If you must, test that patches
+            # work after your changes.
+            #
+            # target_file is the git repo in the current working directory
+            rel = os.path.relpath(dirpath, source_dir)
+            target_file = os.path.join(rel, fname)
+            # source_file is the run file loaded in a tmp directory
+            source_file = os.path.join(dirpath, fname)
+
+            if sys.stdout.isatty() and not output:
+                color = ["--color"]
+            else:
+                color = ["--no-color"]
+
+            if os.path.exists(os.path.join(target_dir, target_file)):
+                cmd = (
+                    ["git", "diff", "--no-index", "--exit-code"]
+                    + color
+                    + [
+                        target_file,
+                        source_file,
+                    ]
+                )
+                result: CompletedProcess = run(
+                    cmd, text=True, stdout=PIPE, cwd=target_dir
+                )
+                if result.returncode == 0:
+                    if not output:
+                        echo(f"✅  {target_file} is identical, skipping")
+                    continue
+
+                if output:
+                    diffs.append(result.stdout)
+                else:
+                    run(["less", "-R"], input=result.stdout, text=True)
+            else:
+                if not output:
+                    echo(f"❗  {target_file} not in the target directory, skipping")
+    return diffs if output else None
+
+
+def run_op(
+    runspec: str, op: Callable[..., Optional[List[str]]], **op_args: Mapping[str, Any]
+) -> Optional[List[str]]:
+    tmp = None
+    try:
+        tmp = extract_code_package(runspec)
+        return op(tmp.name, **op_args)
+    finally:
+        if tmp and os.path.exists(tmp.name):
+            shutil.rmtree(tmp.name)
+
+
+def run_op_diff_runs(
+    source_run_pathspec: str, target_run_pathspec: str, **op_args: Mapping[str, Any]
+) -> Optional[List[str]]:
+    source_tmp = None
+    target_tmp = None
+    try:
+        source_tmp = extract_code_package(source_run_pathspec)
+        target_tmp = extract_code_package(target_run_pathspec)
+        return perform_diff(source_tmp.name, target_tmp.name, **op_args)
+    finally:
+        for d in [source_tmp, target_tmp]:
+            if d and os.path.exists(d.name):
+                shutil.rmtree(d.name)
+
+
+def op_diff(tmpdir: str, **kwargs: Mapping[str, Any]) -> Optional[List[str]]:
+    kwargs_dict = dict(kwargs)
+    target_dir = cast(Optional[str], kwargs_dict.pop("target_dir", None))
+    output: bool = bool(kwargs_dict.pop("output", False))
+    op_args: Mapping[str, Any] = {**kwargs_dict}
+    return perform_diff(tmpdir, target_dir=target_dir, output=output, **op_args)
+
+
+def op_pull(tmpdir: str, dst: str, **op_args: Mapping[str, Any]) -> None:
+    if os.path.exists(dst):
+        echo(f"❌  Directory *{dst}* already exists")
+    else:
+        shutil.move(tmpdir, dst)
+        echo(f"Code downloaded to *{dst}*")
+
+
+def op_patch(tmpdir: str, dst: str, **kwargs: Mapping[str, Any]) -> None:
+    diffs = perform_diff(tmpdir, output=True) or []
+    with open(dst, "w", encoding="utf-8") as f:
+        for out in diffs:
+            out = out.replace(tmpdir, "/.")
+            out = out.replace("+++ b/./", "+++ b/")
+            out = out.replace("--- b/./", "--- b/")
+            out = out.replace("--- a/./", "--- a/")
+            out = out.replace("+++ a/./", "+++ a/")
+            f.write(out)
+    echo(f"Patch saved in *{dst}*")
+    path = run(
+        ["git", "rev-parse", "--show-prefix"], text=True, stdout=PIPE
+    ).stdout.strip()
+    if path:
+        diropt = f" --directory={path.rstrip('/')}"
+    else:
+        diropt = ""
+    echo("Apply the patch by running:")
+    echo_always(
+        f"git apply --verbose{diropt} {dst}", highlight=True, bold=True, err=True
+    )
+
+
+@code.command()
+@click.argument("run_pathspec")
+def diff(run_pathspec: str, **kwargs: Mapping[str, Any]) -> None:
+    """
+    Do a 'git diff' of the current directory and a Metaflow run.
+    """
+    _ = run_op(run_pathspec, op_diff, **kwargs)
+
+
+@code.command()
+@click.argument("source_run_pathspec")
+@click.argument("target_run_pathspec")
+def diff_runs(
+    source_run_pathspec: str, target_run_pathspec: str, **kwargs: Mapping[str, Any]
+) -> None:
+    """
+    Do a 'git diff' between two Metaflow runs.
+    """
+    _ = run_op_diff_runs(source_run_pathspec, target_run_pathspec, **kwargs)
+
+
+@code.command()
+@click.argument("run_pathspec")
+@click.option(
+    "--dir", help="Destination directory (default: {run_pathspec}_code)", default=None
+)
+def pull(
+    run_pathspec: str, dir: Optional[str] = None, **kwargs: Mapping[str, Any]
+) -> None:
+    """
+    Pull the code of a Metaflow run.
+    """
+    if dir is None:
+        dir = run_pathspec.lower().replace("/", "_") + "_code"
+    op_args: Mapping[str, Any] = {**kwargs, "dst": dir}
+    run_op(run_pathspec, op_pull, **op_args)
+
+
+@code.command()
+@click.argument("run_pathspec")
+@click.option(
+    "--file_path",
+    help="Patch file name. If not provided, defaults to a sanitized version of RUN_PATHSPEC "
+    "with slashes replaced by underscores, plus '.patch'.",
+    show_default=False,
+)
+@click.option(
+    "--overwrite", is_flag=True, help="Overwrite the patch file if it exists."
+)
+def patch(
+    run_pathspec: str,
+    file_path: Optional[str] = None,
+    overwrite: bool = False,
+    **kwargs: Mapping[str, Any],
+) -> None:
+    """
+    Create a patch by comparing current dir with a Metaflow run.
+    """
+    if file_path is None:
+        file_path = run_pathspec.lower().replace("/", "_") + ".patch"
+    if os.path.exists(file_path) and not overwrite:
+        echo(f"File *{file_path}* already exists. To overwrite, specify --overwrite.")
+        return
+    op_args: Mapping[str, Any] = {**kwargs, "dst": file_path}
+    run_op(run_pathspec, op_patch, **op_args)