wandb 0.19.1__py3-none-musllinux_1_2_aarch64.whl

wandb/filesync/step_prepare.py

@@ -0,0 +1,179 @@
+"""Batching file prepare requests to our API."""
+
+import queue
+import threading
+import time
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
+
+if TYPE_CHECKING:
+    from wandb.sdk.internal.internal_api import (
+        Api,
+        CreateArtifactFileSpecInput,
+        CreateArtifactFilesResponseFile,
+    )
+
+
+# Request for a file to be prepared.
+class RequestPrepare(NamedTuple):
+    file_spec: "CreateArtifactFileSpecInput"
+    response_channel: "queue.Queue[ResponsePrepare]"
+
+
+class RequestFinish(NamedTuple):
+    pass
+
+
+class ResponsePrepare(NamedTuple):
+    birth_artifact_id: str
+    upload_url: Optional[str]
+    upload_headers: Sequence[str]
+    upload_id: Optional[str]
+    storage_path: Optional[str]
+    multipart_upload_urls: Optional[Dict[int, str]]
+
+
+Request = Union[RequestPrepare, RequestFinish]
+
+
+def _clamp(x: float, low: float, high: float) -> float:
+    return max(low, min(x, high))
+
+
+def gather_batch(
+    request_queue: "queue.Queue[Request]",
+    batch_time: float,
+    inter_event_time: float,
+    max_batch_size: int,
+    clock: Callable[[], float] = time.monotonic,
+) -> Tuple[bool, Sequence[RequestPrepare]]:
+    batch_start_time = clock()
+    remaining_time = batch_time
+
+    first_request = request_queue.get()
+    if isinstance(first_request, RequestFinish):
+        return True, []
+
+    batch: List[RequestPrepare] = [first_request]
+
+    while remaining_time > 0 and len(batch) < max_batch_size:
+        try:
+            request = request_queue.get(
+                timeout=_clamp(
+                    x=inter_event_time,
+                    low=1e-12,  # 0 = "block forever", so just use something tiny
+                    high=remaining_time,
+                ),
+            )
+            if isinstance(request, RequestFinish):
+                return True, batch
+
+            batch.append(request)
+            remaining_time = batch_time - (clock() - batch_start_time)
+
+        except queue.Empty:
+            break
+
+    return False, batch
+
+
+def prepare_response(response: "CreateArtifactFilesResponseFile") -> ResponsePrepare:
+    multipart_resp = response.get("uploadMultipartUrls")
+    part_list = multipart_resp["uploadUrlParts"] if multipart_resp else []
+    multipart_parts = {u["partNumber"]: u["uploadUrl"] for u in part_list} or None
+
+    return ResponsePrepare(
+        birth_artifact_id=response["artifact"]["id"],
+        upload_url=response["uploadUrl"],
+        upload_headers=response["uploadHeaders"],
+        upload_id=multipart_resp and multipart_resp.get("uploadID"),
+        storage_path=response.get("storagePath"),
+        multipart_upload_urls=multipart_parts,
+    )
+
+
+class StepPrepare:
+    """A thread that batches requests to our file prepare API.
+
+    Any number of threads may call prepare() in parallel. The PrepareBatcher thread
+    will batch requests up and send them all to the backend at once.
+    """
+
+    def __init__(
+        self,
+        api: "Api",
+        batch_time: float,
+        inter_event_time: float,
+        max_batch_size: int,
+        request_queue: Optional["queue.Queue[Request]"] = None,
+    ) -> None:
+        self._api = api
+        self._inter_event_time = inter_event_time
+        self._batch_time = batch_time
+        self._max_batch_size = max_batch_size
+        self._request_queue: queue.Queue[Request] = request_queue or queue.Queue()
+        self._thread = threading.Thread(target=self._thread_body)
+        self._thread.daemon = True
+
+    def _thread_body(self) -> None:
+        while True:
+            finish, batch = gather_batch(
+                request_queue=self._request_queue,
+                batch_time=self._batch_time,
+                inter_event_time=self._inter_event_time,
+                max_batch_size=self._max_batch_size,
+            )
+            if batch:
+                batch_response = self._prepare_batch(batch)
+                # send responses
+                for prepare_request in batch:
+                    name = prepare_request.file_spec["name"]
+                    response_file = batch_response[name]
+                    response = prepare_response(response_file)
+                    prepare_request.response_channel.put(response)
+            if finish:
+                break
+
+    def _prepare_batch(
+        self, batch: Sequence[RequestPrepare]
+    ) -> Mapping[str, "CreateArtifactFilesResponseFile"]:
+        """Execute the prepareFiles API call.
+
+        Args:
+            batch: List of RequestPrepare objects
+        Returns:
+            dict of (save_name: ResponseFile) pairs where ResponseFile is a dict with
+                an uploadUrl key. The value of the uploadUrl key is None if the file
+                already exists, or a url string if the file should be uploaded.
+        """
+        return self._api.create_artifact_files([req.file_spec for req in batch])
+
+    def prepare(
+        self, file_spec: "CreateArtifactFileSpecInput"
+    ) -> "queue.Queue[ResponsePrepare]":
+        response_queue: queue.Queue[ResponsePrepare] = queue.Queue()
+        self._request_queue.put(RequestPrepare(file_spec, response_queue))
+        return response_queue
+
+    def start(self) -> None:
+        self._thread.start()
+
+    def finish(self) -> None:
+        self._request_queue.put(RequestFinish())
+
+    def is_alive(self) -> bool:
+        return self._thread.is_alive()
+
+    def shutdown(self) -> None:
+        self.finish()
+        self._thread.join()

