wandb 0.19.6__py3-none-macosx_11_0_arm64.whl → 0.19.7__py3-none-macosx_11_0_arm64.whl

wandb/__init__.py

@@ -10,7 +10,7 @@ For reference documentation, see https://docs.wandb.com/ref/python.
 """
 from __future__ import annotations
 
-__version__ = "0.19.6"
+__version__ = "0.19.7"
 
 
 from wandb.errors import Error

wandb/__init__.pyi

@@ -63,6 +63,7 @@ from typing import (
     Any,
     Callable,
     Dict,
+    Iterable,
     List,
     Literal,
     Optional,
@@ -103,7 +104,7 @@ if TYPE_CHECKING:
     import wandb
     from wandb.plot import CustomChart
 
-__version__: str = "0.19.6"
+__version__: str = "0.19.7"
 
 run: Run | None
 config: wandb_config.Config
@@ -114,6 +115,25 @@ _sentry: Sentry
 api: InternalApi
 patched: Dict[str, List[Callable]]
 
+def require(
+    requirement: str | Iterable[str] | None = None,
+    experiment: str | Iterable[str] | None = None,
+) -> None:
+    """Indicate which experimental features are used by the script.
+
+    This should be called before any other `wandb` functions, ideally right
+    after importing `wandb`.
+
+    Args:
+        requirement: The name of a feature to require or an iterable of
+            feature names.
+        experiment: An alias for `requirement`.
+
+    Raises:
+        wandb.errors.UnsupportedError: If a feature name is unknown.
+    """
+    ...
+
 def setup(settings: Settings | None = None) -> _WandbSetup:
     """Prepares W&B for use in the current process and its children.
 
@@ -271,10 +291,10 @@ def init(
             on the system, such as checking the git root or the current program
             file. If we can't infer the project name, the project will default to
             `"uncategorized"`.
-        dir: An absolute path to the directory where metadata and downloaded
-            files will be stored. When calling `download()` on an artifact, files
-            will be saved to this directory. If not specified, this defaults to
-            the `./wandb` directory.
+        dir: The absolute path to the directory where experiment logs and
+            metadata files are stored. If not specified, this defaults
+            to the `./wandb` directory. Note that this does not affect the
+            location where artifacts are stored when calling `download()`.
         id: A unique identifier for this run, used for resuming. It must be unique
             within the project and cannot be reused once a run is deleted. The
             identifier must not contain any of the following special characters:

wandb/apis/public/_generated/__init__.py

@@ -0,0 +1,21 @@
+# 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,
+    ServerFeaturesQueryServerInfo,
+    ServerFeaturesQueryServerInfoFeatures,
+)
+
+__all__ = [
+    "Base",
+    "GQLBase",
+    "GQLId",
+    "SerializedToJson",
+    "Typename",
+    "SERVER_FEATURES_QUERY_GQL",
+    "ServerFeaturesQuery",
+    "ServerFeaturesQueryServerInfo",
+    "ServerFeaturesQueryServerInfoFeatures",
+]

wandb/apis/public/_generated/base.py

@@ -0,0 +1,128 @@
+# Generated by ariadne-codegen
+
+from __future__ import annotations
+
+from typing import Any, Literal, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field, Json, ValidationError, WrapValidator
+from pydantic.alias_generators import to_camel
+from pydantic.main import IncEx
+from pydantic_core import to_json
+from pydantic_core.core_schema import ValidatorFunctionWrapHandler
+
+from .typing_compat import Annotated, override
+
+
+# Base class for all automation classes/types.
+# Omitted from docstring to avoid inclusion in generated docs.
+class Base(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+        validate_assignment=True,
+        validate_default=True,
+        extra="forbid",
+        alias_generator=to_camel,
+        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.
+# Omitted from docstring to avoid 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),
+]
+
+
+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/_generated/enums.py

@@ -0,0 +1,4 @@
+# Generated by ariadne-codegen
+# Source: core/api/graphql/schemas/schema-latest.graphql
+
+from __future__ import annotations

wandb/apis/public/_generated/input_types.py

@@ -0,0 +1,4 @@
+# Generated by ariadne-codegen
+# Source: core/api/graphql/schemas/schema-latest.graphql
+
+from __future__ import annotations

wandb/apis/public/_generated/operations.py

