wandb 0.17.4__py3-none-win32.whl → 0.17.5__py3-none-win32.whl

wandb/sdk/artifacts/artifact.py

@@ -1300,7 +1300,8 @@ class Artifact:
                 automatic integrity validation. Disabling checksumming will speed up
                 artifact creation but reference directories will not iterated through so the
                 objects in the directory will not be saved to the artifact. We recommend
-                adding reference objects in the case checksumming is false.
+                setting `checksum=False` when adding reference objects, in which case
+                a new version will only be created if the reference URI changes.
             max_objects: The maximum number of objects to consider when adding a
                 reference that points to directory or bucket store prefix. By default,
                 the maximum number of objects allowed for Amazon S3,
@@ -1627,16 +1628,19 @@ class Artifact:
     ) -> FilePathStr:
         """Download the contents of the artifact to the specified root directory.
 
-        Existing files located within `root` are not modified. Explicitly delete
-        `root` before you call `download` if you want the contents of `root` to exactly
-        match the artifact.
+        Existing files located within `root` are not modified. Explicitly delete `root`
+        before you call `download` if you want the contents of `root` to exactly match
+        the artifact.
 
         Arguments:
             root: The directory W&B stores the artifact's files.
             allow_missing_references: If set to `True`, any invalid reference paths
                 will be ignored while downloading referenced files.
-            skip_cache: If set to `True`, the artifact cache will be skipped when downloading
-                and W&B will download each file into the default root or specified download directory.
+            skip_cache: If set to `True`, the artifact cache will be skipped when
+                downloading and W&B will download each file into the default root or
+                specified download directory.
+            path_prefix: If specified, only files with a path that starts with the given
+                prefix will be downloaded. Uses unix format (forward slashes).
 
         Returns:
             The path to the downloaded contents.
@@ -1663,23 +1667,6 @@ class Artifact:
             path_prefix=path_prefix,
         )
 
-    @classmethod
-    def path_contains_dir_prefix(cls, path: StrPath, dir_path: StrPath) -> bool:
-        """Returns true if `path` contains `dir_path` as a prefix."""
-        if not dir_path:
-            return True
-        path_parts = PurePosixPath(path).parts
-        dir_parts = PurePosixPath(dir_path).parts
-        return path_parts[: len(dir_parts)] == dir_parts
-
-    @classmethod
-    def should_download_entry(
-        cls, entry: ArtifactManifestEntry, prefix: Optional[StrPath]
-    ) -> bool:
-        if prefix is None:
-            return True
-        return cls.path_contains_dir_prefix(entry.path, prefix)
-
     def _download_using_core(
         self,
         root: str,
@@ -1816,7 +1803,7 @@ class Artifact:
                         # Handled by core
                         continue
                     entry._download_url = edge["node"]["directUrl"]
-                    if self.should_download_entry(entry, prefix=path_prefix):
+                    if (not path_prefix) or entry.path.startswith(str(path_prefix)):
                         active_futures.add(executor.submit(download_entry, entry))
                 # Wait for download threads to catch up.
                 max_backlog = fetch_url_batch_size

wandb/sdk/interface/interface.py

@@ -358,7 +358,7 @@ class InterfaceBase:
                 proto_extra.value_json = json.dumps(v)
         return proto_manifest
 
-    def publish_link_artifact(
+    def deliver_link_artifact(
         self,
         run: "Run",
         artifact: "Artifact",
@@ -366,8 +366,8 @@ class InterfaceBase:
         aliases: Iterable[str],
         entity: Optional[str] = None,
         project: Optional[str] = None,
-    ) -> None:
-        link_artifact = pb.LinkArtifactRecord()
+    ) -> MailboxHandle:
+        link_artifact = pb.LinkArtifactRequest()
         if artifact.is_draft():
             link_artifact.client_id = artifact._client_id
         else:
@@ -377,10 +377,12 @@ class InterfaceBase:
         link_artifact.portfolio_project = project or run.project
         link_artifact.portfolio_aliases.extend(aliases)
 
-        self._publish_link_artifact(link_artifact)
+        return self._deliver_link_artifact(link_artifact)
 
     @abstractmethod
-    def _publish_link_artifact(self, link_artifact: pb.LinkArtifactRecord) -> None:
+    def _deliver_link_artifact(
+        self, link_artifact: pb.LinkArtifactRequest
+    ) -> MailboxHandle:
         raise NotImplementedError
 
     @staticmethod
@@ -749,6 +751,7 @@ class InterfaceBase:
         self,
         include_paths: List[List[str]],
         exclude_paths: List[List[str]],
+        input_schema: Optional[dict],
         run_config: bool = False,
         file_path: str = "",
     ):
@@ -766,6 +769,8 @@ class InterfaceBase:
         Args:
             include_paths: paths within config to include as job inputs.
             exclude_paths: paths within config to exclude as job inputs.
+            input_schema: A JSON Schema describing which attributes will be
+                editable from the Launch drawer.
             run_config: bool indicating whether wandb.config is the input source.
             file_path: path to file to include as a job input.
         """
@@ -788,6 +793,8 @@ class InterfaceBase:
                 pb.JobInputSource.ConfigFileSource(path=file_path),
             )
         request.input_source.CopyFrom(source)
+        if input_schema:
+            request.input_schema = json_dumps_safer(input_schema)
 
         return self._publish_job_input(request)
 

wandb/sdk/interface/interface_shared.py

@@ -137,6 +137,7 @@ class InterfaceShared(InterfaceBase):
         check_version: Optional[pb.CheckVersionRequest] = None,
         log_artifact: Optional[pb.LogArtifactRequest] = None,
         download_artifact: Optional[pb.DownloadArtifactRequest] = None,
+        link_artifact: Optional[pb.LinkArtifactRequest] = None,
         defer: Optional[pb.DeferRequest] = None,
         attach: Optional[pb.AttachRequest] = None,
         server_info: Optional[pb.ServerInfoRequest] = None,
@@ -184,6 +185,8 @@ class InterfaceShared(InterfaceBase):
             request.log_artifact.CopyFrom(log_artifact)
         elif download_artifact:
             request.download_artifact.CopyFrom(download_artifact)
+        elif link_artifact:
+            request.link_artifact.CopyFrom(link_artifact)
         elif defer:
             request.defer.CopyFrom(defer)
         elif attach:
@@ -242,7 +245,6 @@ class InterfaceShared(InterfaceBase):
         request: Optional[pb.Request] = None,
         telemetry: Optional[tpb.TelemetryRecord] = None,
         preempting: Optional[pb.RunPreemptingRecord] = None,
-        link_artifact: Optional[pb.LinkArtifactRecord] = None,
         use_artifact: Optional[pb.UseArtifactRecord] = None,
         output: Optional[pb.OutputRecord] = None,
         output_raw: Optional[pb.OutputRawRecord] = None,
@@ -282,8 +284,6 @@ class InterfaceShared(InterfaceBase):
             record.metric.CopyFrom(metric)
         elif preempting:
             record.preempting.CopyFrom(preempting)
-        elif link_artifact:
-            record.link_artifact.CopyFrom(link_artifact)
         elif use_artifact:
             record.use_artifact.CopyFrom(use_artifact)
         elif output:
@@ -393,10 +393,6 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(files=files)
         self._publish(rec)
 
-    def _publish_link_artifact(self, link_artifact: pb.LinkArtifactRecord) -> Any:
-        rec = self._make_record(link_artifact=link_artifact)
-        self._publish(rec)
-
     def _publish_use_artifact(self, use_artifact: pb.UseArtifactRecord) -> Any:
         rec = self._make_record(use_artifact=use_artifact)
         self._publish(rec)
@@ -411,6 +407,12 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_request(download_artifact=download_artifact)
         return self._deliver_record(rec)
 
+    def _deliver_link_artifact(
+        self, link_artifact: pb.LinkArtifactRequest
+    ) -> MailboxHandle:
+        rec = self._make_request(link_artifact=link_artifact)
+        return self._deliver_record(rec)
+
     def _publish_artifact(self, proto_artifact: pb.ArtifactRecord) -> None:
         rec = self._make_record(artifact=proto_artifact)
         self._publish(rec)

wandb/sdk/internal/handler.py

@@ -230,7 +230,7 @@ class HandleManager:
     def handle_files(self, record: Record) -> None:
         self._dispatch_record(record)
 
-    def handle_link_artifact(self, record: Record) -> None:
+    def handle_request_link_artifact(self, record: Record) -> None:
         self._dispatch_record(record)
 
     def handle_use_artifact(self, record: Record) -> None:

wandb/sdk/internal/internal_api.py

@@ -232,14 +232,14 @@ class Api:
 
         # todo: remove these hacky hacks after settings refactor is complete
         #  keeping this code here to limit scope and so that it is easy to remove later
