wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -0,0 +1,378 @@
+"""WandB storage policy."""
+
+from __future__ import annotations
+
+import hashlib
+import math
+import os
+import shutil
+from typing import TYPE_CHECKING, Any, Sequence
+from urllib.parse import quote
+
+import requests
+import urllib3
+
+from wandb.errors.term import termwarn
+from wandb.sdk.artifacts.artifact_file_cache import (
+    ArtifactFileCache,
+    get_artifact_file_cache,
+)
+from wandb.sdk.artifacts.staging import get_staging_dir
+from wandb.sdk.artifacts.storage_handlers.azure_handler import AzureHandler
+from wandb.sdk.artifacts.storage_handlers.gcs_handler import GCSHandler
+from wandb.sdk.artifacts.storage_handlers.http_handler import HTTPHandler
+from wandb.sdk.artifacts.storage_handlers.local_file_handler import LocalFileHandler
+from wandb.sdk.artifacts.storage_handlers.multi_handler import MultiHandler
+from wandb.sdk.artifacts.storage_handlers.s3_handler import S3Handler
+from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandler
+from wandb.sdk.artifacts.storage_handlers.wb_artifact_handler import WBArtifactHandler
+from wandb.sdk.artifacts.storage_handlers.wb_local_artifact_handler import (
+    WBLocalArtifactHandler,
+)
+from wandb.sdk.artifacts.storage_layout import StorageLayout
+from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY
+from wandb.sdk.artifacts.storage_policy import StoragePolicy
+from wandb.sdk.internal.internal_api import Api as InternalApi
+from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
+from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, hex_to_b64_id
+from wandb.sdk.lib.paths import FilePathStr, URIStr
+
+if TYPE_CHECKING:
+    from wandb.filesync.step_prepare import StepPrepare
+    from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+    from wandb.sdk.internal import progress
+
+# Sleep length: 0, 2, 4, 8, 16, 32, 64, 120, 120, 120, 120, 120, 120, 120, 120, 120
+# seconds, i.e. a total of 20min 6s.
+_REQUEST_RETRY_STRATEGY = urllib3.util.retry.Retry(
+    backoff_factor=1,
+    total=16,
+    status_forcelist=(308, 408, 409, 429, 500, 502, 503, 504),
+)
+_REQUEST_POOL_CONNECTIONS = 64
+_REQUEST_POOL_MAXSIZE = 64
+
+# AWS S3 max upload parts without having to make additional requests for extra parts
+S3_MAX_PART_NUMBERS = 1000
+S3_MIN_MULTI_UPLOAD_SIZE = 2 * 1024**3
+S3_MAX_MULTI_UPLOAD_SIZE = 5 * 1024**4
+
+
+class WandbStoragePolicy(StoragePolicy):
+    @classmethod
+    def name(cls) -> str:
+        return WANDB_STORAGE_POLICY
+
+    @classmethod
+    def from_config(
+        cls, config: dict, api: InternalApi | None = None
+    ) -> WandbStoragePolicy:
+        return cls(config=config, api=api)
+
+    def __init__(
+        self,
+        config: dict | None = None,
+        cache: ArtifactFileCache | None = None,
+        api: InternalApi | None = None,
+    ) -> None:
+        self._cache = cache or get_artifact_file_cache()
+        self._config = config or {}
+        self._session = requests.Session()
+        adapter = requests.adapters.HTTPAdapter(
+            max_retries=_REQUEST_RETRY_STRATEGY,
+            pool_connections=_REQUEST_POOL_CONNECTIONS,
+            pool_maxsize=_REQUEST_POOL_MAXSIZE,
+        )
+        self._session.mount("http://", adapter)
+        self._session.mount("https://", adapter)
+
+        s3 = S3Handler()
+        gcs = GCSHandler()
+        azure = AzureHandler()
+        http = HTTPHandler(self._session)
+        https = HTTPHandler(self._session, scheme="https")
+        artifact = WBArtifactHandler()
+        local_artifact = WBLocalArtifactHandler()
+        file_handler = LocalFileHandler()
+
+        self._api = api or InternalApi()
+        self._handler = MultiHandler(
+            handlers=[
+                s3,
+                gcs,
+                azure,
+                http,
+                https,
+                artifact,
+                local_artifact,
+                file_handler,
+            ],
+            default_handler=TrackingHandler(),
+        )
+
+    def config(self) -> dict:
+        return self._config
+
+    def load_file(
+        self,
+        artifact: Artifact,
+        manifest_entry: ArtifactManifestEntry,
+        dest_path: str | None = None,
+    ) -> FilePathStr:
+        if dest_path is not None:
+            self._cache._override_cache_path = dest_path
+
+        path, hit, cache_open = self._cache.check_md5_obj_path(
+            B64MD5(manifest_entry.digest),  # TODO(spencerpearson): unsafe cast
+            manifest_entry.size if manifest_entry.size is not None else 0,
+        )
+        if hit:
+            return path
+
+        if manifest_entry._download_url is not None:
+            response = self._session.get(manifest_entry._download_url, stream=True)
+            try:
+                response.raise_for_status()
+            except Exception:
+                # Signed URL might have expired, fall back to fetching it one by one.
+                manifest_entry._download_url = None
+        if manifest_entry._download_url is None:
+            auth = None
+            http_headers = _thread_local_api_settings.headers or {}
+            if self._api.access_token is not None:
+                http_headers["Authorization"] = f"Bearer {self._api.access_token}"
+            elif _thread_local_api_settings.cookies is None:
+                auth = ("api", self._api.api_key or "")
+
+            response = self._session.get(
+                self._file_url(self._api, artifact.entity, manifest_entry),
+                auth=auth,
+                cookies=_thread_local_api_settings.cookies,
+                headers=http_headers,
+                stream=True,
+            )
+            response.raise_for_status()
+
+        with cache_open(mode="wb") as file:
+            for data in response.iter_content(chunk_size=16 * 1024):
+                file.write(data)
+        return path
+
+    def store_reference(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: str | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        return self._handler.store_path(
+            artifact, path, name=name, checksum=checksum, max_objects=max_objects
+        )
+
+    def load_reference(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+        dest_path: str | None = None,
+    ) -> FilePathStr | URIStr:
+        assert manifest_entry.ref is not None
+        used_handler = self._handler._get_handler(manifest_entry.ref)
+        if hasattr(used_handler, "_cache") and (dest_path is not None):
+            used_handler._cache._override_cache_path = dest_path
+        return self._handler.load_path(manifest_entry, local)
+
+    def _file_url(
+        self,
+        api: InternalApi,
+        entity_name: str,
+        manifest_entry: ArtifactManifestEntry,
+    ) -> str:
+        storage_layout = self._config.get("storageLayout", StorageLayout.V1)
+        storage_region = self._config.get("storageRegion", "default")
+        md5_hex = b64_to_hex_id(B64MD5(manifest_entry.digest))
+
+        if storage_layout == StorageLayout.V1:
+            return "{}/artifacts/{}/{}".format(
+                api.settings("base_url"), entity_name, md5_hex
+            )
+        elif storage_layout == StorageLayout.V2:
+            return "{}/artifactsV2/{}/{}/{}/{}".format(
+                api.settings("base_url"),
+                storage_region,
+                entity_name,
+                quote(
+                    manifest_entry.birth_artifact_id
+                    if manifest_entry.birth_artifact_id is not None
+                    else ""
+                ),
+                md5_hex,
+            )
+        else:
+            raise Exception(f"unrecognized storage layout: {storage_layout}")
+
+    def s3_multipart_file_upload(
+        self,
+        file_path: str,
+        chunk_size: int,
+        hex_digests: dict[int, str],
+        multipart_urls: dict[int, str],
+        extra_headers: dict[str, str],
+    ) -> list[dict[str, Any]]:
+        etags = []
+        part_number = 1
+
+        with open(file_path, "rb") as f:
+            while True:
+                data = f.read(chunk_size)
+                if not data:
+                    break
+                md5_b64_str = str(hex_to_b64_id(hex_digests[part_number]))
+                upload_resp = self._api.upload_multipart_file_chunk_retry(
+                    multipart_urls[part_number],
+                    data,
+                    extra_headers={
+                        "content-md5": md5_b64_str,
+                        "content-length": str(len(data)),
+                        "content-type": extra_headers.get("Content-Type", ""),
+                    },
+                )
+                assert upload_resp is not None
+                etags.append(
+                    {"partNumber": part_number, "hexMD5": upload_resp.headers["ETag"]}
+                )
+                part_number += 1
+        return etags
+
+    def default_file_upload(
+        self,
+        upload_url: str,
+        file_path: str,
+        extra_headers: dict[str, Any],
+        progress_callback: progress.ProgressFn | None = None,
+    ) -> None:
+        """Upload a file to the artifact store and write to cache."""
+        with open(file_path, "rb") as file:
+            # This fails if we don't send the first byte before the signed URL expires.
+            self._api.upload_file_retry(
+                upload_url,
+                file,
+                progress_callback,
+                extra_headers=extra_headers,
+            )
+
+    def calc_chunk_size(self, file_size: int) -> int:
+        # Default to chunk size of 100MiB. S3 has cap of 10,000 upload parts.
+        # If file size exceeds the default chunk size, recalculate chunk size.
+        default_chunk_size = 100 * 1024**2
+        if default_chunk_size * S3_MAX_PART_NUMBERS < file_size:
+            return math.ceil(file_size / S3_MAX_PART_NUMBERS)
+        return default_chunk_size
+
+    def store_file(
+        self,
+        artifact_id: str,
+        artifact_manifest_id: str,
+        entry: ArtifactManifestEntry,
+        preparer: StepPrepare,
+        progress_callback: progress.ProgressFn | None = None,
+    ) -> bool:
+        """Upload a file to the artifact store.
+
+        Returns:
+            True if the file was a duplicate (did not need to be uploaded),
+            False if it needed to be uploaded or was a reference (nothing to dedupe).
+        """
+        file_size = entry.size if entry.size is not None else 0
+        chunk_size = self.calc_chunk_size(file_size)
+        upload_parts = []
+        hex_digests = {}
+        file_path = entry.local_path if entry.local_path is not None else ""
+        # Logic for AWS s3 multipart upload.
+        # Only chunk files if larger than 2 GiB. Currently can only support up to 5TiB.
+        if (
+            file_size >= S3_MIN_MULTI_UPLOAD_SIZE
+            and file_size <= S3_MAX_MULTI_UPLOAD_SIZE
+        ):
+            part_number = 1
+            with open(file_path, "rb") as f:
+                while True:
+                    data = f.read(chunk_size)
+                    if not data:
+                        break
+                    hex_digest = hashlib.md5(data).hexdigest()
+                    upload_parts.append(
+                        {"hexMD5": hex_digest, "partNumber": part_number}
+                    )
+                    hex_digests[part_number] = hex_digest
+                    part_number += 1
+
+        resp = preparer.prepare(
+            {
+                "artifactID": artifact_id,
+                "artifactManifestID": artifact_manifest_id,
+                "name": entry.path,
+                "md5": entry.digest,
+                "uploadPartsInput": upload_parts,
+            }
+        ).get()
+
+        entry.birth_artifact_id = resp.birth_artifact_id
+
+        multipart_urls = resp.multipart_upload_urls
+        if resp.upload_url is None:
+            return True
+        if entry.local_path is None:
+            return False
+        extra_headers = {
+            header.split(":", 1)[0]: header.split(":", 1)[1]
+            for header in (resp.upload_headers or {})
+        }
+
+        # This multipart upload isn't available, do a regular single url upload
+        if multipart_urls is None and resp.upload_url:
+            self.default_file_upload(
+                resp.upload_url, file_path, extra_headers, progress_callback
+            )
+        else:
+            if multipart_urls is None:
+                raise ValueError(f"No multipart urls to upload for file: {file_path}")
+            # Upload files using s3 multipart upload urls
+            etags = self.s3_multipart_file_upload(
+                file_path,
+                chunk_size,
+                hex_digests,
+                multipart_urls,
+                extra_headers,
+            )
+            assert resp.storage_path is not None
+            self._api.complete_multipart_upload_artifact(
+                artifact_id, resp.storage_path, etags, resp.upload_id
+            )
+        self._write_cache(entry)
+
+        return False
+
+    def _write_cache(self, entry: ArtifactManifestEntry) -> None:
+        if entry.local_path is None:
+            return
+
+        # Cache upon successful upload.
+        _, hit, cache_open = self._cache.check_md5_obj_path(
+            B64MD5(entry.digest),
+            entry.size if entry.size is not None else 0,
+        )
+
+        staging_dir = get_staging_dir()
+        try:
+            if not entry.skip_cache and not hit:
+                with cache_open("wb") as f, open(entry.local_path, "rb") as src:
+                    shutil.copyfileobj(src, f)
+            if entry.local_path.startswith(staging_dir):
+                # Delete staged files here instead of waiting till
+                # all the files are uploaded
+                os.chmod(entry.local_path, 0o600)
+                os.remove(entry.local_path)
+        except OSError as e:
+            termwarn(f"Failed to cache {entry.local_path}, ignoring {e}")

wandb/sdk/artifacts/storage_policy.py

@@ -0,0 +1,72 @@
+"""Storage policy."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from wandb.sdk.internal.internal_api import Api as InternalApi
+from wandb.sdk.lib.paths import FilePathStr, URIStr
+
+if TYPE_CHECKING:
+    from wandb.filesync.step_prepare import StepPrepare
+    from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+    from wandb.sdk.internal.progress import ProgressFn
+
+
+class StoragePolicy:
+    @classmethod
+    def lookup_by_name(cls, name: str) -> type[StoragePolicy]:
+        import wandb.sdk.artifacts.storage_policies  # noqa: F401
+
+        for sub in cls.__subclasses__():
+            if sub.name() == name:
+                return sub
+        raise NotImplementedError(f"Failed to find storage policy '{name}'")
+
+    @classmethod
+    def name(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def from_config(cls, config: dict, api: InternalApi | None = None) -> StoragePolicy:
+        raise NotImplementedError
+
+    def config(self) -> dict:
+        raise NotImplementedError
+
+    def load_file(
+        self,
+        artifact: Artifact,
+        manifest_entry: ArtifactManifestEntry,
+        dest_path: str | None = None,
+    ) -> FilePathStr:
+        raise NotImplementedError
+
+    def store_file(
+        self,
+        artifact_id: str,
+        artifact_manifest_id: str,
+        entry: ArtifactManifestEntry,
+        preparer: StepPrepare,
+        progress_callback: ProgressFn | None = None,
+    ) -> bool:
+        raise NotImplementedError
+
+    def store_reference(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: str | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        raise NotImplementedError
+
+    def load_reference(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+        dest_path: str | None = None,
+    ) -> FilePathStr | URIStr:
+        raise NotImplementedError

wandb/sdk/backend/backend.py

@@ -0,0 +1,222 @@
+"""Backend - Send to internal process.
+
+Manage backend.
+
+"""
+
+import importlib.machinery
+import logging
+import multiprocessing
+import os
+import queue
+import sys
+import threading
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
+
+import wandb
+from wandb.sdk.interface.interface import InterfaceBase
+from wandb.sdk.interface.interface_queue import InterfaceQueue
+from wandb.sdk.internal.internal import wandb_internal
+from wandb.sdk.internal.settings_static import SettingsStatic
+from wandb.sdk.lib.mailbox import Mailbox
+from wandb.sdk.wandb_settings import Settings
+
+if TYPE_CHECKING:
+    from wandb.proto.wandb_internal_pb2 import Record, Result
+    from wandb.sdk.lib import service_connection
+
+    from ..wandb_run import Run
+
+    RecordQueue = Union["queue.Queue[Record]", multiprocessing.Queue[Record]]
+    ResultQueue = Union["queue.Queue[Result]", multiprocessing.Queue[Result]]
+
+logger = logging.getLogger("wandb")
+
+
+class BackendThread(threading.Thread):
+    """Class to running internal process as a thread."""
+
+    def __init__(self, target: Callable, kwargs: Dict[str, Any]) -> None:
+        threading.Thread.__init__(self)
+        self.name = "BackendThr"
+        self._target = target
+        self._kwargs = kwargs
+        self.daemon = True
+        self.pid = 0
+
+    def run(self) -> None:
+        self._target(**self._kwargs)
+
+
+class Backend:
+    # multiprocessing context or module
+    _multiprocessing: multiprocessing.context.BaseContext
+    interface: Optional[InterfaceBase]
+    _internal_pid: Optional[int]
+    wandb_process: Optional[multiprocessing.process.BaseProcess]
+    _settings: Optional[Settings]
+    record_q: Optional["RecordQueue"]
+    result_q: Optional["ResultQueue"]
+    _mailbox: Mailbox
+
+    def __init__(
+        self,
+        mailbox: Mailbox,
+        settings: Optional[Settings] = None,
+        log_level: Optional[int] = None,
+        service: "Optional[service_connection.ServiceConnection]" = None,
+    ) -> None:
+        self._done = False
+        self.record_q = None
+        self.result_q = None
+        self.wandb_process = None
+        self.interface = None
+        self._internal_pid = None
+        self._settings = settings
+        self._log_level = log_level
+        self._service = service
+        self._mailbox = mailbox
+
+        self._multiprocessing = multiprocessing  # type: ignore
+        self._multiprocessing_setup()
+
+        # for _module_main_* methods
+        self._save_mod_path: Optional[str] = None
+        self._save_mod_spec = None
+
+    def _hack_set_run(self, run: "Run") -> None:
+        assert self.interface
+        self.interface._hack_set_run(run)
+
+    def _multiprocessing_setup(self) -> None:
+        assert self._settings
+        if self._settings.start_method == "thread":
+            return
+
+        # defaulting to spawn for now, fork needs more testing
+        start_method = self._settings.start_method or "spawn"
+
+        # TODO: use fork context if unix and frozen?
+        # if py34+, else fall back
+        if not hasattr(multiprocessing, "get_context"):
+            return
+        all_methods = multiprocessing.get_all_start_methods()
+        logger.info(
+            "multiprocessing start_methods={}, using: {}".format(
+                ",".join(all_methods), start_method
+            )
+        )
+        ctx = multiprocessing.get_context(start_method)
+        self._multiprocessing = ctx
+
+    def _module_main_install(self) -> None:
+        # Support running code without a: __name__ == "__main__"
+        main_module = sys.modules["__main__"]
+        main_mod_spec = getattr(main_module, "__spec__", None)
+        main_mod_path = getattr(main_module, "__file__", None)
+        if main_mod_spec is None:  # hack for pdb
+            # Note: typing has trouble with BuiltinImporter
+            loader: Loader = importlib.machinery.BuiltinImporter  # type: ignore # noqa: F821
+            main_mod_spec = importlib.machinery.ModuleSpec(
+                name="wandb.mpmain", loader=loader
+            )
+            main_module.__spec__ = main_mod_spec
+        else:
+            self._save_mod_spec = main_mod_spec
+
+        if main_mod_path is not None:
+            self._save_mod_path = main_module.__file__
+            fname = os.path.join(
+                os.path.dirname(wandb.__file__), "mpmain", "__main__.py"
+            )
+            main_module.__file__ = fname
+
+    def _module_main_uninstall(self) -> None:
+        main_module = sys.modules["__main__"]
+        # Undo temporary changes from: __name__ == "__main__"
+        main_module.__spec__ = self._save_mod_spec
+        if self._save_mod_path:
+            main_module.__file__ = self._save_mod_path
+
+    def ensure_launched(self) -> None:
+        """Launch backend worker if not running."""
+        if self._service:
+            self.interface = self._service.make_interface(self._mailbox)
+            return
+
+        assert self._settings
+        settings = self._settings.copy()
+        settings.update(_log_level=self._log_level or logging.DEBUG)
+
+        start_method = settings.start_method
+
+        settings_static = SettingsStatic(settings.to_proto())
+        user_pid = os.getpid()
+
+        if start_method == "thread":
+            self.record_q = queue.Queue()
+            self.result_q = queue.Queue()
+            wandb._set_internal_process(disable=True)  # type: ignore
+            wandb_thread = BackendThread(
+                target=wandb_internal,
+                kwargs=dict(
+                    settings=settings_static,
+                    record_q=self.record_q,
+                    result_q=self.result_q,
+                    user_pid=user_pid,
+                ),
+            )
+            # TODO: risky cast, assumes BackendThread Process duck typing
+            self.wandb_process = wandb_thread  # type: ignore
+        else:
+            self.record_q = self._multiprocessing.Queue()
+            self.result_q = self._multiprocessing.Queue()
+            self.wandb_process = self._multiprocessing.Process(  # type: ignore
+                target=wandb_internal,
+                kwargs=dict(
+                    settings=settings_static,
+                    record_q=self.record_q,
+                    result_q=self.result_q,
+                    user_pid=user_pid,
+                ),
+            )
+            assert self.wandb_process
+            self.wandb_process.name = "wandb_internal"
+
+        self._module_main_install()
+
+        logger.info("starting backend process...")
+        # Start the process with __name__ == "__main__" workarounds
+        assert self.wandb_process
+        self.wandb_process.start()
+        self._internal_pid = self.wandb_process.pid
+        logger.info(f"started backend process with pid: {self.wandb_process.pid}")
+
+        self._module_main_uninstall()
+
+        self.interface = InterfaceQueue(
+            process=self.wandb_process,
+            record_q=self.record_q,  # type: ignore
+            result_q=self.result_q,  # type: ignore
+            mailbox=self._mailbox,
+        )
+
+    def server_status(self) -> None:
+        """Report server status."""
+        pass
+
+    def cleanup(self) -> None:
+        # TODO: make _done atomic
+        if self._done:
+            return
+        self._done = True
+        if self.interface:
+            self.interface.join()
+        if self.wandb_process:
+            self.wandb_process.join()
+
+        if self.record_q and hasattr(self.record_q, "close"):
+            self.record_q.close()
+        if self.result_q and hasattr(self.result_q, "close"):
+            self.result_q.close()
+        # No printing allowed from here until redirect restore!!!