wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/artifacts/storage_policies/s3_bucket_policy.py

@@ -0,0 +1,61 @@
+"""S3 bucket storage policy."""
+from typing import TYPE_CHECKING, Dict, Optional, Sequence, Union
+
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.storage_handlers.local_file_handler import LocalFileHandler
+from wandb.sdk.artifacts.storage_handlers.multi_handler import MultiHandler
+from wandb.sdk.artifacts.storage_handlers.s3_handler import S3Handler
+from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandler
+from wandb.sdk.artifacts.storage_policy import StoragePolicy
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.lib.paths import FilePathStr, URIStr
+
+
+# Don't use this yet!
+class __S3BucketPolicy(StoragePolicy):  # noqa: N801
+    @classmethod
+    def name(cls) -> str:
+        return "wandb-s3-bucket-policy-v1"
+
+    @classmethod
+    def from_config(cls, config: Dict[str, str]) -> "__S3BucketPolicy":
+        if "bucket" not in config:
+            raise ValueError("Bucket name not found in config")
+        return cls(config["bucket"])
+
+    def __init__(self, bucket: str) -> None:
+        self._bucket = bucket
+        s3 = S3Handler(bucket)
+        local = LocalFileHandler()
+
+        self._handler = MultiHandler(
+            handlers=[
+                s3,
+                local,
+            ],
+            default_handler=TrackingHandler(),
+        )
+
+    def config(self) -> Dict[str, str]:
+        return {"bucket": self._bucket}
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> Union[URIStr, FilePathStr]:
+        return self._handler.load_path(manifest_entry, local=local)
+
+    def store_path(
+        self,
+        artifact: "Artifact",
+        path: Union[URIStr, FilePathStr],
+        name: Optional[str] = None,
+        checksum: bool = True,
+        max_objects: Optional[int] = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        return self._handler.store_path(
+            artifact, path, name=name, checksum=checksum, max_objects=max_objects
+        )

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

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

wandb/sdk/{interface/artifacts/artifact_storage.py → artifacts/storage_policy.py}

@@ -1,20 +1,15 @@
+"""Storage policy."""
 from typing import TYPE_CHECKING, Dict, Optional, Sequence, Type, Union
 
 from wandb.sdk.lib.paths import FilePathStr, URIStr
 
 if TYPE_CHECKING:
-    from urllib.parse import ParseResult
-
     from wandb.filesync.step_prepare import StepPrepare
-    from wandb.sdk.interface.artifacts import Artifact, ArtifactManifestEntry
+    from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
     from wandb.sdk.internal.progress import ProgressFn
 
 
-class StorageLayout:
-    V1 = "V1"
-    V2 = "V2"
-
-
 class StoragePolicy:
     @classmethod
     def lookup_by_name(cls, name: str) -> Optional[Type["StoragePolicy"]]:
@@ -36,7 +31,7 @@ class StoragePolicy:
 
     def load_file(
         self, artifact: "Artifact", manifest_entry: "ArtifactManifestEntry"
-    ) -> str:
+    ) -> FilePathStr:
         raise NotImplementedError
 
     def store_file_sync(
@@ -74,52 +69,5 @@ class StoragePolicy:
         self,
         manifest_entry: "ArtifactManifestEntry",
         local: bool = False,
-    ) -> str:
-        raise NotImplementedError
-
-
-class StorageHandler:
-    def can_handle(self, parsed_url: "ParseResult") -> bool:
-        """Checks whether this handler can handle the given url.
-
-        Returns:
-            Whether this handler can handle the given url.
-        """
-        raise NotImplementedError
-
-    def load_path(
-        self,
-        manifest_entry: "ArtifactManifestEntry",
-        local: bool = False,
-    ) -> Union[URIStr, FilePathStr]:
-        """Load a file or directory given the corresponding index entry.
-
-        Args:
-            manifest_entry: The index entry to load
-            local: Whether to load the file locally or not
-
-        Returns:
-            A path to the file represented by `index_entry`
-        """
-        raise NotImplementedError
-
-    def store_path(
-        self,
-        artifact: "Artifact",
-        path: Union[URIStr, FilePathStr],
-        name: Optional[str] = None,
-        checksum: bool = True,
-        max_objects: Optional[int] = None,
-    ) -> Sequence["ArtifactManifestEntry"]:
-        """Store the file or directory at the given path to the specified artifact.
-
-        Args:
-            path: The path to store
-            name: If specified, the logical name that should map to `path`
-            checksum: Whether to compute the checksum of the file
-            max_objects: The maximum number of objects to store
-
-        Returns:
-            A list of manifest entries to store within the artifact
-        """
+    ) -> Union[FilePathStr, URIStr]:
         raise NotImplementedError

wandb/sdk/data_types/_dtypes.py

@@ -13,8 +13,7 @@ from wandb.util import (
 np = get_module("numpy")  # intentionally not required
 
 if t.TYPE_CHECKING:
-    from wandb.apis.public import Artifact as DownloadedArtifact
-    from wandb.sdk.wandb_artifacts import Artifact as ArtifactInCreation
+    from wandb.sdk.artifacts.artifact import Artifact
 
 _TYPES_STRIPPED = not (sys.version_info.major == 3 and sys.version_info.minor >= 6)
 if not _TYPES_STRIPPED:
@@ -77,7 +76,7 @@ class TypeRegistry:
 
     @staticmethod
     def type_from_dict(
-        json_dict: t.Dict[str, t.Any], artifact: t.Optional["DownloadedArtifact"] = None
+        json_dict: t.Dict[str, t.Any], artifact: t.Optional["Artifact"] = None
     ) -> "Type":
         wb_type = json_dict.get("wb_type")
         if wb_type is None:
@@ -135,7 +134,7 @@ class TypeRegistry:
 
 def _params_obj_to_json_obj(
     params_obj: t.Any,
-    artifact: t.Optional["ArtifactInCreation"] = None,
+    artifact: t.Optional["Artifact"] = None,
 ) -> t.Any:
     """Helper method."""
     if params_obj.__class__ == dict:
@@ -152,7 +151,7 @@ def _params_obj_to_json_obj(
 
 
 def _json_obj_to_params_obj(
-    json_obj: t.Any, artifact: t.Optional["DownloadedArtifact"] = None
+    json_obj: t.Any, artifact: t.Optional["Artifact"] = None
 ) -> t.Any:
     """Helper method."""
     if json_obj.__class__ == dict:
@@ -222,9 +221,7 @@ class Type:
         else:
             return InvalidType()
 
-    def to_json(
-        self, artifact: t.Optional["ArtifactInCreation"] = None
-    ) -> t.Dict[str, t.Any]:
+    def to_json(self, artifact: t.Optional["Artifact"] = None) -> t.Dict[str, t.Any]:
         """Generate a jsonable dictionary serialization the type.
 
         If overridden by subclass, ensure that `from_json` is equivalently overridden.
@@ -249,7 +246,7 @@ class Type:
     def from_json(
         cls,
         json_dict: t.Dict[str, t.Any],
-        artifact: t.Optional["DownloadedArtifact"] = None,
+        artifact: t.Optional["Artifact"] = None,
     ) -> "Type":
         """Construct a new instance of the type using a JSON dictionary.
 
@@ -756,9 +753,7 @@ class NDArrayType(Type):
 
         return InvalidType()
 
-    def to_json(
-        self, artifact: t.Optional["ArtifactInCreation"] = None
-    ) -> t.Dict[str, t.Any]:
+    def to_json(self, artifact: t.Optional["Artifact"] = None) -> t.Dict[str, t.Any]:
         # custom override to support serialization path outside of params internal dict
         res = {
             "wb_type": self.name,

wandb/sdk/data_types/base_types/json_metadata.py

@@ -9,7 +9,8 @@ from .._private import MEDIA_TMP
 from .media import Media
 
 if TYPE_CHECKING:  # pragma: no cover
-    from ...wandb_artifacts import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
+
     from ...wandb_run import Run as LocalRun
 
 
@@ -39,7 +40,7 @@ class JSONMetadata(Media):
     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", "LocalArtifact"]) -> dict:
+    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()
 

wandb/sdk/data_types/base_types/media.py

@@ -9,15 +9,15 @@ 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  # type: ignore
+    import numpy as np
 
-    from wandb.apis.public import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
-    from ...wandb_artifacts import Artifact as LocalArtifact
     from ...wandb_run import Run as LocalRun
 
 
@@ -143,7 +143,7 @@ class Media(WBValue):
             self._path = new_path
             _datatypes_callback(media_path)
 
-    def to_json(self, run: Union["LocalRun", "LocalArtifact"]) -> dict:
+    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
@@ -193,11 +193,11 @@ class Media(WBValue):
                 # by definition is_bound, but are needed for
                 # mypy to understand that these are strings below.
                 assert isinstance(self._path, str)
-                json_obj["path"] = util.to_forward_slash_path(
+                json_obj["path"] = LogicalPath(
                     os.path.relpath(self._path, self._run.dir)
                 )
 
-        elif isinstance(run, wandb.wandb_sdk.wandb_artifacts.Artifact):
+        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
@@ -253,7 +253,7 @@ class Media(WBValue):
 
     @classmethod
     def from_json(
-        cls: Type["Media"], json_obj: dict, source_artifact: "PublicArtifact"
+        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_path(json_obj["path"]).download())
@@ -315,4 +315,4 @@ def _numpy_arrays_to_lists(
     # Protects against logging non serializable objects
     elif isinstance(payload, Media):
         return str(payload.__class__.__name__)
-    return payload
+    return payload  # type: ignore