wandb 0.19.8__py3-none-any.whl → 0.19.10__py3-none-any.whl

wandb/sdk/data_types/video.py

@@ -1,11 +1,12 @@
+import functools
 import logging
 import os
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type, Union
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Type, Union
 
 import wandb
 from wandb import util
-from wandb.sdk.lib import filesystem, runid
+from wandb.sdk.lib import filesystem, printer, printer_asyncio, runid
 
 from . import _dtypes
 from ._private import MEDIA_TMP
@@ -90,13 +91,12 @@ class Video(BatchableMedia):
         fps: Optional[int] = None,
         format: Optional[str] = None,
     ):
-        super().__init__()
+        super().__init__(caption=caption)
 
         self._format = format or "gif"
         self._width = None
         self._height = None
         self._channels = None
-        self._caption = caption
         if self._format not in Video.EXTS:
             raise ValueError(
                 "wandb.Video accepts {} formats".format(", ".join(Video.EXTS))
@@ -135,7 +135,11 @@ class Video(BatchableMedia):
                     "wandb.Video accepts a file path or numpy like data as input"
                 )
             fps = fps or 4
-            self.encode(fps=fps)
+            printer_asyncio.run_async_with_spinner(
+                printer.new_printer(),
+                "Encoding video...",
+                functools.partial(self.encode, fps=fps),
+            )
 
     def encode(self, fps: int = 4) -> None:
         # import ImageSequenceClip from the appropriate MoviePy module
@@ -153,29 +157,12 @@ class Video(BatchableMedia):
         filename = os.path.join(
             MEDIA_TMP.name, runid.generate_id() + "." + self._format
         )
-        if TYPE_CHECKING:
-            kwargs: Dict[str, Optional[bool]] = {}
-        try:  # older versions of moviepy do not support logger argument
-            kwargs = {"logger": None}
-            if self._format == "gif":
-                write_gif_with_image_io(clip, filename)
-            else:
-                clip.write_videofile(filename, **kwargs)
-        except TypeError:
-            try:  # even older versions of moviepy do not support progress_bar argument
-                kwargs = {"verbose": False, "progress_bar": False}
-                if self._format == "gif":
-                    clip.write_gif(filename, **kwargs)
-                else:
-                    clip.write_videofile(filename, **kwargs)
-            except TypeError:
-                kwargs = {
-                    "verbose": False,
-                }
-                if self._format == "gif":
-                    clip.write_gif(filename, **kwargs)
-                else:
-                    clip.write_videofile(filename, **kwargs)
+
+        if self._format == "gif":
+            write_gif_with_image_io(clip, filename)
+        else:
+            clip.write_videofile(filename, logger=None)
+
         self._set_file(filename, is_tmp=True)
 
     @classmethod
@@ -190,8 +177,6 @@ class Video(BatchableMedia):
             json_dict["width"] = self._width
         if self._height is not None:
             json_dict["height"] = self._height
-        if self._caption:
-            json_dict["caption"] = self._caption
 
         return json_dict
 

wandb/sdk/interface/interface.py

@@ -173,6 +173,8 @@ class InterfaceBase:
             self._make_config(data=config_dict, obj=proto_run.config)
         if run._telemetry_obj:
             proto_run.telemetry.MergeFrom(run._telemetry_obj)
+        if run._start_runtime:
+            proto_run.runtime = run._start_runtime
         return proto_run
 
     def publish_run(self, run: "Run") -> None:

wandb/{apis/public → sdk/internal}/_generated/__init__.py

@@ -1,6 +1,5 @@
 # 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,
@@ -9,11 +8,6 @@ from .server_features_query import (
 )
 
 __all__ = [
-    "Base",
-    "GQLBase",
-    "GQLId",
-    "SerializedToJson",
-    "Typename",
     "SERVER_FEATURES_QUERY_GQL",
     "ServerFeaturesQuery",
     "ServerFeaturesQueryServerInfo",

wandb/{apis/public → sdk/internal}/_generated/server_features_query.py

@@ -7,15 +7,15 @@ from typing import List, Optional
 
 from pydantic import Field
 
-from .base import GQLBase
+from wandb._pydantic import GQLBase
 
 
 class ServerFeaturesQuery(GQLBase):
-    server_info: Optional["ServerFeaturesQueryServerInfo"] = Field(alias="serverInfo")
+    server_info: Optional[ServerFeaturesQueryServerInfo] = Field(alias="serverInfo")
 
 
 class ServerFeaturesQueryServerInfo(GQLBase):
-    features: List[Optional["ServerFeaturesQueryServerInfoFeatures"]]
+    features: List[Optional[ServerFeaturesQueryServerInfoFeatures]]
 
 
 class ServerFeaturesQueryServerInfoFeatures(GQLBase):

wandb/sdk/internal/internal_api.py

@@ -36,6 +36,7 @@ import requests
 import yaml
 from wandb_gql import Client, gql
 from wandb_gql.client import RetryError
+from wandb_graphql.language.ast import Document
 
 import wandb
 from wandb import env, util
@@ -43,7 +44,9 @@ from wandb.apis.normalize import normalize_exceptions, parse_backend_error_messa
 from wandb.errors import AuthenticationError, CommError, UnsupportedError, UsageError
 from wandb.integration.sagemaker import parse_sm_secrets
 from wandb.old.settings import Settings
+from wandb.proto.wandb_internal_pb2 import ServerFeature
 from wandb.sdk.artifacts._validators import is_artifact_registry_project
+from wandb.sdk.internal._generated import SERVER_FEATURES_QUERY_GQL, ServerFeaturesQuery
 from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
 from wandb.sdk.lib.gql_request import GraphQLSession
 from wandb.sdk.lib.hashutil import B64MD5, md5_file_b64
@@ -365,6 +368,7 @@ class Api:
         self.server_create_run_queue_supports_priority: Optional[bool] = None
         self.server_supports_template_variables: Optional[bool] = None
         self.server_push_to_run_queue_supports_priority: Optional[bool] = None
+        self._server_features_cache: Optional[dict[str, bool]] = None
 
     def gql(self, *args: Any, **kwargs: Any) -> Any:
         ret = self._retry_gql(
@@ -869,6 +873,52 @@ class Api:
         _, _, mutations = self.server_info_introspection()
         return "updateRunQueueItemWarning" in mutations
 
+    def _check_server_feature(self, feature_value: ServerFeature) -> bool:
+        """Check if a server feature is enabled.
+
+        Args:
+            feature_value (ServerFeature): The enum value of the feature to check.
+
+        Returns:
+            bool: True if the feature is enabled, False otherwise.
+
+        Raises:
+            Exception: If server doesn't support feature queries or other errors occur
+        """
+        if self._server_features_cache is None:
+            query = gql(SERVER_FEATURES_QUERY_GQL)
+            response = self.gql(query)
+            server_info = ServerFeaturesQuery.model_validate(response).server_info
+            if server_info and (features := server_info.features):
+                self._server_features_cache = {
+                    f.name: f.is_enabled for f in features if f
+                }
+            else:
+                self._server_features_cache = {}
+
+        return self._server_features_cache.get(ServerFeature.Name(feature_value), False)
+
+    def _check_server_feature_with_fallback(self, feature_value: ServerFeature) -> bool:
+        """Wrapper around check_server_feature that warns and returns False for older unsupported servers.
+
+        Good to use for features that have a fallback mechanism for older servers.
+
+        Args:
+            feature_value (ServerFeature): The enum value of the feature to check.
+
+        Returns:
+            bool: True if the feature is enabled, False otherwise.
+
+        Exceptions:
+            Exception: If an error other than the server not supporting feature queries occurs.
+        """
+        try:
+            return self._check_server_feature(feature_value)
+        except Exception as e:
+            if 'Cannot query field "features" on type "ServerInfo".' in str(e):
+                return False
+            raise e
+
     @normalize_exceptions
     def update_run_queue_item_warning(
         self,
@@ -3703,67 +3753,108 @@ class Api:
         else:
             raise ValueError(f"Unable to find an organization under entity {entity!r}.")
 
-    def use_artifact(
+    def _construct_use_artifact_query(
         self,
         artifact_id: str,
         entity_name: Optional[str] = None,
         project_name: Optional[str] = None,
         run_name: Optional[str] = None,
         use_as: Optional[str] = None,
-    ) -> Optional[Dict[str, Any]]:
-        query_template = """
-        mutation UseArtifact(
-            $entityName: String!,
-            $projectName: String!,
-            $runName: String!,
-            $artifactID: ID!,
-            _USED_AS_TYPE_
-        ) {
-            useArtifact(input: {
-                entityName: $entityName,
-                projectName: $projectName,
-                runName: $runName,
-                artifactID: $artifactID,
-                _USED_AS_VALUE_
-            }) {
-                artifact {
-                    id
-                    digest
-                    description
-                    state
-                    createdAt
-                    metadata
-                }
-            }
-        }
-        """
+        artifact_entity_name: Optional[str] = None,
+        artifact_project_name: Optional[str] = None,
+    ) -> Tuple[Document, Dict[str, Any]]:
+        query_vars = [
+            "$entityName: String!",
+            "$projectName: String!",
+            "$runName: String!",
+            "$artifactID: ID!",
+        ]
+        query_args = [
+            "entityName: $entityName",
+            "projectName: $projectName",
+            "runName: $runName",
+            "artifactID: $artifactID",
+        ]
 
         artifact_types = self.server_use_artifact_input_introspection()
-        if "usedAs" in artifact_types:
-            query_template = query_template.replace(
-                "_USED_AS_TYPE_", "$usedAs: String"
-            ).replace("_USED_AS_VALUE_", "usedAs: $usedAs")
-        else:
-            query_template = query_template.replace("_USED_AS_TYPE_", "").replace(
-                "_USED_AS_VALUE_", ""
-            )
-
-        query = gql(query_template)
+        if "usedAs" in artifact_types and use_as:
+            query_vars.append("$usedAs: String")
+            query_args.append("usedAs: $usedAs")
 
         entity_name = entity_name or self.settings("entity")
         project_name = project_name or self.settings("project")
         run_name = run_name or self.current_run_id
 
-        response = self.gql(
-            query,
-            variable_values={
-                "entityName": entity_name,
-                "projectName": project_name,
-                "runName": run_name,
-                "artifactID": artifact_id,
-                "usedAs": use_as,
-            },
+        variable_values: Dict[str, Any] = {
+            "entityName": entity_name,
+            "projectName": project_name,
+            "runName": run_name,
+            "artifactID": artifact_id,
+            "usedAs": use_as,
+        }
+
+        server_allows_entity_project_information = (
+            self._check_server_feature_with_fallback(
+                ServerFeature.USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION  # type: ignore
+            )
+        )
+        if server_allows_entity_project_information:
+            query_vars.extend(
+                [
+                    "$artifactEntityName: String",
+                    "$artifactProjectName: String",
+                ]
+            )
+            query_args.extend(
+                [
+                    "artifactEntityName: $artifactEntityName",
+                    "artifactProjectName: $artifactProjectName",
+                ]
+            )
+            variable_values["artifactEntityName"] = artifact_entity_name
+            variable_values["artifactProjectName"] = artifact_project_name
+
+        vars_str = ", ".join(query_vars)
+        args_str = ", ".join(query_args)
+
+        query = gql(
+            f"""
+            mutation UseArtifact({vars_str}) {{
+                useArtifact(input: {{{args_str}}}) {{
+                    artifact {{
+                        id
+                        digest
+                        description
+                        state
+                        createdAt
+                        metadata
+                    }}
+                }}
+            }}
+            """
+        )
+        return query, variable_values
+
+    def use_artifact(
+        self,
+        artifact_id: str,
+        entity_name: Optional[str] = None,
+        project_name: Optional[str] = None,
+        run_name: Optional[str] = None,
+        artifact_entity_name: Optional[str] = None,
+        artifact_project_name: Optional[str] = None,
+        use_as: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        query, variable_values = self._construct_use_artifact_query(
+            artifact_id,
+            entity_name,
+            project_name,
+            run_name,
+            use_as,
+            artifact_entity_name,
+            artifact_project_name,
         )
+        response = self.gql(query, variable_values)
 
         if response["useArtifact"]["artifact"]:
             artifact: Dict[str, Any] = response["useArtifact"]["artifact"]

wandb/sdk/internal/profiler.py

@@ -18,12 +18,13 @@ def torch_trace_handler():
      torch.profiler.profile(..., on_trace_ready=wandb.profiler.torch_trace_handler())
      ```
 
-    Calling this function ensures that profiler charts & tables can be viewed in your run dashboard
-    on wandb.ai.
+    Calling this function ensures that profiler charts & tables can be viewed in
+    your run dashboard on wandb.ai.
 
-    Please note that `wandb.init()` must be called before this function is invoked.
-    The PyTorch (torch) version must also be at least 1.9, in order to ensure stability
-    of their Profiler API.
+    Please note that `wandb.init()` must be called before this function is
+    invoked, and the reinit setting must not be set to "create_new". The PyTorch
+    (torch) version must also be at least 1.9, in order to ensure stability of
+    their Profiler API.
 
     Args:
         None

wandb/sdk/internal/run.py

@@ -5,21 +5,28 @@ Semi-stubbed run for internal process use.
 
 """
 
-from wandb._globals import _datatypes_set_callback
+import sys
 
-from .. import wandb_run
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+from wandb.sdk import wandb_run
 
 
 class InternalRun(wandb_run.Run):
     def __init__(self, run_obj, settings, datatypes_cb):
         super().__init__(settings=settings)
         self._run_obj = run_obj
+        self._datatypes_cb = datatypes_cb
 
-        # TODO: This overwrites what's done in the constructor of wandb_run.Run.
-        # We really want a common interface for wandb_run.Run and InternalRun.
-        _datatypes_set_callback(datatypes_cb)
-
+    @override
     def _set_backend(self, backend):
         # This type of run object can't have a backend
         # or do any writes.
         pass
+
+    @override
+    def _publish_file(self, fname: str) -> None:
+        self._datatypes_cb(fname)

wandb/sdk/internal/sender.py

@@ -1531,6 +1531,8 @@ class SendManager:
 
         metadata = json.loads(artifact.metadata) if artifact.metadata else None
         res = saver.save(
+            entity=artifact.entity,
+            project=artifact.project,
             type=artifact.type,
             name=artifact.name,
             client_id=artifact.client_id,

wandb/sdk/internal/sender_config.py

@@ -47,17 +47,14 @@ class ConfigState:
         # Add any top-level keys that aren't already set.
         self._add_unset_keys_from_subtree(old_config_tree, [])
 
-        # Unfortunately, when a user logs visualizations, we store them in the
-        # run's config. When resuming a run, we want to avoid erasing previously
-        # logged visualizations, hence this special handling:
-        self._add_unset_keys_from_subtree(
-            old_config_tree,
-            [_WANDB_INTERNAL_KEY, "visualize"],
-        )
-        self._add_unset_keys_from_subtree(
-            old_config_tree,
-            [_WANDB_INTERNAL_KEY, "viz"],
-        )
+        # When resuming a run, we want to ensure the some of the old configs keys
+        # are maintained. So we have this logic here to add back
+        # any keys that were in the old config but not in the new config
+        for key in ["viz", "visualize", "mask/class_labels"]:
+            self._add_unset_keys_from_subtree(
+                old_config_tree,
+                [_WANDB_INTERNAL_KEY, key],
+            )
 
     def _add_unset_keys_from_subtree(
         self,

wandb/sdk/internal/settings_static.py

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

wandb/sdk/lib/apikey.py

@@ -11,7 +11,7 @@ import textwrap
 from functools import partial
 
 # import Literal
-from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional, Union
+from typing import TYPE_CHECKING, Callable, Literal
 from urllib.parse import urlparse
 
 import click
@@ -23,6 +23,9 @@ from wandb.errors import term
 from wandb.errors.links import url_registry
 from wandb.util import _is_databricks, isatty, prompt_choices
 
+if TYPE_CHECKING:
+    from wandb.sdk.wandb_settings import Settings
+
 LOGIN_CHOICE_ANON = "Private W&B dashboard, no account required"
 LOGIN_CHOICE_NEW = "Create a W&B account"
 LOGIN_CHOICE_EXISTS = "Use an existing W&B account"
@@ -50,18 +53,14 @@ class WriteNetrcError(Exception):
 Mode = Literal["allow", "must", "never", "false", "true"]
 
 
-if TYPE_CHECKING:
-    from wandb.sdk.wandb_settings import Settings
-
-
 getpass = partial(click.prompt, hide_input=True, err=True)
 
 
-def _fixup_anon_mode(default: Optional[Mode]) -> Optional[Mode]:
+def _fixup_anon_mode(default: Mode | None) -> Mode | None:
     # Convert weird anonymode values from legacy settings files
     # into one of our expected values.
     anon_mode = default or "never"
-    mapping: Dict[Mode, Mode] = {"true": "allow", "false": "never"}
+    mapping: dict[Mode, Mode] = {"true": "allow", "false": "never"}
     return mapping.get(anon_mode, anon_mode)
 
 
@@ -84,15 +83,35 @@ def get_netrc_file_path() -> str:
     return os.path.join(os.path.expanduser("~"), netrc_file)
 
 
+def _api_key_prompt_str(app_url: str, referrer: str | None = None) -> str:
+    """Generate a prompt string for API key authorization.
+
+    Creates a URL string that directs users to the authorization page where they
+    can find their API key.
+
+    Args:
+        app_url: The base URL of the W&B application.
+        referrer: Optional referrer parameter to include in the URL.
+
+    Returns:
+        A formatted string with instructions and the authorization URL.
+    """
+    ref = ""
+    if referrer:
+        ref = f"?ref={referrer}"
+    return f"You can find your API key in your browser here: {app_url}/authorize{ref}"
+
+
 def prompt_api_key(  # noqa: C901
-    settings: "Settings",
-    api: Optional[InternalApi] = None,
-    input_callback: Optional[Callable] = None,
-    browser_callback: Optional[Callable] = None,
+    settings: Settings,
+    api: InternalApi | None = None,
+    input_callback: Callable | None = None,
+    browser_callback: Callable | None = None,
     no_offline: bool = False,
     no_create: bool = False,
     local: bool = False,
-) -> Union[str, bool, None]:
+    referrer: str | None = None,
+) -> str | bool | None:
     """Prompt for api key.
 
     Returns:
@@ -150,7 +169,10 @@ def prompt_api_key(  # noqa: C901
         key = browser_callback(signup=True) if browser_callback else None
 
         if not key:
-            wandb.termlog(f"Create an account here: {app_url}/authorize?signup=true")
+            ref = f"&ref={referrer}" if referrer else ""
+            wandb.termlog(
+                f"Create an account here: {app_url}/authorize?signup=true{ref}"
+            )
             key = input_callback(api_ask).strip()
     elif result == LOGIN_CHOICE_EXISTS:
         key = browser_callback() if browser_callback else None
@@ -165,9 +187,7 @@ def prompt_api_key(  # noqa: C901
                     f"Logging into {host}. (Learn how to deploy a W&B server "
                     f"locally: {url_registry.url('wandb-server')})"
                 )
-            wandb.termlog(
-                f"You can find your API key in your browser here: {app_url}/authorize"
-            )
+            wandb.termlog(_api_key_prompt_str(app_url, referrer))
             key = input_callback(api_ask).strip()
     elif result == LOGIN_CHOICE_NOTTY:
         # TODO: Needs refactor as this needs to be handled by caller
@@ -276,9 +296,9 @@ def write_netrc(host: str, entity: str, key: str):
 
 
 def write_key(
-    settings: "Settings",
-    key: Optional[str],
-    api: Optional["InternalApi"] = None,
+    settings: Settings,
+    key: str | None,
+    api: InternalApi | None = None,
 ) -> None:
     if not key:
         raise ValueError("No API key specified.")
@@ -298,7 +318,7 @@ def write_key(
     write_netrc(settings.base_url, "user", key)
 
 
-def api_key(settings: Optional["Settings"] = None) -> Optional[str]:
+def api_key(settings: Settings | None = None) -> str | None:
     if settings is None:
         settings = wandb.setup().settings
     if settings.api_key:

wandb/sdk/lib/asyncio_compat.py

@@ -53,7 +53,7 @@ class _Runner:
     """
 
     def __init__(self) -> None:
-        self._lock = threading.Condition()
+        self._lock = threading.Lock()
 
         self._is_cancelled = False
         self._started = False