ob-metaflow-stubs 6.0.3.117__py2.py3-none-any.whl → 6.0.3.119__py2.py3-none-any.whl

metaflow-stubs/{metadata → metadata_provider}/metadata.pyi

@@ -1,31 +1,19 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.27.1+obcheckpoint(0.1.2);ob(v1)                               #
-# Generated on 2024-11-13T19:05:29.394809                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.818374                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
+import metaflow
 import typing
 if typing.TYPE_CHECKING:
-    import metaflow.metadata.metadata
-    import metaflow.exception
+    import metaflow.metadata_provider.metadata
 
-class MetaflowInternalError(metaflow.exception.MetaflowException, metaclass=type):
-    ...
-
-class MetaflowTaggingError(metaflow.exception.MetaflowException, metaclass=type):
-    ...
-
-def validate_tag(tag):
-    """
-    - Tag must be either of bytes-type or unicode-type.
-    - If tag is of bytes-type, it must be UTF-8 decodable
-    - If tag is of unicode-type, it must be UTF-8 encodable
-    - Tag may not be empty string.
-    - Tag cannot be too long (500 chars)
-    """
-    ...
+from ..exception import MetaflowInternalError as MetaflowInternalError
+from ..exception import MetaflowTaggingError as MetaflowTaggingError
+from ..tagging_util import validate_tag as validate_tag
 
 class DataArtifact(tuple, metaclass=type):
     """

metaflow-stubs/metadata_provider/util.pyi

@@ -0,0 +1,19 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.818617                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+
+
+def copy_tree(src, dst, update = False):
+    ...
+
+def sync_local_metadata_to_datastore(metadata_local_dir, task_ds):
+    ...
+
+def sync_local_metadata_from_datastore(metadata_local_dir, task_ds):
+    ...
+

metaflow-stubs/metaflow_config.pyi

@@ -1,18 +1,13 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.27.1+obcheckpoint(0.1.2);ob(v1)                               #
-# Generated on 2024-11-13T19:05:29.343920                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.799422                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
 
-class MetaflowException(Exception, metaclass=type):
-    def __init__(self, msg = "", lineno = None):
-        ...
-    def __str__(self):
-        ...
-    ...
+from .exception import MetaflowException as MetaflowException
 
 DATASTORE_LOCAL_DIR: str
 
@@ -40,6 +35,8 @@ DEFAULT_SECRETS_BACKEND_TYPE: None
 
 DEFAULT_SECRETS_ROLE: None
 
+DEFAULT_FROM_DEPLOYMENT_IMPL: str
+
 USER: None
 
 DATASTORE_SYSROOT_LOCAL: None
@@ -240,8 +237,6 @@ AIRFLOW_KUBERNETES_KUBECONFIG_FILE: None
 
 AIRFLOW_KUBERNETES_KUBECONFIG_CONTEXT: None
 
-FROM_DEPLOYMENT_IMPL: str
-
 CONDA_PACKAGE_S3ROOT: None
 
 CONDA_PACKAGE_AZUREROOT: None

metaflow-stubs/metaflow_current.pyi

@@ -1,21 +1,22 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.27.1+obcheckpoint(0.1.2);ob(v1)                               #
-# Generated on 2024-11-13T19:05:29.440634                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.899445                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
 import typing
 if typing.TYPE_CHECKING:
-    import typing
-    import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.hf_hub.decorator
-    import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.modeling_utils.core
     import metaflow.plugins.cards.component_serializer
-    import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.decorator
+    import metaflow
+    import typing
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.modeling_utils.core
     import metaflow.events
     import metaflow.metaflow_current
-    import metaflow
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.hf_hub.decorator
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.decorator
+
 
 TYPE_CHECKING: bool
 
@@ -224,10 +225,10 @@ class Current(object, metaclass=type):
         """
         ...
     @property
-    def graph(_):
+    def graph(self):
         ...
     @property
-    def model(self) -> "metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.modeling_utils.core.ModelSerializer":
+    def model(self) -> "metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.modeling_utils.core.ModelSerializer":
         """
         (only in the presence of the @model decorator)
         
@@ -306,36 +307,7 @@ class Current(object, metaclass=type):
         """
         ...
     @property
-    def parallel(self) -> "metaflow.metaflow_current.Parallel":
-        """
-        (only in the presence of the @parallel decorator)
-        
-        Returns a namedtuple with relevant information about the parallel task.
-        
-        Returns
-        -------
-        Parallel
-            `namedtuple` with the following fields:
-                - main_ip (`str`)
-                    The IP address of the control task.
-                - num_nodes (`int`)
-                    The total number of tasks created by @parallel
-                - node_index (`int`)
-                    The index of the current task in all the @parallel tasks.
-                - control_task_id (`Optional[str]`)
-                    The task ID of the control task. Available to all tasks.
-        """
-        ...
-    @property
-    def is_parallel(self) -> bool:
-        """
-        (only in the presence of the @parallel decorator)
-        
-        True if the current step is a @parallel step.
-        """
-        ...
-    @property
-    def checkpoint(self) -> "metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.decorator.CurrentCheckpointer":
+    def checkpoint(self) -> "metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.decorator.CurrentCheckpointer":
         """
         (only in the presence of the @checkpoint decorator)
         
@@ -396,7 +368,7 @@ class Current(object, metaclass=type):
         """
         ...
     @property
-    def huggingface_hub(self) -> "metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.hf_hub.decorator.HuggingfaceRegistry":
+    def huggingface_hub(self) -> "metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.hf_hub.decorator.HuggingfaceRegistry":
         """
         (only in the presence of the @huggingface_hub decorator)
         
@@ -465,9 +437,38 @@ class Current(object, metaclass=type):
         """
         ...
     @property
+    def parallel(self) -> "metaflow.metaflow_current.Parallel":
+        """
+        (only in the presence of the @parallel decorator)
+        
+        Returns a namedtuple with relevant information about the parallel task.
+        
+        Returns
+        -------
+        Parallel
+            `namedtuple` with the following fields:
+                - main_ip (`str`)
+                    The IP address of the control task.
+                - num_nodes (`int`)
+                    The total number of tasks created by @parallel
+                - node_index (`int`)
+                    The index of the current task in all the @parallel tasks.
+                - control_task_id (`Optional[str]`)
+                    The task ID of the control task. Available to all tasks.
+        """
+        ...
+    @property
+    def is_parallel(self) -> bool:
+        """
+        (only in the presence of the @parallel decorator)
+        
+        True if the current step is a @parallel step.
+        """
+        ...
+    @property
     def trigger(self) -> "metaflow.events.Trigger":
         """
-        (only in the presence of the @trigger_on_finish, or @trigger decorators)
+        (only in the presence of the @trigger, or @trigger_on_finish decorators)
         
         Returns `Trigger` if the current run is triggered by an event
         

metaflow-stubs/mf_extensions/__init__.pyi

@@ -0,0 +1,6 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.804706                                                            #
+######################################################################################################
+

metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi

@@ -0,0 +1,6 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.813004                                                            #
+######################################################################################################
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi

@@ -0,0 +1,6 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.813069                                                            #
+######################################################################################################
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi

@@ -0,0 +1,6 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.813128                                                            #
+######################################################################################################
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi

@@ -0,0 +1,6 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.813185                                                            #
+######################################################################################################
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi

@@ -0,0 +1,129 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.876821                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import metaflow
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.core
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.checkpoint_storage
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
+    import metaflow.datastore.datastore_storage
+
+from ..exceptions import KeyNotCompatibleWithObjectException as KeyNotCompatibleWithObjectException
+from ..utils.identity_utils import pathspec_hash as pathspec_hash
+from ..utils.general import replace_start_and_end_slash as replace_start_and_end_slash
+from ..datastore.core import allow_safe as allow_safe
+from ..datastore.core import DatastoreInterface as DatastoreInterface
+from ..datastore.core import ObjectStorage as ObjectStorage
+from ..datastore.core import STORAGE_FORMATS as STORAGE_FORMATS
+from ..datastore.exceptions import DatastoreReadInitException as DatastoreReadInitException
+from ..datastore.exceptions import DatastoreWriteInitException as DatastoreWriteInitException
+from ..datastore.exceptions import DatastoreNotReadyException as DatastoreNotReadyException
+from ..datastructures import CheckpointArtifact as CheckpointArtifact
+from ..datastore.utils import safe_serialize as safe_serialize
+
+DEFAULT_NAME: str
+
+CHECKPOINTS_STORAGE_PREFIX: str
+
+DEFAULT_STORAGE_FORMAT: str
+
+ARTIFACT_STORE_NAME: str
+
+METADATA_STORE_NAME: str
+
+ARTIFACT_METADATA_STORE_NAME: str
+
+class CheckpointsPathComponents(tuple, metaclass=type):
+    """
+    CheckpointsPathComponents(flow_name, step_name, scope, task_identifier, pathspec_hash, attempt, name, version_id, is_metadata, key_name, root_prefix)
+    """
+    @staticmethod
+    def __new__(_cls, flow_name, step_name, scope, task_identifier, pathspec_hash, attempt, name, version_id, is_metadata, key_name, root_prefix):
+        """
+        Create new instance of CheckpointsPathComponents(flow_name, step_name, scope, task_identifier, pathspec_hash, attempt, name, version_id, is_metadata, key_name, root_prefix)
+        """
+        ...
+    def __repr__(self):
+        """
+        Return a nicely formatted representation string
+        """
+        ...
+    def __getnewargs__(self):
+        """
+        Return self as a plain tuple.  Used by copy and pickle.
+        """
+        ...
+    ...
+
+def decompose_key_artifact_metadata_store(key) -> CheckpointsPathComponents:
+    ...
+
+def decompose_key_artifact_store(key) -> CheckpointsPathComponents:
+    """
+    Convert Key into Path Components.
+    PATH COMPONENTS: mf.checkpoints/artifacts/<flow_name>/<step_name>/<scope>/<task_identifier>/<pathspec_hash>.<attempt>.<name>.<version_id>
+    """
+    ...
+
+def decompose_key_metadata_store(key) -> CheckpointsPathComponents:
+    """
+    Convert Key into Path Components.
+    PATH COMPONENTS: mf.checkpoints/artifacts/<flow_name>/<step_name>/<scope>/<task_identifier>/<pathspec_hash>.<attempt>.<name>.<version_id>
+    """
+    ...
+
+class CheckpointDatastore(metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.core.DatastoreInterface, metaclass=type):
+    """
+    Consisits of 3 main components:
+    - `artifact_store`: This is where the checkpoint artifacts are stored.
+        - This key to the checkpoint in this store becomes the "key for the checkpoint"
+    - `metadata_store`: This is where the metadata of the checkpoint is stored it based on the currently executing task (path structure resembles that of a metaflow pathspec).
+        - This store helps retrieve information about all the checkpoints stored for a task during the execution.
+    - `artifact_metadatastore`: This is similar to the metadata store but holds a pathstructure similar to the artifact store.
+        - this store helps reverse lookup the Checkpoint metadata object from the checkpoint key.
+    """
+    @property
+    def metadata_ready(self):
+        ...
+    @property
+    def artifact_ready(self):
+        ...
+    def set_root_prefix(self, root_prefix):
+        ...
+    @classmethod
+    def init_read_store(cls, storage_backend: metaflow.datastore.datastore_storage.DataStoreStorage, pathspec = None, checkpoint_key = None):
+        """
+        This will initialize the datastore for reading.
+        
+        - If there is only the pathspec that's provided then it can mean the user is doing a list operations
+        - if only the checkpoint_key is provided then it can mean the user is trying to load a specific checkpoint
+        """
+        ...
+    def create_key_name(self, *args):
+        ...
+    @classmethod
+    def init_write_store(cls, storage_backend: metaflow.datastore.datastore_storage.DataStoreStorage, pathspec, scope, task_identifier):
+        ...
+    def save(self, local_path: str, attempt, version_id, name = 'mfchckpt', metadata = {}, set_latest = True, storage_format = 'files') -> metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures.CheckpointArtifact:
+        ...
+    def latest(*args, **kwargs):
+        ...
+    def load(self, local_path, version_id, attempt, name, storage_format = 'files'):
+        ...
+    def load_metadata(self, attempt, version_id, name = 'mfchckpt') -> dict:
+        ...
+    def list(self, name = None, attempt = None, within_task = True):
+        ...
+    @classmethod
+    def decompose_key(cls, key) -> CheckpointsPathComponents:
+        ...
+    def __init__(self, artifact_store: metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.core.ObjectStorage, metadata_store: metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.core.ObjectStorage, artifact_metadatastore: metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.core.ObjectStorage):
+        ...
+    ...
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi

@@ -0,0 +1,26 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.835340                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
+
+from .core import Checkpointer as Checkpointer
+from .core import WriteResolver as WriteResolver
+from .core import ReadResolver as ReadResolver
+from ..datastructures import CheckpointArtifact as CheckpointArtifact
+
+TYPE_CHECKING: bool
+
+CHECKPOINT_UID_ENV_VAR_NAME: str
+
+DEFAULT_NAME: str
+
+def load_checkpoint(checkpoint: typing.Union[metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures.CheckpointArtifact, dict, str], local_path: str):
+    ...
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi

@@ -0,0 +1,156 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.865978                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.checkpoint_storage
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.core
+
+from ......metaflow_current import current as current
+from ......exception import MetaflowException as MetaflowException
+from .exceptions import CheckpointNotAvailableException as CheckpointNotAvailableException
+from .exceptions import CheckpointException as CheckpointException
+from ..utils import flowspec_utils as flowspec_utils
+from .checkpoint_storage import CheckpointDatastore as CheckpointDatastore
+from ..datastructures import CheckpointArtifact as CheckpointArtifact
+from ..datastore.task_utils import init_datastorage_object as init_datastorage_object
+from ..datastore.task_utils import resolve_storage_backend as resolve_task_storage_backend
+
+TYPE_CHECKING: bool
+
+CHECKPOINT_TAG_PREFIX: str
+
+MAX_HASH_LEN: int
+
+DEFAULT_NAME: str
+
+CHECKPOINT_UID_ENV_VAR_NAME: str
+
+DEFAULT_STORAGE_FORMAT: str
+
+class Checkpointer(object, metaclass=type):
+    @property
+    def current_version(self):
+        ...
+    def set_root_prefix(self, root_prefix):
+        ...
+    def __init__(self, datastore: metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.checkpoint_storage.CheckpointDatastore, attempt: int = 0):
+        ...
+    def load_metadata(self, version_id: int = None, name = 'mfchckpt') -> dict:
+        ...
+    def save(self, path: str, metadata = {}, latest = True, name = 'mfchckpt', storage_format = 'files') -> metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures.CheckpointArtifact:
+        ...
+    ...
+
+class CheckpointLoadPolicy(object, metaclass=type):
+    @classmethod
+    def fresh(cls, datastore: metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.checkpoint_storage.CheckpointDatastore, flow: "metaflow.FlowSpec") -> typing.Optional[metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures.CheckpointArtifact]:
+        """
+        ```python
+        @checkpoint(load_policy="fresh")
+        ```
+        While in `fresh` mode, we want to load the "latest" checkpoint from
+        what ever task is executing at the current memoment.
+        
+        The behavior is will be such that 1st attempt of any task will not load
+        any checkpoint and there after it will load the checkpoint from the previous
+        attempt of the task (ala the lastest checkpoint within the task).
+        """
+        ...
+    @classmethod
+    def eager(cls, datastore: metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.checkpoint_storage.CheckpointDatastore, flow: "metaflow.FlowSpec") -> typing.Optional[metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures.CheckpointArtifact]:
+        """
+        ```python
+        @checkpoint(load_policy="eager")
+        ```
+        While in `eager` mode, we want to load the "latest" checkpoint ever
+        written at the "scope" level based on the kind of task that is executing.
+        
+        Setting this mode helps "checkpoints leak across executions" for the same task
+        there by allowing a way to reboot the state when new executions start.
+        """
+        ...
+    ...
+
+class ScopeResolver(object, metaclass=type):
+    @classmethod
+    def from_namespace(cls):
+        ...
+    @classmethod
+    def from_tags(cls, tags):
+        ...
+    ...
+
+class ReadResolver(object, metaclass=type):
+    """
+    Responsible for instantiating the `CheckpointDatastore` during read operations
+    based on different context's.
+    """
+    @classmethod
+    def from_key(cls, checkpoint_key):
+        """
+        """
+        ...
+    @classmethod
+    def from_key_and_pathspec(cls, pathspec, checkpoint_key):
+        """
+        """
+        ...
+    @classmethod
+    def from_key_and_run(cls, run, checkpoint_key) -> metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.checkpoint_storage.CheckpointDatastore:
+        """
+        """
+        ...
+    ...
+
+class WriteResolver(object, metaclass=type):
+    """
+    Responsible for instantiating the `CheckpointDatastore` and the subsequent
+    `_checkpointer_uid` which can instantiate the `Checkpointer` object outside
+    of the metaflow context.
+    """
+    @classmethod
+    def can_resolve_from_envionment(cls):
+        ...
+    @classmethod
+    def decompose_checkpoint_id(cls, checkpoint_id):
+        ...
+    @classmethod
+    def resolver_info(cls, flow, run, step, taskid, taskidf, scope, attempt):
+        ...
+    @classmethod
+    def construct_checkpoint_id(cls, resolver_info: ResolverInfo):
+        ...
+    @classmethod
+    def from_environment(cls):
+        ...
+    @classmethod
+    def from_run(cls, run: "metaflow.FlowSpec", scope: str, task_identifier: typing.Optional[str] = None, gang_scheduled_task = False):
+        """
+        The task-identifier gets computed in the Metaflow main process with the
+        i.e. in the decorator. The pathspec we choose to write the metadata store
+        depends on if the task is being gang scheduled or not.
+        """
+        ...
+    ...
+
+class CheckpointReferenceResolver(object, metaclass=type):
+    """
+    Resolve the Metaflow checkpoint object based on the flow artifact reference
+    or key; Used for lineage derivation.
+    """
+    @classmethod
+    def from_key(cls, flow, checkpoint_key):
+        """
+        Used by lineage derivation
+        """
+        ...
+    ...
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi

@@ -0,0 +1,25 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.835614                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import metaflow
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow.exception
+
+from ......exception import MetaflowException as MetaflowException
+
+class CheckpointNotAvailableException(metaflow.exception.MetaflowException, metaclass=type):
+    def __init__(self, message):
+        ...
+    ...
+
+class CheckpointException(metaflow.exception.MetaflowException, metaclass=type):
+    def __init__(self, message):
+        ...
+    ...
+