-        extra_http_headers = self.settings("_extra_http_headers") or json.loads(
+        self._extra_http_headers = self.settings("_extra_http_headers") or json.loads(
             self._environ.get("WANDB__EXTRA_HTTP_HEADERS", "{}")
         )
-        extra_http_headers.update(_thread_local_api_settings.headers or {})
+        self._extra_http_headers.update(_thread_local_api_settings.headers or {})
 
         auth = None
         if self.access_token is not None:
-            extra_http_headers["Authorization"] = f"Bearer {self.access_token}"
+            self._extra_http_headers["Authorization"] = f"Bearer {self.access_token}"
         elif _thread_local_api_settings.cookies is None:
             auth = ("api", self.api_key or "")
 
@@ -253,7 +253,7 @@ class Api:
                     "User-Agent": self.user_agent,
                     "X-WANDB-USERNAME": env.get_username(env=self._environ),
                     "X-WANDB-USER-EMAIL": env.get_user_email(env=self._environ),
-                    **extra_http_headers,
+                    **self._extra_http_headers,
                 },
                 use_json=True,
                 # this timeout won't apply when the DNS lookup fails. in that case, it will be 60s

wandb/sdk/internal/sender.py

@@ -1473,8 +1473,13 @@ class SendManager:
         # tbrecord watching threads are handled by handler.py
         pass
 
-    def send_link_artifact(self, record: "Record") -> None:
-        link = record.link_artifact
+    def send_request_link_artifact(self, record: "Record") -> None:
+        if not (record.control.req_resp or record.control.mailbox_slot):
+            raise ValueError(
+                f"Expected either `req_resp` or `mailbox_slot`, got: {record.control!r}"
+            )
+        result = proto_util._result_from_record(record)
+        link = record.request.link_artifact
         client_id = link.client_id
         server_id = link.server_id
         portfolio_name = link.portfolio_name
@@ -1490,7 +1495,9 @@ class SendManager:
                     client_id, server_id, portfolio_name, entity, project, aliases
                 )
             except Exception as e:
+                result.response.log_artifact_response.error_message = f'error linking artifact to "{entity}/{project}/{portfolio_name}"; error: {e}'
                 logger.warning("Failed to link artifact to portfolio: %s", e)
+        self._respond_result(result)
 
     def send_use_artifact(self, record: "Record") -> None:
         """Pretend to send a used artifact.

wandb/sdk/launch/builder/kaniko_builder.py

@@ -263,11 +263,17 @@ class KanikoBuilder(AbstractBuilder):
         repo_uri = await self.registry.get_repo_uri()
         image_uri = repo_uri + ":" + image_tag
 
-        if (
-            not launch_project.build_required()
-            and await self.registry.check_image_exists(image_uri)
-        ):
-            return image_uri
+        # The DOCKER_CONFIG_SECRET option is mutually exclusive with the
+        # registry classes, so we must skip the check for image existence in
+        # that case.
+        if not launch_project.build_required():
+            if DOCKER_CONFIG_SECRET:
+                wandb.termlog(
+                    f"Skipping check for existing image {image_uri} due to custom dockerconfig."
+                )
+            else:
+                if await self.registry.check_image_exists(image_uri):
+                    return image_uri
 
         _logger.info(f"Building image {image_uri}...")
         _, api_client = await get_kube_context_and_api_client(
@@ -286,7 +292,12 @@ class KanikoBuilder(AbstractBuilder):
         wandb.termlog(f"{LOG_PREFIX}Created kaniko job {build_job_name}")
 
         try:
-            if isinstance(self.registry, AzureContainerRegistry):
+            # DOCKER_CONFIG_SECRET is a user provided dockerconfigjson. Skip our
+            # dockerconfig handling if it's set.
+            if (
+                isinstance(self.registry, AzureContainerRegistry)
+                and not DOCKER_CONFIG_SECRET
+            ):
                 dockerfile_config_map = client.V1ConfigMap(
                     metadata=client.V1ObjectMeta(
                         name=f"docker-config-{build_job_name}"
@@ -344,7 +355,10 @@ class KanikoBuilder(AbstractBuilder):
         finally:
             wandb.termlog(f"{LOG_PREFIX}Cleaning up resources")
             try:
-                if isinstance(self.registry, AzureContainerRegistry):
+                if (
+                    isinstance(self.registry, AzureContainerRegistry)
+                    and not DOCKER_CONFIG_SECRET
+                ):
                     await core_v1.delete_namespaced_config_map(
                         f"docker-config-{build_job_name}", "wandb"
                     )
@@ -498,7 +512,10 @@ class KanikoBuilder(AbstractBuilder):
                     "readOnly": True,
                 }
             )
-        if isinstance(self.registry, AzureContainerRegistry):
+        if (
+            isinstance(self.registry, AzureContainerRegistry)
+            and not DOCKER_CONFIG_SECRET
+        ):
             # Add the docker config map
             volumes.append(
                 {
@@ -533,7 +550,11 @@ class KanikoBuilder(AbstractBuilder):
         # Apply the rest of our defaults
         pod_labels["wandb"] = "launch"
         # This annotation is required to enable azure workload identity.
-        if isinstance(self.registry, AzureContainerRegistry):
+        # Don't add this label if using a docker config secret for auth.
+        if (
+            isinstance(self.registry, AzureContainerRegistry)
+            and not DOCKER_CONFIG_SECRET
+        ):
             pod_labels["azure.workload.identity/use"] = "true"
         pod_spec["restartPolicy"] = pod_spec.get("restartPolicy", "Never")
         pod_spec["activeDeadlineSeconds"] = pod_spec.get(

wandb/sdk/launch/inputs/internal.py

@@ -11,7 +11,7 @@ import os
 import pathlib
 import shutil
 import tempfile
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 import wandb
 import wandb.data_types
@@ -62,11 +62,13 @@ class JobInputArguments:
         self,
         include: Optional[List[str]] = None,
         exclude: Optional[List[str]] = None,
+        schema: Optional[dict] = None,
         file_path: Optional[str] = None,
         run_config: Optional[bool] = None,
     ):
         self.include = include
         self.exclude = exclude
+        self.schema = schema
         self.file_path = file_path
         self.run_config = run_config
 
@@ -121,15 +123,66 @@ def _publish_job_input(
         exclude_paths=[_split_on_unesc_dot(path) for path in input.exclude]
         if input.exclude
         else [],
+        input_schema=input.schema,
         run_config=input.run_config,
         file_path=input.file_path or "",
     )
 
 
+def _replace_refs_and_allofs(schema: dict, defs: dict) -> dict:
+    """Recursively fix JSON schemas with common issues.
+
+    1. Replaces any instances of $ref with their associated definition in defs
+    2. Removes any "allOf" lists that only have one item, "lifting" the item up
+    See test_internal.py for examples
+    """
+    ret: Dict[str, Any] = {}
+    if "$ref" in schema:
+        # Reference found, replace it with its definition
+        def_key = schema["$ref"].split("#/$defs/")[1]
+        # Also run recursive replacement in case a ref contains more refs
+        return _replace_refs_and_allofs(defs.pop(def_key), defs)
+    for key, val in schema.items():
+        if isinstance(val, dict):
+            # Step into dicts recursively
+            new_val_dict = _replace_refs_and_allofs(val, defs)
+            ret[key] = new_val_dict
+        elif isinstance(val, list):
+            # Step into each item in the list
+            new_val_list = []
+            for item in val:
+                if isinstance(item, dict):
+                    new_val_list.append(_replace_refs_and_allofs(item, defs))
+                else:
+                    new_val_list.append(item)
+            # Lift up allOf blocks with only one item
+            if (
+                key == "allOf"
+                and len(new_val_list) == 1
+                and isinstance(new_val_list[0], dict)
+            ):
+                ret.update(new_val_list[0])
+            else:
+                ret[key] = new_val_list
+        else:
+            # For anything else (str, int, etc) keep it as-is
+            ret[key] = val
+    return ret
+
+
+def _convert_pydantic_model_to_jsonschema(model: Any) -> dict:
+    schema = model.model_json_schema()
+    defs = schema.pop("$defs")
+    if not defs:
+        return schema
+    return _replace_refs_and_allofs(schema, defs)
+
+
 def handle_config_file_input(
     path: str,
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     """Declare an overridable configuration file for a launch job.
 
@@ -151,9 +204,20 @@ def handle_config_file_input(
         path,
         dest,
     )
+    # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+    # or the BaseModel class itself (e.g. schema=MySchema)
+    if hasattr(schema, "model_json_schema") and callable(
+        schema.model_json_schema  # type: ignore
+    ):
+        schema = _convert_pydantic_model_to_jsonschema(schema)
+    if schema and not isinstance(schema, dict):
+        raise LaunchError(
+            "schema must be a dict, Pydantic model instance, or Pydantic model class."
+        )
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,
+        schema=schema,
         file_path=path,
         run_config=False,
     )
