wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/sdk/internal/internal_api.py

@@ -58,7 +58,7 @@ if TYPE_CHECKING:
     from .progress import ProgressFn
 
     class CreateArtifactFileSpecInput(TypedDict, total=False):
-        """Corresponds to `type CreateArtifactFileSpecInput` in schema.graphql"""
+        """Corresponds to `type CreateArtifactFileSpecInput` in schema.graphql."""
 
         artifactID: str  # noqa: N815
         name: str
@@ -106,7 +106,7 @@ class _ThreadLocalData(threading.local):
 
 
 class Api:
-    """W&B Internal Api wrapper
+    """W&B Internal Api wrapper.
 
     Note:
         Settings are automatically overridden by looking for
@@ -205,6 +205,7 @@ class Api:
         self._azure_blob_module = util.get_module("azure.storage.blob")
 
         self.query_types: Optional[List[str]] = None
+        self.mutation_types: Optional[List[str]] = None
         self.server_info_types: Optional[List[str]] = None
         self.server_use_artifact_input_info: Optional[List[str]] = None
         self._max_cli_version: Optional[str] = None
@@ -229,11 +230,11 @@ class Api:
         return self._local_data.context or self._global_context
 
     def reauth(self) -> None:
-        """Ensures the current api key is set in the transport"""
+        """Ensure the current api key is set in the transport."""
         self.client.transport.auth = ("api", self.api_key or "")
 
     def relocate(self) -> None:
-        """Ensures the current api points to the right server"""
+        """Ensure the current api points to the right server."""
         self.client.transport.url = "%s/graphql" % self.settings("base_url")
 
     def execute(self, *args: Any, **kwargs: Any) -> "_Response":
@@ -384,7 +385,7 @@ class Api:
     def parse_slug(
         self, slug: str, project: Optional[str] = None, run: Optional[str] = None
     ) -> Tuple[str, str]:
-        """Parses a slug into a project and run.
+        """Parse a slug into a project and run.
 
         Arguments:
             slug (str): The slug to parse
@@ -409,12 +410,15 @@ class Api:
         return project, run
 
     @normalize_exceptions
-    def server_info_introspection(self) -> Tuple[List[str], List[str]]:
+    def server_info_introspection(self) -> Tuple[List[str], List[str], List[str]]:
         query_string = """
            query ProbeServerCapabilities {
                QueryType: __type(name: "Query") {
                    ...fieldData
                 }
+                MutationType: __type(name: "Mutation") {
+                   ...fieldData
+                }
                ServerInfoType: __type(name: "ServerInfo") {
                    ...fieldData
                 }
@@ -426,7 +430,11 @@ class Api:
                 }
             }
         """
-        if self.query_types is None or self.server_info_types is None:
+        if (
+            self.query_types is None
+            or self.mutation_types is None
+            or self.server_info_types is None
+        ):
             query = gql(query_string)
             res = self.gql(query)
 
@@ -434,11 +442,15 @@ class Api:
                 field.get("name", "")
                 for field in res.get("QueryType", {}).get("fields", [{}])
             ]
+            self.mutation_types = [
+                field.get("name", "")
+                for field in res.get("MutationType", {}).get("fields", [{}])
+            ]
             self.server_info_types = [
                 field.get("name", "")
                 for field in res.get("ServerInfoType", {}).get("fields", [{}])
             ]
-        return self.query_types, self.server_info_types
+        return self.query_types, self.server_info_types, self.mutation_types
 
     @normalize_exceptions
     def server_settings_introspection(self) -> None:
@@ -505,6 +517,35 @@ class Api:
         res = self.gql(query)
         return res.get("LaunchAgentType") or None
 
+    @normalize_exceptions
+    def fail_run_queue_item_introspection(self) -> bool:
+        _, _, mutations = self.server_info_introspection()
+        return "failRunQueueItem" in mutations
+
+    @normalize_exceptions
+    def fail_run_queue_item(self, run_queue_item_id: str) -> bool:
+        mutation = gql(
+            """
+        mutation failRunQueueItem($runQueueItemId: ID!) {
+            failRunQueueItem(
+                input: {
+                    runQueueItemId: $runQueueItemId
+                }
+            ) {
+                success
+            }
+        }
+        """
+        )
+        response = self.gql(
+            mutation,
+            variable_values={
+                "runQueueItemId": run_queue_item_id,
+            },
+        )
+        result: bool = response["failRunQueueItem"]["success"]
+        return result
+
     @normalize_exceptions
     def viewer(self) -> Dict[str, Any]:
         query = gql(
@@ -530,11 +571,10 @@ class Api:
 
     @normalize_exceptions
     def max_cli_version(self) -> Optional[str]:
-
         if self._max_cli_version is not None:
             return self._max_cli_version
 
-        query_types, server_info_types = self.server_info_introspection()
+        query_types, server_info_types, _ = self.server_info_introspection()
         cli_version_exists = (
             "serverInfo" in query_types and "cliVersionInfo" in server_info_types
         )
@@ -580,7 +620,7 @@ class Api:
             _CLI_QUERY_
         }
         """
-        query_types, server_info_types = self.server_info_introspection()
+        query_types, server_info_types, _ = self.server_info_introspection()
 
         cli_version_exists = (
             "serverInfo" in query_types and "cliVersionInfo" in server_info_types
@@ -603,7 +643,7 @@ class Api:
 
     @normalize_exceptions
     def list_projects(self, entity: Optional[str] = None) -> List[Dict[str, str]]:
-        """Lists projects in W&B scoped by entity.
+        """List projects in W&B scoped by entity.
 
         Arguments:
             entity (str, optional): The entity to scope this project to.
@@ -635,7 +675,7 @@ class Api:
 
     @normalize_exceptions
     def project(self, project: str, entity: Optional[str] = None) -> "_Response":
-        """Retrieve project
+        """Retrieve project.
 
         Arguments:
             project (str): The project to get details for
@@ -743,7 +783,7 @@ class Api:
     def list_runs(
         self, project: str, entity: Optional[str] = None
     ) -> List[Dict[str, str]]:
-        """Lists runs in W&B scoped by project.
+        """List runs in W&B scoped by project.
 
         Arguments:
             project (str): The project to scope the runs to
@@ -784,7 +824,7 @@ class Api:
     def run_config(
         self, project: str, run: Optional[str] = None, entity: Optional[str] = None
     ) -> Tuple[str, Dict[str, Any], Optional[str], Dict[str, Any]]:
-        """Get the relevant configs for a run
+        """Get the relevant configs for a run.
 
         Arguments:
             project (str): The project to download, (can include bucket)
@@ -973,7 +1013,7 @@ class Api:
         description: Optional[str] = None,
         entity: Optional[str] = None,
     ) -> Dict[str, Any]:
-        """Create a new project
+        """Create a new project.
 
         Arguments:
             project (str): The project to create
@@ -1006,6 +1046,32 @@ class Api:
         result: Dict[str, Any] = response["upsertModel"]["model"]
         return result
 
+    @normalize_exceptions
+    def entity_is_team(self, entity: str) -> bool:
+        query = gql(
+            """
+            query EntityIsTeam($entity: String!) {
+                entity(name: $entity) {
+                    id
+                    isTeam
+                }
+            }
+            """
+        )
+        variable_values = {
+            "entity": entity,
+        }
+
+        res = self.gql(query, variable_values)
+        if res.get("entity") is None:
+            raise Exception(
+                f"Error fetching entity {entity} "
+                "check that you have access to this entity"
+            )
+
+        is_team: bool = res["entity"]["isTeam"]
+        return is_team
+
     @normalize_exceptions
     def get_project_run_queues(self, entity: str, project: str) -> List[Dict[str, str]]:
         query = gql(
@@ -1029,10 +1095,19 @@ class Api:
 
         res = self.gql(query, variable_values)
         if res.get("project") is None:
-            raise Exception(
-                f"Error fetching run queues for {entity}/{project} "
-                "check that you have access to this entity and project"
-            )
+            # circular dependency: (LAUNCH_DEFAULT_PROJECT = model-registry)
+            if project == "model-registry":
+                msg = (
+                    f"Error fetching run queues for {entity} "
+                    "check that you have access to this entity and project"
+                )
+            else:
+                msg = (
+                    f"Error fetching run queues for {entity}/{project} "
+                    "check that you have access to this entity and project"
+                )
+
+            raise Exception(msg)
 
         project_run_queues: List[Dict[str, str]] = res["project"]["runQueues"]
         return project_run_queues
@@ -1073,10 +1148,7 @@ class Api:
     def push_to_run_queue_by_name(
         self, entity: str, project: str, queue_name: str, run_spec: str
     ) -> Optional[Dict[str, Any]]:
-        """
-        Queryless mutation, should be used before legacy fallback method
-        """
-
+        """Queryless mutation, should be used before legacy fallback method."""
         mutation = gql(
             """
         mutation pushToRunQueueByName(
@@ -1445,7 +1517,7 @@ class Api:
         summary_metrics: Optional[str] = None,
         num_retries: Optional[int] = None,
     ) -> Tuple[dict, bool, Optional[List]]:
-        """Update a run
+        """Update a run.
 
         Arguments:
             id (str, optional): The existing run to update
@@ -1468,7 +1540,6 @@ class Api:
             summary_metrics (str, optional): The JSON summary metrics
             num_retries (int, optional): Number of retries
         """
-
         query_string = """
         mutation UpsertBucket(
             $id: String,
@@ -1706,7 +1777,7 @@ class Api:
         entity: Optional[str] = None,
         description: Optional[str] = None,
     ) -> Tuple[str, List[str], Dict[str, Dict[str, Any]]]:
-        """Generate temporary resumable upload urls
+        """Generate temporary resumable upload urls.
 
         Arguments:
             project (str): The project to download
@@ -1775,7 +1846,7 @@ class Api:
         run: Optional[str] = None,
         entity: Optional[str] = None,
     ) -> Dict[str, Dict[str, str]]:
-        """Generate download urls
+        """Generate download urls.
 
         Arguments:
             project (str): The project to download
@@ -1834,7 +1905,7 @@ class Api:
         run: Optional[str] = None,
         entity: Optional[str] = None,
     ) -> Optional[Dict[str, str]]:
-        """Generate download urls
+        """Generate download urls.
 
         Arguments:
             project (str): The project to download
@@ -1887,7 +1958,7 @@ class Api:
 
     @normalize_exceptions
     def download_file(self, url: str) -> Tuple[int, requests.Response]:
-        """Initiate a streaming download
+        """Initiate a streaming download.
 
         Arguments:
             url (str): The url to download
@@ -1905,7 +1976,7 @@ class Api:
         metadata: Dict[str, str],
         out_dir: Optional[str] = None,
     ) -> Tuple[str, Optional[requests.Response]]:
-        """Download a file from a run and write it to wandb/
+        """Download a file from a run and write it to wandb/.
 
         Arguments:
             metadata (obj): The metadata object for the file to download. Comes from Api.download_urls().
@@ -1931,9 +2002,7 @@ class Api:
     def upload_file_azure(
         self, url: str, file: Any, extra_headers: Dict[str, str]
     ) -> None:
-        """
-        Upload a file to azure
-        """
+        """Upload a file to azure."""
         from azure.core.exceptions import AzureError  # type: ignore
 
         # Configure the client without retries so our existing logic can handle them
@@ -1972,7 +2041,7 @@ class Api:
         callback: Optional["ProgressFn"] = None,
         extra_headers: Optional[Dict[str, str]] = None,
     ) -> Optional[requests.Response]:
-        """Uploads a file to W&B with failure resumption
+        """Upload a file to W&B with failure resumption.
 
         Arguments:
             url: The url to download
@@ -2039,7 +2108,7 @@ class Api:
         project_name: Optional[str] = None,
         entity: Optional[str] = None,
     ) -> dict:
-        """Register a new agent
+        """Register a new agent.
 
         Arguments:
             host (str): hostname
@@ -2318,7 +2387,7 @@ class Api:
 
     @normalize_exceptions
     def create_anonymous_api_key(self) -> str:
-        """Creates a new API key belonging to a new anonymous user."""
+        """Create a new API key belonging to a new anonymous user."""
         mutation = gql(
             """
         mutation CreateAnonymousApiKey {
@@ -2337,14 +2406,14 @@ class Api:
 
     @staticmethod
     def file_current(fname: str, md5: B64MD5) -> bool:
-        """Checksum a file and compare the md5 with the known md5"""
+        """Checksum a file and compare the md5 with the known md5."""
         return os.path.isfile(fname) and md5_file_b64(fname) == md5
 
     @normalize_exceptions
     def pull(
         self, project: str, run: Optional[str] = None, entity: Optional[str] = None
     ) -> "List[requests.Response]":
-        """Download files from W&B
+        """Download files from W&B.
 
         Arguments:
             project (str): The project to download
@@ -2379,7 +2448,7 @@ class Api:
         force: bool = True,
         progress: Union[TextIO, bool] = False,
     ) -> "List[Optional[requests.Response]]":
-        """Uploads multiple files to W&B
+        """Uploads multiple files to W&B.
 
         Arguments:
             files (list or dict): The filenames to upload, when dict the values are open files
@@ -2949,7 +3018,6 @@ class Api:
         self,
         client_id: str,
     ) -> Optional[str]:
-
         if client_id in self._client_id_mapping:
             return self._client_id_mapping[client_id]
 
@@ -3142,9 +3210,7 @@ class Api:
         entity: Optional[str] = None,
         project: Optional[str] = None,
     ) -> None:
-        """
-        Finish the sweep to stop running new runs and let currently running runs finish.
-        """
+        """Finish the sweep to stop running new runs and let currently running runs finish."""
         self.set_sweep_state(
             sweep=sweep, state="FINISHED", entity=entity, project=project
         )
@@ -3155,9 +3221,7 @@ class Api:
         entity: Optional[str] = None,
         project: Optional[str] = None,
     ) -> None:
-        """
-        Cancel the sweep to kill all running runs and stop running new runs.
-        """
+        """Cancel the sweep to kill all running runs and stop running new runs."""
         self.set_sweep_state(
             sweep=sweep, state="CANCELED", entity=entity, project=project
         )
@@ -3168,9 +3232,7 @@ class Api:
         entity: Optional[str] = None,
         project: Optional[str] = None,
     ) -> None:
-        """
-        Pause the sweep to temporarily stop running new runs.
-        """
+        """Pause the sweep to temporarily stop running new runs."""
         self.set_sweep_state(
             sweep=sweep, state="PAUSED", entity=entity, project=project
         )
@@ -3181,20 +3243,18 @@ class Api:
         entity: Optional[str] = None,
         project: Optional[str] = None,
     ) -> None:
-        """
-        Resume the sweep to continue running new runs.
-        """
+        """Resume the sweep to continue running new runs."""
         self.set_sweep_state(
             sweep=sweep, state="RUNNING", entity=entity, project=project
         )
 
     def _status_request(self, url: str, length: int) -> requests.Response:
-        """Ask google how much we've uploaded"""
+        """Ask google how much we've uploaded."""
         return requests.put(
             url=url,
             headers={"Content-Length": "0", "Content-Range": "bytes */%i" % length},
         )
 
     def _flatten_edges(self, response: "_Response") -> List[Dict]:
-        """Return an array from the nested graphql relay structure"""
+        """Return an array from the nested graphql relay structure."""
         return [node["node"] for node in response["edges"]]

wandb/sdk/internal/job_builder.py

@@ -1,6 +1,4 @@
-"""
-job builder.
-"""
+"""job builder."""
 import json
 import os
 import sys

wandb/sdk/internal/profiler.py

@@ -10,7 +10,7 @@ PYTORCH_PROFILER_MODULE = "torch.profiler"
 
 
 def torch_trace_handler():
-    """Creates a trace handler for traces generated by the profiler.
+    """Create a trace handler for traces generated by the profiler.
 
      Provide as an argument to `torch.profiler.profile`:
      ```python

wandb/sdk/internal/progress.py

@@ -1,6 +1,4 @@
-"""
-progress.
-"""
+"""progress."""
 
 import os
 import sys
@@ -20,7 +18,7 @@ if TYPE_CHECKING:
 
 
 class Progress:
-    """A helper class for displaying progress"""
+    """A helper class for displaying progress."""
 
     ITER_BYTES = 1024 * 1024
 
@@ -40,7 +38,7 @@ class Progress:
         self.len = os.fstat(file.fileno()).st_size
 
     def read(self, size=-1):
-        """Read bytes and call the callback"""
+        """Read bytes and call the callback."""
         bites = self.file.read(size)
         self.bytes_read += len(bites)
         if not bites and self.bytes_read < self.len:
@@ -64,7 +62,7 @@ class Progress:
         self.file.seek(0)
 
     def __getattr__(self, name):
-        """Fallback to the file object for attrs not defined here"""
+        """Fallback to the file object for attrs not defined here."""
         if hasattr(self.file, name):
             return getattr(self.file, name)
         else:

wandb/sdk/internal/sample.py

@@ -1,6 +1,4 @@
-"""
-sample.
-"""
+"""sample."""
 
 import math
 

wandb/sdk/internal/sender.py

@@ -1,6 +1,4 @@
-"""
-sender.
-"""
+"""sender."""
 
 
 import json
@@ -29,10 +27,11 @@ import requests
 
 import wandb
 from wandb import util
-from wandb.errors import ContextCancelledError
+from wandb.errors import CommError
 from wandb.filesync.dir_watcher import DirWatcher
 from wandb.proto import wandb_internal_pb2
 from wandb.sdk.lib import redirect
+from wandb.sdk.lib.mailbox import ContextCancelledError
 
 from ..interface import interface
 from ..interface.interface_queue import InterfaceQueue
@@ -271,7 +270,8 @@ class SendManager:
 
     @classmethod
     def setup(cls, root_dir: str, resume: Union[None, bool, str]) -> "SendManager":
-        """This is a helper class method to set up a standalone SendManager.
+        """Set up a standalone SendManager.
+
         Currently, we're using this primarily for `sync.py`.
         """
         files_dir = os.path.join(root_dir, "files")
@@ -710,8 +710,7 @@ class SendManager:
     def _maybe_setup_resume(
         self, run: "RunRecord"
     ) -> Optional["wandb_internal_pb2.ErrorInfo"]:
-        """This maybe queries the backend for a run and fails if the settings are
-        incompatible."""
+        """Queries the backend for a run; fail if the settings are incompatible."""
         if not self._settings.resume:
             return None
 
@@ -850,7 +849,7 @@ class SendManager:
         config_util.save_config_file_from_dict(config_path, config_value_dict)
 
     def _sync_spell(self) -> None:
-        """Syncs this run with spell"""
+        """Sync this run with spell."""
         if not self._run:
             return
         try:
@@ -926,7 +925,19 @@ class SendManager:
             config_value_dict = self._config_format(None)
             self._config_save(config_value_dict)
 
-        self._init_run(run, config_value_dict)
+        try:
+            self._init_run(run, config_value_dict)
+        except CommError as e:
+            logger.error(e, exc_info=True)
+            if record.control.req_resp or record.control.mailbox_slot:
+                result = proto_util._result_from_record(record)
+                result.run_result.run.CopyFrom(run)
+                error = wandb_internal_pb2.ErrorInfo()
+                error.message = str(e)
+                result.run_result.error.CopyFrom(error)
+                self._respond_result(result)
+            return
+
         assert self._run  # self._run is configured in _init_run()
 
         if record.control.req_resp or record.control.mailbox_slot:
@@ -1360,9 +1371,9 @@ class SendManager:
                 logger.warning("Failed to link artifact to portfolio: %s", e)
 
     def send_use_artifact(self, record: "Record") -> None:
-        """
-        This function doesn't actually send anything, it is just used
-        internally
+        """Pretend to send a used artifact.
+
+        This function doesn't actually send anything, it is just used internally.
         """
         use = record.use_artifact
         if use.type == "job":
@@ -1530,10 +1541,11 @@ class SendManager:
         return self._cached_server_info
 
     def get_local_info(self) -> "LocalInfo":
-        """
-        This is a helper function that queries the server to get the local version information.
-        First, we perform an introspection, if it returns empty we deduce that the docker image is
-        out-of-date. Otherwise, we use the returned values to deduce the state of the local server.
+        """Queries the server to get the local version information.
+
+        First, we perform an introspection, if it returns empty we deduce that the
+        docker image is out-of-date. Otherwise, we use the returned values to deduce the
+        state of the local server.
         """
         local_info = wandb_internal_pb2.LocalInfo()
         if self._settings._offline:

wandb/sdk/internal/settings_static.py

@@ -1,7 +1,5 @@
-"""
-static settings.
-"""
-from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
+"""static settings."""
+from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Union
 
 SettingsDict = Dict[str, Union[str, float, Tuple, None]]
 
@@ -20,6 +18,8 @@ class SettingsStatic:
     _stats_samples_to_average: int
     _stats_join_assets: bool
     _stats_neuron_monitor_config_path: Optional[str]
+    _stats_open_metrics_endpoints: Mapping[str, str]
+    _stats_open_metrics_filters: Mapping[str, Mapping[str, str]]
     files_dir: str
     program_relpath: Optional[str]
     log_internal: str
@@ -32,7 +32,7 @@ class SettingsStatic:
     _jupyter_name: Optional[str]
     _jupyter_root: Optional[str]
     _network_buffer: Optional[int]
-    _require_service: Optional[str]
+    _disable_service: Optional[bool]
     _live_policy_rate_limit: Optional[int]
     resume: Optional[str]
     program: Optional[str]

wandb/sdk/internal/system/assets/__init__.py

@@ -6,5 +6,6 @@ from .gpu_apple import GPUApple  # noqa: F401
 from .ipu import IPU  # noqa: F401
 from .memory import Memory  # noqa: F401
 from .network import Network  # noqa: F401
+from .open_metrics import OpenMetrics  # noqa: F401
 from .tpu import TPU  # noqa: F401
 from .trainium import Trainium  # noqa: F401

wandb/sdk/internal/system/assets/cpu.py

@@ -20,9 +20,7 @@ if TYPE_CHECKING:
 
 
 class ProcessCpuPercent:
-    """
-    CPU usage of the process in percent normalized by the number of CPUs.
-    """
+    """CPU usage of the process in percent normalized by the number of CPUs."""
 
     # name = "process_cpu_percent"
     name = "cpu"
@@ -58,9 +56,7 @@ class ProcessCpuPercent:
 
 
 class CpuPercent:
-    """
-    CPU usage of the system in percent per core.
-    """
+    """CPU usage of the system in percent per core."""
 
     name = "cpu.{i}.cpu_percent"
 
@@ -87,9 +83,7 @@ class CpuPercent:
 
 
 class ProcessCpuThreads:
-    """
-    Number of threads used by the process.
-    """
+    """Number of threads used by the process."""
 
     name = "proc.cpu.threads"