wandb 0.21.4__py3-none-musllinux_1_2_aarch64.whl → 0.22.1__py3-none-musllinux_1_2_aarch64.whl

wandb/apis/public/runs.py

@@ -85,13 +85,92 @@ RUN_FRAGMENT = """fragment RunFragment on Run {
     historyKeys
 }"""
 
+# Lightweight fragment for listing operations - excludes heavy fields
+LIGHTWEIGHT_RUN_FRAGMENT = """fragment LightweightRunFragment on Run {
+    id
+    tags
+    name
+    displayName
+    sweepName
+    state
+    group
+    jobType
+    commit
+    readOnly
+    createdAt
+    heartbeatAt
+    description
+    notes
+    historyLineCount
+    user {
+        name
+        username
+    }
+}"""
+
+# Fragment name constants to avoid string parsing
+RUN_FRAGMENT_NAME = "RunFragment"
+LIGHTWEIGHT_RUN_FRAGMENT_NAME = "LightweightRunFragment"
+
+
+def _create_runs_query(
+    *, lazy: bool, with_internal_id: bool, with_project_id: bool
+) -> gql:
+    """Create GraphQL query for runs with appropriate fragment."""
+    fragment = LIGHTWEIGHT_RUN_FRAGMENT if lazy else RUN_FRAGMENT
+    fragment_name = LIGHTWEIGHT_RUN_FRAGMENT_NAME if lazy else RUN_FRAGMENT_NAME
+
+    return gql(
+        f"""#graphql
+        query Runs($project: String!, $entity: String!, $cursor: String, $perPage: Int = 50, $order: String, $filters: JSONString) {{
+            project(name: $project, entityName: $entity) {{
+                {"internalId" if with_internal_id else ""}
+                runCount(filters: $filters)
+                readOnly
+                runs(filters: $filters, after: $cursor, first: $perPage, order: $order) {{
+                    edges {{
+                        node {{
+                            {"projectId" if with_project_id else ""}
+                            ...{fragment_name}
+                        }}
+                        cursor
+                    }}
+                    pageInfo {{
+                        endCursor
+                        hasNextPage
+                    }}
+                }}
+            }}
+        }}
+        {fragment}
+        """
+    )
+
 
 @normalize_exceptions
 def _server_provides_internal_id_for_project(client) -> bool:
-    """Returns True if the server allows us to query the internalId field for a project.
-
-    This check is done by utilizing GraphQL introspection in the available fields on the Project type.
+    """Returns True if the server allows us to query the internalId field for a project."""
+    query_string = """
+       query ProbeProjectInput {
+            ProjectType: __type(name:"Project") {
+                fields {
+                    name
+                }
+            }
+        }
     """
+
+    # Only perform the query once to avoid extra network calls
+    query = gql(query_string)
+    res = client.execute(query)
+    return "internalId" in [
+        x["name"] for x in (res.get("ProjectType", {}).get("fields", [{}]))
+    ]
+
+
+@normalize_exceptions
+def _server_provides_project_id_for_run(client) -> bool:
+    """Returns True if the server allows us to query the projectId field for a run."""
     query_string = """
        query ProbeRunInput {
             RunType: __type(name:"Run") {
@@ -201,34 +280,15 @@ class Runs(SizedPaginator["Run"]):
         order: str = "+created_at",
         per_page: int = 50,
         include_sweeps: bool = True,
+        lazy: bool = True,
     ):
         if not order:
             order = "+created_at"
 
-        self.QUERY = gql(
-            f"""#graphql
-            query Runs($project: String!, $entity: String!, $cursor: String, $perPage: Int = 50, $order: String, $filters: JSONString) {{
-                project(name: $project, entityName: $entity) {{
-                    internalId
-                    runCount(filters: $filters)
-                    readOnly
-                    runs(filters: $filters, after: $cursor, first: $perPage, order: $order) {{
-                        edges {{
-                            node {{
-                                {"projectId" if _server_provides_internal_id_for_project(client) else ""}
-                                ...RunFragment
-                            }}
-                            cursor
-                        }}
-                        pageInfo {{
-                            endCursor
-                            hasNextPage
-                        }}
-                    }}
-                }}
-            }}
-            {RUN_FRAGMENT}
-            """
+        self.QUERY = _create_runs_query(
+            lazy=lazy,
+            with_internal_id=_server_provides_internal_id_for_project(client),
+            with_project_id=_server_provides_project_id_for_run(client),
         )
 
         self.entity = entity
@@ -238,6 +298,7 @@ class Runs(SizedPaginator["Run"]):
         self.order = order
         self._sweeps = {}
         self._include_sweeps = include_sweeps
+        self._lazy = lazy
         variables = {
             "project": self.project,
             "entity": self.entity,
@@ -296,6 +357,7 @@ class Runs(SizedPaginator["Run"]):
                 run_response["node"]["name"],
                 run_response["node"],
                 include_sweeps=self._include_sweeps,
+                lazy=self._lazy,
             )
             objs.append(run)
 
@@ -422,6 +484,39 @@ class Runs(SizedPaginator["Run"]):
     def __repr__(self):
         return f"<Runs {self.entity}/{self.project}>"
 
+    def upgrade_to_full(self):
+        """Upgrade this Runs collection from lazy to full mode.
+
+        This switches to fetching full run data and
+        upgrades any already-loaded Run objects to have full data.
+        Uses parallel loading for better performance when upgrading multiple runs.
+        """
+        if not self._lazy:
+            return  # Already in full mode
+
+        # Switch to full mode
+        self._lazy = False
+
+        # Regenerate query with full fragment
+        self.QUERY = _create_runs_query(
+            lazy=False,
+            with_internal_id=_server_provides_internal_id_for_project(self.client),
+            with_project_id=_server_provides_project_id_for_run(self.client),
+        )
+
+        # Upgrade any existing runs that have been loaded - use parallel loading for performance
+        lazy_runs = [run for run in self.objects if run._lazy]
+        if lazy_runs:
+            from concurrent.futures import ThreadPoolExecutor
+
+            # Limit workers to avoid overwhelming the server
+            max_workers = min(len(lazy_runs), 10)
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                futures = [executor.submit(run.load_full_data) for run in lazy_runs]
+                # Wait for all to complete
+                for future in futures:
+                    future.result()
+
 
 class Run(Attrs):
     """A single run associated with an entity and project.
@@ -465,6 +560,7 @@ class Run(Attrs):
         run_id: str,
         attrs: Mapping | None = None,
         include_sweeps: bool = True,
+        lazy: bool = True,
     ):
         """Initialize a Run object.
 
@@ -481,6 +577,8 @@ class Run(Attrs):
         self.id = run_id
         self.sweep = None
         self._include_sweeps = include_sweeps
+        self._lazy = lazy
+        self._full_data_loaded = False  # Track if we've loaded full data
         self.dir = os.path.join(self._base_dir, *self.path)
         try:
             os.makedirs(self.dir)
@@ -490,6 +588,7 @@ class Run(Attrs):
         self._metadata: dict[str, Any] | None = None
         self._state = _attrs.get("state", "not found")
         self.server_provides_internal_id_field: bool | None = None
+        self._server_provides_project_id_field: bool | None = None
         self._is_loaded: bool = False
 
         self.load(force=not _attrs)
@@ -594,23 +693,34 @@ class Run(Attrs):
                 "notes": None,
                 "state": state,
             },
+            lazy=False,  # Created runs should have full data available immediately
         )
 
-    def load(self, force=False):
-        if force or not self._attrs:
-            self._is_loaded = False
-            query = gql(f"""#graphql
-            query Run($project: String!, $entity: String!, $name: String!) {{
-                project(name: $project, entityName: $entity) {{
-                    run(name: $name) {{
-                        {"projectId" if _server_provides_internal_id_for_project(self.client) else ""}
-                        ...RunFragment
-                    }}
+    def _load_with_fragment(
+        self, fragment: str, fragment_name: str, force: bool = False
+    ):
+        """Load run data using specified GraphQL fragment."""
+        # Cache the server capability check to avoid repeated network calls
+        if self._server_provides_project_id_field is None:
+            self._server_provides_project_id_field = (
+                _server_provides_project_id_for_run(self.client)
+            )
+
+        query = gql(
+            f"""
+        query Run($project: String!, $entity: String!, $name: String!) {{
+            project(name: $project, entityName: $entity) {{
+                run(name: $name) {{
+                    {"projectId" if self._server_provides_project_id_field else ""}
+                    ...{fragment_name}
                 }}
             }}
-            {RUN_FRAGMENT}
-            """)
+        }}
+        {fragment}
+        """
+        )
 
+        if force or not self._attrs:
             response = self._exec(query)
             if (
                 response is None
@@ -620,6 +730,10 @@ class Run(Attrs):
                 raise ValueError("Could not find run {}".format(self))
             self._attrs = response["project"]["run"]
 
+            self._state = self._attrs["state"]
+            if self._attrs.get("user"):
+                self.user = public.User(self.client, self._attrs["user"])
+
             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.
@@ -632,39 +746,78 @@ class Run(Attrs):
                 )
 
         if not self._is_loaded:
-            self._load_from_attrs()
+            # Always set _project_internal_id if projectId is available, regardless of fragment type
+            if "projectId" in self._attrs:
+                self._project_internal_id = int(self._attrs["projectId"])
+            else:
+                self._project_internal_id = None
+
+            # Only call _load_from_attrs when using the full fragment or when the fields are actually present
+            if fragment_name == RUN_FRAGMENT_NAME or (
+                "config" in self._attrs
+                or "summaryMetrics" in self._attrs
+                or "systemMetrics" in self._attrs
+            ):
+                self._load_from_attrs()
             self._is_loaded = True
 
         return self._attrs
 
     def _load_from_attrs(self):
         self._state = self._attrs.get("state", None)
-        self._attrs["config"] = _convert_to_dict(self._attrs.get("config"))
-        self._attrs["summaryMetrics"] = _convert_to_dict(
-            self._attrs.get("summaryMetrics")
-        )
-        self._attrs["systemMetrics"] = _convert_to_dict(
-            self._attrs.get("systemMetrics")
-        )
 
-        if "projectId" in self._attrs:
-            self._project_internal_id = int(self._attrs["projectId"])
-        else:
-            self._project_internal_id = None
+        # Only convert fields if they exist in _attrs
+        if "config" in self._attrs:
+            self._attrs["config"] = _convert_to_dict(self._attrs.get("config"))
+        if "summaryMetrics" in self._attrs:
+            self._attrs["summaryMetrics"] = _convert_to_dict(
+                self._attrs.get("summaryMetrics")
+            )
+        if "systemMetrics" in self._attrs:
+            self._attrs["systemMetrics"] = _convert_to_dict(
+                self._attrs.get("systemMetrics")
+            )
+
+        # Only check for sweeps if sweep_name is available (not in lazy mode or if it exists)
+        if self._include_sweeps and self._attrs.get("sweepName") and not self.sweep:
+            # There may be a lot of runs. Don't bother pulling them all
+            self.sweep = public.Sweep(
+                self.client,
+                self.entity,
+                self.project,
+                self._attrs["sweepName"],
+                withRuns=False,
+            )
 
-        if self._attrs.get("user"):
-            self.user = public.User(self.client, self._attrs["user"])
         config_user, config_raw = {}, {}
-        for key, value in self._attrs.get("config").items():
-            config = config_raw if key in WANDB_INTERNAL_KEYS else config_user
-            if isinstance(value, dict) and "value" in value:
-                config[key] = value["value"]
-            else:
-                config[key] = value
+        if self._attrs.get("config"):
+            try:
+                # config is already converted to dict by _convert_to_dict
+                for key, value in self._attrs.get("config", {}).items():
+                    config = config_raw if key in WANDB_INTERNAL_KEYS else config_user
+                    if isinstance(value, dict) and "value" in value:
+                        config[key] = value["value"]
+                    else:
+                        config[key] = value
+            except (TypeError, AttributeError):
+                # Handle case where config is malformed or not a dict
+                pass
+
         config_raw.update(config_user)
         self._attrs["config"] = config_user
         self._attrs["rawconfig"] = config_raw
 
+        return self._attrs
+
+    def load(self, force=False):
+        """Load run data using appropriate fragment based on lazy mode."""
+        if self._lazy:
+            return self._load_with_fragment(
+                LIGHTWEIGHT_RUN_FRAGMENT, LIGHTWEIGHT_RUN_FRAGMENT_NAME, force
+            )
+        else:
+            return self._load_with_fragment(RUN_FRAGMENT, RUN_FRAGMENT_NAME, force)
+
     @normalize_exceptions
     def wait_until_finished(self):
         """Check the state of the run until it is finished."""
@@ -1126,9 +1279,43 @@ class Run(Attrs):
         )
         return artifact
 
+    def load_full_data(self, force: bool = False) -> dict[str, Any]:
+        """Load full run data including heavy fields like config, systemMetrics, summaryMetrics.
+
+        This method is useful when you initially used lazy=True for listing runs,
+        but need access to the full data for specific runs.
+
+        Args:
+            force: Force reload even if data is already loaded
+
+        Returns:
+            The loaded run attributes
+        """
+        if not self._lazy and not force:
+            # Already in full mode, no need to reload
+            return self._attrs
+
+        # Load full data and mark as loaded
+        result = self._load_with_fragment(RUN_FRAGMENT, RUN_FRAGMENT_NAME, force=True)
+        self._full_data_loaded = True
+        return result
+
+    @property
+    def config(self):
+        """Get run config. Auto-loads full data if in lazy mode."""
+        if self._lazy and not self._full_data_loaded and "config" not in self._attrs:
+            self.load_full_data()
+        return self._attrs.get("config", {})
+
     @property
     def summary(self):
-        """A mutable dict-like property that holds summary values associated with the run."""
+        """Get run summary metrics. Auto-loads full data if in lazy mode."""
+        if (
+            self._lazy
+            and not self._full_data_loaded
+            and "summaryMetrics" not in self._attrs
+        ):
+            self.load_full_data()
         if self._summary is None:
             from wandb.old.summary import HTTPSummary
 
@@ -1136,6 +1323,41 @@ class Run(Attrs):
             self._summary = HTTPSummary(self, self.client, summary=self.summary_metrics)
         return self._summary
 
+    @property
+    def system_metrics(self):
+        """Get run system metrics. Auto-loads full data if in lazy mode."""
+        if (
+            self._lazy
+            and not self._full_data_loaded
+            and "systemMetrics" not in self._attrs
+        ):
+            self.load_full_data()
+        return self._attrs.get("systemMetrics", {})
+
+    @property
+    def summary_metrics(self):
+        """Get run summary metrics. Auto-loads full data if in lazy mode."""
+        if (
+            self._lazy
+            and not self._full_data_loaded
+            and "summaryMetrics" not in self._attrs
+        ):
+            self.load_full_data()
+        return self._attrs.get("summaryMetrics", {})
+
+    @property
+    def rawconfig(self):
+        """Get raw run config including internal keys. Auto-loads full data if in lazy mode."""
+        if self._lazy and not self._full_data_loaded and "rawconfig" not in self._attrs:
+            self.load_full_data()
+        return self._attrs.get("rawconfig", {})
+
+    @property
+    def sweep_name(self):
+        """Get sweep name. Always available since sweepName is in lightweight fragment."""
+        # sweepName is included in lightweight fragment, so no need to load full data
+        return self._attrs.get("sweepName")
+
     @property
     def path(self):
         """The path of the run. The path is a list containing the entity, project, and run_id."""

wandb/apis/public/sweeps.py

@@ -27,8 +27,9 @@ Note:
     and wandb.agent() functions from the main wandb package.
 """
 
+from __future__ import annotations
+
 import urllib
-from typing import Optional
 
 from wandb_gql import gql
 
@@ -103,7 +104,7 @@ class Sweeps(SizedPaginator["Sweep"]):
         entity: str,
         project: str,
         per_page: int = 50,
-    ) -> "Sweeps":
+    ) -> Sweeps:
         """An iterable collection of `Sweep` objects.
 
         Args:
@@ -317,7 +318,7 @@ class Sweep(Attrs):
             return None
 
     @property
-    def expected_run_count(self) -> Optional[int]:
+    def expected_run_count(self) -> int | None:
         """Return the number of expected runs in the sweep or None for infinite runs."""
         return self._attrs.get("runCountExpected")
 
@@ -360,12 +361,12 @@ class Sweep(Attrs):
     @classmethod
     def get(
         cls,
-        client: "RetryingClient",
-        entity: Optional[str] = None,
-        project: Optional[str] = None,
-        sid: Optional[str] = None,
-        order: Optional[str] = None,
-        query: Optional[str] = None,
+        client: RetryingClient,
+        entity: str | None = None,
+        project: str | None = None,
+        sid: str | None = None,
+        order: str | None = None,
+        query: str | None = None,
         **kwargs,
     ):
         """Execute a query against the cloud backend.

wandb/apis/public/teams.py

@@ -8,6 +8,8 @@ Note:
     permissions.
 """
 
+from __future__ import annotations
+
 import requests
 from wandb_gql import gql
 

wandb/apis/public/users.py

@@ -7,6 +7,8 @@ Note:
     users and their authentication. Some operations require admin privileges.
 """
 
+from __future__ import annotations
+
 import requests
 from wandb_gql import gql
 

wandb/apis/public/utils.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import re
 from enum import Enum
-from typing import Any, Dict, Iterable, Mapping, Optional, Set
+from typing import Any, Iterable, Mapping
 from urllib.parse import urlparse
 
 from wandb_gql import gql
@@ -75,7 +77,7 @@ def parse_org_from_registry_path(path: str, path_type: PathType) -> str:
 
 
 def fetch_org_from_settings_or_entity(
-    settings: dict, default_entity: Optional[str] = None
+    settings: dict, default_entity: str | None = None
 ) -> str:
     """Fetch the org from either the settings or deriving it from the entity.
 
@@ -110,17 +112,17 @@ def fetch_org_from_settings_or_entity(
 class _GQLCompatRewriter(visitor.Visitor):
     """GraphQL AST visitor to rewrite queries/mutations to be compatible with older server versions."""
 
-    omit_variables: Set[str]
-    omit_fragments: Set[str]
-    omit_fields: Set[str]
-    rename_fields: Dict[str, str]
+    omit_variables: set[str]
+    omit_fragments: set[str]
+    omit_fields: set[str]
+    rename_fields: dict[str, str]
 
     def __init__(
         self,
-        omit_variables: Optional[Iterable[str]] = None,
-        omit_fragments: Optional[Iterable[str]] = None,
-        omit_fields: Optional[Iterable[str]] = None,
-        rename_fields: Optional[Mapping[str, str]] = None,
+        omit_variables: Iterable[str] | None = None,
+        omit_fragments: Iterable[str] | None = None,
+        omit_fields: Iterable[str] | None = None,
+        rename_fields: Mapping[str, str] | None = None,
     ):
         self.omit_variables = set(omit_variables or ())
         self.omit_fragments = set(omit_fragments or ())
@@ -130,7 +132,6 @@ class _GQLCompatRewriter(visitor.Visitor):
     def enter_VariableDefinition(self, node: ast.VariableDefinition, *_, **__) -> Any:  # noqa: N802
         if node.variable.name.value in self.omit_variables:
             return visitor.REMOVE
-        # return node
 
     def enter_ObjectField(self, node: ast.ObjectField, *_, **__) -> Any:  # noqa: N802
         # For context, note that e.g.:
@@ -176,10 +177,10 @@ class _GQLCompatRewriter(visitor.Visitor):
 
 def gql_compat(
     request_string: str,
-    omit_variables: Optional[Iterable[str]] = None,
-    omit_fragments: Optional[Iterable[str]] = None,
-    omit_fields: Optional[Iterable[str]] = None,
-    rename_fields: Optional[Mapping[str, str]] = None,
+    omit_variables: Iterable[str] | None = None,
+    omit_fragments: Iterable[str] | None = None,
+    omit_fields: Iterable[str] | None = None,
+    rename_fields: Mapping[str, str] | None = None,
 ) -> ast.Document:
     """Rewrite a GraphQL request string to ensure compatibility with older server versions.