wandb 0.18.0__py3-none-macosx_10_13_x86_64.whl

wandb/filesync/step_checksum.py

@@ -0,0 +1,142 @@
+"""Batching file prepare requests to our API."""
+
+import concurrent.futures
+import functools
+import os
+import queue
+import shutil
+import threading
+from typing import TYPE_CHECKING, NamedTuple, Optional, Union, cast
+
+from wandb.filesync import step_upload
+from wandb.sdk.lib import filesystem, runid
+from wandb.sdk.lib.paths import LogicalPath
+
+if TYPE_CHECKING:
+    import tempfile
+
+    from wandb.filesync import stats
+    from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
+    from wandb.sdk.artifacts.artifact_saver import SaveFn
+    from wandb.sdk.internal import internal_api
+
+
+class RequestUpload(NamedTuple):
+    path: str
+    save_name: LogicalPath
+    copy: bool
+
+
+class RequestStoreManifestFiles(NamedTuple):
+    manifest: "ArtifactManifest"
+    artifact_id: str
+    save_fn: "SaveFn"
+
+
+class RequestCommitArtifact(NamedTuple):
+    artifact_id: str
+    finalize: bool
+    before_commit: step_upload.PreCommitFn
+    result_future: "concurrent.futures.Future[None]"
+
+
+class RequestFinish(NamedTuple):
+    callback: Optional[step_upload.OnRequestFinishFn]
+
+
+Event = Union[
+    RequestUpload, RequestStoreManifestFiles, RequestCommitArtifact, RequestFinish
+]
+
+
+class StepChecksum:
+    def __init__(
+        self,
+        api: "internal_api.Api",
+        tempdir: "tempfile.TemporaryDirectory",
+        request_queue: "queue.Queue[Event]",
+        output_queue: "queue.Queue[step_upload.Event]",
+        stats: "stats.Stats",
+    ) -> None:
+        self._api = api
+        self._tempdir = tempdir
+        self._request_queue = request_queue
+        self._output_queue = output_queue
+        self._stats = stats
+
+        self._thread = threading.Thread(target=self._thread_body)
+        self._thread.daemon = True
+
+    def _thread_body(self) -> None:
+        while True:
+            req = self._request_queue.get()
+            if isinstance(req, RequestUpload):
+                path = req.path
+                if req.copy:
+                    path = os.path.join(
+                        self._tempdir.name,
+                        f"{runid.generate_id()}-{req.save_name}",
+                    )
+                    filesystem.mkdir_exists_ok(os.path.dirname(path))
+                    try:
+                        # certain linux distros throw an exception when copying
+                        # large files: https://bugs.python.org/issue43743
+                        shutil.copy2(req.path, path)
+                    except OSError:
+                        shutil._USE_CP_SENDFILE = False  # type: ignore[attr-defined]
+                        shutil.copy2(req.path, path)
+                self._stats.init_file(req.save_name, os.path.getsize(path))
+                self._output_queue.put(
+                    step_upload.RequestUpload(
+                        path,
+                        req.save_name,
+                        None,
+                        None,
+                        req.copy,
+                        None,
+                        None,
+                    )
+                )
+            elif isinstance(req, RequestStoreManifestFiles):
+                for entry in req.manifest.entries.values():
+                    if entry.local_path:
+                        self._stats.init_file(
+                            entry.local_path,
+                            cast(int, entry.size),
+                            is_artifact_file=True,
+                        )
+                        self._output_queue.put(
+                            step_upload.RequestUpload(
+                                entry.local_path,
+                                entry.path,
+                                req.artifact_id,
+                                entry.digest,
+                                False,
+                                functools.partial(req.save_fn, entry),
+                                entry.digest,
+                            )
+                        )
+            elif isinstance(req, RequestCommitArtifact):
+                self._output_queue.put(
+                    step_upload.RequestCommitArtifact(
+                        req.artifact_id,
+                        req.finalize,
+                        req.before_commit,
+                        req.result_future,
+                    )
+                )
+            elif isinstance(req, RequestFinish):
+                break
+            else:
+                raise Exception("internal error")
+
+        self._output_queue.put(step_upload.RequestFinish(req.callback))
+
+    def start(self) -> None:
+        self._thread.start()
+
+    def is_alive(self) -> bool:
+        return self._thread.is_alive()
+
+    def finish(self) -> None:
+        self._request_queue.put(RequestFinish(None))

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.
+
+        Arguments:
+            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,290 @@
+"""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 wandb.filesync import stats
+    from wandb.sdk.internal import file_stream, internal_api, progress
+    from wandb.sdk.internal.settings_static import SettingsStatic
+
+    if sys.version_info >= (3, 8):
+        from typing import TypedDict
+    else:
+        from typing_extensions import TypedDict
+
+    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()