mxm-scaffold 0.19.0.dev5259284__tar.gz → 0.19.0.dev58203836__tar.gz

{mxm_scaffold-0.19.0.dev5259284 → mxm_scaffold-0.19.0.dev58203836}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: mxm-scaffold
-Version: 0.19.0.dev5259284
+Version: 0.19.0.dev58203836
 Summary: Merantix Momentum AI Platform client library
 Author: Merantix Momentum GmbH
 License: MIT License
@@ -10,6 +10,9 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Typing :: Typed
 Requires-Dist: matplotlib
 Requires-Dist: hydra-zen
+Requires-Dist: aim>=3.29.1 ; extra == 'aim'
+Requires-Dist: fsspec>=2024.12.0 ; extra == 'aim'
+Requires-Dist: gcsfs>=2024.12.0 ; extra == 'aim'
 Requires-Dist: tqdm ; extra == 'data'
 Requires-Dist: random-name ; extra == 'data'
 Requires-Dist: docarray ; extra == 'data'
@@ -36,6 +39,7 @@ Requires-Python: >=3.11, <3.13
 Project-URL: Documentation, https://docs.scaffold.merantix-momentum.cloud/
 Project-URL: Homepage, https://merantix-momentum.com
 Project-URL: Repository, https://github.com/merantix-momentum/scaffold-core/
+Provides-Extra: aim
 Provides-Extra: data
 Provides-Extra: flyte
 Provides-Extra: monitoring

