wandb 0.18.0__py3-none-macosx_10_13_x86_64.whl

wandb/sdk/data_types/_private.py

@@ -0,0 +1,10 @@
+import atexit
+import tempfile
+
+# Staging directory, so we can encode raw data into files, then hash them before
+# we put them into the Run directory to be uploaded.
+MEDIA_TMP = tempfile.TemporaryDirectory("wandb-media")
+
+
+def _cleanup_media_tmp_dir() -> None:
+    atexit.register(MEDIA_TMP.cleanup)

wandb/sdk/data_types/base_types/json_metadata.py

@@ -0,0 +1,55 @@
+import codecs
+import os
+from typing import TYPE_CHECKING, Type, Union
+
+from wandb import util
+from wandb.sdk.lib import runid
+
+from .._private import MEDIA_TMP
+from .media import Media
+
+if TYPE_CHECKING:  # pragma: no cover
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ...wandb_run import Run as LocalRun
+
+
+# Allows encoding of arbitrary JSON structures
+# as a file
+#
+# This class should be used as an abstract class
+# extended to have validation methods
+
+
+class JSONMetadata(Media):
+    """JSONMetadata is a type for encoding arbitrary metadata as files."""
+
+    def __init__(self, val: dict) -> None:
+        super().__init__()
+
+        self.validate(val)
+        self._val = val
+
+        ext = "." + self.type_name() + ".json"
+        tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ext)
+        with codecs.open(tmp_path, "w", encoding="utf-8") as fp:
+            util.json_dump_uncompressed(self._val, fp)
+        self._set_file(tmp_path, is_tmp=True, extension=ext)
+
+    @classmethod
+    def get_media_subdir(cls: Type["JSONMetadata"]) -> str:
+        return os.path.join("media", "metadata", cls.type_name())
+
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        json_dict = super().to_json(run_or_artifact)
+        json_dict["_type"] = self.type_name()
+
+        return json_dict
+
+    # These methods should be overridden in the child class
+    @classmethod
+    def type_name(cls) -> str:
+        return "metadata"
+
+    def validate(self, val: dict) -> bool:
+        return True

wandb/sdk/data_types/base_types/media.py

@@ -0,0 +1,315 @@
+import hashlib
+import os
+import platform
+import re
+import shutil
+from typing import TYPE_CHECKING, Optional, Sequence, Type, Union, cast
+
+import wandb
+from wandb import util
+from wandb._globals import _datatypes_callback
+from wandb.sdk.lib import filesystem
+from wandb.sdk.lib.paths import LogicalPath
+
+from .wb_value import WBValue
+
+if TYPE_CHECKING:  # pragma: no cover
+    import numpy as np
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ...wandb_run import Run as LocalRun
+
+
+SYS_PLATFORM = platform.system()
+
+
+def _wb_filename(
+    key: Union[str, int], step: Union[str, int], id: Union[str, int], extension: str
+) -> str:
+    return f"{str(key)}_{str(step)}_{str(id)}{extension}"
+
+
+class Media(WBValue):
+    """A WBValue stored as a file outside JSON that can be rendered in a media panel.
+
+    If necessary, we move or copy the file into the Run's media directory so that it
+    gets uploaded.
+    """
+
+    _path: Optional[str]
+    _run: Optional["LocalRun"]
+    _caption: Optional[str]
+    _is_tmp: Optional[bool]
+    _extension: Optional[str]
+    _sha256: Optional[str]
+    _size: Optional[int]
+
+    def __init__(self, caption: Optional[str] = None) -> None:
+        super().__init__()
+        self._path = None
+        # The run under which this object is bound, if any.
+        self._run = None
+        self._caption = caption
+
+    def _set_file(
+        self, path: str, is_tmp: bool = False, extension: Optional[str] = None
+    ) -> None:
+        self._path = path
+        self._is_tmp = is_tmp
+        self._extension = extension
+        assert extension is None or path.endswith(
+            extension
+        ), f'Media file extension "{extension}" must occur at the end of path "{path}".'
+
+        with open(self._path, "rb") as f:
+            self._sha256 = hashlib.sha256(f.read()).hexdigest()
+        self._size = os.path.getsize(self._path)
+
+    @classmethod
+    def get_media_subdir(cls: Type["Media"]) -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    def captions(
+        media_items: Sequence["Media"],
+    ) -> Union[bool, Sequence[Optional[str]]]:
+        if media_items[0]._caption is not None:
+            return [m._caption for m in media_items]
+        else:
+            return False
+
+    def is_bound(self) -> bool:
+        return self._run is not None
+
+    def file_is_set(self) -> bool:
+        return self._path is not None and self._sha256 is not None
+
+    def bind_to_run(
+        self,
+        run: "LocalRun",
+        key: Union[int, str],
+        step: Union[int, str],
+        id_: Optional[Union[int, str]] = None,
+        ignore_copy_err: Optional[bool] = None,
+    ) -> None:
+        """Bind this object to a particular Run.
+
+        Calling this function is necessary so that we have somewhere specific to put the
+        file associated with this object, from which other Runs can refer to it.
+        """
+        assert self.file_is_set(), "bind_to_run called before _set_file"
+
+        if SYS_PLATFORM == "Windows" and not util.check_windows_valid_filename(key):
+            raise ValueError(
+                f"Media {key} is invalid. Please remove invalid filename characters"
+            )
+
+        # The following two assertions are guaranteed to pass
+        # by definition file_is_set, but are needed for
+        # mypy to understand that these are strings below.
+        assert isinstance(self._path, str)
+        assert isinstance(self._sha256, str)
+
+        assert run is not None, 'Argument "run" must not be None.'
+        self._run = run
+
+        if self._extension is None:
+            _, extension = os.path.splitext(os.path.basename(self._path))
+        else:
+            extension = self._extension
+
+        if id_ is None:
+            id_ = self._sha256[:20]
+
+        file_path = _wb_filename(key, step, id_, extension)
+        media_path = os.path.join(self.get_media_subdir(), file_path)
+        new_path = os.path.join(self._run.dir, media_path)
+        filesystem.mkdir_exists_ok(os.path.dirname(new_path))
+
+        if self._is_tmp:
+            shutil.move(self._path, new_path)
+            self._path = new_path
+            self._is_tmp = False
+            _datatypes_callback(media_path)
+        else:
+            try:
+                shutil.copy(self._path, new_path)
+            except shutil.SameFileError as e:
+                if not ignore_copy_err:
+                    raise e
+            self._path = new_path
+            _datatypes_callback(media_path)
+
+    def to_json(self, run: Union["LocalRun", "Artifact"]) -> dict:
+        """Serialize the object into a JSON blob.
+
+        Uses run or artifact to store additional data. If `run_or_artifact` is a
+        wandb.Run then `self.bind_to_run()` must have been previously been called.
+
+        Args:
+            run_or_artifact (wandb.Run | wandb.Artifact): the Run or Artifact for which
+                this object should be generating JSON for - this is useful to store
+                additional data if needed.
+
+        Returns:
+            dict: JSON representation
+        """
+        # NOTE: uses of Audio in this class are a temporary hack -- when Ref support moves up
+        # into Media itself we should get rid of them
+        from wandb import Image
+        from wandb.data_types import Audio
+        from wandb.sdk.wandb_run import Run
+
+        json_obj = {}
+
+        if isinstance(run, Run):
+            json_obj.update(
+                {
+                    "_type": "file",  # TODO(adrian): This isn't (yet) a real media type we support on the frontend.
+                    "sha256": self._sha256,
+                    "size": self._size,
+                }
+            )
+            artifact_entry_url = self._get_artifact_entry_ref_url()
+            if artifact_entry_url is not None:
+                json_obj["artifact_path"] = artifact_entry_url
+            artifact_entry_latest_url = self._get_artifact_entry_latest_ref_url()
+            if artifact_entry_latest_url is not None:
+                json_obj["_latest_artifact_path"] = artifact_entry_latest_url
+
+            if artifact_entry_url is None or self.is_bound():
+                assert self.is_bound(), "Value of type {} must be bound to a run with bind_to_run() before being serialized to JSON.".format(
+                    type(self).__name__
+                )
+
+                assert (
+                    self._run is run
+                ), "We don't support referring to media files across runs."
+
+                # The following two assertions are guaranteed to pass
+                # by definition is_bound, but are needed for
+                # mypy to understand that these are strings below.
+                assert isinstance(self._path, str)
+                json_obj["path"] = LogicalPath(
+                    os.path.relpath(self._path, self._run.dir)
+                )
+
+        elif isinstance(run, wandb.Artifact):
+            if self.file_is_set():
+                # The following two assertions are guaranteed to pass
+                # by definition of the call above, but are needed for
+                # mypy to understand that these are strings below.
+                assert isinstance(self._path, str)
+                assert isinstance(self._sha256, str)
+                artifact = run  # Checks if the concrete image has already been added to this artifact
+                name = artifact.get_added_local_path_name(self._path)
+                if name is None:
+                    if self._is_tmp:
+                        name = os.path.join(
+                            self.get_media_subdir(), os.path.basename(self._path)
+                        )
+                    else:
+                        # If the files is not temporary, include the first 8 characters of the file's SHA256 to
+                        # avoid name collisions. This way, if there are two images `dir1/img.png` and `dir2/img.png`
+                        # we end up with a unique path for each.
+                        name = os.path.join(
+                            self.get_media_subdir(),
+                            self._sha256[:20],
+                            os.path.basename(self._path),
+                        )
+
+                    # if not, check to see if there is a source artifact for this object
+                    if (
+                        self._artifact_source is not None
+                        # and self._artifact_source.artifact != artifact
+                    ):
+                        default_root = self._artifact_source.artifact._default_root()
+                        # if there is, get the name of the entry (this might make sense to move to a helper off artifact)
+                        if self._path.startswith(default_root):
+                            name = self._path[len(default_root) :]
+                            name = name.lstrip(os.sep)
+
+                        # Add this image as a reference
+                        path = self._artifact_source.artifact.get_entry(name)
+                        artifact.add_reference(path.ref_url(), name=name)
+                    elif (
+                        isinstance(self, Audio) or isinstance(self, Image)
+                    ) and self.path_is_reference(self._path):
+                        artifact.add_reference(self._path, name=name)
+                    else:
+                        entry = artifact.add_file(
+                            self._path, name=name, is_tmp=self._is_tmp
+                        )
+                        name = entry.path
+
+                json_obj["path"] = name
+                json_obj["sha256"] = self._sha256
+            json_obj["_type"] = self._log_type
+        return json_obj
+
+    @classmethod
+    def from_json(
+        cls: Type["Media"], json_obj: dict, source_artifact: "Artifact"
+    ) -> "Media":
+        """Likely will need to override for any more complicated media objects."""
+        return cls(source_artifact.get_entry(json_obj["path"]).download())
+
+    def __eq__(self, other: object) -> bool:
+        """Likely will need to override for any more complicated media objects."""
+        return (
+            isinstance(other, self.__class__)
+            and hasattr(self, "_sha256")
+            and hasattr(other, "_sha256")
+            and self._sha256 == other._sha256
+        )
+
+    @staticmethod
+    def path_is_reference(path: Optional[str]) -> bool:
+        return bool(path and re.match(r"^(gs|s3|https?)://", path))
+
+
+class BatchableMedia(Media):
+    """Media that is treated in batches.
+
+    E.g. images and thumbnails. Apart from images, we just use these batches to help
+    organize files by name in the media directory.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @classmethod
+    def seq_to_json(
+        cls: Type["BatchableMedia"],
+        seq: Sequence["BatchableMedia"],
+        run: "LocalRun",
+        key: str,
+        step: Union[int, str],
+    ) -> dict:
+        raise NotImplementedError
+
+
+def _numpy_arrays_to_lists(
+    payload: Union[dict, Sequence, "np.ndarray"],
+) -> Union[Sequence, dict, str, int, float, bool]:
+    # Casts all numpy arrays to lists so we don't convert them to histograms, primarily for Plotly
+
+    if isinstance(payload, dict):
+        res = {}
+        for key, val in payload.items():
+            res[key] = _numpy_arrays_to_lists(val)
+        return res
+    elif isinstance(payload, Sequence) and not isinstance(payload, str):
+        return [_numpy_arrays_to_lists(v) for v in payload]
+    elif util.is_numpy_array(payload):
+        if TYPE_CHECKING:
+            payload = cast("np.ndarray", payload)
+        return [
+            _numpy_arrays_to_lists(v)
+            for v in (payload.tolist() if payload.ndim > 0 else [payload.tolist()])
+        ]
+    # Protects against logging non serializable objects
+    elif isinstance(payload, Media):
+        return str(payload.__class__.__name__)
+    return payload  # type: ignore

wandb/sdk/data_types/base_types/wb_value.py

@@ -0,0 +1,274 @@
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Type, Union
+
+from wandb import util
+
+if TYPE_CHECKING:  # pragma: no cover
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ...wandb_run import Run as LocalRun
+
+    TypeMappingType = Dict[str, Type["WBValue"]]
+
+
+def _server_accepts_client_ids() -> bool:
+    from wandb.util import parse_version
+
+    # First, if we are offline, assume the backend server cannot
+    # accept client IDs. Unfortunately, this is the best we can do
+    # until we are sure that all local versions are > "0.11.0" max_cli_version.
+    # The practical implication is that tables logged in offline mode
+    # will not show up in the workspace (but will still show up in artifacts). This
+    # means we never lose data, and we can still view using weave. If we decided
+    # to use client ids in offline mode, then the manifests and artifact data
+    # would never be resolvable and would lead to failed uploads. Our position
+    # is to never lose data - and instead take the tradeoff in the UI.
+    if util._is_offline():
+        return False
+
+    # If the script is online, request the max_cli_version and ensure the server
+    # is of a high enough version.
+    max_cli_version = util._get_max_cli_version()
+    if max_cli_version is None:
+        return False
+    accepts_client_ids: bool = parse_version("0.11.0") <= parse_version(max_cli_version)
+    return accepts_client_ids
+
+
+class _WBValueArtifactSource:
+    artifact: "Artifact"
+    name: Optional[str]
+
+    def __init__(self, artifact: "Artifact", name: Optional[str] = None) -> None:
+        self.artifact = artifact
+        self.name = name
+
+
+class _WBValueArtifactTarget:
+    artifact: "Artifact"
+    name: Optional[str]
+
+    def __init__(self, artifact: "Artifact", name: Optional[str] = None) -> None:
+        self.artifact = artifact
+        self.name = name
+
+
+class WBValue:
+    """Typed objects that can be logged with `wandb.log()` and visualized by wandb.
+
+    The objects will be serialized as JSON and always have a _type attribute that
+    indicates how to interpret the other fields.
+    """
+
+    # Class Attributes
+    _type_mapping: ClassVar[Optional["TypeMappingType"]] = None
+    # override _log_type to indicate the type which the subclass deserializes
+    _log_type: ClassVar[Optional[str]] = None
+
+    # Instance Attributes
+    _artifact_source: Optional[_WBValueArtifactSource]
+    _artifact_target: Optional[_WBValueArtifactTarget]
+
+    def __init__(self) -> None:
+        self._artifact_source = None
+        self._artifact_target = None
+
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        """Serialize the object into a JSON blob.
+
+        Uses current run or artifact to store additional data.
+
+        Args:
+            run_or_artifact (wandb.Run | wandb.Artifact): the Run or Artifact for which
+                this object should be generating JSON for - this is useful to to store
+                additional data if needed.
+
+        Returns:
+            dict: JSON representation
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def from_json(
+        cls: Type["WBValue"], json_obj: dict, source_artifact: "Artifact"
+    ) -> "WBValue":
+        """Deserialize a `json_obj` into it's class representation.
+
+        If additional resources were stored in the `run_or_artifact` artifact during the
+        `to_json` call, then those resources should be in the `source_artifact`.
+
+        Args:
+            json_obj (dict): A JSON dictionary to deserialize source_artifact
+            (wandb.Artifact): An artifact which will hold any additional
+                resources which were stored during the `to_json` function.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def with_suffix(cls: Type["WBValue"], name: str, filetype: str = "json") -> str:
+        """Get the name with the appropriate suffix.
+
+        Args:
+            name (str): the name of the file
+            filetype (str, optional): the filetype to use. Defaults to "json".
+
+        Returns:
+            str: a filename which is suffixed with it's `_log_type` followed by the
+                filetype.
+        """
+        if cls._log_type is not None:
+            suffix = cls._log_type + "." + filetype
+        else:
+            suffix = filetype
+        if not name.endswith(suffix):
+            return name + "." + suffix
+        return name
+
+    @staticmethod
+    def init_from_json(
+        json_obj: dict, source_artifact: "Artifact"
+    ) -> Optional["WBValue"]:
+        """Initialize a `WBValue` from a JSON blob based on the class that creatd it.
+
+        Looks through all subclasses and tries to match the json obj with the class
+        which created it. It will then call that subclass' `from_json` method.
+        Importantly, this function will set the return object's `source_artifact`
+        attribute to the passed in source artifact. This is critical for artifact
+        bookkeeping. If you choose to create a wandb.Value via it's `from_json` method,
+        make sure to properly set this `artifact_source` to avoid data duplication.
+
+        Args:
+            json_obj (dict): A JSON dictionary to deserialize. It must contain a `_type`
+                key. This is used to lookup the correct subclass to use.
+            source_artifact (wandb.Artifact): An artifact which will hold any additional
+                resources which were stored during the `to_json` function.
+
+        Returns:
+            wandb.Value: a newly created instance of a subclass of wandb.Value
+        """
+        class_option = WBValue.type_mapping().get(json_obj["_type"])
+        if class_option is not None:
+            obj = class_option.from_json(json_obj, source_artifact)
+            obj._set_artifact_source(source_artifact)
+            return obj
+
+        return None
+
+    @staticmethod
+    def type_mapping() -> "TypeMappingType":
+        """Return a map from `_log_type` to subclass. Used to lookup correct types for deserialization.
+
+        Returns:
+            dict: dictionary of str:class
+        """
+        if WBValue._type_mapping is None:
+            WBValue._type_mapping = {}
+            frontier = [WBValue]
+            explored = set()
+            while len(frontier) > 0:
+                class_option = frontier.pop()
+                explored.add(class_option)
+                if class_option._log_type is not None:
+                    WBValue._type_mapping[class_option._log_type] = class_option
+                for subclass in class_option.__subclasses__():
+                    if subclass not in explored:
+                        frontier.append(subclass)
+        return WBValue._type_mapping
+
+    def __eq__(self, other: object) -> bool:
+        return id(self) == id(other)
+
+    def __ne__(self, other: object) -> bool:
+        return not self.__eq__(other)
+
+    def to_data_array(self) -> List[Any]:
+        """Convert the object to a list of primitives representing the underlying data."""
+        raise NotImplementedError
+
+    def _set_artifact_source(
+        self, artifact: "Artifact", name: Optional[str] = None
+    ) -> None:
+        assert (
+            self._artifact_source is None
+        ), "Cannot update artifact_source. Existing source: {}/{}".format(
+            self._artifact_source.artifact, self._artifact_source.name
+        )
+        self._artifact_source = _WBValueArtifactSource(artifact, name)
+
+    def _set_artifact_target(
+        self, artifact: "Artifact", name: Optional[str] = None
+    ) -> None:
+        assert (
+            self._artifact_target is None
+        ), "Cannot update artifact_target. Existing target: {}/{}".format(
+            self._artifact_target.artifact, self._artifact_target.name
+        )
+        self._artifact_target = _WBValueArtifactTarget(artifact, name)
+
+    def _get_artifact_entry_ref_url(self) -> Optional[str]:
+        # If the object is coming from another artifact
+        if self._artifact_source and self._artifact_source.name:
+            ref_entry = self._artifact_source.artifact.get_entry(
+                type(self).with_suffix(self._artifact_source.name)
+            )
+            return str(ref_entry.ref_url())
+        # Else, if the object is destined for another artifact and we support client IDs
+        elif (
+            self._artifact_target
+            and self._artifact_target.name
+            and self._artifact_target.artifact._client_id is not None
+            and self._artifact_target.artifact._final
+            and _server_accepts_client_ids()
+        ):
+            return "wandb-client-artifact://{}/{}".format(
+                self._artifact_target.artifact._client_id,
+                type(self).with_suffix(self._artifact_target.name),
+            )
+        # Else if we do not support client IDs, but online, then block on upload
+        # Note: this is old behavior just to stay backwards compatible
+        # with older server versions. This code path should be removed
+        # once those versions are no longer supported. This path uses a .wait
+        # which blocks the user process on artifact upload.
+        elif (
+            self._artifact_target
+            and self._artifact_target.name
+            and self._artifact_target.artifact._is_draft_save_started()
+            and not util._is_offline()
+            and not _server_accepts_client_ids()
+        ):
+            self._artifact_target.artifact.wait()
+            ref_entry = self._artifact_target.artifact.get_entry(
+                type(self).with_suffix(self._artifact_target.name)
+            )
+            return str(ref_entry.ref_url())
+        return None
+
+    def _get_artifact_entry_latest_ref_url(self) -> Optional[str]:
+        if (
+            self._artifact_target
+            and self._artifact_target.name
+            and self._artifact_target.artifact._client_id is not None
+            and self._artifact_target.artifact._final
+            and _server_accepts_client_ids()
+        ):
+            return "wandb-client-artifact://{}:latest/{}".format(
+                self._artifact_target.artifact._sequence_client_id,
+                type(self).with_suffix(self._artifact_target.name),
+            )
+        # Else if we do not support client IDs, then block on upload
+        # Note: this is old behavior just to stay backwards compatible
+        # with older server versions. This code path should be removed
+        # once those versions are no longer supported. This path uses a .wait
+        # which blocks the user process on artifact upload.
+        elif (
+            self._artifact_target
+            and self._artifact_target.name
+            and self._artifact_target.artifact._is_draft_save_started()
+            and not util._is_offline()
+            and not _server_accepts_client_ids()
+        ):
+            self._artifact_target.artifact.wait()
+            ref_entry = self._artifact_target.artifact.get_entry(
+                type(self).with_suffix(self._artifact_target.name)
+            )
+            return str(ref_entry.ref_url())
+        return None