@@ -0,0 +1,15 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/utils/
+
+__all__ = ["SERVER_FEATURES_QUERY_GQL"]
+
+SERVER_FEATURES_QUERY_GQL = """
+query ServerFeaturesQuery {
+  serverInfo {
+    features {
+      name
+      isEnabled
+    }
+  }
+}
+"""

wandb/apis/public/_generated/server_features_query.py

@@ -0,0 +1,27 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/utils/
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import Field
+
+from .base import GQLBase
+
+
+class ServerFeaturesQuery(GQLBase):
+    server_info: Optional["ServerFeaturesQueryServerInfo"] = Field(alias="serverInfo")
+
+
+class ServerFeaturesQueryServerInfo(GQLBase):
+    features: List[Optional["ServerFeaturesQueryServerInfoFeatures"]]
+
+
+class ServerFeaturesQueryServerInfoFeatures(GQLBase):
+    name: str
+    is_enabled: bool = Field(alias="isEnabled")
+
+
+ServerFeaturesQuery.model_rebuild()
+ServerFeaturesQueryServerInfo.model_rebuild()

wandb/apis/public/_generated/typing_compat.py

@@ -0,0 +1,14 @@
+# Generated by ariadne-codegen
+
+"""Definitions to ensure compatibility with all supported python versions."""
+
+import sys
+
+if sys.version_info >= (3, 12):
+    from typing import Annotated, override
+else:
+    from typing_extensions import Annotated, override
+
+
+Annnotated = Annotated
+override = override

wandb/apis/public/api.py

@@ -25,8 +25,15 @@ import wandb
 from wandb import env, util
 from wandb.apis import public
 from wandb.apis.normalize import normalize_exceptions
+from wandb.apis.public._generated import SERVER_FEATURES_QUERY_GQL, ServerFeaturesQuery
 from wandb.apis.public.const import RETRY_TIMEDELTA
-from wandb.apis.public.utils import PathType, parse_org_from_registry_path
+from wandb.apis.public.registries import Registries
+from wandb.apis.public.utils import (
+    PathType,
+    fetch_org_from_settings_or_entity,
+    parse_org_from_registry_path,
+)
+from wandb.proto.wandb_internal_pb2 import ServerFeature
 from wandb.sdk.artifacts._validators import is_artifact_registry_project
 from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
@@ -248,6 +255,9 @@ class Api:
             self.settings["entity"] = _overrides["username"]
         self.settings["base_url"] = self.settings["base_url"].rstrip("/")
 
+        if "organization" in _overrides:
+            self.settings["organization"] = _overrides["organization"]
+
         self._viewer = None
         self._projects = {}
         self._runs = {}
@@ -279,6 +289,7 @@ class Api:
             )
         )
         self._client = RetryingClient(self._base_client)
+        self._server_features_cache = None
 
     def create_project(self, name: str, entity: str) -> None:
         """Create a new project.
@@ -900,12 +911,46 @@ class Api:
     ):
         """Return a set of runs from a project that match the filters provided.
 
-        You can filter by `config.*`, `summary_metrics.*`, `tags`, `state`, `entity`, `createdAt`, etc.
+        Fields you can filter by include:
+        - `createdAt`: The timestamp when the run was created. (in ISO 8601 format, e.g. "2023-01-01T12:00:00Z")
+        - `displayName`: The human-readable display name of the run. (e.g. "eager-fox-1")
+        - `duration`: The total runtime of the run in seconds.
+        - `group`: The group name used to organize related runs together.
+        - `host`: The hostname where the run was executed.
+        - `jobType`: The type of job or purpose of the run.
+        - `name`: The unique identifier of the run. (e.g. "a1b2cdef")
+        - `state`: The current state of the run.
+        - `tags`: The tags associated with the run.
+        - `username`: The username of the user who initiated the run
+
+        Additionally, you can filter by items in the run config or summary metrics.
+        Such as `config.experiment_name`, `summary_metrics.loss`, etc.
+
+        For more complex filtering, you can use MongoDB query operators.
+        For details, see: https://docs.mongodb.com/manual/reference/operator/query
+        The following operations are supported:
+        - `$and`
+        - `$or`
+        - `$nor`
+        - `$eq`
+        - `$ne`
+        - `$gt`
+        - `$gte`
+        - `$lt`
+        - `$lte`
+        - `$in`
+        - `$nin`
+        - `$exists`
+        - `$regex`
+
 
         Examples:
             Find runs in my_project where config.experiment_name has been set to "foo"
             ```
-            api.runs(path="my_entity/my_project", filters={"config.experiment_name": "foo"})
+            api.runs(
+                path="my_entity/my_project",
+                filters={"config.experiment_name": "foo"},
+            )
             ```
 
             Find runs in my_project where config.experiment_name has been set to "foo" or "bar"
@@ -932,7 +977,24 @@ class Api:
             Find runs in my_project where the run name matches a regex (anchors are not supported)
             ```
             api.runs(
-                path="my_entity/my_project", filters={"display_name": {"$regex": "^foo.*"}}
+                path="my_entity/my_project",
+                filters={"display_name": {"$regex": "^foo.*"}},
+            )
+            ```
+
+            Find runs in my_project where config.experiment contains a nested field "category" with value "testing"
+            ```
+            api.runs(
+                path="my_entity/my_project",
+                filters={"config.experiment.category": "testing"},
+            )
+            ```
+
+            Find runs in my_project with a loss value of 0.5 nested in a dictionary under model1 in the summary metrics
+            ```
+            api.runs(
+                path="my_entity/my_project",
+                filters={"summary_metrics.model1.loss": 0.5},
             )
             ```
 
@@ -947,8 +1009,6 @@ class Api:
                 You can filter by run properties such as config.key, summary_metrics.key, state, entity, createdAt, etc.
                 For example: `{"config.experiment_name": "foo"}` would find runs with a config entry
                     of experiment name set to "foo"
-                You can compose operations to make more complicated queries,
-                    see Reference for the language is at  https://docs.mongodb.com/manual/reference/operator/query
             order: (str) Order can be `created_at`, `heartbeat_at`, `config.*.value`, or `summary_metrics.*`.
                 If you prepend order with a + order is ascending.
                 If you prepend order with a - order is descending (default).
@@ -1139,6 +1199,12 @@ class Api:
             entity = InternalApi()._resolve_org_entity_name(
                 entity=settings_entity, organization=org
             )
+
+        if entity is None:
+            raise ValueError(
+                "Could not determine entity. Please include the entity as part of the collection name path."
+            )
+
         return public.ArtifactCollection(
             self.client, entity, project, collection_name, type_name
         )
@@ -1211,6 +1277,12 @@ class Api:
             entity = InternalApi()._resolve_org_entity_name(
                 entity=settings_entity, organization=organization
             )
