wandb 0.21.0__py3-none-win_amd64.whl → 0.21.2__py3-none-win_amd64.whl

wandb/sdk/interface/interface_sock.py

@@ -1,19 +1,18 @@
-"""InterfaceSock - Derived from InterfaceShared using a socket to send to internal thread.
-
-See interface.py for how interface classes relate to each other.
-
-"""
+from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import override
 
-from wandb.sdk.mailbox import Mailbox
+from wandb.proto import wandb_server_pb2 as spb
 
-from ..lib.sock_client import SockClient
 from .interface_shared import InterfaceShared
 
 if TYPE_CHECKING:
     from wandb.proto import wandb_internal_pb2 as pb
+    from wandb.sdk.lib.service.service_client import ServiceClient
+    from wandb.sdk.mailbox import MailboxHandle
 
 
 logger = logging.getLogger("wandb")
@@ -22,18 +21,29 @@ logger = logging.getLogger("wandb")
 class InterfaceSock(InterfaceShared):
     def __init__(
         self,
-        sock_client: SockClient,
-        mailbox: Mailbox,
+        client: ServiceClient,
         stream_id: str,
     ) -> None:
-        super().__init__(mailbox=mailbox)
-        self._sock_client = sock_client
+        super().__init__()
+        self._client = client
         self._stream_id = stream_id
 
     def _assign(self, record: Any) -> None:
         assert self._stream_id
         record._info.stream_id = self._stream_id
 
-    def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None:
+    @override
+    def _publish(self, record: pb.Record, local: bool | None = None) -> None:
         self._assign(record)
-        self._sock_client.send_record_publish(record)
+        request = spb.ServerRequest()
+        request.record_publish.CopyFrom(record)
+        self._client.publish(request)
+
+    @override
+    def _deliver(self, record: pb.Record) -> MailboxHandle[pb.Result]:
+        self._assign(record)
+        request = spb.ServerRequest()
+        request.record_publish.CopyFrom(record)
+
+        handle = self._client.deliver(request)
+        return handle.map(lambda response: response.result_communicate)

wandb/sdk/internal/internal_api.py

@@ -284,6 +284,7 @@ class Api:
         self._extra_http_headers.update(_thread_local_api_settings.headers or {})
 
         auth = None
+        api_key = api_key or self.default_settings.get("api_key")
         if api_key:
             auth = ("api", api_key)
         elif self.access_token is not None:
@@ -3234,6 +3235,7 @@ class Api:
         entity: Optional[str] = None,
         state: Optional[str] = None,
         prior_runs: Optional[List[str]] = None,
