metaflow-checkpoint 0.0.1__tar.gz

metaflow_checkpoint-0.0.1/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.1
+Name: metaflow-checkpoint
+Version: 0.0.1
+Summary: An EXPERIMENTAL checkpoint decorator for Metaflow
+Author: Valay Dave
+Author-email: help@outerbounds.com
+Description-Content-Type: text/markdown
+
+# Metaflow Checkpoint
+
+TODO : Write readme

metaflow_checkpoint-0.0.1/README.md

@@ -0,0 +1,3 @@
+# Metaflow Checkpoint
+
+TODO : Write readme

metaflow_checkpoint-0.0.1/metaflow_checkpoint.egg-info/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.1
+Name: metaflow-checkpoint
+Version: 0.0.1
+Summary: An EXPERIMENTAL checkpoint decorator for Metaflow
+Author: Valay Dave
+Author-email: help@outerbounds.com
+Description-Content-Type: text/markdown
+
+# Metaflow Checkpoint
+
+TODO : Write readme

metaflow_checkpoint-0.0.1/metaflow_checkpoint.egg-info/SOURCES.txt

@@ -0,0 +1,49 @@
+README.md
+setup.py
+metaflow_checkpoint.egg-info/PKG-INFO
+metaflow_checkpoint.egg-info/SOURCES.txt
+metaflow_checkpoint.egg-info/dependency_links.txt
+metaflow_checkpoint.egg-info/top_level.txt
+metaflow_extensions/obcheckpoint/plugins/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constants.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/task_utils.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/diskspace_usage_card.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/storage_injections.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/decorator.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/cards/model_card.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/constants.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/disk_usage.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.py
+metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.py
+metaflow_extensions/obcheckpoint/toplevel/__init__.py
+metaflow_extensions/obcheckpoint/toplevel/toplevel.py

metaflow_checkpoint-0.0.1/metaflow_checkpoint.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

metaflow_checkpoint-0.0.1/metaflow_checkpoint.egg-info/top_level.txt

@@ -0,0 +1 @@
+metaflow_extensions

metaflow_checkpoint-0.0.1/metaflow_extensions/obcheckpoint/plugins/__init__.py

@@ -0,0 +1,24 @@
+STEP_DECORATORS_DESC = [
+    (
+        "checkpoint",
+        ".machine_learning_utilities.checkpoints.decorator.CheckpointDecorator",
+    ),
+    (
+        "model",
+        ".machine_learning_utilities.modeling_utils.decorator.ModelDecorator",
+    ),
+    (
+        "huggingface_hub",
+        ".machine_learning_utilities.hf_hub.decorator.HuggingfaceHubDecorator",
+    ),
+    (
+        "disk_profiler",
+        ".machine_learning_utilities.checkpoints.cards.diskspace_usage_card.DiskUsageProfilerDecorator",
+    ),
+]
+
+###
+# CONFIGURE: Similar to datatools, you can make visible under metaflow.plugins.* other
+#            submodules not referenced in this file
+###
+__mf_promote_submodules__ = []

metaflow_checkpoint-0.0.1/metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.py

@@ -0,0 +1,2 @@
+from .datastructures import CheckpointArtifact
+from .checkpoints.core import Checkpointer

metaflow_checkpoint-0.0.1/metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.py

@@ -0,0 +1,3 @@
+from .deco_injection_mixin import CardDecoratorInjector
+from .async_cards import CardRefresher, AsyncPeriodicRefresher
+from .extra_components import LineChart, UpadateableTable

metaflow_checkpoint-0.0.1/metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.py

@@ -0,0 +1,135 @@
+from metaflow.metaflow_current import current
+import json
+from threading import Thread, Event
+import time
+
+
+class InfoCollectorThread(Thread):
+    def __init__(
+        self,
+        interval=1,
+        file_name=None,
+    ):
+        super().__init__()
+        self._exit_event = Event()
+        self._interval = interval
+        assert file_name is not None, "file_name must be provided"
+        self._file_name = file_name
+        self.daemon = True
+        self._data = {}
+        self._has_errored = False
+        self._current_error = None
+
+    def read(self):
+        return self._data
+
+    def has_errored(self):
+        return self._has_errored
+
+    def get_error(self):
+        return self._current_error
+
+    def _safely_load(self):
+        try:
+            with open(self._file_name, "r") as f:
+                return json.load(f), None
+        except FileNotFoundError as e:
+            return {}, e
+        except json.JSONDecodeError as e:
+            return {}, e
+        except Exception as e:
+            return {}, e
+
+    def run(self):
+        while self._exit_event.is_set() is False:
+            data, self._current_error = self._safely_load()
+            if not self._current_error:
+                self._data = data
+            self._has_errored = True if self._current_error else False
+            time.sleep(self._interval)
+
+    def stop(self):
+        self._exit_event.set()
+        self.join()
+
+
+class CardRefresher:
+
+    CARD_ID = None
+
+    REFRESHABLE = True
+
+    def on_startup(self, current_card):
+        raise NotImplementedError("make_card method must be implemented")
+
+    def on_error(self, current_card, error_message):
+        raise NotImplementedError("error_card method must be implemented")
+
+    def on_update(self, current_card, data_object):
+        raise NotImplementedError("update_card method must be implemented")
+
+
+class CardUpdaterThread(Thread):
+    def __init__(
+        self,
+        card_refresher: CardRefresher,
+        interval=1,
+        file_name=None,
+        collector_thread: InfoCollectorThread = None,
+    ):
+        super().__init__()
+        self._exit_event = Event()
+        self._interval = interval
+        self._refresher = card_refresher
+        self._file_name = file_name
+        self._collector_thread = collector_thread
+        self.daemon = True
+
+    def run(self):
+        if self._refresher.CARD_ID is None:
+            raise ValueError("CARD_ID must be defined")
+        current_card = current.card[self._refresher.CARD_ID]
+        self._refresher.on_startup(current_card)
+        while self._exit_event.is_set() is False:
+            data = self._collector_thread.read()
+            if self._collector_thread.has_errored():
+                self._refresher.on_error(
+                    current_card, self._collector_thread.get_error()
+                )
+            self._refresher.on_update(current_card, data)
+            time.sleep(self._interval)
+
+    def stop(self):
+        self._exit_event.set()
+        self._collector_thread.stop()
+        self.join()
+
+
+class AsyncPeriodicRefresher:
+    def __init__(
+        self,
+        card_referesher: CardRefresher,
+        updater_interval=1,
+        collector_interval=1,
+        file_name=None,
+    ):
+        assert card_referesher.CARD_ID is not None, "CARD_ID must be defined"
+        self._collector_thread = None
+        self._collector_thread = InfoCollectorThread(
+            interval=collector_interval, file_name=file_name
+        )
+        self._collector_thread.start()
+
+        self._updater_thread = CardUpdaterThread(
+            card_refresher=card_referesher,
+            interval=updater_interval,
+            file_name=file_name,
+            collector_thread=self._collector_thread,
+        )
+
+    def start(self):
+        self._updater_thread.start()
+
+    def stop(self):
+        self._updater_thread.stop()
+        self._collector_thread.stop()

metaflow_checkpoint-0.0.1/metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.py

@@ -0,0 +1,70 @@
+from metaflow.exception import MetaflowException
+from collections import defaultdict
+
+
+class CardDecoratorInjector:
+    """
+    Mixin Useful for injecting @card decorators from other first class Metaflow decorators.
+    """
+
+    _first_time_init = defaultdict(dict)
+
+    @classmethod
+    def _get_first_time_init_cached_value(cls, step_name, card_id):
+        return cls._first_time_init.get(step_name, {}).get(card_id, None)
+
+    @classmethod
+    def _set_first_time_init_cached_value(cls, step_name, card_id, value):
+        cls._first_time_init[step_name][card_id] = value
+
+    def _card_deco_already_attached(self, step, card_id):
+        for decorator in step.decorators:
+            if decorator.name == "card":
+                if decorator.attributes["id"] and card_id in decorator.attributes["id"]:
+                    return True
+        return False
+
+    def _get_step(self, flow, step_name):
+        for step in flow:
+            if step.name == step_name:
+                return step
+        return None
+
+    def _first_time_init_check(self, step_dag_node, card_id):
+        """ """
+        return not self._card_deco_already_attached(step_dag_node, card_id)
+
+    def attach_card_decorator(
+        self,
+        flow,
+        step_name,
+        card_id,
+        card_type,
+        refresh_interval=5,
+    ):
+        """
+        This method is called `step_init` in your StepDecorator code since
+        this class is used as a Mixin
+        """
+        from metaflow import decorators as _decorators
+
+        if not all([card_id, card_type]):
+            raise MetaflowException(
+                "`card_id` and `card_type` must be set in the `CardDecoratorInjector` Mixin"
+            )
+
+        step_dag_node = self._get_step(flow, step_name)
+        if (
+            self._get_first_time_init_cached_value(step_name, card_id) is None
+        ):  # First check class level setting.
+            if self._first_time_init_check(step_dag_node, card_id):
+                self._set_first_time_init_cached_value(step_name, card_id, True)
+                _decorators._attach_decorators_to_step(
+                    step_dag_node,
+                    [
+                        "card:type=%s,id=%s,refresh_interval=%s"
+                        % (card_type, card_id, str(refresh_interval))
+                    ],
+                )
+            else:
+                self._set_first_time_init_cached_value(step_name, card_id, False)

metaflow_checkpoint-0.0.1/metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.py

@@ -0,0 +1,250 @@
+import os
+from metaflow.cards import VegaChart, MetaflowCardComponent, Artifact, Table
+from typing import Optional, List, Union
+import math
+from metaflow.plugins.cards.card_modules.components import (
+    with_default_component_id,
+    TaskToDict,
+    ArtifactsComponent,
+    UserComponent,
+    TableComponent,
+    SectionComponent,
+    render_safely,
+)
+import datetime
+from metaflow.metaflow_current import current
+import json
+from functools import wraps
+from collections import defaultdict
+from threading import Thread, Event
+import time
+
+
+def update_spec_data(spec, data):
+    spec["data"]["values"] = data
+    return spec
+
+
+def update_data_object(data_object, data):
+    data_object["values"].append(data)
+    return data_object
+
+
+def line_chart_spec(
+    title=None,
+    x_name="u",
+    y_name="v",
+    xtitle=None,
+    ytitle=None,
+    width=600,
+    height=400,
+    with_params=True,
+    x_axis_temporal=False,
+):
+    parameters = [
+        {
+            "name": "interpolate",
+            "value": "linear",
+            "bind": {
+                "input": "select",
+                "options": [
+                    "basis",
+                    "cardinal",
+                    "catmull-rom",
+                    "linear",
+                    "monotone",
+                    "natural",
+                    "step",
+                    "step-after",
+                    "step-before",
+                ],
+            },
+        },
+        {
+            "name": "tension",
+            "value": 0,
+            "bind": {"input": "range", "min": 0, "max": 1, "step": 0.05},
+        },
+        {
+            "name": "strokeWidth",
+            "value": 2,
+            "bind": {"input": "range", "min": 0, "max": 10, "step": 0.5},
+        },
+        {
+            "name": "strokeCap",
+            "value": "butt",
+            "bind": {"input": "select", "options": ["butt", "round", "square"]},
+        },
+        {
+            "name": "strokeDash",
+            "value": [1, 0],
+            "bind": {
+                "input": "select",
+                "options": [[1, 0], [8, 8], [8, 4], [4, 4], [4, 2], [2, 1], [1, 1]],
+            },
+        },
+    ]
+    parameter_marks = {
+        "interpolate": {"expr": "interpolate"},
+        "tension": {"expr": "tension"},
+        "strokeWidth": {"expr": "strokeWidth"},
+        "strokeDash": {"expr": "strokeDash"},
+        "strokeCap": {"expr": "strokeCap"},
+    }
+    spec = {
+        "title": title if title else "Line Chart",
+        "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
+        # "width": width,
+        # "height": height,
+        "params": parameters if with_params else [],
+        "data": {"name": "values", "values": []},
+        "mark": {
+            "type": "line",
+            "tooltip": True,
+            **(parameter_marks if with_params else {}),
+        },
+        "selection": {"grid": {"type": "interval", "bind": "scales"}},
+        "encoding": {
+            "x": {
+                "field": x_name,
+                "title": xtitle if xtitle else x_name,
+                **({"timeUnit": "seconds"} if x_axis_temporal else {}),
+                **({"type": "quantitative"} if not x_axis_temporal else {}),
+            },
+            "y": {
+                "field": y_name,
+                "type": "quantitative",
+                "title": ytitle if ytitle else y_name,
+            },
+        },
+        **({"width": width, "height": height} if width and height else {}),
+    }
+    data = {"values": []}
+    return spec, data
+
+
+class LineChart(UserComponent):
+    REALTIME_UPDATABLE = True
+
+    def __init__(
+        self,
+        title,
+        xtitle,
+        ytitle,
+        x_name,
+        y_name,
+        with_params=False,
+        x_axis_temporal=False,
+        width=None,
+        height=None,
+    ):
+        super().__init__()
+
+        self.spec, _ = line_chart_spec(
+            title=title,
+            xtitle=xtitle,
+            ytitle=ytitle,
+            x_name=x_name,
+            y_name=y_name,
+            with_params=with_params,
+            x_axis_temporal=x_axis_temporal,
+            width=width,
+            height=height,
+        )
+
+    def update(self, data):  # Can take a diff
+        self.spec = update_spec_data(self.spec, data)
+
+    @with_default_component_id
+    def render(self):
+        vega_chart = VegaChart(self.spec, show_controls=True)
+        vega_chart.component_id = self.component_id
+        return vega_chart.render()
+
+
+class ArtifactTable(Artifact):
+    def __init__(self, data_dict):
+        self._data = data_dict
+        self._task_to_dict = TaskToDict(only_repr=True)
+
+    @with_default_component_id
+    @render_safely
+    def render(self):
+        _art_list = []
+        for k, v in self._data.items():
+            _art = self._task_to_dict.infer_object(v)
+            _art["name"] = k
+            _art_list.append(_art)
+
+        af_component = ArtifactsComponent(data=_art_list)
+        af_component.component_id = self.component_id
+        return af_component.render()
+
+
+class UpadateableTable(UserComponent):
+    """
+
+    Parameters
+    ----------
+    data : List[List[str or MetaflowCardComponent]], optional
+        List (rows) of lists (columns). Each item can be a string or a `MetaflowCardComponent`.
+    headers : List[str], optional
+        Optional header row for the table.
+    """
+
+    REALTIME_UPDATABLE = True
+
+    def update(self, row: List[Union[str, MetaflowCardComponent]]):
+        self._data.append(row)
+
+    def __init__(
+        self,
+        data: Optional[List[List[Union[str, MetaflowCardComponent]]]] = None,
+        headers: Optional[List[str]] = None,
+        disable_updates: bool = False,
+    ):
+        data = data or [[]]
+        headers = headers or []
+        header_bool, data_bool = TableComponent.validate(headers, data)
+        self._headers = []
+        self._data = [[]]
+        if header_bool:
+            self._headers = headers
+        if data_bool:
+            self._data = data
+
+        if disable_updates:
+            self.REALTIME_UPDATABLE = False
+
+    def _render_subcomponents(self):
+        for row in self._data:
+            for col in row:
+                if isinstance(col, VegaChart):
+                    col._chart_inside_table = True
+
+        return [
+            SectionComponent.render_subcomponents(
+                row,
+                additional_allowed_types=[
+                    str,
+                    bool,
+                    int,
+                    float,
+                    dict,
+                    list,
+                    tuple,
+                    type(None),
+                ],
+                allow_unknowns=True,
+            )
+            for row in self._data
+        ]
+
+    @with_default_component_id
+    @render_safely
+    def render(self):
+        table_component = TableComponent(
+            headers=self._headers, data=self._render_subcomponents()
+        )
+        table_component.component_id = self.component_id
+        return table_component.render()

metaflow_checkpoint-0.0.1/metaflow_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.py

@@ -0,0 +1 @@
+from .lineage_card import create_checkpoint_card, null_card