+
+        if entity is None:
+            raise ValueError(
+                "Could not determine entity. Please include the entity as part of the artifact name path."
+            )
+
         artifact = wandb.Artifact._from_name(
             entity=entity,
             project=project,
@@ -1385,3 +1457,117 @@ class Api:
             return True
         except wandb.errors.CommError:
             return False
+
+    def registries(
+        self,
+        organization: Optional[str] = None,
+        filter: Optional[Dict[str, Any]] = None,
+    ) -> Registries:
+        """Returns a Registry iterator.
+
+        Use the iterator to search and filter registries, collections,
+        or artifact versions across your organization's registry.
+
+        Examples:
+            Find all registries with the names that contain "model"
+            ```python
+            import wandb
+
+            api = wandb.Api()  # specify an org if your entity belongs to multiple orgs
+            api.registries(filter={"name": {"$regex": "model"}})
+            ```
+
+            Find all collections in the registries with the name "my_collection" and the tag "my_tag"
+            ```python
+            api.registries().collections(filter={"name": "my_collection", "tag": "my_tag"})
+            ```
+
+            Find all artifact versions in the registries with a collection name that contains "my_collection" and a version that has the alias "best"
+            ```python
+            api.registries().collections(
+                filter={"name": {"$regex": "my_collection"}}
+            ).versions(filter={"alias": "best"})
+            ```
+
+            Find all artifact versions in the registries that contain "model" and have the tag "prod" or alias "best"
+            ```python
+            api.registries(filter={"name": {"$regex": "model"}}).versions(
+                filter={"$or": [{"tag": "prod"}, {"alias": "best"}]}
+            )
+            ```
+
+        Args:
+            organization: (str, optional) The organization of the registry to fetch.
+                If not specified, use the organization specified in the user's settings.
+            filter: (dict, optional) MongoDB-style filter to apply to each object in the registry iterator.
+                Fields available to filter for collections are
+                    `name`, `description`, `created_at`, `updated_at`.
+                Fields available to filter for collections are
+                    `name`, `tag`, `description`, `created_at`, `updated_at`
+                Fields available to filter for versions are
+                    `tag`, `alias`, `created_at`, `updated_at`, `metadata`
+
+        Returns:
+            A registry iterator.
+        """
+        if not self._check_server_feature_with_fallback(
+            ServerFeature.ARTIFACT_REGISTRY_SEARCH
+        ):
+            raise RuntimeError(
+                "Registry search API is not enabled on this wandb server version. "
+                "Please upgrade your server version or contact support at support@wandb.com."
+            )
+
+        organization = organization or fetch_org_from_settings_or_entity(
+            self.settings, self.default_entity
+        )
+        return Registries(self.client, organization, filter)
+
+    def _check_server_feature(self, feature: ServerFeature) -> bool:
+        """Check if a server feature is enabled.
+
+        Args:
+            feature (ServerFeature): 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:
+            response = self.client.execute(gql(SERVER_FEATURES_QUERY_GQL))
+            self._server_features_cache = ServerFeaturesQuery.model_validate(response)
+
+        feature_name = ServerFeature.Name(feature)
+        if (
+            self._server_features_cache
+            and self._server_features_cache.server_info
+            and self._server_features_cache.server_info.features
+        ):
+            for feature_info in self._server_features_cache.server_info.features:
+                if feature_info and feature_info.name == feature_name:
+                    return feature_info.is_enabled
+
+        return False
+
+    def _check_server_feature_with_fallback(self, feature: 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 (ServerFeature): 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)
+        except Exception as e:
+            if 'Cannot query field "features" on type "ServerInfo".' in str(e):
+                return False
+            raise e

wandb/apis/public/artifacts.py

@@ -12,56 +12,16 @@ from wandb.apis import public
 from wandb.apis.normalize import normalize_exceptions
 from wandb.apis.paginator import Paginator
 from wandb.errors.term import termlog
+from wandb.sdk.artifacts._graphql_fragments import (
+    ARTIFACT_FILES_FRAGMENT,
+    ARTIFACTS_TYPES_FRAGMENT,
+)
 from wandb.sdk.lib import deprecate
 
 if TYPE_CHECKING:
     from wandb.apis.public import RetryingClient, Run
 
 
-ARTIFACTS_TYPES_FRAGMENT = """
-fragment ArtifactTypesFragment on ArtifactTypeConnection {
-    edges {
-         node {
-             id
-             name
-             description
-             createdAt
-         }
-         cursor
-    }
-    pageInfo {
-        endCursor
-        hasNextPage
-    }
-}
-"""
-
-# TODO, factor out common file fragment
-ARTIFACT_FILES_FRAGMENT = """fragment ArtifactFilesFragment on Artifact {
-    files(names: $fileNames, after: $fileCursor, first: $fileLimit) {
-        edges {
-            node {
-                id
-                name: displayName
-                url
-                sizeBytes
-                storagePath
-                mimetype
-                updatedAt
-                digest
-                md5
-                directUrl
-            }
-            cursor
-        }
-        pageInfo {
-            endCursor
-            hasNextPage
-        }
-    }
-}"""
-
-
 class ArtifactTypes(Paginator):
     QUERY = gql(
         """
@@ -317,6 +277,7 @@ class ArtifactCollection:
         project: str,
         name: str,
         type: str,
+        organization: Optional[str] = None,
         attrs: Optional[Mapping[str, Any]] = None,
     ):
         self.client = client
@@ -327,11 +288,14 @@ class ArtifactCollection:
         self._type = type
         self._saved_type = type
         self._attrs = attrs
-        self.load()
+        if self._attrs is None:
+            self.load()
         self._aliases = [a["node"]["alias"] for a in self._attrs["aliases"]["edges"]]
         self._description = self._attrs["description"]
+        self._created_at = self._attrs["createdAt"]
         self._tags = [a["node"]["name"] for a in self._attrs["tags"]["edges"]]
         self._saved_tags = copy(self._tags)
+        self.organization = organization
 
     @property
     def id(self):
@@ -354,6 +318,10 @@ class ArtifactCollection:
         """Artifact Collection Aliases."""
         return self._aliases
 
+    @property
+    def created_at(self):
+        return self._created_at
+
     def load(self):
         query = gql(
             """