+        display_name: Optional[str] = None,
         template_variable_values: Optional[Dict[str, Any]] = None,
     ) -> Tuple[str, List[str]]:
         """Upsert a sweep object.
@@ -3248,6 +3250,7 @@ class Api:
             entity (str): entity to use
             state (str): state
             prior_runs (list): IDs of existing runs to add to the sweep
+            display_name (str): display name for the sweep
             template_variable_values (dict): template variable values
         """
         project_query = """
@@ -3271,6 +3274,7 @@ class Api:
             $scheduler: JSONString,
             $state: String,
             $priorRunsFilters: JSONString,
+            $displayName: String,
         ) {
             upsertSweep(input: {
                 id: $id,
@@ -3282,6 +3286,7 @@ class Api:
                 scheduler: $scheduler,
                 state: $state,
                 priorRunsFilters: $priorRunsFilters,
+                displayName: $displayName,
             }) {
                 sweep {
                     name
@@ -3358,6 +3363,7 @@ class Api:
                     "templateVariableValues": json.dumps(template_variable_values),
                     "scheduler": scheduler,
                     "priorRunsFilters": filters,
+                    "displayName": display_name,
                 }
                 if state:
                     variables["state"] = state

wandb/sdk/internal/job_builder.py

@@ -109,6 +109,7 @@ class JobSourceDict(TypedDict, total=False):
     input_types: Dict[str, Any]
     output_types: Dict[str, Any]
     runtime: Optional[str]
+    services: Dict[str, str]
 
 
 class ArtifactInfoForJob(TypedDict):
@@ -143,6 +144,7 @@ class JobBuilder:
     _job_version_alias: Optional[str]
     _is_notebook_run: bool
     _verbose: bool
+    _services: Dict[str, str]
 
     def __init__(self, settings: SettingsStatic, verbose: bool = False):
         self._settings = settings
@@ -162,6 +164,7 @@ class JobBuilder:
         self._is_notebook_run = self._get_is_notebook_run()
         self._verbose = verbose
         self._partial = False
+        self._services = {}
 
     def set_config(self, config: Dict[str, Any]) -> None:
         self._config = config
@@ -544,6 +547,9 @@ class JobBuilder:
             "runtime": runtime,
         }
 
+        if self._services:
+            source_info["services"] = self._services
+
         assert source_info is not None
         assert name is not None
 

wandb/sdk/internal/settings_static.py

@@ -4,7 +4,7 @@ from typing import Any, Iterable
 
 from wandb.proto import wandb_settings_pb2
 from wandb.sdk.lib import RunMoment
-from wandb.sdk.wandb_settings import Settings
+from wandb.sdk.wandb_settings import CLIENT_ONLY_SETTINGS, Settings
 
 
 class SettingsStatic(Settings):
@@ -41,8 +41,7 @@ class SettingsStatic(Settings):
 
         forks_specified: list[str] = []
         for key in fields:
-            # Skip Python-only keys that do not exist on the proto.
-            if key in ("reinit",):
+            if key in CLIENT_ONLY_SETTINGS:
                 continue
 
             value: Any = None

wandb/sdk/launch/agent/agent.py

@@ -1,6 +1,7 @@
 """Implementation of launch agent."""
 
 import asyncio
+import copy
 import logging
 import os
 import pprint
@@ -421,6 +422,7 @@ class LaunchAgent:
         """Removes the job from our list for now."""
         with self._jobs_lock:
             job_and_run_status = self._jobs[thread_id]
+
         if (
             job_and_run_status.entity is not None
             and job_and_run_status.entity != self._entity
@@ -516,7 +518,11 @@ class LaunchAgent:
         Arguments:
             job: Job to run.
         """
-        _msg = f"{LOG_PREFIX}Launch agent received job:\n{pprint.pformat(job)}\n"
+        job_copy = copy.deepcopy(job)
+        if "runSpec" in job_copy and "_wandb_api_key" in job_copy["runSpec"]:
+            job_copy["runSpec"]["_wandb_api_key"] = "<redacted>"
+
+        _msg = f"{LOG_PREFIX}Launch agent received job:\n{pprint.pformat(job_copy)}\n"
         wandb.termlog(_msg)
         _logger.info(_msg)
         # update agent status
@@ -727,6 +733,7 @@ class LaunchAgent:
         backend = loader.runner_from_config(
             resource, api, backend_config, environment, registry
         )