{mxm_scaffold-0.19.0.dev5259284 → mxm_scaffold-0.19.0.dev58203836}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mxm-scaffold"
-version = "0.19.0.dev5259284"
+version = "0.19.0.dev58203836"
 description = "Merantix Momentum AI Platform client library"
 authors = [
     {name = "Merantix Momentum GmbH"}
@@ -68,12 +68,18 @@ wandb = [
     "flatten-dict",
 ]
 
+aim = [
+    "aim>=3.29.1",
+    "fsspec>=2024.12.0",
+    "gcsfs>=2024.12.0",
+]
+
 [dependency-groups]
 # development dependencies
 dev = [
     "ipykernel>=6.31.0",
     "autopep8==1.7.0",
-    "black==24.3.0",
+    "black==26.3.1",
     "mockito",
     "pre-commit>=2.16.0",
     "py",
@@ -90,6 +96,9 @@ dev = [
     "wheel>=0.46.3",
     "build>=0.10,<1.3",
     "multiprocess>=0.70.14; sys_platform=='darwin'",
+    "aim>=3.29.1",
+    "fsspec>=2024.12.0",
+    "gcsfs>=2024.12.0",
 ]
 
 # docs dependencies

mxm_scaffold-0.19.0.dev58203836/src/scaffold/aim_utils/__init__.py

@@ -0,0 +1,3 @@
+from .core import register_fsspec_to_aim
+
+__all__ = ["register_fsspec_to_aim"]

mxm_scaffold-0.19.0.dev58203836/src/scaffold/aim_utils/core.py

@@ -0,0 +1,27 @@
+import logging
+
+
+def register_fsspec_to_aim():
+    logger = logging.getLogger(__name__)
+    try:
+        import fsspec
+        from aim.storage.artifacts.artifact_registry import registry
+    except ImportError:
+        logger.warning(
+            "It seems that aim_utils or fsspec are not installed. "
+            "Skipping registering fsspec to aim_utils's artifact manager."
+        )
+        return  # if aim_utils or fsspec are not installed, do nothing
+    from scaffold.aim_utils.fsspec_artifact_storage import FsspecArtifactStorage
+
+    for protocol in fsspec.available_protocols():
+        if protocol in registry.registry:  # only register if not already supported
+            # since this functionality was added to explicitly support Google Cloud Storage, warn when aim supports GCS.
+            if protocol == "gs":
+                logger.warning(
+                    "Aim now supports GCS. "
+                    "Calling register_fsspec_to_aim is no longer needed for GCS support. "
+                    "Skipping registering."
+                )
+            continue
+        registry.register(protocol, FsspecArtifactStorage)

mxm_scaffold-0.19.0.dev58203836/src/scaffold/aim_utils/fsspec_artifact_storage.py

@@ -0,0 +1,90 @@
+import pathlib
+import tempfile
+from concurrent.futures import ThreadPoolExecutor
+from typing import Optional
+from urllib.parse import urlparse
+
+from aim.ext.cleanup import AutoClean
+from aim.storage.artifacts.artifact_storage import AbstractArtifactStorage
+
+from scaffold.data.fs import get_fs_from_url
+
+
+class FsspecArtifactStorage(AbstractArtifactStorage):
+    """Fsspec storage backend for Aim."""
+
+    def __init__(self, storage_root: str, max_upload_workers: int = 4):
+        """
+        Args:
+            storage_root: URL with an optional prefix (e.g., 'gs://my-bucket/path/prefix').
+            fs: Optional filesystem object to be used. It will be created based on the storage root if not specified.
+
+        Note: Make sure to have the correct fsspec dependencies
+        installed for the protocol you want to use (e.g., 'gcsfs' for 'gs://').
+        """
+        super().__init__(storage_root)
+        self._storage_root_str = storage_root
+        parsed_url = urlparse(storage_root)
+        self._storage_root = pathlib.Path(parsed_url.netloc) / pathlib.Path(parsed_url.path.lstrip("/"))
+        self.scheme = parsed_url.scheme + "://"
+
+        self.thread_pool = ThreadPoolExecutor(max_workers=max_upload_workers, thread_name_prefix="fs-upload")
+
+        # make sure all threads finish uploading before shutting down
+        self._resources = FsspecArtifactStorageAutoClean(self)
+
+    def upload_artifact(self, file_path: str, artifact_path: str, block: bool = False):
+        """
+        Using fsspec, upload a local file to the artifact storage in a non-blocking manner by default.
+
+        Args:
+            file_path: Local path to the file to upload.
+            artifact_path: Path to the artifact in storage, relative to the storage root specified in the constructor.
+            block: If True, the method will block until the upload is complete.
+        """
+
+        dest_path = pathlib.Path(self._storage_root) / artifact_path.lstrip("/")
+        fs = get_fs_from_url(self._storage_root_str)
+        fs.makedirs(self.scheme + dest_path.parent.as_posix(), exist_ok=True)
+        future = self.thread_pool.submit(fs.put_file, file_path, self.scheme + dest_path.as_posix())
+
+        if block:
+            future.result()
+
+    def download_artifact(self, artifact_path: str, dest_dir: Optional[str] = None) -> str:
+        """
+        Using fsspec, download an artifact to a local directory.
+
+        Args:
+            artifact_path: path to the artifact in storage, relative to the storage root specified in the constructor.
+            dest_dir: Local directory to download the artifact to. If None, a temporary directory is created.
+        Returns:
+              The local path to the downloaded artifact.
+        """
+
+        dest_dir = pathlib.Path(tempfile.mkdtemp() if not dest_dir else dest_dir)
+
+        source_path = pathlib.Path(self._storage_root) / artifact_path.lstrip("/")
+        dest_path = dest_dir / source_path.name
+
+        fs = get_fs_from_url(self._storage_root_str)
+        fs.get_file(self.scheme + source_path.as_posix(), dest_path.as_posix())
+
+        return dest_path.as_posix()
+
+    def delete_artifact(self, artifact_path: str):
+        path = pathlib.Path(self._storage_root) / artifact_path.lstrip("/")
+        fs = get_fs_from_url(self._storage_root_str)
+        fs.rm_file(self.scheme + path.as_posix())
+
+
+class FsspecArtifactStorageAutoClean(AutoClean["FsspecArtifactStorage"]):
+    """Makes sure all upload threads finish before the FsspecArtifactStorage is destroyed."""
+
+    def __init__(self, instance: "FsspecArtifactStorage"):
+        super().__init__(instance)
+        self._thread_pool = instance.thread_pool
+
+    def _close(self):
+        # wait for all threads to finish
+        self._thread_pool.shutdown(wait=True)