wandb 0.19.8__py3-none-win_amd64.whl → 0.19.9__py3-none-win_amd64.whl

wandb/sdk/data_types/molecule.py

@@ -53,9 +53,7 @@ class Molecule(BatchableMedia):
         caption: Optional[str] = None,
         **kwargs: str,
     ) -> None:
-        super().__init__()
-
-        self._caption = caption
+        super().__init__(caption=caption)
 
         if hasattr(data_or_path, "name"):
             # if the file has a path, we just detect the type and copy it from there
@@ -208,8 +206,6 @@ class Molecule(BatchableMedia):
     def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = self._log_type
-        if self._caption:
-            json_dict["caption"] = self._caption
         return json_dict
 
     @classmethod

wandb/sdk/data_types/object_3d.py

@@ -215,9 +215,10 @@ class Object3D(BatchableMedia):
     def __init__(
         self,
         data_or_path: Union["np.ndarray", str, "TextIO", dict],
+        caption: Optional[str] = None,
         **kwargs: Optional[Union[str, "FileFormat3D"]],
     ) -> None:
-        super().__init__()
+        super().__init__(caption=caption)
 
         if hasattr(data_or_path, "name"):
             # if the file has a path, we just detect the type and copy it from there

wandb/sdk/data_types/saved_model.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import os
 import shutil
 import sys
+from types import ModuleType
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
 
 import wandb
@@ -15,9 +16,6 @@ from ._private import MEDIA_TMP
 from .base_types.wb_value import WBValue
 
 if TYPE_CHECKING:
-    from types import ModuleType
-
-    import cloudpickle  # type: ignore
     import sklearn  # type: ignore
     import tensorflow  # type: ignore
     import torch  # type: ignore
@@ -264,9 +262,9 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         self._serialize(self._model_obj, target_path)
 
 
-def _get_cloudpickle() -> "cloudpickle":
+def _get_cloudpickle() -> ModuleType:
     return cast(
-        "cloudpickle",
+        ModuleType,
         util.get_module("cloudpickle", "ModelAdapter requires `cloudpickle`"),
     )
 
@@ -338,9 +336,9 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
         return json_obj
 
 
-def _get_torch() -> "torch":
+def _get_torch() -> ModuleType:
     return cast(
-        "torch",
+        ModuleType,
         util.get_module("torch", "ModelAdapter requires `torch`"),
     )
 
@@ -366,9 +364,9 @@ class _PytorchSavedModel(_PicklingSavedModel["torch.nn.Module"]):
         )
 
 
-def _get_sklearn() -> "sklearn":
+def _get_sklearn() -> ModuleType:
     return cast(
-        "sklearn",
+        ModuleType,
         util.get_module("sklearn", "ModelAdapter requires `sklearn`"),
     )
 

wandb/sdk/data_types/video.py

@@ -90,13 +90,12 @@ class Video(BatchableMedia):
         fps: Optional[int] = None,
         format: Optional[str] = None,
     ):
-        super().__init__()
+        super().__init__(caption=caption)
 
         self._format = format or "gif"
         self._width = None
         self._height = None
         self._channels = None
-        self._caption = caption
         if self._format not in Video.EXTS:
             raise ValueError(
                 "wandb.Video accepts {} formats".format(", ".join(Video.EXTS))
@@ -190,8 +189,6 @@ class Video(BatchableMedia):
             json_dict["width"] = self._width
         if self._height is not None:
             json_dict["height"] = self._height
-        if self._caption:
-            json_dict["caption"] = self._caption
 
         return json_dict
 

wandb/{apis/public → sdk/internal}/_generated/__init__.py

@@ -1,6 +1,5 @@
 # Generated by ariadne-codegen
 
-from .base import Base, GQLBase, GQLId, SerializedToJson, Typename
 from .operations import SERVER_FEATURES_QUERY_GQL
 from .server_features_query import (
     ServerFeaturesQuery,
@@ -9,11 +8,6 @@ from .server_features_query import (
 )
 
 __all__ = [
-    "Base",
-    "GQLBase",
-    "GQLId",
-    "SerializedToJson",
-    "Typename",
     "SERVER_FEATURES_QUERY_GQL",
     "ServerFeaturesQuery",
     "ServerFeaturesQueryServerInfo",

wandb/sdk/internal/_generated/base.py

@@ -0,0 +1,226 @@
+# Generated by ariadne-codegen
+
+"""This module defines base classes for generated types, including partial support for compatibility with Pydantic v1."""
+
+from __future__ import annotations
+
+from contextlib import suppress
+from importlib.metadata import version
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
+
+import pydantic
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic import main as pydantic_main
+from typing_extensions import Annotated, override
+
+if TYPE_CHECKING:
+    from pydantic.main import IncEx
+
+pydantic_major_version, *_ = version(pydantic.__name__).split(".")
+IS_PYDANTIC_V2: bool = int(pydantic_major_version) >= 2
+
+# Maps {v2 -> v1} model config keys that were renamed in v2.
+# See: https://docs.pydantic.dev/latest/migration/#changes-to-config
+_V1_CONFIG_KEYS = {
+    "populate_by_name": "allow_population_by_field_name",
+    "str_to_lower": "anystr_lower",
+    "str_strip_whitespace": "anystr_strip_whitespace",
+    "str_to_upper": "anystr_upper",
+    "ignored_types": "keep_untouched",
+    "str_max_length": "max_anystr_length",
+    "str_min_length": "min_anystr_length",
+    "from_attributes": "orm_mode",
+    "json_schema_extra": "schema_extra",
+    "validate_default": "validate_all",
+}
+
+
+def _convert_v2_config(v2_config: dict[str, Any]) -> dict[str, Any]:
+    """Return a copy of the v2 ConfigDict with renamed v1 keys."""
+    return {_V1_CONFIG_KEYS.get(k, k): v for k, v in v2_config.items()}
+
+
+if IS_PYDANTIC_V2:
+    PydanticModelMetaclass = type  # placeholder
+else:
+    PydanticModelMetaclass = pydantic_main.ModelMetaclass
+
+
+class _V1MixinMetaclass(PydanticModelMetaclass):
+    def __new__(
+        cls,
+        name: str,
+        bases: tuple[type, ...],
+        namespace: dict[str, Any],
+        **kwargs: Any,
+    ):
+        # Converts a model config in a v2 class definition, e.g.:
+        #
+        #     class MyModel(BaseModel):
+        #         model_config = ConfigDict(populate_by_name=True)
+        #
+        # ...to a Config class in a v1 class definition, e.g.:
+        #
+        #     class MyModel(BaseModel):
+        #         class Config:
+        #             populate_by_name = True
+        #
+        if config_dict := namespace.pop("model_config", None):
+            namespace["Config"] = type("Config", (), _convert_v2_config(config_dict))
+        return super().__new__(cls, name, bases, namespace, **kwargs)
+
+
+class PydanticV1Mixin(metaclass=_V1MixinMetaclass):
+    @classmethod
+    def __try_update_forward_refs__(cls, **localns: Any) -> None:
+        with suppress(AttributeError):
+            super().__try_update_forward_refs__(**localns)
+
+    @classmethod
+    def model_rebuild(cls, *args: Any, **kwargs: Any) -> None:
+        return cls.update_forward_refs(*args, **kwargs)
+
+    @classmethod
+    def model_construct(cls, *args: Any, **kwargs: Any) -> Any:
+        return cls.construct(*args, **kwargs)
+
+    @classmethod
+    def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
+        return cls.parse_obj(*args, **kwargs)
+
+    @classmethod
+    def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
+        return cls.parse_raw(*args, **kwargs)
+
+    def model_dump(self, *args: Any, **kwargs: Any) -> Any:
+        return self.dict(*args, **kwargs)
+
+    def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
+        return self.json(*args, **kwargs)
+
+    def model_copy(self, *args: Any, **kwargs: Any) -> Any:
+        return self.copy(*args, **kwargs)
+
+
+class PydanticV2Mixin:
+    # Placeholder: Pydantic v2 is already compatible with itself, so no need for extra mixins.
+    pass
+
+
+# Pick the mixin type based on the detected Pydantic version.
+PydanticCompatMixin = PydanticV2Mixin if IS_PYDANTIC_V2 else PydanticV1Mixin
+
+
+# Base class for all generated classes/types.
+# This is deliberately not a docstring to prevent inclusion in generated docs.
+class Base(BaseModel, PydanticCompatMixin):
+    model_config = ConfigDict(
+        populate_by_name=True,
+        validate_assignment=True,
+        validate_default=True,
+        extra="forbid",
+        use_attribute_docstrings=True,
+        from_attributes=True,
+        revalidate_instances="always",
+    )
+
+    @override
+    def model_dump(
+        self,
+        *,
+        mode: Literal["json", "python"] | str = "json",  # NOTE: changed default
+        include: IncEx | None = None,
+        exclude: IncEx | None = None,
+        context: dict[str, Any] | None = None,
+        by_alias: bool = True,  # NOTE: changed default
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        round_trip: bool = True,  # NOTE: changed default
+        warnings: bool | Literal["none", "warn", "error"] = True,
+        serialize_as_any: bool = False,
+    ) -> dict[str, Any]:
+        return super().model_dump(
+            mode=mode,
+            include=include,
+            exclude=exclude,
+            context=context,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+            serialize_as_any=serialize_as_any,
+        )
+
+    @override
+    def model_dump_json(
+        self,
+        *,
+        indent: int | None = None,
+        include: IncEx | None = None,
+        exclude: IncEx | None = None,
+        context: dict[str, Any] | None = None,
+        by_alias: bool = True,  # NOTE: changed default
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        round_trip: bool = True,  # NOTE: changed default
+        warnings: bool | Literal["none", "warn", "error"] = True,
+        serialize_as_any: bool = False,
+    ) -> str:
+        return super().model_dump_json(
+            indent=indent,
+            include=include,
+            exclude=exclude,
+            context=context,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+            serialize_as_any=serialize_as_any,
+        )
+
+
+# Base class with extra customization for GQL generated types.
+# This is deliberately not a docstring to prevent inclusion in generated docs.
+class GQLBase(Base):
+    model_config = ConfigDict(
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+
+# ------------------------------------------------------------------------------
+# Reusable annotations for field types
+T = TypeVar("T")
+
+GQLId = Annotated[
+    str,
+    Field(repr=False, strict=True, frozen=True),
+]
+
+Typename = Annotated[
+    T,
+    Field(repr=False, alias="__typename", frozen=True),
+]
+
+
+# FIXME: Restore or modify this after ensuring pydantic v1 compatibility.
+# def validate_maybe_json(v: Any, handler: ValidatorFunctionWrapHandler) -> Any:
+#     """Wraps default Json[...] field validator to allow instantiation with an already-decoded value."""
+#     try:
+#         return handler(v)
+#     except ValidationError:
+#         # Try revalidating after properly jsonifying the value
+#         return handler(to_json(v, by_alias=True, round_trip=True))
+
+
+# SerializedToJson = Annotated[
+#     Json[T],
+#     # Allow lenient instantiation/validation: incoming data may already be deserialized.
+#     WrapValidator(validate_maybe_json),
+# ]

wandb/{apis/public → sdk/internal}/_generated/server_features_query.py

@@ -7,15 +7,15 @@ from typing import List, Optional
 
 from pydantic import Field
 
-from .base import GQLBase
+from wandb._pydantic import GQLBase
 
 
 class ServerFeaturesQuery(GQLBase):
-    server_info: Optional["ServerFeaturesQueryServerInfo"] = Field(alias="serverInfo")
+    server_info: Optional[ServerFeaturesQueryServerInfo] = Field(alias="serverInfo")
 
 
 class ServerFeaturesQueryServerInfo(GQLBase):
-    features: List[Optional["ServerFeaturesQueryServerInfoFeatures"]]
+    features: List[Optional[ServerFeaturesQueryServerInfoFeatures]]
 
 
 class ServerFeaturesQueryServerInfoFeatures(GQLBase):

wandb/{apis/public → sdk/internal}/_generated/typing_compat.py

@@ -10,5 +10,5 @@ else:
     from typing_extensions import Annotated, override
 
 
-Annnotated = Annotated
+Annotated = Annotated
 override = override

wandb/sdk/internal/internal_api.py

@@ -36,6 +36,7 @@ import requests
 import yaml
 from wandb_gql import Client, gql
 from wandb_gql.client import RetryError
+from wandb_graphql.language.ast import Document
 
 import wandb
 from wandb import env, util
@@ -43,7 +44,9 @@ from wandb.apis.normalize import normalize_exceptions, parse_backend_error_messa
 from wandb.errors import AuthenticationError, CommError, UnsupportedError, UsageError
 from wandb.integration.sagemaker import parse_sm_secrets
 from wandb.old.settings import Settings
+from wandb.proto.wandb_internal_pb2 import ServerFeature
 from wandb.sdk.artifacts._validators import is_artifact_registry_project
+from wandb.sdk.internal._generated import SERVER_FEATURES_QUERY_GQL, ServerFeaturesQuery
 from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
 from wandb.sdk.lib.gql_request import GraphQLSession
 from wandb.sdk.lib.hashutil import B64MD5, md5_file_b64
@@ -365,6 +368,7 @@ class Api:
         self.server_create_run_queue_supports_priority: Optional[bool] = None
         self.server_supports_template_variables: Optional[bool] = None
         self.server_push_to_run_queue_supports_priority: Optional[bool] = None
+        self._server_features_cache: Optional[dict[str, bool]] = None
 
     def gql(self, *args: Any, **kwargs: Any) -> Any:
         ret = self._retry_gql(
@@ -869,6 +873,52 @@ class Api:
         _, _, mutations = self.server_info_introspection()
         return "updateRunQueueItemWarning" in mutations
 
+    def _check_server_feature(self, feature_value: ServerFeature) -> bool:
+        """Check if a server feature is enabled.
+
+        Args:
+            feature_value (ServerFeature): The enum value of the feature to check.
+
+        Returns:
+            bool: True if the feature is enabled, False otherwise.
+
+        Raises:
+            Exception: If server doesn't support feature queries or other errors occur
+        """
+        if self._server_features_cache is None:
+            query = gql(SERVER_FEATURES_QUERY_GQL)
+            response = self.gql(query)
+            server_info = ServerFeaturesQuery.model_validate(response).server_info
+            if server_info and (features := server_info.features):
+                self._server_features_cache = {
+                    f.name: f.is_enabled for f in features if f
+                }
+            else:
+                self._server_features_cache = {}
+
+        return self._server_features_cache.get(ServerFeature.Name(feature_value), False)
+
+    def _check_server_feature_with_fallback(self, feature_value: ServerFeature) -> bool:
+        """Wrapper around check_server_feature that warns and returns False for older unsupported servers.
+
+        Good to use for features that have a fallback mechanism for older servers.
+
+        Args:
+            feature_value (ServerFeature): The enum value of the feature to check.
+
+        Returns:
+            bool: True if the feature is enabled, False otherwise.
+
+        Exceptions:
+            Exception: If an error other than the server not supporting feature queries occurs.
+        """
+        try:
+            return self._check_server_feature(feature_value)
+        except Exception as e:
+            if 'Cannot query field "features" on type "ServerInfo".' in str(e):
+                return False
+            raise e
+
     @normalize_exceptions
     def update_run_queue_item_warning(
         self,
@@ -3703,67 +3753,108 @@ class Api:
         else:
             raise ValueError(f"Unable to find an organization under entity {entity!r}.")
 
-    def use_artifact(
+    def _construct_use_artifact_query(
         self,
         artifact_id: str,
         entity_name: Optional[str] = None,
         project_name: Optional[str] = None,
         run_name: Optional[str] = None,
         use_as: Optional[str] = None,
-    ) -> Optional[Dict[str, Any]]:
-        query_template = """
-        mutation UseArtifact(
-            $entityName: String!,
-            $projectName: String!,
-            $runName: String!,
-            $artifactID: ID!,
-            _USED_AS_TYPE_
-        ) {
-            useArtifact(input: {
-                entityName: $entityName,
-                projectName: $projectName,
-                runName: $runName,
-                artifactID: $artifactID,
-                _USED_AS_VALUE_
-            }) {
-                artifact {
-                    id
-                    digest
-                    description
-                    state
-                    createdAt
-                    metadata
-                }
-            }
-        }
-        """
+        artifact_entity_name: Optional[str] = None,
+        artifact_project_name: Optional[str] = None,
+    ) -> Tuple[Document, Dict[str, Any]]:
+        query_vars = [
+            "$entityName: String!",
+            "$projectName: String!",
+            "$runName: String!",
+            "$artifactID: ID!",
+        ]
+        query_args = [
+            "entityName: $entityName",
+            "projectName: $projectName",
+            "runName: $runName",
+            "artifactID: $artifactID",
+        ]
 
         artifact_types = self.server_use_artifact_input_introspection()
-        if "usedAs" in artifact_types:
-            query_template = query_template.replace(
-                "_USED_AS_TYPE_", "$usedAs: String"
-            ).replace("_USED_AS_VALUE_", "usedAs: $usedAs")
-        else:
-            query_template = query_template.replace("_USED_AS_TYPE_", "").replace(
-                "_USED_AS_VALUE_", ""
-            )
-
-        query = gql(query_template)
+        if "usedAs" in artifact_types and use_as:
+            query_vars.append("$usedAs: String")
+            query_args.append("usedAs: $usedAs")
 
         entity_name = entity_name or self.settings("entity")
         project_name = project_name or self.settings("project")
         run_name = run_name or self.current_run_id
 
-        response = self.gql(
-            query,
-            variable_values={
-                "entityName": entity_name,
-                "projectName": project_name,
-                "runName": run_name,
-                "artifactID": artifact_id,
-                "usedAs": use_as,
-            },
+        variable_values: Dict[str, Any] = {
+            "entityName": entity_name,
+            "projectName": project_name,
+            "runName": run_name,
+            "artifactID": artifact_id,
+            "usedAs": use_as,
+        }
+
+        server_allows_entity_project_information = (
+            self._check_server_feature_with_fallback(
+                ServerFeature.USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION  # type: ignore
+            )
+        )
+        if server_allows_entity_project_information:
+            query_vars.extend(
+                [
+                    "$artifactEntityName: String",
+                    "$artifactProjectName: String",
+                ]
+            )
+            query_args.extend(
+                [
+                    "artifactEntityName: $artifactEntityName",
+                    "artifactProjectName: $artifactProjectName",
+                ]
+            )
+            variable_values["artifactEntityName"] = artifact_entity_name
+            variable_values["artifactProjectName"] = artifact_project_name
+
+        vars_str = ", ".join(query_vars)
+        args_str = ", ".join(query_args)
+
+        query = gql(
+            f"""
+            mutation UseArtifact({vars_str}) {{
+                useArtifact(input: {{{args_str}}}) {{
+                    artifact {{
+                        id
+                        digest
+                        description
+                        state
+                        createdAt
+                        metadata
+                    }}
+                }}
+            }}
+            """
+        )
+        return query, variable_values
+
+    def use_artifact(
+        self,
+        artifact_id: str,
+        entity_name: Optional[str] = None,
+        project_name: Optional[str] = None,
+        run_name: Optional[str] = None,
+        artifact_entity_name: Optional[str] = None,
+        artifact_project_name: Optional[str] = None,
+        use_as: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        query, variable_values = self._construct_use_artifact_query(
+            artifact_id,
+            entity_name,
+            project_name,
+            run_name,
+            use_as,
+            artifact_entity_name,
+            artifact_project_name,
         )
+        response = self.gql(query, variable_values)
 
         if response["useArtifact"]["artifact"]:
             artifact: Dict[str, Any] = response["useArtifact"]["artifact"]

wandb/sdk/internal/sender.py

@@ -1531,6 +1531,8 @@ class SendManager:
 
         metadata = json.loads(artifact.metadata) if artifact.metadata else None
         res = saver.save(
+            entity=artifact.entity,
+            project=artifact.project,
             type=artifact.type,
             name=artifact.name,
             client_id=artifact.client_id,

wandb/sdk/internal/sender_config.py

@@ -47,17 +47,14 @@ class ConfigState:
         # Add any top-level keys that aren't already set.
         self._add_unset_keys_from_subtree(old_config_tree, [])
 
-        # Unfortunately, when a user logs visualizations, we store them in the
-        # run's config. When resuming a run, we want to avoid erasing previously
-        # logged visualizations, hence this special handling:
-        self._add_unset_keys_from_subtree(
-            old_config_tree,
-            [_WANDB_INTERNAL_KEY, "visualize"],
-        )
-        self._add_unset_keys_from_subtree(
-            old_config_tree,
-            [_WANDB_INTERNAL_KEY, "viz"],
-        )
+        # When resuming a run, we want to ensure the some of the old configs keys
+        # are maintained. So we have this logic here to add back
+        # any keys that were in the old config but not in the new config
+        for key in ["viz", "visualize", "mask/class_labels"]:
+            self._add_unset_keys_from_subtree(
+                old_config_tree,
+                [_WANDB_INTERNAL_KEY, key],
+            )
 
     def _add_unset_keys_from_subtree(
         self,