wandb 0.19.7__py3-none-win32.whl → 0.19.9__py3-none-win32.whl

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

@@ -1,6 +1,7 @@
 """sender."""
 
 import contextlib
+import glob
 import gzip
 import json
 import logging
@@ -1408,7 +1409,8 @@ class SendManager:
         for k in files.files:
             # TODO(jhr): fix paths with directories
             self._save_file(
-                interface.GlobStr(k.path), interface.file_enum_to_policy(k.policy)
+                interface.GlobStr(glob.escape(k.path)),
+                interface.file_enum_to_policy(k.policy),
             )
 
     def send_header(self, record: "Record") -> None:
@@ -1529,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,

wandb/sdk/internal/settings_static.py

@@ -21,11 +21,33 @@ class SettingsStatic(Settings):
     def _proto_to_dict(self, proto: wandb_settings_pb2.Settings) -> dict:
         data = {}
 
+        exclude_fields = {
+            "model_config",
+            "model_fields",
+            "model_fields_set",
+            "__fields__",
+            "__model_fields_set",
+            "__pydantic_self__",
+            "__pydantic_initialised__",
+        }
+
+        fields = (
+            Settings.model_fields
+            if hasattr(Settings, "model_fields")
+            else Settings.__fields__
+        )  # type: ignore [attr-defined]
+
+        fields = {k: v for k, v in fields.items() if k not in exclude_fields}  # type: ignore [union-attr]
+
         forks_specified: list[str] = []
-        for key in Settings.model_fields:  # type: ignore [attr-defined]
+        for key in fields:
+            # Skip Python-only keys that do not exist on the proto.
+            if key in ("reinit",):
+                continue
+
             value: Any = None
 
-            field_info = Settings.model_fields[key]
+            field_info = fields[key]
             annotation = str(field_info.annotation)
 
             if key == "_stats_open_metrics_filters":

wandb/sdk/lib/apikey.py

@@ -11,7 +11,7 @@ import textwrap
 from functools import partial
 
 # import Literal
-from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional, Union
+from typing import TYPE_CHECKING, Callable, Literal
 from urllib.parse import urlparse
 
 import click
@@ -23,6 +23,9 @@ from wandb.errors import term
 from wandb.errors.links import url_registry
 from wandb.util import _is_databricks, isatty, prompt_choices
 
+if TYPE_CHECKING:
+    from wandb.sdk.wandb_settings import Settings
+
 LOGIN_CHOICE_ANON = "Private W&B dashboard, no account required"
 LOGIN_CHOICE_NEW = "Create a W&B account"
 LOGIN_CHOICE_EXISTS = "Use an existing W&B account"
@@ -50,18 +53,14 @@ class WriteNetrcError(Exception):
 Mode = Literal["allow", "must", "never", "false", "true"]
 
 
-if TYPE_CHECKING:
-    from wandb.sdk.wandb_settings import Settings
-
-
 getpass = partial(click.prompt, hide_input=True, err=True)
 
 
-def _fixup_anon_mode(default: Optional[Mode]) -> Optional[Mode]:
+def _fixup_anon_mode(default: Mode | None) -> Mode | None:
     # Convert weird anonymode values from legacy settings files
     # into one of our expected values.
     anon_mode = default or "never"
-    mapping: Dict[Mode, Mode] = {"true": "allow", "false": "never"}
+    mapping: dict[Mode, Mode] = {"true": "allow", "false": "never"}
     return mapping.get(anon_mode, anon_mode)
 
 
@@ -85,14 +84,14 @@ def get_netrc_file_path() -> str:
 
 
 def prompt_api_key(  # noqa: C901
-    settings: "Settings",
-    api: Optional[InternalApi] = None,
-    input_callback: Optional[Callable] = None,
-    browser_callback: Optional[Callable] = None,
+    settings: Settings,
+    api: InternalApi | None = None,
+    input_callback: Callable | None = None,
+    browser_callback: Callable | None = None,
     no_offline: bool = False,
     no_create: bool = False,
     local: bool = False,
-) -> Union[str, bool, None]:
+) -> str | bool | None:
     """Prompt for api key.
 
     Returns:
@@ -276,9 +275,9 @@ def write_netrc(host: str, entity: str, key: str):
 
 
 def write_key(
-    settings: "Settings",
-    key: Optional[str],
-    api: Optional["InternalApi"] = None,
+    settings: Settings,
+    key: str | None,
+    api: InternalApi | None = None,
 ) -> None:
     if not key:
         raise ValueError("No API key specified.")
@@ -298,7 +297,7 @@ def write_key(
     write_netrc(settings.base_url, "user", key)
 
 
-def api_key(settings: Optional["Settings"] = None) -> Optional[str]:
+def api_key(settings: Settings | None = None) -> str | None:
     if settings is None:
         settings = wandb.setup().settings
     if settings.api_key: