wandb 0.20.2rc20250616__py3-none-macosx_11_0_arm64.whl → 0.21.1__py3-none-macosx_11_0_arm64.whl

wandb/sdk/interface/interface_shared.py

@@ -59,6 +59,10 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(telemetry=telem)
         self._publish(rec)
 
+    def _publish_environment(self, environment: pb.EnvironmentRecord) -> None:
+        rec = self._make_record(environment=environment)
+        self._publish(rec)
+
     def _publish_job_input(
         self, job_input: pb.JobInputRequest
     ) -> MailboxHandle[pb.Result]:
@@ -106,11 +110,9 @@ class InterfaceShared(InterfaceBase):
         summary_record: Optional[pb.SummaryRecordRequest] = None,
         telemetry_record: Optional[pb.TelemetryRecordRequest] = None,
         get_system_metrics: Optional[pb.GetSystemMetricsRequest] = None,
-        get_system_metadata: Optional[pb.GetSystemMetadataRequest] = None,
         python_packages: Optional[pb.PythonPackagesRequest] = None,
         job_input: Optional[pb.JobInputRequest] = None,
         run_finish_without_exit: Optional[pb.RunFinishWithoutExitRequest] = None,
-        metadata: Optional[pb.MetadataRequest] = None,
     ) -> pb.Record:
         request = pb.Request()
         if get_summary:
@@ -169,8 +171,6 @@ class InterfaceShared(InterfaceBase):
             request.telemetry_record.CopyFrom(telemetry_record)
         elif get_system_metrics:
             request.get_system_metrics.CopyFrom(get_system_metrics)
-        elif get_system_metadata:
-            request.get_system_metadata.CopyFrom(get_system_metadata)
         elif sync_finish:
             request.sync_finish.CopyFrom(sync_finish)
         elif python_packages:
@@ -179,8 +179,6 @@ class InterfaceShared(InterfaceBase):
             request.job_input.CopyFrom(job_input)
         elif run_finish_without_exit:
             request.run_finish_without_exit.CopyFrom(run_finish_without_exit)
-        elif metadata:
-            request.metadata.CopyFrom(metadata)
         else:
             raise Exception("Invalid request")
         record = self._make_record(request=request)
@@ -212,6 +210,7 @@ class InterfaceShared(InterfaceBase):
         use_artifact: Optional[pb.UseArtifactRecord] = None,
         output: Optional[pb.OutputRecord] = None,
         output_raw: Optional[pb.OutputRawRecord] = None,
+        environment: Optional[pb.EnvironmentRecord] = None,
     ) -> pb.Record:
         record = pb.Record()
         if run:
@@ -254,6 +253,8 @@ class InterfaceShared(InterfaceBase):
             record.output.CopyFrom(output)
         elif output_raw:
             record.output_raw.CopyFrom(output_raw)
+        elif environment:
+            record.environment.CopyFrom(environment)
         else:
             raise Exception("Invalid record")
         return record
@@ -304,10 +305,6 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(summary=summary)
         self._publish(rec)
 
-    def _publish_metadata(self, metadata: pb.MetadataRequest) -> None:
-        rec = self._make_request(metadata=metadata)
-        self._publish(rec)
-
     def _publish_metric(self, metric: pb.MetricRecord) -> None:
         rec = self._make_record(metric=metric)
         self._publish(rec)
@@ -422,12 +419,6 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(get_system_metrics=get_system_metrics)
         return self._deliver_record(record)
 
-    def _deliver_get_system_metadata(
-        self, get_system_metadata: pb.GetSystemMetadataRequest
-    ) -> MailboxHandle[pb.Result]:
-        record = self._make_request(get_system_metadata=get_system_metadata)
-        return self._deliver_record(record)
-
     def _deliver_exit(
         self,
         exit_data: pb.RunExitRecord,

wandb/sdk/internal/datastore.py

@@ -124,18 +124,18 @@ class DataStore:
         header = self._fp.read(LEVELDBLOG_HEADER_LEN)
         if len(header) == 0:
             return None
-        assert (
-            len(header) == LEVELDBLOG_HEADER_LEN
-        ), f"record header is {len(header)} bytes instead of the expected {LEVELDBLOG_HEADER_LEN}"
+        assert len(header) == LEVELDBLOG_HEADER_LEN, (
+            f"record header is {len(header)} bytes instead of the expected {LEVELDBLOG_HEADER_LEN}"
+        )
         fields = struct.unpack("<IHB", header)
         checksum, dlength, dtype = fields
         # check len, better fit in the block
         self._index += LEVELDBLOG_HEADER_LEN
         data = self._fp.read(dlength)
         checksum_computed = zlib.crc32(data, self._crc[dtype]) & 0xFFFFFFFF
-        assert (
-            checksum == checksum_computed
-        ), "record checksum is invalid, data may be corrupt"
+        assert checksum == checksum_computed, (
+            "record checksum is invalid, data may be corrupt"
+        )
         self._index += dlength
         return dtype, data
 
@@ -158,9 +158,9 @@ class DataStore:
         if dtype == LEVELDBLOG_FULL:
             return data
 
-        assert (
-            dtype == LEVELDBLOG_FIRST
-        ), f"expected record to be type {LEVELDBLOG_FIRST} but found {dtype}"
+        assert dtype == LEVELDBLOG_FIRST, (
+            f"expected record to be type {LEVELDBLOG_FIRST} but found {dtype}"
+        )
         while True:
             offset = self._index % LEVELDBLOG_BLOCK_LEN
             record = self.scan_record()
@@ -170,9 +170,9 @@ class DataStore:
             if dtype == LEVELDBLOG_LAST:
                 data += new_data
                 break
-            assert (
-                dtype == LEVELDBLOG_MIDDLE
-            ), f"expected record to be type {LEVELDBLOG_MIDDLE} but found {dtype}"
+            assert dtype == LEVELDBLOG_MIDDLE, (
+                f"expected record to be type {LEVELDBLOG_MIDDLE} but found {dtype}"
+            )
             data += new_data
         return data
 
@@ -183,17 +183,17 @@ class DataStore:
             LEVELDBLOG_HEADER_MAGIC,
             LEVELDBLOG_HEADER_VERSION,
         )
-        assert (
-            len(data) == LEVELDBLOG_HEADER_LEN
-        ), f"header size is {len(data)} bytes, expected {LEVELDBLOG_HEADER_LEN}"
+        assert len(data) == LEVELDBLOG_HEADER_LEN, (
+            f"header size is {len(data)} bytes, expected {LEVELDBLOG_HEADER_LEN}"
+        )
         self._fp.write(data)
         self._index += len(data)
 
     def _read_header(self):
         header = self._fp.read(LEVELDBLOG_HEADER_LEN)
-        assert (
-            len(header) == LEVELDBLOG_HEADER_LEN
-        ), f"header is {len(header)} bytes instead of the expected {LEVELDBLOG_HEADER_LEN}"
+        assert len(header) == LEVELDBLOG_HEADER_LEN, (
+            f"header is {len(header)} bytes instead of the expected {LEVELDBLOG_HEADER_LEN}"
+        )
         ident, magic, version = struct.unpack("<4sHB", header)
         if ident != strtobytes(LEVELDBLOG_HEADER_IDENT):
             raise Exception("Invalid header")

wandb/sdk/internal/handler.py

