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

metaflow/plugins/kubernetes/spot_metadata_cli.py

@@ -0,0 +1,69 @@
+from metaflow._vendor import click
+from datetime import datetime, timezone
+from metaflow.tagging_util import validate_tags
+from metaflow.metadata_provider import MetaDatum
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Commands related to spot metadata.")
+def spot_metadata():
+    pass
+
+
+@spot_metadata.command(help="Record spot termination metadata for a task.")
+@click.option(
+    "--run-id",
+    required=True,
+    help="Run ID for which metadata is to be recorded.",
+)
+@click.option(
+    "--step-name",
+    required=True,
+    help="Step Name for which metadata is to be recorded.",
+)
+@click.option(
+    "--task-id",
+    required=True,
+    help="Task ID for which metadata is to be recorded.",
+)
+@click.option(
+    "--termination-notice-time",
+    required=True,
+    help="Spot termination notice time.",
+)
+@click.option(
+    "--tag",
+    "tags",
+    multiple=True,
+    required=False,
+    default=None,
+    help="List of tags.",
+)
+@click.pass_obj
+def record(obj, run_id, step_name, task_id, termination_notice_time, tags=None):
+    validate_tags(tags)
+
+    tag_list = list(tags) if tags else []
+
+    entries = [
+        MetaDatum(
+            field="spot-termination-received-at",
+            value=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+            type="spot-termination-received-at",
+            tags=tag_list,
+        ),
+        MetaDatum(
+            field="spot-termination-time",
+            value=termination_notice_time,
+            type="spot-termination-time",
+            tags=tag_list,
+        ),
+    ]
+
+    obj.metadata.register_metadata(
+        run_id=run_id, step_name=step_name, task_id=task_id, metadata=entries
+    )

metaflow/plugins/kubernetes/spot_monitor_sidecar.py

@@ -0,0 +1,109 @@
+import os
+import sys
+import time
+import signal
+import requests
+import subprocess
+from multiprocessing import Process
+from datetime import datetime, timezone
+from metaflow.sidecar import MessageTypes
+
+
+class SpotTerminationMonitorSidecar(object):
+    EC2_TYPE_URL = "http://169.254.169.254/latest/meta-data/instance-life-cycle"
+    METADATA_URL = "http://169.254.169.254/latest/meta-data/spot/termination-time"
+    TOKEN_URL = "http://169.254.169.254/latest/api/token"
+    POLL_INTERVAL = 5  # seconds
+
+    def __init__(self):
+        self.is_alive = True
+        self._process = None
+        self._token = None
+        self._token_expiry = 0
+
+        if self._is_aws_spot_instance():
+            self._process = Process(target=self._monitor_loop)
+            self._process.start()
+
+    def process_message(self, msg):
+        if msg.msg_type == MessageTypes.SHUTDOWN:
+            self.is_alive = False
+            if self._process:
+                self._process.terminate()
+
+    @classmethod
+    def get_worker(cls):
+        return cls
+
+    def _get_imds_token(self):
+        current_time = time.time()
+        if current_time >= self._token_expiry - 60:  # Refresh 60s before expiry
+            try:
+                response = requests.put(
+                    url=self.TOKEN_URL,
+                    headers={"X-aws-ec2-metadata-token-ttl-seconds": "300"},
+                    timeout=1,
+                )
+                if response.status_code == 200:
+                    self._token = response.text
+                    self._token_expiry = current_time + 240  # Slightly less than TTL
+            except requests.exceptions.RequestException:
+                pass
+        return self._token
+
+    def _make_ec2_request(self, url, timeout):
+        token = self._get_imds_token()
+        headers = {"X-aws-ec2-metadata-token": token} if token else {}
+        response = requests.get(url=url, headers=headers, timeout=timeout)
+        return response
+
+    def _is_aws_spot_instance(self):
+        try:
+            response = self._make_ec2_request(url=self.EC2_TYPE_URL, timeout=1)
+            return response.status_code == 200 and response.text == "spot"
+        except (requests.exceptions.RequestException, requests.exceptions.Timeout):
+            return False
+
+    def _monitor_loop(self):
+        while self.is_alive:
+            try:
+                response = self._make_ec2_request(url=self.METADATA_URL, timeout=1)
+                if response.status_code == 200:
+                    termination_time = response.text
+                    self._emit_termination_metadata(termination_time)
+                    os.kill(os.getppid(), signal.SIGTERM)
+                    break
+            except (requests.exceptions.RequestException, requests.exceptions.Timeout):
+                pass
+            time.sleep(self.POLL_INTERVAL)
+
+    def _emit_termination_metadata(self, termination_time):
+        flow_filename = os.getenv("METAFLOW_FLOW_FILENAME")
+        pathspec = os.getenv("MF_PATHSPEC")
+        _, run_id, step_name, task_id = pathspec.split("/")
+        retry_count = os.getenv("MF_ATTEMPT")
+
+        with open("/tmp/spot_termination_notice", "w") as fp:
+            fp.write(termination_time)
+
+        command = [
+            sys.executable,
+            f"/metaflow/{flow_filename}",
+            "spot-metadata",
+            "record",
+            "--run-id",
+            run_id,
+            "--step-name",
+            step_name,
+            "--task-id",
+            task_id,
+            "--termination-notice-time",
+            termination_time,
+            "--tag",
+            "attempt_id:{}".format(retry_count),
+        ]
+
+        result = subprocess.run(command, capture_output=True, text=True)
+
+        if result.returncode != 0:
+            print(f"Failed to record spot termination metadata: {result.stderr}")

metaflow/plugins/logs_cli.py

@@ -0,0 +1,359 @@
+from metaflow._vendor import click
+from metaflow.cli import LOGGER_TIMESTAMP
+
+from ..exception import CommandException
+from ..datastore import TaskDataStoreSet, TaskDataStore
+
+
+from ..mflog import mflog, LOG_SOURCES
+
+
+# main motivation from https://github.com/pallets/click/issues/430
+# in order to support a default command being called for a Click group.
+#
+# NOTE: We need this in order to not introduce breaking changes to existing CLI, as we wanted to
+# nest both existing `logs` and the new `logs scrub` under a shared group, but `logs` already has
+# a well defined behavior of showing the logs.
+class CustomGroup(click.Group):
+    def __init__(self, name=None, commands=None, default_cmd=None, **attrs):
+        super(CustomGroup, self).__init__(name, commands, **attrs)
+        self.default_cmd = default_cmd
+
+    def get_command(self, ctx, cmd_name):
+        if cmd_name not in self.list_commands(ctx):
+            # input from the CLI does not match a command, so we pass that
+            # as the args to the default command instead.
+            ctx.passed_cmd = cmd_name
+            cmd_name = self.default_cmd
+        return super(CustomGroup, self).get_command(ctx, cmd_name)
+
+    def parse_args(self, ctx, args):
+        # We first try to parse args as is, to determine whether we need to fall back to the default commmand
+        # if any options are supplied, the parse will fail, as the group does not support the options.
+        # In this case we fallback to the default command, inserting that as the first arg and parsing again.
+        # copy args as trying to parse will destroy them.
+        original_args = list(args)
+        try:
+            super().parse_args(ctx, args)
+            args_parseable = True
+        except Exception:
+            args_parseable = False
+        if not args or not args_parseable:
+            original_args.insert(0, self.default_cmd)
+        return super().parse_args(ctx, original_args)
+
+    def resolve_command(self, ctx, args):
+        cmd_name, cmd_obj, args = super(CustomGroup, self).resolve_command(ctx, args)
+        passed_cmd = getattr(ctx, "passed_cmd", None)
+        if passed_cmd is not None:
+            args.insert(0, passed_cmd)
+
+        return cmd_name, cmd_obj, args
+
+    def format_commands(self, ctx, formatter):
+        formatter = CustomFormatter(self.default_cmd, formatter)
+        return super(CustomGroup, self).format_commands(ctx, formatter)
+
+
+class CustomFormatter:
+    def __init__(self, default_cmd, original_formatter) -> None:
+        self.default_cmd = default_cmd
+        self.formatter = original_formatter
+
+    def __getattr__(self, name):
+        return getattr(self.formatter, name)
+
+    def write_dl(self, rows):
+        def _format(dup):
+            cmd, help = dup
+            if cmd == self.default_cmd:
+                cmd = cmd + " [Default]"
+            return (cmd, help)
+
+        rows = [_format(dup) for dup in rows]
+
+        return self.formatter.write_dl(rows)
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(cls=CustomGroup, help="Commands related to logs", default_cmd="show")
+@click.pass_context
+def logs(ctx):
+    # the logger is configured in cli.py
+    global echo
+    echo = ctx.obj.echo
+
+
+@logs.command(
+    help="Show stdout/stderr produced by a task or all tasks in a step. "
+    "The format for input-path is either <run_id>/<step_name> or "
+    "<run_id>/<step_name>/<task_id>."
+)
+@click.argument("input-path")
+@click.option(
+    "--stdout/--no-stdout",
+    default=False,
+    show_default=True,
+    help="Show stdout of the task.",
+)
+@click.option(
+    "--stderr/--no-stderr",
+    default=False,
+    show_default=True,
+    help="Show stderr of the task.",
+)
+@click.option(
+    "--both/--no-both",
+    default=True,
+    show_default=True,
+    help="Show both stdout and stderr of the task.",
+)
+@click.option(
+    "--timestamps/--no-timestamps",
+    default=False,
+    show_default=True,
+    help="Show timestamps.",
+)
+@click.option(
+    "--attempt",
+    default=None,
+    type=int,
+    show_default=False,
+    help="Attempt number of a task to show, defaults to the latest attempt.",
+)
+@click.pass_obj
+def show(
+    obj, input_path, stdout=None, stderr=None, both=None, timestamps=False, attempt=None
+):
+    types = set()
+    if stdout:
+        types.add("stdout")
+        both = False
+    if stderr:
+        types.add("stderr")
+        both = False
+    if both:
+        types.update(("stdout", "stderr"))
+
+    streams = list(sorted(types, reverse=True))
+
+    # Pathspec can either be run_id/step_name or run_id/step_name/task_id.
+    parts = input_path.split("/")
+    if len(parts) == 2:
+        run_id, step_name = parts
+        task_id = None
+    elif len(parts) == 3:
+        run_id, step_name, task_id = parts
+    else:
+        raise CommandException(
+            "input_path should either be run_id/step_name "
+            "or run_id/step_name/task_id"
+        )
+
+    datastore_set = TaskDataStoreSet(
+        obj.flow_datastore, run_id, steps=[step_name], allow_not_done=True
+    )
+    if task_id:
+        ds_list = [
+            TaskDataStore(
+                obj.flow_datastore,
+                run_id=run_id,
+                step_name=step_name,
+                task_id=task_id,
+                mode="r",
+                allow_not_done=True,
+            )
+        ]
+    else:
+        ds_list = list(datastore_set)  # get all tasks
+
+    if ds_list:
+
+        def echo_unicode(line, **kwargs):
+            click.secho(line.decode("UTF-8", errors="replace"), **kwargs)
+
+        # old style logs are non mflog-style logs
+        maybe_old_style = True
+        for ds in ds_list:
+            echo(
+                "Dumping logs of run_id=*{run_id}* "
+                "step=*{step}* task_id=*{task_id}*".format(
+                    run_id=ds.run_id, step=ds.step_name, task_id=ds.task_id
+                ),
+                fg="magenta",
+            )
+
+            for stream in streams:
+                echo(stream, bold=True)
+                logs = ds.load_logs(LOG_SOURCES, stream, attempt_override=attempt)
+                if any(data for _, data in logs):
+                    # attempt to read new, mflog-style logs
+                    for line in mflog.merge_logs([blob for _, blob in logs]):
+                        if timestamps:
+                            ts = mflog.utc_to_local(line.utc_tstamp)
+                            tstamp = ts.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+                            click.secho(tstamp + " ", fg=LOGGER_TIMESTAMP, nl=False)
+                        echo_unicode(line.msg)
+                    maybe_old_style = False
+                elif maybe_old_style:
+                    # if they are not available, we may be looking at
+                    # a legacy run (unless we have seen new-style data already
+                    # for another stream). This return an empty string if
+                    # nothing is found
+                    log = ds.load_log_legacy(stream, attempt_override=attempt)
+                    if log and timestamps:
+                        raise CommandException(
+                            "We can't show --timestamps for old runs. Sorry!"
+                        )
+                    echo_unicode(log, nl=False)
+    else:
+        raise CommandException(
+            "No Tasks found at the given path -- "
+            "either none exist or none have started yet"
+        )
+
+
+@logs.command(
+    help="Scrub stdout/stderr produced by a task or all tasks in a step. "
+    "The format for input-path is either <run_id>/<step_name> or "
+    "<run_id>/<step_name>/<task_id>."
+)
+@click.argument("input-path")
+@click.option(
+    "--stdout/--no-stdout",
+    default=False,
+    show_default=True,
+    help="Scrub stdout of the step or task.",
+)
+@click.option(
+    "--stderr/--no-stderr",
+    default=False,
+    show_default=True,
+    help="Scrub stderr of the step or task.",
+)
+@click.option(
+    "--both/--no-both",
+    default=True,
+    show_default=True,
+    help="Scrub both stdout and stderr of the step or task.",
+)
+@click.option(
+    "--attempt",
+    default=None,
+    type=int,
+    show_default=False,
+    help="Attempt number of a task to scrub, defaults to the latest attempt.",
+)
+@click.option(
+    "--latest/--all",
+    default=True,
+    show_default=False,
+    help="Scrub latest/all attempts of a step or task",
+)
+@click.option(
+    "--include-not-done",
+    default=False,
+    show_default=False,
+    is_flag=True,
+    help="Also scrub steps or tasks that are not done. Use this for tasks that did not finish correctly, and could not otherwise be scrubbed.",
+)
+@click.pass_obj
+def scrub(
+    obj,
+    input_path,
+    stdout=None,
+    stderr=None,
+    both=None,
+    attempt=None,
+    latest=None,
+    include_not_done=None,
+):
+    types = set()
+    if stdout:
+        types.add("stdout")
+        both = False
+    if stderr:
+        types.add("stderr")
+        both = False
+    if both:
+        types.update(("stdout", "stderr"))
+
+    streams = list(sorted(types, reverse=True))
+
+    # Pathspec can either be run_id/step_name or run_id/step_name/task_id.
+    parts = input_path.split("/")
+    if len(parts) == 2:
+        run_id, step_name = parts
+        task_id = None
+    elif len(parts) == 3:
+        run_id, step_name, task_id = parts
+    else:
+        raise CommandException(
+            "input_path should either be run_id/step_name "
+            "or run_id/step_name/task_id"
+        )
+
+    if task_id:
+        if latest:
+            ds_list = obj.flow_datastore.get_task_datastores(
+                pathspecs=[input_path],
+                attempt=attempt,
+                mode="d",
+                allow_not_done=include_not_done,
+            )
+        else:
+            ds_list = obj.flow_datastore.get_task_datastores(
+                pathspecs=[input_path],
+                attempt=attempt,
+                mode="d",
+                allow_not_done=include_not_done,
+                include_prior=True,
+            )
+    else:
+        if latest:
+            ds_list = obj.flow_datastore.get_task_datastores(
+                run_id=run_id,
+                steps=[step_name],
+                attempt=attempt,
+                mode="d",
+                allow_not_done=include_not_done,
+            )
+        else:
+            ds_list = obj.flow_datastore.get_task_datastores(
+                run_id=run_id,
+                steps=[step_name],
+                attempt=attempt,
+                mode="d",
+                allow_not_done=include_not_done,
+                include_prior=True,
+            )
+
+    if ds_list:
+        for ds in ds_list:
+            failures = []
+            for stream in streams:
+                try:
+                    ds.scrub_logs(LOG_SOURCES, stream)
+                except Exception:
+                    failures.append(stream)
+            if failures:
+                obj.echo_always(
+                    "Failed to scrub %s - attempt %s : *%s*"
+                    % (ds.pathspec, ds.attempt, ",".join(failures))
+                )
+            else:
+                echo(
+                    "Logs have been scrubbed for %s - attempt %s"
+                    % (ds.pathspec, ds.attempt)
+                )
+
+    else:
+        raise CommandException(
+            "No Tasks found at the given path -- "
+            "either none exist or they have not finished yet.\n"
+            "If you know the task has finished, you can supply --include-not-done to force scrub it."
+        )