wandb 0.21.2__py3-none-macosx_12_0_arm64.whl

wandb/filesync/dir_watcher.py

@@ -0,0 +1,404 @@
+import abc
+import fnmatch
+import glob
+import logging
+import os
+import queue
+import time
+from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, MutableSet, Optional
+
+from wandb import util
+from wandb.sdk.interface.interface import GlobStr
+from wandb.sdk.lib.paths import LogicalPath
+
+if TYPE_CHECKING:
+    import wandb.vendor.watchdog_0_9_0.observers.api as wd_api
+    import wandb.vendor.watchdog_0_9_0.observers.polling as wd_polling
+    import wandb.vendor.watchdog_0_9_0.watchdog.events as wd_events
+    from wandb.sdk.interface.interface import PolicyName
+    from wandb.sdk.internal.file_pusher import FilePusher
+    from wandb.sdk.internal.settings_static import SettingsStatic
+else:
+    wd_polling = util.vendor_import("wandb_watchdog.observers.polling")
+    wd_events = util.vendor_import("wandb_watchdog.events")
+
+PathStr = str  # TODO(spencerpearson): would be nice to use Path here
+
+
+logger = logging.getLogger(__name__)
+
+
+class FileEventHandler(abc.ABC):
+    def __init__(
+        self,
+        file_path: PathStr,
+        save_name: LogicalPath,
+        file_pusher: "FilePusher",
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.file_path = file_path
+        # Convert windows paths to unix paths
+        self.save_name = LogicalPath(save_name)
+        self._file_pusher = file_pusher
+        self._last_sync: Optional[float] = None
+
+    @property
+    @abc.abstractmethod
+    def policy(self) -> "PolicyName":
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def on_modified(self, force: bool = False) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def finish(self) -> None:
+        raise NotImplementedError
+
+    def on_renamed(self, new_path: PathStr, new_name: LogicalPath) -> None:
+        self.file_path = new_path
+        self.save_name = new_name
+        self.on_modified()
+
+
+class PolicyNow(FileEventHandler):
+    """This policy only uploads files now."""
+
+    def on_modified(self, force: bool = False) -> None:
+        # only upload if we've never uploaded or when .save is called
+        if self._last_sync is None or force:
+            self._file_pusher.file_changed(self.save_name, self.file_path)
+            self._last_sync = os.path.getmtime(self.file_path)
+
+    def finish(self) -> None:
+        pass
+
+    @property
+    def policy(self) -> "PolicyName":
+        return "now"
+
+
+class PolicyEnd(FileEventHandler):
+    """This policy only updates at the end of the run."""
+
+    def on_modified(self, force: bool = False) -> None:
+        pass
+
+    # TODO: make sure we call this
+    def finish(self) -> None:
+        # We use copy=False to avoid possibly expensive copies, and because
+        # user files shouldn't still be changing at the end of the run.
+        self._last_sync = os.path.getmtime(self.file_path)
+        self._file_pusher.file_changed(self.save_name, self.file_path, copy=False)
+
+    @property
+    def policy(self) -> "PolicyName":
+        return "end"
+
+
+class PolicyLive(FileEventHandler):
+    """Event handler that uploads respecting throttling.
+
+    Uploads files every RATE_LIMIT_SECONDS, which changes as the size increases to deal
+    with throttling.
+    """
+
+    RATE_LIMIT_SECONDS = 15
+    unit_dict = dict(util.POW_10_BYTES)
+    # Wait to upload until size has increased 20% from last upload
+    RATE_LIMIT_SIZE_INCREASE = 1.2
+
+    def __init__(
+        self,
+        file_path: PathStr,
+        save_name: LogicalPath,
+        file_pusher: "FilePusher",
+        settings: Optional["SettingsStatic"] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(file_path, save_name, file_pusher, *args, **kwargs)
+        self._last_uploaded_time: Optional[float] = None
+        self._last_uploaded_size: int = 0
+        if settings is not None:
+            if settings.x_live_policy_rate_limit is not None:
+                self.RATE_LIMIT_SECONDS = settings.x_live_policy_rate_limit
+            self._min_wait_time: Optional[float] = settings.x_live_policy_wait_time
+        else:
+            self._min_wait_time = None
+
+    @property
+    def current_size(self) -> int:
+        return os.path.getsize(self.file_path)
+
+    @classmethod
+    def min_wait_for_size(cls, size: int) -> float:
+        if size < 10 * cls.unit_dict["MB"]:
+            return 60
+        elif size < 100 * cls.unit_dict["MB"]:
+            return 5 * 60
+        elif size < cls.unit_dict["GB"]:
+            return 10 * 60
+        else:
+            return 20 * 60
+
+    def should_update(self) -> bool:
+        if self._last_uploaded_time is not None:
+            # Check rate limit by time elapsed
+            time_elapsed = time.time() - self._last_uploaded_time
+            # if more than 15 seconds has passed potentially upload it
+            if time_elapsed < self.RATE_LIMIT_SECONDS:
+                return False
+
+            # Check rate limit by size increase
+            if float(self._last_uploaded_size) > 0:
+                size_increase = self.current_size / float(self._last_uploaded_size)
+                if size_increase < self.RATE_LIMIT_SIZE_INCREASE:
+                    return False
+            return time_elapsed > (
+                self._min_wait_time or self.min_wait_for_size(self.current_size)
+            )
+
+        # if the file has never been uploaded, we'll upload it
+        return True
+
+    def on_modified(self, force: bool = False) -> None:
+        if self.current_size == 0:
+            return
+        if self._last_sync == os.path.getmtime(self.file_path):
+            return
+        if force or self.should_update():
+            self.save_file()
+
+    def save_file(self) -> None:
+        self._last_sync = os.path.getmtime(self.file_path)
+        self._last_uploaded_time = time.time()
+        self._last_uploaded_size = self.current_size
+        self._file_pusher.file_changed(self.save_name, self.file_path)
+
+    def finish(self) -> None:
+        self.on_modified(force=True)
+
+    @property
+    def policy(self) -> "PolicyName":
+        return "live"
+
+
+class DirWatcher:
+    def __init__(
+        self,
+        settings: "SettingsStatic",
+        file_pusher: "FilePusher",
+        file_dir: Optional[PathStr] = None,
+    ) -> None:
+        self._file_count = 0
+        self._dir = file_dir or settings.files_dir
+        self._settings = settings
+        self._savename_file_policies: MutableMapping[LogicalPath, PolicyName] = {}
+        self._user_file_policies: Mapping[PolicyName, MutableSet[GlobStr]] = {
+            "end": set(),
+            "live": set(),
+            "now": set(),
+        }
+        self._file_pusher = file_pusher
+        self._file_event_handlers: MutableMapping[LogicalPath, FileEventHandler] = {}
+        self._file_observer = wd_polling.PollingObserver()
+        self._file_observer.schedule(
+            self._per_file_event_handler(), self._dir, recursive=True
+        )
+        self._file_observer.start()
+        logger.info("watching files in: %s", settings.files_dir)
+
+    @property
+    def emitter(self) -> Optional["wd_api.EventEmitter"]:
+        try:
+            return next(iter(self._file_observer.emitters))
+        except StopIteration:
+            return None
+
+    def update_policy(self, path: GlobStr, policy: "PolicyName") -> None:
+        # When we're dealing with one of our own media files, there's no need
+        # to store the policy in memory.  _get_file_event_handler will always
+        # return PolicyNow.  Using the path makes syncing historic runs much
+        # faster if the name happens to include glob escapable characters.  In
+        # the future we may add a flag to "files" records that indicates it's
+        # policy is not dynamic and doesn't need to be stored / checked.
+        save_name = LogicalPath(
+            os.path.relpath(os.path.join(self._dir, path), self._dir)
+        )
+        if save_name.startswith("media/"):
+            pass
+        elif path == glob.escape(path):
+            self._savename_file_policies[save_name] = policy
+        else:
+            self._user_file_policies[policy].add(path)
+
+        for src_path in glob.glob(os.path.join(self._dir, path)):
+            save_name = LogicalPath(os.path.relpath(src_path, self._dir))
+            feh = self._get_file_event_handler(src_path, save_name)
+            # handle the case where the policy changed
+            if feh.policy != policy:
+                try:
+                    del self._file_event_handlers[save_name]
+                except KeyError:
+                    # TODO: probably should do locking, but this handles moved files for now
+                    pass
+                feh = self._get_file_event_handler(src_path, save_name)
+            feh.on_modified(force=True)
+
+    def _per_file_event_handler(self) -> "wd_events.FileSystemEventHandler":
+        """Create a Watchdog file event handler that does different things for every file."""
+        file_event_handler = wd_events.PatternMatchingEventHandler()
+        file_event_handler.on_created = self._on_file_created
+        file_event_handler.on_modified = self._on_file_modified
+        file_event_handler.on_moved = self._on_file_moved
+        file_event_handler._patterns = [os.path.join(self._dir, os.path.normpath("*"))]
+        # Ignore hidden files/folders
+        #  TODO: what other files should we skip?
+        file_event_handler._ignore_patterns = [
+            "*.tmp",
+            "*.wandb",
+            "wandb-summary.json",
+            os.path.join(self._dir, ".*"),
+            os.path.join(self._dir, "*/.*"),
+        ]
+        for glb in self._settings.ignore_globs:
+            file_event_handler._ignore_patterns.append(os.path.join(self._dir, glb))
+
+        return file_event_handler
+
+    def _on_file_created(self, event: "wd_events.FileCreatedEvent") -> None:
+        logger.info("file/dir created: %s", event.src_path)
+        if os.path.isdir(event.src_path):
+            return None
+        self._file_count += 1
+        # We do the directory scan less often as it grows
+        if self._file_count % 100 == 0:
+            emitter = self.emitter
+            if emitter:
+                emitter._timeout = int(self._file_count / 100) + 1
+        save_name = LogicalPath(os.path.relpath(event.src_path, self._dir))
+        self._get_file_event_handler(event.src_path, save_name).on_modified()
+
+    # TODO(spencerpearson): this pattern repeats so many times we should have a method/function for it
+    # def _save_name(self, path: PathStr) -> LogicalPath:
+    #     return LogicalPath(os.path.relpath(path, self._dir))
+
+    def _on_file_modified(self, event: "wd_events.FileModifiedEvent") -> None:
+        logger.info(f"file/dir modified: {event.src_path}")
+        if os.path.isdir(event.src_path):
+            return None
+        save_name = LogicalPath(os.path.relpath(event.src_path, self._dir))
+        self._get_file_event_handler(event.src_path, save_name).on_modified()
+
+    def _on_file_moved(self, event: "wd_events.FileMovedEvent") -> None:
+        # TODO: test me...
+        logger.info(f"file/dir moved: {event.src_path} -> {event.dest_path}")
+        if os.path.isdir(event.dest_path):
+            return None
+        old_save_name = LogicalPath(os.path.relpath(event.src_path, self._dir))
+        new_save_name = LogicalPath(os.path.relpath(event.dest_path, self._dir))
+
+        # We have to move the existing file handler to the new name
+        handler = self._get_file_event_handler(event.src_path, old_save_name)
+        self._file_event_handlers[new_save_name] = handler
+        del self._file_event_handlers[old_save_name]
+
+        handler.on_renamed(event.dest_path, new_save_name)
+
+    def _get_file_event_handler(
+        self, file_path: PathStr, save_name: LogicalPath
+    ) -> FileEventHandler:
+        """Get or create an event handler for a particular file.
+
+        file_path: the file's actual path
+        save_name: its path relative to the run directory (aka the watch directory)
+        """
+        # Always return PolicyNow for any of our media files.
+        if save_name.startswith("media/"):
+            return PolicyNow(file_path, save_name, self._file_pusher, self._settings)
+        if save_name not in self._file_event_handlers:
+            # TODO: we can use PolicyIgnore if there are files we never want to sync
+            if "tfevents" in save_name or "graph.pbtxt" in save_name:
+                self._file_event_handlers[save_name] = PolicyLive(
+                    file_path, save_name, self._file_pusher, self._settings
+                )
+            elif save_name in self._savename_file_policies:
+                policy_name = self._savename_file_policies[save_name]
+                make_handler = (
+                    PolicyLive
+                    if policy_name == "live"
+                    else PolicyNow
+                    if policy_name == "now"
+                    else PolicyEnd
+                )
+                self._file_event_handlers[save_name] = make_handler(
+                    file_path, save_name, self._file_pusher, self._settings
+                )
+            else:
+                make_handler = PolicyEnd
+                for policy, globs in self._user_file_policies.items():
+                    if policy == "end":
+                        continue
+                    # Convert set to list to avoid RuntimeError's
+                    # TODO: we may need to add locks
+                    for g in list(globs):
+                        paths = glob.glob(os.path.join(self._dir, g))
+                        if any(save_name in p for p in paths):
+                            if policy == "live":
+                                make_handler = PolicyLive
+                            elif policy == "now":
+                                make_handler = PolicyNow
+                self._file_event_handlers[save_name] = make_handler(
+                    file_path, save_name, self._file_pusher, self._settings
+                )
+        return self._file_event_handlers[save_name]
+
+    def finish(self) -> None:
+        logger.info("shutting down directory watcher")
+        try:
+            # avoid hanging if we crashed before the observer was started
+            if self._file_observer.is_alive():
+                # rather unfortunately we need to manually do a final scan of the dir
+                # with `queue_events`, then iterate through all events before stopping
+                # the observer to catch all files written.  First we need to prevent the
+                # existing thread from consuming our final events, then we process them
+                self._file_observer._timeout = 0
+                self._file_observer._stopped_event.set()
+                self._file_observer.join()
+                self.emitter.queue_events(0)  # type: ignore[union-attr]
+                while True:
+                    try:
+                        self._file_observer.dispatch_events(
+                            self._file_observer.event_queue, 0
+                        )
+                    except queue.Empty:
+                        break
+                # Calling stop unschedules any inflight events so we handled them above
+                self._file_observer.stop()
+        # TODO: py2 TypeError: PyCObject_AsVoidPtr called with null pointer
+        except TypeError:
+            pass
+        # TODO: py3 SystemError: <built-in function stop> returned an error
+        except SystemError:
+            pass
+
+        # Ensure we've at least noticed every file in the run directory. Sometimes
+        # we miss things because asynchronously watching filesystems isn't reliable.
+        logger.info("scan: %s", self._dir)
+
+        for dirpath, _, filenames in os.walk(self._dir):
+            for fname in filenames:
+                file_path = os.path.join(dirpath, fname)
+                save_name = LogicalPath(os.path.relpath(file_path, self._dir))
+                ignored = False
+                for glb in self._settings.ignore_globs:
+                    if len(fnmatch.filter([save_name], glb)) > 0:
+                        ignored = True
+                        logger.info("ignored: %s matching glob %s", save_name, glb)
+                        break
+                if ignored:
+                    continue
+                logger.info("scan save: %s %s", file_path, save_name)
+                self._get_file_event_handler(file_path, save_name).finish()

wandb/filesync/stats.py

@@ -0,0 +1,100 @@
+import threading
+from typing import MutableMapping, NamedTuple
+
+import wandb
+
+
+class FileStats(NamedTuple):
+    deduped: bool
+    total: int
+    uploaded: int
+    failed: bool
+    artifact_file: bool
+
+
+class Summary(NamedTuple):
+    uploaded_bytes: int
+    total_bytes: int
+    deduped_bytes: int
+
+
+class FileCountsByCategory(NamedTuple):
+    artifact: int
+    wandb: int
+    media: int
+    other: int
+
+
+class Stats:
+    def __init__(self) -> None:
+        self._stats: MutableMapping[str, FileStats] = {}
+        self._lock = threading.Lock()
+
+    def init_file(
+        self, save_name: str, size: int, is_artifact_file: bool = False
+    ) -> None:
+        with self._lock:
+            self._stats[save_name] = FileStats(
+                deduped=False,
+                total=size,
+                uploaded=0,
+                failed=False,
+                artifact_file=is_artifact_file,
+            )
+
+    def set_file_deduped(self, save_name: str) -> None:
+        with self._lock:
+            orig = self._stats[save_name]
+            self._stats[save_name] = orig._replace(
+                deduped=True,
+                uploaded=orig.total,
+            )
+
+    def update_uploaded_file(self, save_name: str, total_uploaded: int) -> None:
+        with self._lock:
+            self._stats[save_name] = self._stats[save_name]._replace(
+                uploaded=total_uploaded,
+            )
+
+    def update_failed_file(self, save_name: str) -> None:
+        with self._lock:
+            self._stats[save_name] = self._stats[save_name]._replace(
+                uploaded=0,
+                failed=True,
+            )
+
+    def summary(self) -> Summary:
+        # Need to use list to ensure we get a copy, since other threads may
+        # modify this while we iterate
+        with self._lock:
+            stats = list(self._stats.values())
+        return Summary(
+            uploaded_bytes=sum(f.uploaded for f in stats),
+            total_bytes=sum(f.total for f in stats),
+            deduped_bytes=sum(f.total for f in stats if f.deduped),
+        )
+
+    def file_counts_by_category(self) -> FileCountsByCategory:
+        artifact_files = 0
+        wandb_files = 0
+        media_files = 0
+        other_files = 0
+        # Need to use list to ensure we get a copy, since other threads may
+        # modify this while we iterate
+        with self._lock:
+            file_stats = list(self._stats.items())
+        for save_name, stats in file_stats:
+            if stats.artifact_file:
+                artifact_files += 1
+            elif wandb.wandb_lib.filenames.is_wandb_file(save_name):  # type: ignore[attr-defined]  # TODO(spencerpearson): this is probably synonymous with wandb.sdk.lib.filenames...?
+                wandb_files += 1
+            elif save_name.startswith("media"):
+                media_files += 1
+            else:
+                other_files += 1
+        return FileCountsByCategory(
+            artifact=artifact_files,
+            wandb=wandb_files,
+            media=media_files,
+            other=other_files,
+        )

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 TypeError
+
+        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))