@@ -37,7 +37,6 @@ from wandb.proto.wandb_internal_pb2 import (
 
 from ..interface.interface_queue import InterfaceQueue
 from ..lib import handler_util, proto_util
-from ..wandb_metadata import Metadata
 from . import context, sample, tb_watcher
 from .settings_static import SettingsStatic
 
@@ -115,7 +114,6 @@ class HandleManager:
         self._context_keeper = context_keeper
 
         self._tb_watcher = None
-        self._metadata: Optional[Metadata] = None
         self._step = 0
 
         self._track_time = None
@@ -173,9 +171,6 @@ class HandleManager:
     def handle_request_cancel(self, record: Record) -> None:
         self._dispatch_record(record)
 
-    def handle_request_metadata(self, record: Record) -> None:
-        logger.warning("Metadata updates are ignored when using the legacy service.")
-
     def handle_request_defer(self, record: Record) -> None:
         defer = record.request.defer
         state = defer.state
@@ -649,6 +644,9 @@ class HandleManager:
     def handle_footer(self, record: Record) -> None:
         self._dispatch_record(record)
 
+    def handle_metadata(self, record: Record) -> None:
+        self._dispatch_record(record)
+
     def handle_request_attach(self, record: Record) -> None:
         result = proto_util._result_from_record(record)
         attach_id = record.request.attach.attach_id

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:
@@ -362,6 +363,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:
@@ -3233,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.
@@ -3247,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 = """
@@ -3270,6 +3274,7 @@ class Api:
             $scheduler: JSONString,
             $state: String,
             $priorRunsFilters: JSONString,
+            $displayName: String,
         ) {
             upsertSweep(input: {
                 id: $id,
@@ -3281,6 +3286,7 @@ class Api:
                 scheduler: $scheduler,
                 state: $state,
                 priorRunsFilters: $priorRunsFilters,
+                displayName: $displayName,
             }) {
                 sweep {
                     name
@@ -3357,6 +3363,7 @@ class Api:
                     "templateVariableValues": json.dumps(template_variable_values),
                     "scheduler": scheduler,
                     "priorRunsFilters": filters,
+                    "displayName": display_name,
                 }
                 if state:
                     variables["state"] = state
@@ -4661,3 +4668,56 @@ class Api:
         success: bool = response["stopRun"].get("success")
 
         return success
+
+    @normalize_exceptions
+    def create_custom_chart(
+        self,
+        entity: str,
+        name: str,
+        display_name: str,
+        spec_type: str,
+        access: str,
+        spec: Union[str, Mapping[str, Any]],
+    ) -> Optional[Dict[str, Any]]:
+        if not isinstance(spec, str):
+            spec = json.dumps(spec)
+
+        mutation = gql(
+            """
+            mutation CreateCustomChart(
+                $entity: String!
+                $name: String!
+                $displayName: String!
+                $type: String!
+                $access: String!
+                $spec: JSONString!
+            ) {
+                createCustomChart(
+                    input: {
+                        entity: $entity
+                        name: $name
+                        displayName: $displayName
+                        type: $type
+                        access: $access
+                        spec: $spec
+                    }
+                ) {
+                    chart { id }
+                }
+            }
+            """
+        )
+
+        variable_values = {
+            "entity": entity,
+            "name": name,
+            "displayName": display_name,
+            "type": spec_type,
+            "access": access,
+            "spec": spec,
+        }
+
+        result: Optional[Dict[str, Any]] = self.gql(mutation, variable_values)[
+            "createCustomChart"
+        ]
+        return result

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/sender.py

@@ -63,6 +63,7 @@ if TYPE_CHECKING:
         ArtifactManifest,
         ArtifactManifestEntry,
         ArtifactRecord,
+        EnvironmentRecord,
         HttpResponse,
         LocalInfo,
         Record,
@@ -212,6 +213,7 @@ class SendManager:
     _context_keeper: context.ContextKeeper
 
     _telemetry_obj: telemetry.TelemetryRecord
+    _environment_obj: "EnvironmentRecord"
     _fs: Optional["file_stream.FileStreamApi"]
     _run: Optional["RunRecord"]
     _entity: Optional[str]
@@ -268,6 +270,7 @@ class SendManager:
 
         self._start_time: int = 0
         self._telemetry_obj = telemetry.TelemetryRecord()
+        self._environment_obj = wandb_internal_pb2.EnvironmentRecord()
         self._config_metric_pbdict_list: List[Dict[int, Any]] = []
         self._metadata_summary: Dict[str, Any] = defaultdict()
         self._cached_summary: Dict[str, Any] = dict()
@@ -790,12 +793,12 @@ class SendManager:
 
     def _config_backend_dict(self) -> sender_config.BackendConfigDict:
         config = self._consolidated_config or sender_config.ConfigState()
-
         return config.to_backend_dict(
             telemetry_record=self._telemetry_obj,
             framework=self._telemetry_get_framework(),
             start_time_millis=self._start_time,
             metric_pbdicts=self._config_metric_pbdict_list,
+            environment_record=self._environment_obj,
         )
 
     def _config_save(
@@ -1379,11 +1382,11 @@ class SendManager:
             next_idx = len(self._config_metric_pbdict_list)
             self._config_metric_pbdict_list.append(md)
             self._config_metric_index_dict[metric.name] = next_idx
-        self._update_config()
+        self._debounce_config()
 
     def _update_telemetry_record(self, telemetry: telemetry.TelemetryRecord) -> None:
         self._telemetry_obj.MergeFrom(telemetry)
-        self._update_config()
+        self._debounce_config()
 
     def send_telemetry(self, record: "Record") -> None:
         self._update_telemetry_record(record.telemetry)
@@ -1417,6 +1420,23 @@ class SendManager:
         # tbrecord watching threads are handled by handler.py
         pass
 
+    def _update_environment_record(self, environment: "EnvironmentRecord") -> None:
+        self._environment_obj.MergeFrom(environment)
+        self._debounce_config()
+
+    def send_environment(self, record: "Record") -> None:
+        """Inject environment info into config and upload as a JSON file."""
+        self._update_environment_record(record.environment)
+
+        environment_json = json.dumps(proto_util.message_to_dict(self._environment_obj))
+
+        with open(
+            os.path.join(self._settings.files_dir, filenames.METADATA_FNAME), "w"
+        ) as f:
+            f.write(environment_json)
+
+        self._save_file(interface.GlobStr(filenames.METADATA_FNAME), policy="now")
+
     def send_request_link_artifact(self, record: "Record") -> None:
         if not (record.control.req_resp or record.control.mailbox_slot):
             raise ValueError(

wandb/sdk/internal/sender_config.py

@@ -79,6 +79,7 @@ class ConfigState:
         framework: Optional[str],
         start_time_millis: int,
         metric_pbdicts: Sequence[Dict[int, Any]],
+        environment_record: wandb_internal_pb2.EnvironmentRecord,
     ) -> BackendConfigDict:
         """Returns a dictionary representation expected by the backend.
 
@@ -125,6 +126,14 @@ class ConfigState:
         if metric_pbdicts:
             wandb_internal["m"] = metric_pbdicts
 
+        ###################################################
+        # Environment
+        ###################################################
+        writer_id = environment_record.writer_id
+        if writer_id:
+            environment_dict = proto_util.message_to_dict(environment_record)
+            wandb_internal["e"] = {writer_id: environment_dict}
+
         return BackendConfigDict(
             {
                 key: {

wandb/sdk/launch/_project_spec.py

@@ -370,9 +370,9 @@ class LaunchProject:
 
     def set_job_entry_point(self, command: List[str]) -> "EntryPoint":
         """Set job entrypoint for the project."""
-        assert (
-            self._entry_point is None
-        ), "Cannot set entry point twice. Use LaunchProject.override_entrypoint"
+        assert self._entry_point is None, (
+            "Cannot set entry point twice. Use LaunchProject.override_entrypoint"
+        )
         new_entrypoint = EntryPoint(name=command[-1], command=command)
         self._entry_point = new_entrypoint
         return new_entrypoint

wandb/sdk/launch/agent/agent.py

@@ -1,6 +1,7 @@
 """Implementation of launch agent."""
 
 import asyncio
+import copy
 import logging
 import os
 import pprint
@@ -76,9 +77,9 @@ class JobSpecAndQueue:
 def _convert_access(access: str) -> str:
     """Convert access string to a value accepted by wandb."""
     access = access.upper()
-    assert (
-        access == "PROJECT" or access == "USER"
-    ), "Queue access must be either project or user"
+    assert access == "PROJECT" or access == "USER", (
+        "Queue access must be either project or user"
+    )
     return access
 
 
@@ -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/job_status_tracker.py

@@ -44,7 +44,9 @@ class JobAndRunStatusTracker:
             self.run_id is not None
             and self.project is not None
             and self.entity is not None
-        ), "Job tracker does not contain run info. Update with run info before checking if run stopped"
+        ), (
+            "Job tracker does not contain run info. Update with run info before checking if run stopped"
+        )
         check_stop = event_loop_thread_exec(api.api.check_stop_requested)
         try:
             return bool(await check_stop(self.project, self.entity, self.run_id))

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

@@ -115,6 +115,7 @@ def _create_job(
     build_context: Optional[str] = None,
     dockerfile: Optional[str] = None,
     base_image: Optional[str] = None,
+    services: Optional[Dict[str, str]] = None,
 ) -> Tuple[Optional[Artifact], str, List[str]]:
     wandb.termlog(f"Creating launch job of type: {job_type}...")
 
@@ -169,6 +170,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(
@@ -421,7 +423,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: