ob-metaflow 2.15.0.1__py2.py3-none-any.whl → 2.15.5.1__py2.py3-none-any.whl

metaflow/__init__.py

@@ -119,6 +119,12 @@ from .includefile import IncludeFile
 # Decorators
 from .decorators import step, _import_plugin_decorators
 
+
+# Parsers (for configs) for now
+from .plugins import _import_tl_plugins
+
+_import_tl_plugins(globals())
+
 # this auto-generates decorator functions from Decorator objects
 # in the top-level metaflow namespace
 _import_plugin_decorators(globals())

metaflow/cmd/code/__init__.py

@@ -0,0 +1,230 @@
+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._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, dirnames, filenames in os.walk(source_dir, followlinks=True):
+        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)

metaflow/cmd/develop/stub_generator.py

@@ -1133,13 +1133,16 @@ class StubGenerator:
             result = result[1:]
         # Add doc to first and last overloads. Jedi uses the last one and pycharm
         # the first one. Go figure.
+        result_docstring = docs["func_doc"]
+        if docs["param_doc"]:
+            result_docstring += "\nParameters\n----------\n" + docs["param_doc"]
         result[0] = (
             result[0][0],
-            docs["func_doc"] + "\nParameters\n----------\n" + docs["param_doc"],
+            result_docstring,
         )
         result[-1] = (
             result[-1][0],
-            docs["func_doc"] + "\nParameters\n----------\n" + docs["param_doc"],
+            result_docstring,
         )
         return result
 

metaflow/cmd/main_cli.py

@@ -67,6 +67,7 @@ CMDS_DESC = [
     ("configure", ".configure_cmd.cli"),
     ("tutorials", ".tutorials_cmd.cli"),
     ("develop", ".develop.cli"),
+    ("code", ".code.cli"),
 ]
 
 process_cmds(globals())

metaflow/cmd/make_wrapper.py

@@ -2,15 +2,47 @@ import sys
 import subprocess
 from pathlib import Path
 import sysconfig
+import site
+
+
+def find_makefile():
+    possible_dirs = []
+
+    # 1) The standard sysconfig-based location
+    data_dir = sysconfig.get_paths()["data"]
+    possible_dirs.append(Path(data_dir) / "share" / "metaflow" / "devtools")
+
+    # 2) The user base (e.g. ~/.local on many systems)
+    user_base = site.getuserbase()  # e.g. /home/runner/.local
+    possible_dirs.append(Path(user_base) / "share" / "metaflow" / "devtools")
+
+    # 3) site-packages can vary, we can guess share/.. near each site-packages
+    # (Works if pip actually placed devtools near site-packages.)
+    for p in site.getsitepackages():
+        possible_dirs.append(Path(p).parent / "share" / "metaflow" / "devtools")
+    user_site = site.getusersitepackages()
+    possible_dirs.append(Path(user_site).parent / "share" / "metaflow" / "devtools")
+
+    for candidate_dir in possible_dirs:
+        makefile_candidate = candidate_dir / "Makefile"
+        if makefile_candidate.is_file():
+            return makefile_candidate
+
+    return None
 
 
 def main():
-    share_dir = Path(sysconfig.get_paths()["data"]) / "share" / "metaflow" / "devtools"
-    makefile_path = share_dir / "Makefile"
+    makefile_path = find_makefile()
+    if not makefile_path:
+        print("ERROR: Could not find executable in any known location.")
+        sys.exit(1)
     cmd = ["make", "-f", str(makefile_path)] + sys.argv[1:]
-    # subprocess.run(cmd, check=True)
+
     try:
         completed = subprocess.run(cmd, check=True)
         sys.exit(completed.returncode)
     except subprocess.CalledProcessError as ex:
         sys.exit(ex.returncode)
+    except KeyboardInterrupt:
+        print("Process interrupted by user. Exiting cleanly.")
+        sys.exit(1)

metaflow/extension_support/plugins.py

@@ -198,6 +198,7 @@ _plugin_categories = {
         list(x.commands)[0] if len(x.commands) == 1 else "too many commands"
     ),
     "runner_cli": lambda x: x.name,
+    "tl_plugin": None,
 }
 
 

metaflow/metaflow_config.py

@@ -382,6 +382,8 @@ KUBERNETES_DISK = from_conf("KUBERNETES_DISK", None)
 # Default kubernetes QoS class
 KUBERNETES_QOS = from_conf("KUBERNETES_QOS", "burstable")
 
+# Architecture of kubernetes nodes - used for @conda/@pypi in metaflow-dev
+KUBERNETES_CONDA_ARCH = from_conf("KUBERNETES_CONDA_ARCH")
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
 

metaflow/metaflow_environment.py

@@ -6,7 +6,7 @@ from .util import get_username
 from . import metaflow_version
 from metaflow.exception import MetaflowException
 from metaflow.extension_support import dump_module_info
-from metaflow.mflog import BASH_MFLOG
+from metaflow.mflog import BASH_MFLOG, BASH_FLUSH_LOGS
 from . import R
 
 
@@ -159,6 +159,7 @@ class MetaflowEnvironment(object):
     def get_package_commands(self, code_package_url, datastore_type):
         cmds = [
             BASH_MFLOG,
+            BASH_FLUSH_LOGS,
             "mflog 'Setting up task environment.'",
             self._get_install_dependencies_cmd(datastore_type),
             "mkdir metaflow",
@@ -176,6 +177,7 @@ class MetaflowEnvironment(object):
             "fi" % code_package_url,
             "TAR_OPTIONS='--warning=no-timestamp' tar xf job.tar",
             "mflog 'Task is starting.'",
+            "flush_mflogs",
         ]
         return cmds
 

metaflow/mflog/__init__.py

@@ -44,6 +44,8 @@ BASH_MFLOG = (
 BASH_SAVE_LOGS_ARGS = ["python", "-m", "metaflow.mflog.save_logs"]
 BASH_SAVE_LOGS = " ".join(BASH_SAVE_LOGS_ARGS)
 
+BASH_FLUSH_LOGS = "flush_mflogs(){ " f"{BASH_SAVE_LOGS}; " "}"
+
 
 # this function returns a bash expression that redirects stdout
 # and stderr of the given bash expression to mflog.tee
@@ -63,7 +65,7 @@ def bash_capture_logs(bash_expr, var_transform=None):
 # update_delay determines how often logs should be uploaded to S3
 # as a function of the task execution time
 
-MIN_UPDATE_DELAY = 1.0  # the most frequent update interval
+MIN_UPDATE_DELAY = 0.25  # the most frequent update interval
 MAX_UPDATE_DELAY = 30.0  # the least frequent update interval
 
 
@@ -110,7 +112,6 @@ def export_mflog_env_vars(
 
 def tail_logs(prefix, stdout_tail, stderr_tail, echo, has_log_updates):
     def _available_logs(tail, stream, echo, should_persist=False):
-        # print the latest batch of lines
         try:
             for line in tail:
                 if should_persist:
@@ -128,7 +129,7 @@ def tail_logs(prefix, stdout_tail, stderr_tail, echo, has_log_updates):
 
     start_time = time.time()
     next_log_update = start_time
-    log_update_delay = 1
+    log_update_delay = update_delay(0)
     while has_log_updates():
         if time.time() > next_log_update:
             _available_logs(stdout_tail, "stdout", echo)

metaflow/plugins/__init__.py

@@ -167,6 +167,12 @@ DEPLOYER_IMPL_PROVIDERS_DESC = [
     ),
 ]
 
+TL_PLUGINS_DESC = [
+    ("requirements_txt_parser", ".pypi.parsers.requirements_txt_parser"),
+    ("pyproject_toml_parser", ".pypi.parsers.pyproject_toml_parser"),
+    ("conda_environment_yml_parser", ".pypi.parsers.conda_environment_yml_parser"),
+]
+
 process_plugins(globals())
 
 
@@ -207,6 +213,8 @@ GCP_CLIENT_PROVIDERS = resolve_plugins("gcp_client_provider")
 if sys.version_info >= (3, 7):
     DEPLOYER_IMPL_PROVIDERS = resolve_plugins("deployer_impl_provider")
 
+TL_PLUGINS = resolve_plugins("tl_plugin")
+
 from .cards.card_modules import MF_EXTERNAL_CARDS
 
 # Cards; due to the way cards were designed, it is harder to make them fit
@@ -251,3 +259,9 @@ CARDS = [
     TestRefreshComponentCard,
 ]
 merge_lists(CARDS, MF_EXTERNAL_CARDS, "type")
+
+
+def _import_tl_plugins(globals_dict):
+
+    for name, p in TL_PLUGINS.items():
+        globals_dict[name] = p

metaflow/plugins/argo/argo_client.py

@@ -256,12 +256,19 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def trigger_workflow_template(self, name, parameters={}):
+    def trigger_workflow_template(self, name, usertype, username, parameters={}):
         client = self._client.get()
         body = {
             "apiVersion": "argoproj.io/v1alpha1",
             "kind": "Workflow",
-            "metadata": {"generateName": name + "-"},
+            "metadata": {
+                "generateName": name + "-",
+                "annotations": {
+                    "metaflow/triggered_by_user": json.dumps(
+                        {"type": usertype, "name": username}
+                    )
+                },
+            },
             "spec": {
                 "workflowTemplateRef": {"name": name},
                 "arguments": {

metaflow/plugins/argo/argo_workflows.py

@@ -65,6 +65,7 @@ from metaflow.util import (
 )
 
 from .argo_client import ArgoClient
+from metaflow.util import resolve_identity
 
 
 class ArgoWorkflowsException(MetaflowException):
@@ -109,8 +110,7 @@ class ArgoWorkflows(object):
         notify_slack_webhook_url=None,
         notify_pager_duty_integration_key=None,
         notify_incident_io_api_key=None,
-        incident_io_success_severity_id=None,
-        incident_io_error_severity_id=None,
+        incident_io_alert_source_config_id=None,
         enable_heartbeat_daemon=True,
         enable_error_msg_capture=False,
     ):
@@ -161,8 +161,7 @@ class ArgoWorkflows(object):
         self.notify_slack_webhook_url = notify_slack_webhook_url
         self.notify_pager_duty_integration_key = notify_pager_duty_integration_key
         self.notify_incident_io_api_key = notify_incident_io_api_key
-        self.incident_io_success_severity_id = incident_io_success_severity_id
-        self.incident_io_error_severity_id = incident_io_error_severity_id
+        self.incident_io_alert_source_config_id = incident_io_alert_source_config_id
         self.enable_heartbeat_daemon = enable_heartbeat_daemon
         self.enable_error_msg_capture = enable_error_msg_capture
         self.parameters = self._process_parameters()
@@ -316,8 +315,16 @@ class ArgoWorkflows(object):
                     "Workflows before proceeding." % name
                 )
         try:
+            id_parts = resolve_identity().split(":")
+            parts_size = len(id_parts)
+            usertype = id_parts[0] if parts_size > 0 else "unknown"
+            username = id_parts[1] if parts_size > 1 else "unknown"
+
             return ArgoClient(namespace=KUBERNETES_NAMESPACE).trigger_workflow_template(
-                name, parameters
+                name,
+                usertype,
+                username,
+                parameters,
             )
         except Exception as e:
             raise ArgoWorkflowsException(str(e))
@@ -2528,25 +2535,49 @@ class ArgoWorkflows(object):
     def _incident_io_alert_template(self):
         if self.notify_incident_io_api_key is None:
             return None
-        if self.incident_io_error_severity_id is None:
+        if self.incident_io_alert_source_config_id is None:
             raise MetaflowException(
-                "Creating incidents for errors requires a severity id."
+                "Creating alerts for errors requires a alert source config ID."
             )
+        ui_links = self._incident_io_ui_urls_for_run()
         return Template("notify-incident-io-on-error").http(
             Http("POST")
-            .url("https://api.incident.io/v2/incidents")
+            .url(
+                "https://api.incident.io/v2/alert_events/http/%s"
+                % self.incident_io_alert_source_config_id
+            )
             .header("Content-Type", "application/json")
             .header("Authorization", "Bearer %s" % self.notify_incident_io_api_key)
             .body(
                 json.dumps(
                     {
                         "idempotency_key": "argo-{{workflow.name}}",  # use run id to deduplicate alerts.
-                        "visibility": "public",
-                        "severity_id": self.incident_io_error_severity_id,
-                        "name": "Flow %s has failed." % self.flow.name,
-                        "summary": "Metaflow run %s/argo-{{workflow.name}} failed! %s"
-                        % (self.flow.name, self._incident_io_ui_urls_for_run()),
-                        # TODO: Add support for custom field entries.
+                        "status": "firing",
+                        "title": "Flow %s has failed." % self.flow.name,
+                        "description": "Metaflow run {run_pathspec} failed!{urls}".format(
+                            run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
+                            urls=(
+                                "\n\nSee details for the run at:\n\n"
+                                + "\n\n".join(ui_links)
+                                if ui_links
+                                else ""
+                            ),
+                        ),
+                        "source_url": (
+                            "%s/%s/%s"
+                            % (
+                                UI_URL.rstrip("/"),
+                                self.flow.name,
+                                "argo-{{workflow.name}}",
+                            )
+                            if UI_URL
+                            else None
+                        ),
+                        "metadata": {
+                            "run_status": "failed",
+                            "flow_name": self.flow.name,
+                            "run_id": "argo-{{workflow.name}}",
+                        },
                     }
                 )
             )
@@ -2555,27 +2586,49 @@ class ArgoWorkflows(object):
     def _incident_io_change_template(self):
         if self.notify_incident_io_api_key is None:
             return None
-        if self.incident_io_success_severity_id is None:
+        if self.incident_io_alert_source_config_id is None:
             raise MetaflowException(
-                "Creating incidents for successes requires a severity id."
+                "Creating alerts for successes requires an alert source config ID."
             )
+        ui_links = self._incident_io_ui_urls_for_run()
         return Template("notify-incident-io-on-success").http(
             Http("POST")
-            .url("https://api.incident.io/v2/incidents")
+            .url(
+                "https://api.incident.io/v2/alert_events/http/%s"
+                % self.incident_io_alert_source_config_id
+            )
             .header("Content-Type", "application/json")
             .header("Authorization", "Bearer %s" % self.notify_incident_io_api_key)
             .body(
                 json.dumps(
                     {
                         "idempotency_key": "argo-{{workflow.name}}",  # use run id to deduplicate alerts.
-                        "visibility": "public",
-                        "severity_id": self.incident_io_success_severity_id,
-                        # TODO: Do we need to make incident type configurable for successes? otherwise they are created as 'investigating'
-                        # "incident_type_id": ""
-                        "name": "Flow %s has succeeded." % self.flow.name,
-                        "summary": "Metaflow run %s/argo-{{workflow.name}} succeeded!%s"
-                        % (self.flow.name, self._incident_io_ui_urls_for_run()),
-                        # TODO: Add support for custom field entries.
+                        "status": "firing",
+                        "title": "Flow %s has succeeded." % self.flow.name,
+                        "description": "Metaflow run {run_pathspec} succeeded!{urls}".format(
+                            run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
+                            urls=(
+                                "\n\nSee details for the run at:\n\n"
+                                + "\n\n".join(ui_links)
+                                if ui_links
+                                else ""
+                            ),
+                        ),
+                        "source_url": (
+                            "%s/%s/%s"
+                            % (
+                                UI_URL.rstrip("/"),
+                                self.flow.name,
+                                "argo-{{workflow.name}}",
+                            )
+                            if UI_URL
+                            else None
+                        ),
+                        "metadata": {
+                            "run_status": "succeeded",
+                            "flow_name": self.flow.name,
+                            "run_id": "argo-{{workflow.name}}",
+                        },
                     }
                 )
             )
@@ -2597,9 +2650,7 @@ class ArgoWorkflows(object):
                 "{{workflow.name}}",
             )
             links.append(url)
-        if links:
-            links = ["See details for the run at: ", *links]
-        return "\n\n".join(links)
+        return links
 
     def _pager_duty_change_template(self):
         # https://developer.pagerduty.com/docs/ZG9jOjExMDI5NTgy-send-a-change-event