wandb 0.20.1rc20250604__py3-none-any.whl → 0.21.0__py3-none-any.whl

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import os
+import re
 import time
 from pathlib import PurePosixPath
 from typing import TYPE_CHECKING, Sequence
@@ -52,10 +53,20 @@ class S3Handler(StorageHandler):
             required="s3:// references requires the boto3 library, run pip install wandb[aws]",
             lazy=False,
         )
+
+        from botocore.client import Config  # type: ignore
+
+        s3_endpoint = os.getenv("AWS_S3_ENDPOINT_URL")
+        config = (
+            Config(s3={"addressing_style": "virtual"})
+            if s3_endpoint and self._is_coreweave_endpoint(s3_endpoint)
+            else None
+        )
         self._s3 = boto.session.Session().resource(
             "s3",
-            endpoint_url=os.getenv("AWS_S3_ENDPOINT_URL"),
+            endpoint_url=s3_endpoint,
             region_name=os.getenv("AWS_REGION"),
+            config=config,
         )
         self._botocore = util.get_module("botocore")
         return self._s3
@@ -296,3 +307,33 @@ class S3Handler(StorageHandler):
         if hasattr(obj, "version_id") and obj.version_id and obj.version_id != "null":
             extra["versionID"] = obj.version_id
         return extra
+
+    _CW_LEGACY_NETLOC_REGEX: re.Pattern[str] = re.compile(
+        r"""
+        # accelerated endpoints like "accel-object.<region>.coreweave.com"
+        accel-object\.[a-z0-9-]+\.coreweave\.com
+        |
+        # URLs like "object.<region>.coreweave.com"
+        object\.[a-z0-9-]+\.coreweave\.com
+        """,
+        flags=re.VERBOSE,
+    )
+
+    def _is_coreweave_endpoint(self, endpoint_url: str) -> bool:
+        if not (url := endpoint_url.strip().rstrip("/")):
+            return False
+
+        # Only http://cwlota.com is supported using HTTP
+        if url == "http://cwlota.com":
+            return True
+
+        # Enforce HTTPS otherwise
+        https_url = url if url.startswith("https://") else f"https://{url}"
+        netloc = urlparse(https_url).netloc
+        return bool(
+            # Match for https://cwobject.com
+            (netloc == "cwobject.com")
+            or
+            # Check for legacy endpoints
+            self._CW_LEGACY_NETLOC_REGEX.fullmatch(netloc)
+        )

wandb/sdk/backend/backend.py

@@ -11,7 +11,7 @@ from wandb.sdk.interface.interface import InterfaceBase
 from wandb.sdk.wandb_settings import Settings
 
 if TYPE_CHECKING:
-    from wandb.sdk.lib import service_connection
+    from wandb.sdk.lib.service import service_connection
 
 logger = logging.getLogger("wandb")
 

wandb/sdk/data_types/base_types/media.py

@@ -127,9 +127,9 @@ class Media(WBValue):
         self._path = path
         self._is_tmp = is_tmp
         self._extension = extension
-        assert extension is None or path.endswith(
-            extension
-        ), f'Media file extension "{extension}" must occur at the end of path "{path}".'
+        assert extension is None or path.endswith(extension), (
+            f'Media file extension "{extension}" must occur at the end of path "{path}".'
+        )
 
         with open(self._path, "rb") as f:
             self._sha256 = hashlib.sha256(f.read()).hexdigest()
@@ -247,11 +247,13 @@ class Media(WBValue):
                 json_obj["_latest_artifact_path"] = artifact_entry_latest_url
 
             if artifact_entry_url is None or self.is_bound():
-                assert self.is_bound(), f"Value of type {type(self).__name__} must be bound to a run with bind_to_run() before being serialized to JSON."
+                assert self.is_bound(), (
+                    f"Value of type {type(self).__name__} must be bound to a run with bind_to_run() before being serialized to JSON."
+                )
 
-                assert (
-                    self._run is run
-                ), "We don't support referring to media files across runs."
+                assert self._run is run, (
+                    "We don't support referring to media files across runs."
+                )
 
                 # The following two assertions are guaranteed to pass
                 # by definition is_bound, but are needed for

wandb/sdk/data_types/base_types/wb_value.py

@@ -218,17 +218,17 @@ class WBValue:
     def _set_artifact_source(
         self, artifact: "Artifact", name: Optional[str] = None
     ) -> None:
-        assert (
-            self._artifact_source is None
-        ), f"Cannot update artifact_source. Existing source: {self._artifact_source.artifact}/{self._artifact_source.name}"
+        assert self._artifact_source is None, (
+            f"Cannot update artifact_source. Existing source: {self._artifact_source.artifact}/{self._artifact_source.name}"
+        )
         self._artifact_source = _WBValueArtifactSource(artifact, name)
 
     def _set_artifact_target(
         self, artifact: "Artifact", name: Optional[str] = None
     ) -> None:
-        assert (
-            self._artifact_target is None
-        ), f"Cannot update artifact_target. Existing target: {self._artifact_target.artifact}/{self._artifact_target.name}"
+        assert self._artifact_target is None, (
+            f"Cannot update artifact_target. Existing target: {self._artifact_target.artifact}/{self._artifact_target.name}"
+        )
         self._artifact_target = _WBValueArtifactTarget(artifact, name)
 
     def _get_artifact_entry_ref_url(self) -> Optional[str]:

wandb/sdk/data_types/saved_model.py

@@ -257,9 +257,9 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         self._model_obj = None
 
     def _set_obj(self, model_obj: Any) -> None:
-        assert model_obj is not None and self._validate_obj(
-            model_obj
-        ), f"Invalid model object {model_obj}"
+        assert model_obj is not None and self._validate_obj(model_obj), (
+            f"Invalid model object {model_obj}"
+        )
         self._model_obj = model_obj
 
     def _dump(self, target_path: str) -> None:

wandb/sdk/data_types/table.py

@@ -203,8 +203,8 @@ class Table(Media):
     This means you can embed `Images`, `Video`, `Audio`, and other sorts of rich, annotated media
     directly in Tables, alongside other traditional scalar values.
 
-    This class is the primary class used to generate the Table Visualizer
-    in the UI: https://docs.wandb.ai/guides/data-vis/tables.
+    This class is the primary class used to generate the W&B Tables
+    https://docs.wandb.ai/guides/models/tables/.
     """
 
     MAX_ROWS = 10000
@@ -292,9 +292,9 @@ class Table(Media):
                 self._init_from_list([], columns, optional, dtype)
 
     def _validate_log_mode(self, log_mode):
-        assert (
-            log_mode in _SUPPORTED_LOGGING_MODES
-        ), f"Invalid log_mode: {log_mode}. Must be one of {_SUPPORTED_LOGGING_MODES}"
+        assert log_mode in _SUPPORTED_LOGGING_MODES, (
+            f"Invalid log_mode: {log_mode}. Must be one of {_SUPPORTED_LOGGING_MODES}"
+        )
 
     @staticmethod
     def _assert_valid_columns(columns):
@@ -314,9 +314,9 @@ class Table(Media):
             self.add_data(*row)
 
     def _init_from_ndarray(self, ndarray, columns, optional=True, dtype=None):
-        assert util.is_numpy_array(
-            ndarray
-        ), "ndarray argument expects a `numpy.ndarray` object"
+        assert util.is_numpy_array(ndarray), (
+            "ndarray argument expects a `numpy.ndarray` object"
+        )
         self.data = []
         self._assert_valid_columns(columns)
         self.columns = columns
@@ -325,9 +325,9 @@ class Table(Media):
             self.add_data(*row)
 
     def _init_from_dataframe(self, dataframe, columns, optional=True, dtype=None):
-        assert util.is_pandas_data_frame(
-            dataframe
-        ), "dataframe argument expects a `pandas.core.frame.DataFrame` object"
+        assert util.is_pandas_data_frame(dataframe), (
+            "dataframe argument expects a `pandas.core.frame.DataFrame` object"
+        )
         self.data = []
         columns = list(dataframe.columns)
         self._assert_valid_columns(columns)
@@ -440,17 +440,17 @@ class Table(Media):
         is_fk = isinstance(wbtype, _ForeignKeyType)
         is_fi = isinstance(wbtype, _ForeignIndexType)
         if is_pk or is_fk or is_fi:
-            assert (
-                not optional
-            ), "Primary keys, foreign keys, and foreign indexes cannot be optional."
+            assert not optional, (
+                "Primary keys, foreign keys, and foreign indexes cannot be optional."
+            )
 
         if (is_fk or is_fk) and id(wbtype.params["table"]) == id(self):
             raise AssertionError("Cannot set a foreign table reference to same table.")
 
         if is_pk:
-            assert (
-                self._pk_col is None
-            ), f"Cannot have multiple primary keys - {self._pk_col} is already set as the primary key."
+            assert self._pk_col is None, (
+                f"Cannot have multiple primary keys - {self._pk_col} is already set as the primary key."
+            )
 
         # Update the column type
         self._column_types.params["type_map"][col_name] = wbtype
@@ -464,21 +464,21 @@ class Table(Media):
 
     def _eq_debug(self, other, should_assert=False):
         eq = isinstance(other, Table)
-        assert (
-            not should_assert or eq
-        ), f"Found type {other.__class__}, expected {Table}"
+        assert not should_assert or eq, (
+            f"Found type {other.__class__}, expected {Table}"
+        )
         eq = eq and len(self.data) == len(other.data)
-        assert (
-            not should_assert or eq
-        ), f"Found {len(other.data)} rows, expected {len(self.data)}"
+        assert not should_assert or eq, (
+            f"Found {len(other.data)} rows, expected {len(self.data)}"
+        )
         eq = eq and self.columns == other.columns
-        assert (
-            not should_assert or eq
-        ), f"Found columns {other.columns}, expected {self.columns}"
+        assert not should_assert or eq, (
+            f"Found columns {other.columns}, expected {self.columns}"
+        )
         eq = eq and self._column_types == other._column_types
-        assert (
-            not should_assert or eq
-        ), f"Found column type {other._column_types}, expected column type {self._column_types}"
+        assert not should_assert or eq, (
+            f"Found column type {other._column_types}, expected column type {self._column_types}"
+        )
         if eq:
             for row_ndx in range(len(self.data)):
                 for col_ndx in range(len(self.data[row_ndx])):
@@ -487,9 +487,9 @@ class Table(Media):
                     if util.is_numpy_array(_eq):
                         _eq = ((_eq * -1) + 1).sum() == 0
                     eq = eq and _eq
-                    assert (
-                        not should_assert or eq
-                    ), f"Unequal data at row_ndx {row_ndx} col_ndx {col_ndx}: found {other.data[row_ndx][col_ndx]}, expected {self.data[row_ndx][col_ndx]}"
+                    assert not should_assert or eq, (
+                        f"Unequal data at row_ndx {row_ndx} col_ndx {col_ndx}: found {other.data[row_ndx][col_ndx]}, expected {self.data[row_ndx][col_ndx]}"
+                    )
                     if not eq:
                         return eq
         return eq
@@ -932,9 +932,9 @@ class Table(Media):
         assert isinstance(data, list) or is_np
         assert isinstance(optional, bool)
         is_first_col = len(self.columns) == 0
-        assert is_first_col or len(data) == len(
-            self.data
-        ), f"Expected length {len(self.data)}, found {len(data)}"
+        assert is_first_col or len(data) == len(self.data), (
+            f"Expected length {len(self.data)}, found {len(data)}"
+        )
 
         # Add the new data
         for ndx in range(max(len(data), len(self.data))):
@@ -1257,13 +1257,13 @@ class JoinedTable(Media):
 
     def _eq_debug(self, other, should_assert=False):
         eq = isinstance(other, JoinedTable)
-        assert (
-            not should_assert or eq
-        ), f"Found type {other.__class__}, expected {JoinedTable}"
+        assert not should_assert or eq, (
+            f"Found type {other.__class__}, expected {JoinedTable}"
+        )
         eq = eq and self._join_key == other._join_key
-        assert (
-            not should_assert or eq
-        ), f"Found {other._join_key} join key, expected {self._join_key}"
+        assert not should_assert or eq, (
+            f"Found {other._join_key} join key, expected {self._join_key}"
+        )
         eq = eq and self._table1._eq_debug(other._table1, should_assert)
         eq = eq and self._table2._eq_debug(other._table2, should_assert)
         return eq

wandb/sdk/data_types/trace_tree.py

@@ -261,14 +261,14 @@ class Trace:
             A Span object.
         """
         if kind is not None:
-            assert (
-                kind.upper() in SpanKind.__members__
-            ), "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'"
+            assert kind.upper() in SpanKind.__members__, (
+                "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'"
+            )
             kind = SpanKind(kind.upper())
         if status_code is not None:
-            assert (
-                status_code.upper() in StatusCode.__members__
-            ), "Invalid status code, can be one of 'SUCCESS' or 'ERROR'"
+            assert status_code.upper() in StatusCode.__members__, (
+                "Invalid status code, can be one of 'SUCCESS' or 'ERROR'"
+            )
             status_code = StatusCode(status_code.upper())
         if inputs is not None:
             assert isinstance(inputs, dict), "Inputs must be a dictionary"
@@ -419,9 +419,9 @@ class Trace:
         Args:
             value: The kind of the trace to be set.
         """
-        assert (
-            value.upper() in SpanKind.__members__
-        ), "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'"
+        assert value.upper() in SpanKind.__members__, (
+            "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'"
+        )
         self._span.span_kind = SpanKind(value.upper())
 
     def log(self, name: str) -> None:
@@ -433,8 +433,8 @@ class Trace:
         trace_tree = WBTraceTree(self._span, self._model_dict)
         # NOTE: Does not work for reinit="create_new" runs.
         #   This method should be deprecated and users should call run.log().
-        assert (
-            wandb.run is not None
-        ), "You must call wandb.init() before logging a trace"
+        assert wandb.run is not None, (
+            "You must call wandb.init() before logging a trace"
+        )
         assert len(name.strip()) > 0, "You must provide a valid name to log the trace"
         wandb.run.log({name: trace_tree})

wandb/sdk/interface/interface.py

@@ -150,8 +150,7 @@ class InterfaceBase:
         if run._settings.run_notes is not None:
             proto_run.notes = run._settings.run_notes
         if run._settings.run_tags is not None:
-            for tag in run._settings.run_tags:
-                proto_run.tags.append(tag)
+            proto_run.tags.extend(run._settings.run_tags)
         if run._start_time is not None:
             proto_run.start_time.FromMicroseconds(int(run._start_time * 1e6))
         if run._starting_step is not None:
@@ -217,13 +216,6 @@ class InterfaceBase:
     def _publish_config(self, cfg: pb.ConfigRecord) -> None:
         raise NotImplementedError
 
-    def publish_metadata(self, metadata: pb.MetadataRequest) -> None:
-        self._publish_metadata(metadata)
-
-    @abstractmethod
-    def _publish_metadata(self, metadata: pb.MetadataRequest) -> None:
-        raise NotImplementedError
-
     @abstractmethod
     def _publish_metric(self, metric: pb.MetricRecord) -> None:
         raise NotImplementedError
@@ -671,6 +663,13 @@ class InterfaceBase:
     def _publish_telemetry(self, telem: tpb.TelemetryRecord) -> None:
         raise NotImplementedError
 
+    def publish_environment(self, environment: pb.EnvironmentRecord) -> None:
+        self._publish_environment(environment)
+
+    @abstractmethod
+    def _publish_environment(self, environment: pb.EnvironmentRecord) -> None:
+        raise NotImplementedError
+
     def publish_partial_history(
         self,
         run: "Run",
@@ -1000,16 +999,6 @@ class InterfaceBase:
     ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_get_system_metadata(self) -> MailboxHandle[pb.Result]:
-        get_system_metadata = pb.GetSystemMetadataRequest()
-        return self._deliver_get_system_metadata(get_system_metadata)
-
-    @abstractmethod
-    def _deliver_get_system_metadata(
-        self, get_system_metadata: pb.GetSystemMetadataRequest
-    ) -> MailboxHandle[pb.Result]:
-        raise NotImplementedError
-
     def deliver_exit(self, exit_code: Optional[int]) -> MailboxHandle[pb.Result]:
         exit_data = self._make_exit(exit_code)
         return self._deliver_exit(exit_data)

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")