+
         if not (
             project.docker_image
             or project.job_base_image

wandb/sdk/launch/agent/run_queue_item_file_saver.py

@@ -11,7 +11,7 @@ FileSubtypes = Literal["warning", "error"]
 class RunQueueItemFileSaver:
     def __init__(
         self,
-        agent_run: Optional["wandb.sdk.wandb_run.Run"],
+        agent_run: Optional["wandb.Run"],
         run_queue_item_id: str,
     ):
         self.run_queue_item_id = run_queue_item_id
@@ -20,7 +20,7 @@ class RunQueueItemFileSaver:
     def save_contents(
         self, contents: str, fname: str, file_sub_type: FileSubtypes
     ) -> Optional[List[str]]:
-        if not isinstance(self.run, wandb.sdk.wandb_run.Run):
+        if not isinstance(self.run, wandb.Run):
             wandb.termwarn("Not saving file contents because agent has no run")
             return None
         root_dir = self.run._settings.files_dir

wandb/sdk/launch/create_job.py

@@ -11,6 +11,7 @@ from wandb.apis.internal import Api
 from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal.job_builder import JobBuilder
 from wandb.sdk.launch.git_reference import GitReference
+from wandb.sdk.launch.inputs.internal import _validate_schema
 from wandb.sdk.launch.utils import (
     _is_git_uri,
     get_current_python_version,
@@ -115,6 +116,8 @@ def _create_job(
     build_context: Optional[str] = None,
     dockerfile: Optional[str] = None,
     base_image: Optional[str] = None,
+    services: Optional[Dict[str, str]] = None,
+    schema: Optional[Dict[str, Any]] = None,
 ) -> Tuple[Optional[Artifact], str, List[str]]:
     wandb.termlog(f"Creating launch job of type: {job_type}...")
 
@@ -169,6 +172,7 @@ def _create_job(
 
     job_builder = _configure_job_builder_for_partial(tempdir.name, job_source=job_type)
     job_builder._settings.job_name = name
+    job_builder._services = services or {}
     if job_type == "code":
         assert entrypoint is not None
         job_name = _make_code_artifact(
@@ -204,6 +208,15 @@ def _create_job(
     if "latest" not in aliases:
         aliases += ["latest"]
 
+    metadata = {"_partial": True}
+    if schema:
+        _validate_schema(schema)
+        metadata = {
+            "input_schemas": {
+                "@wandb.config": schema,
+            }
+        }
+
     res, _ = api.create_artifact(
         artifact_type_name="job",
         artifact_collection_name=name,
@@ -214,7 +227,7 @@ def _create_job(
         project_name=project,
         run_name=run.id,  # type: ignore # run will be deleted after creation
         description=description,
-        metadata={"_partial": True},
+        metadata=metadata,
         is_user_created=True,
         aliases=[{"artifactCollectionName": name, "alias": a} for a in aliases],
     )
@@ -421,7 +434,7 @@ def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuild
 def _make_code_artifact(
     api: Api,
     job_builder: JobBuilder,
-    run: "wandb.sdk.wandb_run.Run",
+    run: "wandb.Run",
     path: str,
     entrypoint: str,
     entity: Optional[str],

wandb/sdk/launch/inputs/internal.py

@@ -17,7 +17,6 @@ import wandb
 import wandb.data_types
 from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.inputs.schema import META_SCHEMA
-from wandb.sdk.wandb_run import Run
 from wandb.util import get_module
 
 from .files import config_path_is_valid, override_file
@@ -93,7 +92,7 @@ class StagedLaunchInputs:
     ):
         self._staged_inputs.append(input_arguments)
 
-    def apply(self, run: Run):
+    def apply(self, run: wandb.Run):
         """Apply the staged inputs to the given run."""
         for input in self._staged_inputs:
             _publish_job_input(input, run)
@@ -101,13 +100,13 @@ class StagedLaunchInputs:
 
 def _publish_job_input(
     input: JobInputArguments,
-    run: Run,
+    run: wandb.Run,
 ) -> None:
     """Publish a job input to the backend interface of the given run.
 
     Arguments:
         input (JobInputArguments): The arguments for the job input.
-        run (Run): The run to publish the job input to.
+        run (wandb.Run): The run to publish the job input to.
     """
     assert run._backend is not None
     assert run._backend.interface is not None

wandb/sdk/launch/inputs/schema.py

@@ -7,6 +7,7 @@ META_SCHEMA = {
         },
         "title": {"type": "string"},
         "description": {"type": "string"},
+        "format": {"type": "string"},
         "enum": {"type": "array", "items": {"type": ["integer", "number", "string"]}},
         "properties": {"type": "object", "patternProperties": {".*": {"$ref": "#"}}},
         "allOf": {"type": "array", "items": {"$ref": "#"}},

wandb/sdk/launch/runner/kubernetes_monitor.py

@@ -27,6 +27,7 @@ WANDB_K8S_LABEL_NAMESPACE = "wandb.ai"
 WANDB_K8S_RUN_ID = f"{WANDB_K8S_LABEL_NAMESPACE}/run-id"
 WANDB_K8S_LABEL_AGENT = f"{WANDB_K8S_LABEL_NAMESPACE}/agent"
 WANDB_K8S_LABEL_MONITOR = f"{WANDB_K8S_LABEL_NAMESPACE}/monitor"
+WANDB_K8S_LABEL_AUXILIARY_RESOURCE = f"{WANDB_K8S_LABEL_NAMESPACE}/auxiliary-resource"
 
 
 class Resources: