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

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

@@ -0,0 +1,111 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.813236                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
+
+from ..datastructures import CheckpointArtifact as CheckpointArtifact
+from .constructors import load_checkpoint as load_checkpoint
+from .exceptions import CheckpointNotAvailableException as CheckpointNotAvailableException
+
+TYPE_CHECKING: bool
+
+CHECKPOINT_UID_ENV_VAR_NAME: str
+
+DEFAULT_NAME: str
+
+TASK_CHECKPOINTS_ARTIFACT_NAME: str
+
+DEFAULT_STORAGE_FORMAT: str
+
+class Checkpoint(object, metaclass=type):
+    def __init__(self, temp_dir_root = None, init_dir = False):
+        ...
+    @property
+    def directory(self):
+        ...
+    def save(self, path = None, metadata = None, latest = True, name = 'mfchckpt', storage_format = 'files'):
+        """
+        Saves the checkpoint to the datastore
+        
+        Parameters
+        ----------
+        path : Optional[Union[str, os.PathLike]], default: None
+            The path to save the checkpoint. Accepts a file path or a directory path.
+                - If a directory path is provided, all the contents within that directory will be saved.
+                When a checkpoint is reloaded during task retries, `the current.checkpoint.directory` will
+                contain the contents of this directory.
+                - If a file path is provided, the file will be directly saved to the datastore (with the same filename).
+                When the checkpoint is reloaded during task retries, the file with the same name will be available in the
+                `current.checkpoint.directory`.
+                - If no path is provided then the `Checkpoint.directory` will be saved as the checkpoint.
+        
+        name : Optional[str], default: "mfchckpt"
+            The name of the checkpoint.
+        
+        metadata : Optional[Dict], default: {}
+            Any metadata that needs to be saved with the checkpoint.
+        
+        latest : bool, default: True
+            If True, the checkpoint will be marked as the latest checkpoint.
+            This helps determine if the checkpoint gets loaded when the task restarts.
+        
+        storage_format : str, default: files
+            If `tar`, the contents of the directory will be tarred before saving to the datastore.
+            If `files`, saves directory directly to the datastore.
+        """
+        ...
+    def __enter__(self):
+        ...
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        ...
+    def list(self, name: typing.Optional[str] = None, task: typing.Union["metaflow.Task", str, None] = None, attempt: typing.Union[int, str, None] = None, as_dict: bool = True, within_task: bool = True) -> typing.Iterable[typing.Union[typing.Dict, metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures.CheckpointArtifact]]:
+        """
+        lists the checkpoints in the datastore based on the Task.
+        It will always be task scoped.
+        
+        Usage:
+        ------
+        
+        ```python
+        
+        Checkpoint().list(name="best") # lists checkpoints in the current task with the name "best"
+        Checkpoint().list(task="anotherflow/somerunid/somestep/sometask", name="best") # Identical as the above one but
+        Checkpoint().list() # lists all the checkpoints in the current task
+        
+        ```
+        
+        Parameters
+        ----------
+        
+        - `name`:
+            - name of the checkpoint to filter for
+        - `task`:
+            - Task object outside the one that is currently set in the `Checkpoint` object; Can be a pathspec string.
+        - `attempt`:
+            - attempt number of the task (optional filter. If none, then lists all checkpoints from all attempts)
+        """
+        ...
+    def load(self, reference: typing.Union[str, typing.Dict, metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures.CheckpointArtifact], path: typing.Optional[str] = None):
+        """
+        loads a checkpoint reference from the datastore. (resembles a read op)
+        
+        Parameters
+        ----------
+        
+        `reference` :
+            - can be a string, dict or a CheckpointArtifact object:
+                - string: a string reference to the checkpoint (checkpoint key)
+                - dict: a dictionary reference to the checkpoint
+                - CheckpointArtifact: a CheckpointArtifact object reference to the checkpoint
+        """
+        ...
+    ...
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__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.836248                                                            #
+######################################################################################################
+

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

@@ -0,0 +1,188 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.867198                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.core
+    import metaflow.datastore.datastore_storage
+
+from ......exception import MetaflowException as MetaflowException
+from ..utils.tar_utils import create_tarball_on_disk as create_tarball_on_disk
+from ..utils.tar_utils import extract_tarball as extract_tarball
+from ..utils.general import safe_serialize as safe_serialize
+from ..exceptions import KeyNotFoundError as KeyNotFoundError
+
+STORAGE_INJECTIONS_SINGLE_FILE_SAVE: dict
+
+STORAGE_INJECTIONS_MULTIPLE_FILE_SAVE: dict
+
+STORAGE_INJECTIONS_LOAD_FILES: dict
+
+class DatastoreBlob(tuple, metaclass=type):
+    """
+    DatastoreBlob(blob, url, path)
+    """
+    @staticmethod
+    def __new__(_cls, blob, url, path):
+        """
+        Create new instance of DatastoreBlob(blob, url, path)
+        """
+        ...
+    def __repr__(self):
+        """
+        Return a nicely formatted representation string
+        """
+        ...
+    def __getnewargs__(self):
+        """
+        Return self as a plain tuple.  Used by copy and pickle.
+        """
+        ...
+    ...
+
+class ListPathResult(tuple, metaclass=type):
+    """
+    ListPathResult(full_url, key)
+    """
+    @staticmethod
+    def __new__(_cls, full_url, key):
+        """
+        Create new instance of ListPathResult(full_url, key)
+        """
+        ...
+    def __repr__(self):
+        """
+        Return a nicely formatted representation string
+        """
+        ...
+    def __getnewargs__(self):
+        """
+        Return self as a plain tuple.  Used by copy and pickle.
+        """
+        ...
+    ...
+
+COMPRESSION_METHOD: None
+
+class STORAGE_FORMATS(object, metaclass=type):
+    ...
+
+def warning_message(message, logger = None, ts = False, prefix = '[@checkpoint][artifact-store]'):
+    ...
+
+def allow_safe(func):
+    ...
+
+class ObjectStorage(object, metaclass=type):
+    """
+    `ObjectStorage` wraps around the DataStoreStorage object and provides the lower level
+    storage APIs needed by the subsequent classes. This datastore's main function that
+    distingushes it from the DataStoreStorage object is that it manages a generalizable
+    path structure over different storage backends (s3, azure, gs, local etc.).
+    
+    This object will be used to create multiple "Logical" datastores for constructs like
+    `Checkpoints`, `Models` etc.
+    
+    Usage
+    -----
+    ```
+    storage = ObjectStorage(
+        storage_backend,
+        root_prefix = "mf.checkpoints",
+        path_components = ["artifacts", "MyFlow", "step_a", "cd2312rd12d", "x5avasdhtsdfqw"]
+    )
+    ```
+    """
+    def __init__(self, storage_backend: metaflow.datastore.datastore_storage.DataStoreStorage, root_prefix: str, path_components: typing.List):
+        ...
+    @property
+    def path_components(self):
+        ...
+    def set_full_prefix(self, root_prefix):
+        ...
+    @staticmethod
+    def resolve_root(storage_type):
+        ...
+    def full_base_url(self, prefix = None):
+        ...
+    def create_key_name(self, *args):
+        ...
+    @property
+    def datastore_root(self):
+        ...
+    def resolve_key_relative_path(self, key):
+        ...
+    def resolve_key_full_url(self, key):
+        ...
+    def resolve_key_path(self, key):
+        ...
+    def put(self, key: str, obj: typing.Union[str, bytes], overwrite: bool = False) -> str:
+        """
+        TODO : [THIS IS TERRIBLY INEFFICIENT]
+        """
+        ...
+    def put_file(self, key: str, path: str, overwrite = False):
+        ...
+    def put_files(self, key_paths: typing.List[typing.Tuple[str, str]], overwrite = False):
+        ...
+    def get(self, key) -> DatastoreBlob:
+        """
+        TODO : [THIS IS TERRIBLY INEFFICIENT]
+        """
+        ...
+    def get_file(self, key):
+        ...
+    def list_paths(self, keys) -> typing.Iterator[ListPathResult]:
+        """
+        List all objects in the datastore's `keys` index.
+        """
+        ...
+    ...
+
+class DatastoreInterface(object, metaclass=type):
+    """
+    This is the root abstraction used by any underlying datastores like Checkpoint/Model etc.
+    to create the saving/loading mechanics using multiple ObjectStores.
+    
+    The inherited classes require the following implemented:
+        - `ROOT_PREFIX` : The root prefix for the datastore such as `mf.checkpoints` or `mf.models`.
+        - `init_read_store` : The method to initialize the read store; The inheriting class can compose together any number of `BaseDatastore` objects
+        - `init_write_store` : The method to initialize the write store; The inheriting class can compose together any number of `BaseDatastore` objects
+        - `save` : The method to save the artifact data.
+        - `load` : The method to load the artifact data.
+        - `save_metadata` : The method to save the metadata about the artifact.
+        - `load_metadata` : The method to load the metadata about the artifact.
+        - `list` : The method to list all the artifacts.
+    """
+    def set_root_prefix(self, root_prefix):
+        """
+        This function helps ensuring that the root prefix of the datastore
+        and it's underlying ObjectStores can change trivially.
+        """
+        ...
+    def save(self, *args, **kwargs):
+        ...
+    def load(self, *args, **kwargs):
+        ...
+    def save_metadata(self, *args, **kwargs):
+        ...
+    def load_metadata(self, *args, **kwargs):
+        ...
+    def list(self, *args, **kwargs):
+        ...
+    @classmethod
+    def init_read_store(cls, storage_backend: metaflow.datastore.datastore_storage.DataStoreStorage, *args, **kwargs):
+        ...
+    @classmethod
+    def init_write_store(cls, storage_backend: metaflow.datastore.datastore_storage.DataStoreStorage, *args, **kwargs):
+        ...
+    ...
+
+def resolve_root(datastore_type):
+    ...
+

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

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

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi

@@ -0,0 +1,28 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.836299                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import metaflow
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow
+    import metaflow.exception
+
+from ......exception import MetaflowException as MetaflowException
+from .core import resolve_root as resolve_root
+
+TYPE_CHECKING: bool
+
+class UnresolvableDatastoreException(metaflow.exception.MetaflowException, metaclass=type):
+    ...
+
+def init_datastorage_object():
+    ...
+
+def resolve_storage_backend(pathspec: typing.Union[str, "metaflow.Task"] = None):
+    ...
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.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.887615                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+
+
+def is_json_serializable(value):
+    """
+    Check if value is JSON serializable.
+    """
+    ...
+
+def safe_serialize(data):
+    ...
+

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

@@ -0,0 +1,115 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.814163                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import metaflow
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
+
+from .exceptions import KeyNotCompatibleWithObjectException as KeyNotCompatibleWithObjectException
+from .exceptions import KeyNotCompatibleException as KeyNotCompatibleException
+from .exceptions import IncompatibleObjectTypeException as IncompatibleObjectTypeException
+from .datastore.task_utils import init_datastorage_object as init_datastorage_object
+
+class MetaflowDataArtifactReference(object, metaclass=type):
+    @property
+    def size(self):
+        ...
+    @property
+    def url(self):
+        ...
+    @property
+    def key(self):
+        ...
+    @property
+    def pathspec(self):
+        ...
+    @property
+    def attempt(self):
+        ...
+    @property
+    def created_on(self):
+        ...
+    @property
+    def metadata(self):
+        ...
+    def __init__(self, **kwargs):
+        ...
+    def validate(self, data):
+        ...
+    @classmethod
+    def from_dict(cls, data) -> typing.Union["ModelArtifact", "CheckpointArtifact"]:
+        ...
+    @classmethod
+    def hydrate(cls, data: typing.Union["ModelArtifact", "CheckpointArtifact", dict]):
+        ...
+    def to_dict(self):
+        ...
+    ...
+
+class ModelArtifact(MetaflowDataArtifactReference, metaclass=type):
+    def __init__(self, **kwargs):
+        ...
+    @property
+    def blob(self):
+        ...
+    @property
+    def uuid(self):
+        ...
+    @property
+    def serializer(self):
+        ...
+    @property
+    def source(self):
+        ...
+    @property
+    def storage_format(self):
+        ...
+    @classmethod
+    def create(cls, pathspec = None, attempt = None, key = None, url = None, model_uuid = None, metadata = None, storage_format = None, source = None, serializer = None, label = None):
+        ...
+    ...
+
+class CheckpointArtifact(MetaflowDataArtifactReference, metaclass=type):
+    @property
+    def storage_format(self):
+        ...
+    @property
+    def version_id(self):
+        ...
+    @property
+    def name(self):
+        ...
+    def __init__(self, **kwargs):
+        ...
+    ...
+
+class Factory(object, metaclass=type):
+    @classmethod
+    def hydrate(cls, data):
+        ...
+    @classmethod
+    def from_dict(cls, data):
+        ...
+    @classmethod
+    def load(cls, data, local_path, storage_backend):
+        ...
+    @classmethod
+    def object_type_from_key(cls, reference_key):
+        ...
+    @classmethod
+    def load_from_key(cls, key_object, local_path, storage_backend):
+        ...
+    @classmethod
+    def load_metadata_from_key(cls, key_object, storage_backend) -> typing.Union[CheckpointArtifact, ModelArtifact]:
+        ...
+    ...
+
+def load_model(reference: typing.Union[str, MetaflowDataArtifactReference, dict], path: str):
+    ...
+

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

@@ -0,0 +1,40 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.835960                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import metaflow
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow.exception
+
+from .....exception import MetaflowException as MetaflowException
+
+class TODOException(metaflow.exception.MetaflowException, metaclass=type):
+    def __init__(self, message):
+        ...
+    ...
+
+class KeyNotFoundError(metaflow.exception.MetaflowException, metaclass=type):
+    def __init__(self, url):
+        ...
+    ...
+
+class KeyNotCompatibleException(metaflow.exception.MetaflowException, metaclass=type):
+    def __init__(self, key, supported_types):
+        ...
+    ...
+
+class KeyNotCompatibleWithObjectException(metaflow.exception.MetaflowException, metaclass=type):
+    def __init__(self, key, store, message = None):
+        ...
+    ...
+
+class IncompatibleObjectTypeException(metaflow.exception.MetaflowException, metaclass=type):
+    def __init__(self, message):
+        ...
+    ...
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__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.875777                                                            #
+######################################################################################################
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi

@@ -0,0 +1,71 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.875831                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow
+
+from ......metaflow_current import current as current
+from .identity_utils import TaskIdentifier as TaskIdentifier
+from .identity_utils import FlowNotRunningException as FlowNotRunningException
+
+TYPE_CHECKING: bool
+
+class ResolvedTask(tuple, metaclass=type):
+    """
+    ResolvedTask(flow_name, run_id, step_name, task_id, is_foreach, current_attempt, pathspec, control_task_pathspec)
+    """
+    @staticmethod
+    def __new__(_cls, flow_name, run_id, step_name, task_id, is_foreach, current_attempt, pathspec, control_task_pathspec):
+        """
+        Create new instance of ResolvedTask(flow_name, run_id, step_name, task_id, is_foreach, current_attempt, pathspec, control_task_pathspec)
+        """
+        ...
+    def __repr__(self):
+        """
+        Return a nicely formatted representation string
+        """
+        ...
+    def __getnewargs__(self):
+        """
+        Return self as a plain tuple.  Used by copy and pickle.
+        """
+        ...
+    ...
+
+class OriginInfo(tuple, metaclass=type):
+    """
+    OriginInfo(origin_run_id, origin_task_id, origin_attempt, is_foreach)
+    """
+    @staticmethod
+    def __new__(_cls, origin_run_id, origin_task_id, origin_attempt, is_foreach):
+        """
+        Create new instance of OriginInfo(origin_run_id, origin_task_id, origin_attempt, is_foreach)
+        """
+        ...
+    def __repr__(self):
+        """
+        Return a nicely formatted representation string
+        """
+        ...
+    def __getnewargs__(self):
+        """
+        Return self as a plain tuple.  Used by copy and pickle.
+        """
+        ...
+    ...
+
+def resolve_pathspec_for_flowspec(run: "metaflow.FlowSpec" = None):
+    ...
+
+def resolve_storage_backend(run: "metaflow.FlowSpec" = None):
+    ...
+
+def resolve_task_identifier(run: "metaflow.FlowSpec", gang_scheduled_task = False, gang_schedule_task_idf_index = 0):
+    ...
+

metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi

@@ -0,0 +1,56 @@
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.30.2+obcheckpoint(0.1.4);ob(v1)                                                   #
+# Generated on 2024-11-21T22:12:20.877625                                                            #
+######################################################################################################
+
+from __future__ import annotations
+
+import typing
+if typing.TYPE_CHECKING:
+    import metaflow.mf_extensions.obcheckpoint.plugins.machine_learning_utilities.utils.general
+
+
+class PathSize(tuple, metaclass=type):
+    """
+    PathSize(path, calculated_size, early_stopping, not_found)
+    """
+    @staticmethod
+    def __new__(_cls, path, calculated_size, early_stopping, not_found):
+        """
+        Create new instance of PathSize(path, calculated_size, early_stopping, not_found)
+        """
+        ...
+    def __repr__(self):
+        """
+        Return a nicely formatted representation string
+        """
+        ...
+    def __getnewargs__(self):
+        """
+        Return self as a plain tuple.  Used by copy and pickle.
+        """
+        ...
+    ...
+
+def replace_start_and_end_slash(string):
+    ...
+
+def get_path_size(start_path, early_stopping_limit = None) -> PathSize:
+    ...
+
+def warning_message(message, logger = None, ts = False, prefix = '[no-prefix]'):
+    ...
+
+def unit_convert(number, base_unit, convert_unit):
+    ...
+
+def is_json_serializable(value):
+    """
+    Check if value is JSON serializable.
+    """
+    ...
+
+def safe_serialize(data):
+    ...
+