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

wandb/apis/public/api.py

@@ -25,7 +25,6 @@ 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.registries import Registries
 from wandb.apis.public.utils import (
@@ -289,7 +288,7 @@ class Api:
             )
         )
         self._client = RetryingClient(self._base_client)
-        self._server_features_cache = None
+        self._server_features_cache: Optional[dict[str, bool]] = None
 
     def create_project(self, name: str, entity: str) -> None:
         """Create a new project.
@@ -757,15 +756,14 @@ class Api:
         return parts
 
     def projects(
-        self, entity: Optional[str] = None, per_page: Optional[int] = 200
+        self, entity: Optional[str] = None, per_page: int = 200
     ) -> "public.Projects":
         """Get projects for a given entity.
 
         Args:
             entity: (str) Name of the entity requested.  If None, will fall back to the
                 default entity passed to `Api`.  If no default entity, will raise a `ValueError`.
-            per_page: (int) Sets the page size for query pagination.  None will use the default size.
-                Usually there is no reason to change this.
+            per_page: (int) Sets the page size for query pagination.  Usually there is no reason to change this.
 
         Returns:
             A `Projects` object which is an iterable collection of `Project` objects.
@@ -808,7 +806,7 @@ class Api:
         return public.Project(self.client, entity, name, {})
 
     def reports(
-        self, path: str = "", name: Optional[str] = None, per_page: Optional[int] = 50
+        self, path: str = "", name: Optional[str] = None, per_page: int = 50
     ) -> "public.Reports":
         """Get reports for a given project path.
 
@@ -817,8 +815,7 @@ class Api:
         Args:
             path: (str) path to project the report resides in, should be in the form: "entity/project"
             name: (str, optional) optional name of the report requested.
-            per_page: (int) Sets the page size for query pagination.  None will use the default size.
-                Usually there is no reason to change this.
+            per_page: (int) Sets the page size for query pagination.  Usually there is no reason to change this.
 
         Returns:
             A `Reports` object which is an iterable collection of `BetaReport` objects.
@@ -1153,15 +1150,14 @@ class Api:
 
     @normalize_exceptions
     def artifact_collections(
-        self, project_name: str, type_name: str, per_page: Optional[int] = 50
+        self, project_name: str, type_name: str, per_page: int = 50
     ) -> "public.ArtifactCollections":
         """Return a collection of matching artifact collections.
 
         Args:
             project_name: (str) The name of the project to filter on.
             type_name: (str) The name of the artifact type to filter on.
-            per_page: (int, optional) Sets the page size for query pagination.  None will use the default size.
-                Usually there is no reason to change this.
+            per_page: (int) Sets the page size for query pagination.  Usually there is no reason to change this.
 
         Returns:
             An iterable `ArtifactCollections` object.
@@ -1226,7 +1222,7 @@ class Api:
         self,
         type_name: str,
         name: str,
-        per_page: Optional[int] = 50,
+        per_page: int = 50,
         tags: Optional[List[str]] = None,
     ) -> "public.Artifacts":
         """Return an `Artifacts` collection from the given parameters.
@@ -1234,8 +1230,7 @@ class Api:
         Args:
             type_name: (str) The type of artifacts to fetch.
             name: (str) An artifact collection name. May be prefixed with entity/project.
-            per_page: (int, optional) Sets the page size for query pagination.  None will use the default size.
-                Usually there is no reason to change this.
+            per_page: (int) Sets the page size for query pagination.  Usually there is no reason to change this.
             tags: (list[str], optional) Only return artifacts with all of these tags.
 
         Returns:
@@ -1510,7 +1505,7 @@ class Api:
         Returns:
             A registry iterator.
         """
-        if not self._check_server_feature_with_fallback(
+        if not InternalApi()._check_server_feature_with_fallback(
             ServerFeature.ARTIFACT_REGISTRY_SEARCH
         ):
             raise RuntimeError(
@@ -1522,52 +1517,3 @@ class Api:
             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

@@ -10,19 +10,21 @@ from wandb_gql import Client, gql
 import wandb
 from wandb.apis import public
 from wandb.apis.normalize import normalize_exceptions
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import Paginator, SizedPaginator
 from wandb.errors.term import termlog
+from wandb.proto.wandb_internal_pb2 import ServerFeature
 from wandb.sdk.artifacts._graphql_fragments import (
     ARTIFACT_FILES_FRAGMENT,
     ARTIFACTS_TYPES_FRAGMENT,
 )
+from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.lib import deprecate
 
 if TYPE_CHECKING:
     from wandb.apis.public import RetryingClient, Run
 
 
-class ArtifactTypes(Paginator):
+class ArtifactTypes(Paginator["ArtifactType"]):
     QUERY = gql(
         """
         query ProjectArtifacts(
@@ -45,7 +47,7 @@ class ArtifactTypes(Paginator):
         client: Client,
         entity: str,
         project: str,
-        per_page: Optional[int] = 50,
+        per_page: int = 50,
     ):
         self.entity = entity
         self.project = project
@@ -58,7 +60,7 @@ class ArtifactTypes(Paginator):
         super().__init__(client, variable_values, per_page)
 
     @property
-    def length(self):
+    def length(self) -> None:
         # TODO
         return None
 
@@ -167,14 +169,14 @@ class ArtifactType:
         return f"<ArtifactType {self.type}>"
 
 
-class ArtifactCollections(Paginator):
+class ArtifactCollections(SizedPaginator["ArtifactCollection"]):
     def __init__(
         self,
         client: Client,
         entity: str,
         project: str,
         type_name: str,
-        per_page: Optional[int] = 50,
+        per_page: int = 50,
     ):
         self.entity = entity
         self.project = project
@@ -710,7 +712,7 @@ class ArtifactCollection:
         return f"<ArtifactCollection {self._name} ({self._type})>"
 
 
-class Artifacts(Paginator):
+class Artifacts(SizedPaginator["wandb.Artifact"]):
     """An iterable collection of artifact versions associated with a project and optional filter.
 
     This is generally used indirectly via the `Api`.artifact_versions method.
@@ -826,10 +828,8 @@ class Artifacts(Paginator):
         ]
 
 
-class RunArtifacts(Paginator):
-    def __init__(
-        self, client: Client, run: "Run", mode="logged", per_page: Optional[int] = 50
-    ):
+class RunArtifacts(SizedPaginator["wandb.Artifact"]):
+    def __init__(self, client: Client, run: "Run", mode="logged", per_page: int = 50):
         from wandb.sdk.artifacts.artifact import _gql_artifact_fragment
 
         output_query = gql(
@@ -944,9 +944,9 @@ class RunArtifacts(Paginator):
         ]
 
 
-class ArtifactFiles(Paginator):
-    QUERY = gql(
-        """
+class ArtifactFiles(SizedPaginator["public.File"]):
+    ARTIFACT_VERSION_FILES_QUERY = gql(
+        f"""
         query ArtifactFiles(
             $entityName: String!,
             $projectName: String!,
@@ -959,13 +959,40 @@ class ArtifactFiles(Paginator):
             project(name: $projectName, entityName: $entityName) {{
                 artifactType(name: $artifactTypeName) {{
                     artifact(name: $artifactName) {{
-                        ...ArtifactFilesFragment
+                        files(names: $fileNames, after: $fileCursor, first: $fileLimit) {{
+                            ...FilesFragment
+                        }}
                     }}
                 }}
             }}
         }}
-        {}
-    """.format(ARTIFACT_FILES_FRAGMENT)
+        {ARTIFACT_FILES_FRAGMENT}
+    """
+    )
+
+    ARTIFACT_COLLECTION_MEMBERSHIP_FILES_QUERY = gql(
+        f"""
+        query ArtifactCollectionMembershipFiles(
+            $entityName: String!,
+            $projectName: String!,
+            $artifactName: String!,
+            $artifactVersionIndex: String!,
+            $fileNames: [String!],
+            $fileCursor: String,
+            $fileLimit: Int = 50
+        ) {{
+            project(name: $projectName, entityName: $entityName) {{
+                artifactCollection(name: $artifactName) {{
+                    artifactMembership (aliasName: $artifactVersionIndex) {{
+                        files(names: $fileNames, after: $fileCursor, first: $fileLimit) {{
+                            ...FilesFragment
+                        }}
+                    }}
+                }}
+            }}
+        }}
+        {ARTIFACT_FILES_FRAGMENT}
+        """
     )
 
     def __init__(
@@ -975,6 +1002,9 @@ class ArtifactFiles(Paginator):
         names: Optional[Sequence[str]] = None,
         per_page: int = 50,
     ):
+        self.query_via_membership = InternalApi()._check_server_feature_with_fallback(
+            ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILES
+        )
         self.artifact = artifact
         variables = {
             "entityName": artifact.source_entity,
@@ -983,6 +1013,17 @@ class ArtifactFiles(Paginator):
             "artifactName": artifact.source_name,
             "fileNames": names,
         }
+        if self.query_via_membership:
+            self.QUERY = self.ARTIFACT_COLLECTION_MEMBERSHIP_FILES_QUERY
+            variables = {
+                "entityName": artifact.entity,
+                "projectName": artifact.project,
+                "artifactName": artifact.name.split(":")[0],
+                "artifactVersionIndex": artifact.version,
+                "fileNames": names,
+            }
+        else:
+            self.QUERY = self.ARTIFACT_VERSION_FILES_QUERY
         # The server must advertise at least SDK 0.12.21
         # to get storagePath
         if not client.version_supported("0.12.21"):
@@ -1000,6 +1041,10 @@ class ArtifactFiles(Paginator):
     @property
     def more(self):
         if self.last_response:
+            if self.query_via_membership:
+                return self.last_response["project"]["artifactCollection"][
+                    "artifactMembership"
+                ]["files"]["pageInfo"]["hasNextPage"]
             return self.last_response["project"]["artifactType"]["artifact"]["files"][
                 "pageInfo"
             ]["hasNextPage"]
@@ -1009,6 +1054,10 @@ class ArtifactFiles(Paginator):
     @property
     def cursor(self):
         if self.last_response:
+            if self.query_via_membership:
+                return self.last_response["project"]["artifactCollection"][
+                    "artifactMembership"
+                ]["files"]["edges"][-1]["cursor"]
             return self.last_response["project"]["artifactType"]["artifact"]["files"][
                 "edges"
             ][-1]["cursor"]
@@ -1019,6 +1068,13 @@ class ArtifactFiles(Paginator):
         self.variables.update({"fileLimit": self.per_page, "fileCursor": self.cursor})
 
     def convert_objects(self):
+        if self.query_via_membership:
+            return [
+                public.File(self.client, r["node"])
+                for r in self.last_response["project"]["artifactCollection"][
+                    "artifactMembership"
+                ]["files"]["edges"]
+            ]
         return [
             public.File(self.client, r["node"])
             for r in self.last_response["project"]["artifactType"]["artifact"]["files"][

wandb/apis/public/files.py

@@ -12,7 +12,7 @@ import wandb
 from wandb import util
 from wandb.apis.attrs import Attrs
 from wandb.apis.normalize import normalize_exceptions
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import SizedPaginator
 from wandb.apis.public import utils
 from wandb.apis.public.api import Api
 from wandb.apis.public.const import RETRY_TIMEDELTA
@@ -41,7 +41,7 @@ FILE_FRAGMENT = """fragment RunFilesFragment on Run {
 }"""
 
 
-class Files(Paginator):
+class Files(SizedPaginator["File"]):
     """An iterable collection of `File` objects."""
 
     QUERY = gql(

wandb/apis/public/projects.py

@@ -17,7 +17,7 @@ PROJECT_FRAGMENT = """fragment ProjectFragment on Project {
 }"""
 
 
-class Projects(Paginator):
+class Projects(Paginator["Project"]):
     """An iterable collection of `Project` objects."""
 
     QUERY = gql(
@@ -49,7 +49,8 @@ class Projects(Paginator):
         super().__init__(client, variables, per_page)
 
     @property
-    def length(self):
+    def length(self) -> None:
+        # For backwards compatibility, even though this isn't a SizedPaginator
         return None
 
     @property

wandb/apis/public/reports.py

@@ -9,11 +9,11 @@ from wandb_gql import gql
 import wandb
 from wandb.apis import public
 from wandb.apis.attrs import Attrs
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import SizedPaginator
 from wandb.sdk.lib import ipython
 
 
-class Reports(Paginator):
+class Reports(SizedPaginator["BetaReport"]):
     """Reports is an iterable collection of `BetaReport` objects."""
 
     QUERY = gql(

wandb/apis/public/runs.py

@@ -24,7 +24,7 @@ from wandb.apis import public
 from wandb.apis.attrs import Attrs
 from wandb.apis.internal import Api as InternalApi
 from wandb.apis.normalize import normalize_exceptions
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import SizedPaginator
 from wandb.apis.public.const import RETRY_TIMEDELTA
 from wandb.sdk.lib import ipython, json_util, runid
 from wandb.sdk.lib.paths import LogicalPath
@@ -61,7 +61,7 @@ RUN_FRAGMENT = """fragment RunFragment on Run {
 }"""
 
 
-class Runs(Paginator):
+class Runs(SizedPaginator["Run"]):
     """An iterable collection of runs associated with a project and optional filter.
 
     This is generally used indirectly via the `Api`.runs method.
@@ -421,16 +421,15 @@ class Run(Attrs):
             """
         query Run($project: String!, $entity: String!, $name: String!) {{
             project(name: $project, entityName: $entity) {{
-                {}
                 run(name: $name) {{
+                    {}
                     ...RunFragment
                 }}
             }}
         }}
         {}
         """.format(
-                # Only query internalId if the server supports it
-                "internalId" if self._server_provides_internal_id_for_project() else "",
+                "projectId" if self._server_provides_internal_id_for_project() else "",
                 RUN_FRAGMENT,
             )
         )
@@ -444,7 +443,11 @@ class Run(Attrs):
                 raise ValueError("Could not find run {}".format(self))
             self._attrs = response["project"]["run"]
             self._state = self._attrs["state"]
-            self._project_internal_id = response["project"].get("internalId", None)
+
+            self._project_internal_id = (
+                int(self._attrs["projectId"]) if "projectId" in self._attrs else None
+            )
+
             if self._include_sweeps and self.sweep_name and not self.sweep:
                 # There may be a lot of runs. Don't bother pulling them all
                 # just for the sake of this one.
@@ -847,7 +850,12 @@ class Run(Attrs):
         api.set_current_run_id(self.id)
 
         if isinstance(artifact, wandb.Artifact) and not artifact.is_draft():
-            api.use_artifact(artifact.id, use_as=use_as or artifact.name)
+            api.use_artifact(
+                artifact.id,
+                use_as=use_as or artifact.name,
+                artifact_entity_name=artifact.entity,
+                artifact_project_name=artifact.project,
+            )
             return artifact
         elif isinstance(artifact, wandb.Artifact) and artifact.is_draft():
             raise ValueError(
@@ -911,8 +919,8 @@ class Run(Attrs):
         This check is done by utilizing GraphQL introspection in the available fields on the Project type.
         """
         query_string = """
-           query ProbeProjectInput {
-                ProjectType: __type(name:"Project") {
+           query ProbeRunInput {
+                RunType: __type(name:"Run") {
                     fields {
                         name
                     }
@@ -924,8 +932,8 @@ class Run(Attrs):
         if self.server_provides_internal_id_field is None:
             query = gql(query_string)
             res = self.client.execute(query)
-            self.server_provides_internal_id_field = "internalId" in [
-                x["name"] for x in (res.get("ProjectType", {}).get("fields", [{}]))
+            self.server_provides_internal_id_field = "projectId" in [
+                x["name"] for x in (res.get("RunType", {}).get("fields", [{}]))
             ]
 
         return self.server_provides_internal_id_field

wandb/integration/metaflow/metaflow.py

@@ -13,6 +13,7 @@ import inspect
 import pickle
 from functools import wraps
 from pathlib import Path
+from typing import Union
 
 import wandb
 from wandb.sdk.lib import telemetry as wb_telemetry
@@ -25,17 +26,18 @@ except ImportError as e:
     ) from e
 
 try:
-    from fastcore.all import typedispatch
+    from plum import dispatch
 except ImportError as e:
     raise Exception(
-        "Error: `fastcore` not installed >> This integration requires fastcore!  To fix, please `pip install -Uqq fastcore`"
+        "Error: `plum-dispatch` not installed >> "
+        "This integration requires plum-dispatch! To fix, please `pip install -Uqq plum-dispatch`"
     ) from e
 
 
 try:
     import pandas as pd
 
-    @typedispatch  # noqa: F811
+    @dispatch  # noqa: F811
     def _wandb_use(
         name: str,
         data: pd.DataFrame,
@@ -52,7 +54,7 @@ try:
             run.use_artifact(f"{name}:latest")
             wandb.termlog(f"Using artifact: {name} ({type(data)})")
 
-    @typedispatch  # noqa: F811
+    @dispatch  # noqa: F811
     def wandb_track(
         name: str,
         data: pd.DataFrame,
@@ -81,7 +83,7 @@ try:
     import torch
     import torch.nn as nn
 
-    @typedispatch  # noqa: F811
+    @dispatch  # noqa: F811
     def _wandb_use(
         name: str,
         data: nn.Module,
@@ -98,7 +100,7 @@ try:
             run.use_artifact(f"{name}:latest")
             wandb.termlog(f"Using artifact: {name} ({type(data)})")
 
-    @typedispatch  # noqa: F811
+    @dispatch  # noqa: F811
     def wandb_track(
         name: str,
         data: nn.Module,
@@ -126,7 +128,7 @@ except ImportError:
 try:
     from sklearn.base import BaseEstimator
 
-    @typedispatch  # noqa: F811
+    @dispatch  # noqa: F811
     def _wandb_use(
         name: str,
         data: BaseEstimator,
@@ -143,7 +145,7 @@ try:
             run.use_artifact(f"{name}:latest")
             wandb.termlog(f"Using artifact: {name} ({type(data)})")
 
-    @typedispatch  # noqa: F811
+    @dispatch  # noqa: F811
     def wandb_track(
         name: str,
         data: BaseEstimator,
@@ -192,10 +194,10 @@ class ArtifactProxy:
         return getattr(self.flow, key)
 
 
-@typedispatch  # noqa: F811
+@dispatch  # noqa: F811
 def wandb_track(
     name: str,
-    data: (dict, list, set, str, int, float, bool),
+    data: Union[dict, list, set, str, int, float, bool],
     run=None,
     testing=False,
     *args,
@@ -207,7 +209,7 @@ def wandb_track(
     run.log({name: data})
 
 
-@typedispatch  # noqa: F811
+@dispatch  # noqa: F811
 def wandb_track(
     name: str, data: Path, datasets=False, run=None, testing=False, *args, **kwargs
 ):
@@ -225,7 +227,7 @@ def wandb_track(
 
 
 # this is the base case
-@typedispatch  # noqa: F811
+@dispatch  # noqa: F811
 def wandb_track(
     name: str, data, others=False, run=None, testing=False, *args, **kwargs
 ):
@@ -240,7 +242,7 @@ def wandb_track(
         wandb.termlog(f"Logging artifact: {name} ({type(data)})")
 
 
-@typedispatch
+@dispatch  # noqa: F811
 def wandb_use(name: str, data, *args, **kwargs):
     try:
         return _wandb_use(name, data, *args, **kwargs)
@@ -252,14 +254,14 @@ def wandb_use(name: str, data, *args, **kwargs):
         )
 
 
-@typedispatch  # noqa: F811
+@dispatch  # noqa: F811
 def wandb_use(
-    name: str, data: (dict, list, set, str, int, float, bool), *args, **kwargs
+    name: str, data: Union[dict, list, set, str, int, float, bool], *args, **kwargs
 ):  # type: ignore
     pass  # do nothing for these types
 
 
-@typedispatch  # noqa: F811
+@dispatch  # noqa: F811
 def _wandb_use(
     name: str, data: Path, datasets=False, run=None, testing=False, *args, **kwargs
 ):  # type: ignore
@@ -271,7 +273,7 @@ def _wandb_use(
         wandb.termlog(f"Using artifact: {name} ({type(data)})")
 
 
-@typedispatch  # noqa: F811
+@dispatch  # noqa: F811
 def _wandb_use(name: str, data, others=False, run=None, testing=False, *args, **kwargs):  # type: ignore
     if testing:
         return "others" if others else None

wandb/integration/sacred/__init__.py

@@ -23,7 +23,7 @@ class WandbObserver(RunObserver):
         job_type — the type of job you are logging, e.g. eval, worker, ps (default: training)
         save_code — save the main python or notebook file to wandb to enable diffing (default: editable from your settings page)
         group — a string by which to group other runs; see Grouping
-        reinit — whether to allow multiple calls to wandb.init in the same process (default: False)
+        reinit — Shorthand for the reinit setting that defines what to do when `wandb.init()` is called while a run is active. See the setting's documentation.
         id — A unique ID for this run primarily used for Resuming. It must be globally unique, and if you delete a run you can't reuse the ID. Use the name field for a descriptive, useful name for the run. The ID cannot contain special characters.
         resume — if set to True, the run auto resumes; can also be a unique string for manual resuming; see Resuming (default: False)
         anonymous — can be "allow", "never", or "must". This enables or explicitly disables anonymous logging. (default: never)