wandb/filesync/step_upload.py

@@ -0,0 +1,287 @@
+"""Batching file prepare requests to our API."""
+
+import concurrent.futures
+import logging
+import queue
+import sys
+import threading
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    MutableMapping,
+    MutableSequence,
+    MutableSet,
+    NamedTuple,
+    Optional,
+    Union,
+)
+
+from wandb.errors.term import termerror
+from wandb.filesync import upload_job
+from wandb.sdk.lib.paths import LogicalPath
+
+if TYPE_CHECKING:
+    from typing import TypedDict
+
+    from wandb.filesync import stats
+    from wandb.sdk.internal import file_stream, internal_api, progress
+    from wandb.sdk.internal.settings_static import SettingsStatic
+
+    class ArtifactStatus(TypedDict):
+        finalize: bool
+        pending_count: int
+        commit_requested: bool
+        pre_commit_callbacks: MutableSet["PreCommitFn"]
+        result_futures: MutableSet["concurrent.futures.Future[None]"]
+
+
+PreCommitFn = Callable[[], None]
+OnRequestFinishFn = Callable[[], None]
+SaveFn = Callable[["progress.ProgressFn"], bool]
+
+logger = logging.getLogger(__name__)
+
+
+class RequestUpload(NamedTuple):
+    path: str
+    save_name: LogicalPath
+    artifact_id: Optional[str]
+    md5: Optional[str]
+    copied: bool
+    save_fn: Optional[SaveFn]
+    digest: Optional[str]
+
+
+class RequestCommitArtifact(NamedTuple):
+    artifact_id: str
+    finalize: bool
+    before_commit: PreCommitFn
+    result_future: "concurrent.futures.Future[None]"
+
+
+class RequestFinish(NamedTuple):
+    callback: Optional[OnRequestFinishFn]
+
+
+class EventJobDone(NamedTuple):
+    job: RequestUpload
+    exc: Optional[BaseException]
+
+
+Event = Union[RequestUpload, RequestCommitArtifact, RequestFinish, EventJobDone]
+
+
+class StepUpload:
+    def __init__(
+        self,
+        api: "internal_api.Api",
+        stats: "stats.Stats",
+        event_queue: "queue.Queue[Event]",
+        max_threads: int,
+        file_stream: "file_stream.FileStreamApi",
+        settings: Optional["SettingsStatic"] = None,
+    ) -> None:
+        self._api = api
+        self._stats = stats
+        self._event_queue = event_queue
+        self._file_stream = file_stream
+
+        self._thread = threading.Thread(target=self._thread_body)
+        self._thread.daemon = True
+
+        self._pool = concurrent.futures.ThreadPoolExecutor(
+            thread_name_prefix="wandb-upload",
+            max_workers=max_threads,
+        )
+
+        # Indexed by files' `save_name`'s, which are their ID's in the Run.
+        self._running_jobs: MutableMapping[LogicalPath, RequestUpload] = {}
+        self._pending_jobs: MutableSequence[RequestUpload] = []
+
+        self._artifacts: MutableMapping[str, ArtifactStatus] = {}
+
+        self.silent = bool(settings.silent) if settings else False
+
+    def _thread_body(self) -> None:
+        event: Optional[Event]
+        # Wait for event in the queue, and process one by one until a
+        # finish event is received
+        finish_callback = None
+        while True:
+            event = self._event_queue.get()
+            if isinstance(event, RequestFinish):
+                finish_callback = event.callback
+                break
+            self._handle_event(event)
+
+        # We've received a finish event. At this point, further Upload requests
+        # are invalid.
+
+        # After a finish event is received, iterate through the event queue
+        # one by one and process all remaining events.
+        while True:
+            try:
+                event = self._event_queue.get(True, 0.2)
+            except queue.Empty:
+                event = None
+            if event:
+                self._handle_event(event)
+            elif not self._running_jobs:
+                # Queue was empty and no jobs left.
+                self._pool.shutdown(wait=False)
+                if finish_callback:
+                    finish_callback()
+                break
+
+    def _handle_event(self, event: Event) -> None:
+        if isinstance(event, EventJobDone):
+            job = event.job
+
+            if event.exc is not None:
+                logger.exception(
+                    "Failed to upload file: %s", job.path, exc_info=event.exc
+                )
+
+            if job.artifact_id:
+                if event.exc is None:
+                    self._artifacts[job.artifact_id]["pending_count"] -= 1
+                    self._maybe_commit_artifact(job.artifact_id)
+                else:
+                    if not self.silent:
+                        termerror(
+                            "Uploading artifact file failed. Artifact won't be committed."
+                        )
+                    self._fail_artifact_futures(job.artifact_id, event.exc)
+            self._running_jobs.pop(job.save_name)
+            # If we have any pending jobs, start one now
+            if self._pending_jobs:
+                event = self._pending_jobs.pop(0)
+                self._start_upload_job(event)
+        elif isinstance(event, RequestCommitArtifact):
+            if event.artifact_id not in self._artifacts:
+                self._init_artifact(event.artifact_id)
+            self._artifacts[event.artifact_id]["commit_requested"] = True
+            self._artifacts[event.artifact_id]["finalize"] = event.finalize
+            self._artifacts[event.artifact_id]["pre_commit_callbacks"].add(
+                event.before_commit
+            )
+            self._artifacts[event.artifact_id]["result_futures"].add(
+                event.result_future
+            )
+            self._maybe_commit_artifact(event.artifact_id)
+        elif isinstance(event, RequestUpload):
+            if event.artifact_id is not None:
+                if event.artifact_id not in self._artifacts:
+                    self._init_artifact(event.artifact_id)
+                self._artifacts[event.artifact_id]["pending_count"] += 1
+            self._start_upload_job(event)
+        else:
+            raise Exception("Programming error: unhandled event: {}".format(str(event)))
+
+    def _start_upload_job(self, event: RequestUpload) -> None:
+        # Operations on a single backend file must be serialized. if
+        # we're already uploading this file, put the event on the
+        # end of the queue
+        if event.save_name in self._running_jobs:
+            self._pending_jobs.append(event)
+            return
+
+        self._spawn_upload(event)
+
+    def _spawn_upload(self, event: RequestUpload) -> None:
+        """Spawn an upload job, and handles the bookkeeping of `self._running_jobs`.
+
+        Context: it's important that, whenever we add an entry to `self._running_jobs`,
+        we ensure that a corresponding `EventJobDone` message will eventually get handled;
+        otherwise, the `_running_jobs` entry will never get removed, and the StepUpload
+        will never shut down.
+
+        The sole purpose of this function is to make sure that the code that adds an entry
+        to `self._running_jobs` is textually right next to the code that eventually enqueues
+        the `EventJobDone` message. This should help keep them in sync.
+        """
+        # Adding the entry to `self._running_jobs` MUST happen in the main thread,
+        # NOT in the job that gets submitted to the thread-pool, to guard against
+        # this sequence of events:
+        # - StepUpload receives a RequestUpload
+        #     ...and therefore spawns a thread to do the upload
+        # - StepUpload receives a RequestFinish
+        #     ...and checks `self._running_jobs` to see if there are any tasks to wait for...
+        #     ...and there are none, because the addition to `self._running_jobs` happens in
+        #        the background thread, which the scheduler hasn't yet run...
+        #     ...so the StepUpload shuts down. Even though we haven't uploaded the file!
+        #
+        # This would be very bad!
+        # So, this line has to happen _outside_ the `pool.submit()`.
+        self._running_jobs[event.save_name] = event
+
+        def run_and_notify() -> None:
+            try:
+                self._do_upload(event)
+            finally:
+                self._event_queue.put(EventJobDone(event, exc=sys.exc_info()[1]))
+
+        self._pool.submit(run_and_notify)
+
+    def _do_upload(self, event: RequestUpload) -> None:
+        job = upload_job.UploadJob(
+            self._stats,
+            self._api,
+            self._file_stream,
+            self.silent,
+            event.save_name,
+            event.path,
+            event.artifact_id,
+            event.md5,
+            event.copied,
+            event.save_fn,
+            event.digest,
+        )
+        job.run()
+
+    def _init_artifact(self, artifact_id: str) -> None:
+        self._artifacts[artifact_id] = {
+            "finalize": False,
+            "pending_count": 0,
+            "commit_requested": False,
+            "pre_commit_callbacks": set(),
+            "result_futures": set(),
+        }
+
+    def _maybe_commit_artifact(self, artifact_id: str) -> None:
+        artifact_status = self._artifacts[artifact_id]
+        if (
+            artifact_status["pending_count"] == 0
+            and artifact_status["commit_requested"]
+        ):
+            try:
+                for pre_callback in artifact_status["pre_commit_callbacks"]:
+                    pre_callback()
+                if artifact_status["finalize"]:
+                    self._api.commit_artifact(artifact_id)
+            except Exception as exc:
+                termerror(
+                    f"Committing artifact failed. Artifact {artifact_id} won't be finalized."
+                )
+                termerror(str(exc))
+                self._fail_artifact_futures(artifact_id, exc)
+            else:
+                self._resolve_artifact_futures(artifact_id)
+
+    def _fail_artifact_futures(self, artifact_id: str, exc: BaseException) -> None:
+        futures = self._artifacts[artifact_id]["result_futures"]
+        for result_future in futures:
+            result_future.set_exception(exc)
+        futures.clear()
+
+    def _resolve_artifact_futures(self, artifact_id: str) -> None:
+        futures = self._artifacts[artifact_id]["result_futures"]
+        for result_future in futures:
+            result_future.set_result(None)
+        futures.clear()
+
+    def start(self) -> None:
+        self._thread.start()
+
+    def is_alive(self) -> bool:
+        return self._thread.is_alive()

wandb/filesync/upload_job.py

@@ -0,0 +1,142 @@
+import logging
+import os
+from typing import TYPE_CHECKING, Optional
+
+import wandb
+from wandb.sdk.lib.paths import LogicalPath
+
+if TYPE_CHECKING:
+    from wandb.filesync import dir_watcher, stats, step_upload
+    from wandb.sdk.internal import file_stream, internal_api
+
+
+logger = logging.getLogger(__name__)
+
+
+class UploadJob:
+    def __init__(
+        self,
+        stats: "stats.Stats",
+        api: "internal_api.Api",
+        file_stream: "file_stream.FileStreamApi",
+        silent: bool,
+        save_name: LogicalPath,
+        path: "dir_watcher.PathStr",
+        artifact_id: Optional[str],
+        md5: Optional[str],
+        copied: bool,
+        save_fn: Optional["step_upload.SaveFn"],
+        digest: Optional[str],
+    ) -> None:
+        """A file uploader.
+
+        Args:
+            push_function: function(save_name, actual_path) which actually uploads
+                the file.
+            save_name: string logical location of the file relative to the run
+                directory.
+            path: actual string path of the file to upload on the filesystem.
+        """
+        self._stats = stats
+        self._api = api
+        self._file_stream = file_stream
+        self.silent = silent
+        self.save_name = save_name
+        self.save_path = path
+        self.artifact_id = artifact_id
+        self.md5 = md5
+        self.copied = copied
+        self.save_fn = save_fn
+        self.digest = digest
+        super().__init__()
+
+    def run(self) -> None:
+        success = False
+        try:
+            self.push()
+            success = True
+        finally:
+            if self.copied and os.path.isfile(self.save_path):
+                os.remove(self.save_path)
+            if success:
+                self._file_stream.push_success(self.artifact_id, self.save_name)  # type: ignore
+
+    def push(self) -> None:
+        if self.save_fn:
+            # Retry logic must happen in save_fn currently
+            try:
+                deduped = self.save_fn(
+                    lambda _, t: self._stats.update_uploaded_file(self.save_path, t)
+                )
+            except Exception as e:
+                self._stats.update_failed_file(self.save_path)
+                logger.exception("Failed to upload file: %s", self.save_path)
+                wandb._sentry.exception(e)
+                message = str(e)
+                # TODO: this is usually XML, but could be JSON
+                if hasattr(e, "response"):
+                    message = e.response.content
+                wandb.termerror(
+                    f'Error uploading "{self.save_path}": {type(e).__name__}, {message}'
+                )
+                raise
+
+            if deduped:
+                logger.info("Skipped uploading %s", self.save_path)
+                self._stats.set_file_deduped(self.save_path)
+            else:
+                logger.info("Uploaded file %s", self.save_path)
+            return
+
+        if self.md5:
+            # This is the new artifact manifest upload flow, in which we create the
+            # database entry for the manifest file before creating it. This is used for
+            # artifact L0 files. Which now is only artifact_manifest.json
+            _, response = self._api.create_artifact_manifest(
+                self.save_name, self.md5, self.artifact_id
+            )
+            upload_url = response["uploadUrl"]
+            upload_headers = response["uploadHeaders"]
+        else:
+            # The classic file upload flow. We get a signed url and upload the file
+            # then the backend handles the cloud storage metadata callback to create the
+            # file entry. This flow has aged like a fine wine.
+            project = self._api.get_project()
+            _, upload_headers, result = self._api.upload_urls(project, [self.save_name])
+            file_info = result[self.save_name]
+            upload_url = file_info["uploadUrl"]
+
+        if upload_url is None:
+            logger.info("Skipped uploading %s", self.save_path)
+            self._stats.set_file_deduped(self.save_name)
+        else:
+            extra_headers = self._api._extra_http_headers
+            for upload_header in upload_headers:
+                key, val = upload_header.split(":", 1)
+                extra_headers[key] = val
+            # Copied from push TODO(artifacts): clean up
+            # If the upload URL is relative, fill it in with the base URL,
+            # since its a proxied file store like the on-prem VM.
+            if upload_url.startswith("/"):
+                upload_url = f"{self._api.api_url}{upload_url}"
+            try:
+                with open(self.save_path, "rb") as f:
+                    self._api.upload_file_retry(
+                        upload_url,
+                        f,
+                        lambda _, t: self.progress(t),
+                        extra_headers=extra_headers,
+                    )
+                logger.info("Uploaded file %s", self.save_path)
+            except Exception as e:
+                self._stats.update_failed_file(self.save_name)
+                logger.exception("Failed to upload file: %s", self.save_path)
+                wandb._sentry.exception(e)
+                if not self.silent:
+                    wandb.termerror(
+                        f'Error uploading "{self.save_name}": {type(e).__name__}, {e}'
+                    )
+                raise
+
+    def progress(self, total_bytes: int) -> None:
+        self._stats.update_uploaded_file(self.save_name, total_bytes)

wandb/integration/catboost/__init__.py

@@ -0,0 +1,5 @@
+"""W&B callback for CatBoost."""
+
+from .catboost import WandbCallback, log_summary
+
+__all__ = ["log_summary", "WandbCallback"]