@@ -165,7 +229,9 @@ def handle_config_file_input(
 
 
 def handle_run_config_input(
-    include: Optional[List[str]] = None, exclude: Optional[List[str]] = None
+    include: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     """Declare wandb.config as an overridable configuration for a launch job.
 
@@ -175,9 +241,20 @@ def handle_run_config_input(
     If there is no active run, the include and exclude paths are staged and sent
     when a run is created.
     """
+    # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+    # or the BaseModel class itself (e.g. schema=MySchema)
+    if hasattr(schema, "model_json_schema") and callable(
+        schema.model_json_schema  # type: ignore
+    ):
+        schema = _convert_pydantic_model_to_jsonschema(schema)
+    if schema and not isinstance(schema, dict):
+        raise LaunchError(
+            "schema must be a dict, Pydantic model instance, or Pydantic model class."
+        )
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,
+        schema=schema,
         run_config=True,
         file_path=None,
     )

wandb/sdk/launch/inputs/manage.py

@@ -1,12 +1,13 @@
 """Functions for declaring overridable configuration for launch jobs."""
 
-from typing import List, Optional
+from typing import Any, List, Optional
 
 
 def manage_config_file(
     path: str,
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     r"""Declare an overridable configuration file for a launch job.
 
@@ -43,18 +44,27 @@ def manage_config_file(
             relative and must not contain backwards traversal, i.e. `..`.
         include (List[str]): A list of keys to include in the configuration file.
         exclude (List[str]): A list of keys to exclude from the configuration file.
+        schema (dict | Pydantic model): A JSON Schema or Pydantic model describing
+            describing which attributes will be editable from the Launch drawer.
+            Accepts both an instance of a Pydantic BaseModel class or the BaseModel
+            class itself.
 
     Raises:
         LaunchError: If the path is not valid, or if there is no active run.
     """
+    # note: schema's Any type is because in the case where a BaseModel class is
+    # provided, its type is a pydantic internal type that we don't want our typing
+    # to depend on. schema's type should be considered
+    # "Optional[dict | <something with a .model_json_schema() method>]"
     from .internal import handle_config_file_input
 
-    return handle_config_file_input(path, include, exclude)
+    return handle_config_file_input(path, include, exclude, schema)
 
 
 def manage_wandb_config(
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     r"""Declare wandb.config as an overridable configuration for a launch job.
 
@@ -86,10 +96,18 @@ def manage_wandb_config(
     Args:
         include (List[str]): A list of subtrees to include in the configuration.
         exclude (List[str]): A list of subtrees to exclude from the configuration.
+        schema (dict | Pydantic model): A JSON Schema or Pydantic model describing
+            describing which attributes will be editable from the Launch drawer.
+            Accepts both an instance of a Pydantic BaseModel class or the BaseModel
+            class itself.
 
     Raises:
         LaunchError: If there is no active run.
     """
+    # note: schema's Any type is because in the case where a BaseModel class is
+    # provided, its type is a pydantic internal type that we don't want our typing
+    # to depend on. schema's type should be considered
+    # "Optional[dict | <something with a .model_json_schema() method>]"
     from .internal import handle_run_config_input
 
-    handle_run_config_input(include, exclude)
+    handle_run_config_input(include, exclude, schema)

wandb/sdk/lib/tracelog.py

@@ -45,8 +45,8 @@ logger = logging.getLogger(__name__)
 ANNOTATE_QUEUE_NAME = "_DEBUGLOG_QUEUE_NAME"
 
 # capture stdout and stderr before anyone messes with them
-stdout_write = sys.__stdout__.write
-stderr_write = sys.__stderr__.write
+stdout_write = sys.__stdout__.write  # type: ignore
+stderr_write = sys.__stderr__.write  # type: ignore
 
 
 def _log(

wandb/sdk/wandb_manager.py

@@ -114,17 +114,21 @@ class _Manager:
 
         try:
             svc_iface._svc_connect(port=port)
+
         except ConnectionRefusedError as e:
             if not psutil.pid_exists(self._token.pid):
                 message = (
-                    "Connection to wandb service failed "
-                    "since the process is not available. "
+                    "Connection to wandb service failed"
+                    " because the process is not available."
                 )
             else:
-                message = f"Connection to wandb service failed: {e}. "
-            raise ManagerConnectionRefusedError(message)
+                message = "Connection to wandb service failed."
+            raise ManagerConnectionRefusedError(message) from e
+
         except Exception as e:
-            raise ManagerConnectionError(f"Connection to wandb service failed: {e}")
+            raise ManagerConnectionError(
+                "Connection to wandb service failed.",
+            ) from e
 
     def __init__(self, settings: "Settings") -> None:
         """Connects to the internal service, starting it if necessary."""