wandb 0.21.2__py3-none-win32.whl → 0.21.4__py3-none-win32.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- wandb/__init__.py +1 -1
- wandb/__init__.pyi +1 -1
- wandb/_analytics.py +65 -0
- wandb/_iterutils.py +8 -0
- wandb/_pydantic/__init__.py +10 -11
- wandb/_pydantic/base.py +3 -53
- wandb/_pydantic/field_types.py +29 -0
- wandb/_pydantic/v1_compat.py +47 -30
- wandb/_strutils.py +40 -0
- wandb/apis/public/api.py +17 -4
- wandb/apis/public/artifacts.py +5 -4
- wandb/apis/public/automations.py +2 -1
- wandb/apis/public/registries/_freezable_list.py +6 -6
- wandb/apis/public/registries/_utils.py +2 -1
- wandb/apis/public/registries/registries_search.py +4 -0
- wandb/apis/public/registries/registry.py +7 -0
- wandb/automations/_filters/expressions.py +3 -2
- wandb/automations/_filters/operators.py +2 -1
- wandb/automations/_validators.py +20 -0
- wandb/automations/actions.py +4 -2
- wandb/automations/events.py +4 -5
- wandb/bin/gpu_stats.exe +0 -0
- wandb/bin/wandb-core +0 -0
- wandb/cli/beta.py +48 -130
- wandb/cli/beta_sync.py +226 -0
- wandb/cli/cli.py +1 -1
- wandb/integration/dspy/__init__.py +5 -0
- wandb/integration/dspy/dspy.py +422 -0
- wandb/integration/weave/weave.py +55 -0
- wandb/proto/v3/wandb_server_pb2.py +38 -57
- wandb/proto/v3/wandb_sync_pb2.py +87 -0
- wandb/proto/v3/wandb_telemetry_pb2.py +12 -12
- wandb/proto/v4/wandb_server_pb2.py +38 -41
- wandb/proto/v4/wandb_sync_pb2.py +38 -0
- wandb/proto/v4/wandb_telemetry_pb2.py +12 -12
- wandb/proto/v5/wandb_server_pb2.py +38 -41
- wandb/proto/v5/wandb_sync_pb2.py +39 -0
- wandb/proto/v5/wandb_telemetry_pb2.py +12 -12
- wandb/proto/v6/wandb_server_pb2.py +38 -41
- wandb/proto/v6/wandb_sync_pb2.py +49 -0
- wandb/proto/v6/wandb_telemetry_pb2.py +12 -12
- wandb/proto/wandb_generate_proto.py +1 -0
- wandb/proto/wandb_sync_pb2.py +12 -0
- wandb/sdk/artifacts/_validators.py +50 -49
- wandb/sdk/artifacts/artifact.py +11 -11
- wandb/sdk/artifacts/artifact_file_cache.py +1 -1
- wandb/sdk/artifacts/artifact_manifest_entry.py +6 -8
- wandb/sdk/artifacts/exceptions.py +2 -1
- wandb/sdk/artifacts/storage_handlers/gcs_handler.py +1 -1
- wandb/sdk/artifacts/storage_handlers/s3_handler.py +2 -1
- wandb/sdk/launch/inputs/internal.py +25 -24
- wandb/sdk/launch/inputs/schema.py +31 -1
- wandb/sdk/lib/asyncio_compat.py +88 -23
- wandb/sdk/lib/gql_request.py +18 -7
- wandb/sdk/lib/paths.py +23 -21
- wandb/sdk/lib/printer.py +9 -13
- wandb/sdk/lib/progress.py +8 -6
- wandb/sdk/lib/service/service_connection.py +42 -12
- wandb/sdk/mailbox/wait_with_progress.py +1 -1
- wandb/sdk/wandb_init.py +0 -8
- wandb/sdk/wandb_run.py +14 -2
- wandb/sdk/wandb_settings.py +55 -0
- wandb/sdk/wandb_setup.py +2 -2
- {wandb-0.21.2.dist-info → wandb-0.21.4.dist-info}/METADATA +2 -2
- {wandb-0.21.2.dist-info → wandb-0.21.4.dist-info}/RECORD +68 -57
- {wandb-0.21.2.dist-info → wandb-0.21.4.dist-info}/WHEEL +0 -0
- {wandb-0.21.2.dist-info → wandb-0.21.4.dist-info}/entry_points.txt +0 -0
- {wandb-0.21.2.dist-info → wandb-0.21.4.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,422 @@
|
|
1
|
+
"""DSPy ↔ Weights & Biases integration."""
|
2
|
+
|
3
|
+
from __future__ import annotations
|
4
|
+
|
5
|
+
import logging
|
6
|
+
import os
|
7
|
+
from collections.abc import Mapping, Sequence
|
8
|
+
from typing import Any, Literal
|
9
|
+
|
10
|
+
import wandb
|
11
|
+
import wandb.util
|
12
|
+
from wandb.sdk.wandb_run import Run
|
13
|
+
|
14
|
+
dspy = wandb.util.get_module(
|
15
|
+
name="dspy",
|
16
|
+
required=(
|
17
|
+
"To use the W&B DSPy integration you need to have the `dspy` "
|
18
|
+
"python package installed. Install it with `uv pip install dspy`."
|
19
|
+
),
|
20
|
+
lazy=False,
|
21
|
+
)
|
22
|
+
if dspy is not None:
|
23
|
+
assert dspy.__version__ >= "3.0.0", (
|
24
|
+
"DSPy 3.0.0 or higher is required. You have " + dspy.__version__
|
25
|
+
)
|
26
|
+
|
27
|
+
|
28
|
+
logger = logging.getLogger(__name__)
|
29
|
+
|
30
|
+
|
31
|
+
def _flatten_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
32
|
+
"""Flatten a list of nested row dicts into flat key/value dicts.
|
33
|
+
|
34
|
+
Args:
|
35
|
+
rows (list[dict[str, Any]]): List of nested dictionaries to flatten.
|
36
|
+
|
37
|
+
Returns:
|
38
|
+
list[dict[str, Any]]: List of flattened dictionaries.
|
39
|
+
|
40
|
+
"""
|
41
|
+
|
42
|
+
def _flatten(
|
43
|
+
d: dict[str, Any], parent_key: str = "", sep: str = "."
|
44
|
+
) -> dict[str, Any]:
|
45
|
+
items = []
|
46
|
+
for k, v in d.items():
|
47
|
+
new_key = f"{parent_key}{sep}{k}" if parent_key else k
|
48
|
+
if isinstance(v, dict):
|
49
|
+
items.extend(_flatten(v, new_key, sep=sep).items())
|
50
|
+
else:
|
51
|
+
items.append((new_key, v))
|
52
|
+
return dict(items)
|
53
|
+
|
54
|
+
return [_flatten(row) for row in rows]
|
55
|
+
|
56
|
+
|
57
|
+
class WandbDSPyCallback(dspy.utils.BaseCallback):
|
58
|
+
"""W&B callback for tracking DSPy evaluation and optimization.
|
59
|
+
|
60
|
+
This callback logs evaluation scores, per-step predictions (optional), and
|
61
|
+
a table capturing the DSPy program signature over time. It can also save
|
62
|
+
the best program as a W&B Artifact for reproducibility.
|
63
|
+
|
64
|
+
Examples:
|
65
|
+
Basic usage within DSPy settings:
|
66
|
+
|
67
|
+
```python
|
68
|
+
import dspy
|
69
|
+
import wandb
|
70
|
+
from wandb.integration.dspy import WandbDSPyCallback
|
71
|
+
|
72
|
+
with wandb.init(project="dspy-optimization") as run:
|
73
|
+
dspy.settings.callbacks.append(WandbDSPyCallback(run=run))
|
74
|
+
# Run your DSPy optimization/evaluation
|
75
|
+
```
|
76
|
+
"""
|
77
|
+
|
78
|
+
def __init__(self, log_results: bool = True, run: Run | None = None) -> None:
|
79
|
+
"""Initialize the callback.
|
80
|
+
|
81
|
+
Args:
|
82
|
+
log_results (bool): Whether to log per-evaluation prediction tables.
|
83
|
+
run (Run | None): Optional W&B run to use. Defaults to the
|
84
|
+
current global run if available.
|
85
|
+
|
86
|
+
Raises:
|
87
|
+
wandb.Error: If no active run is provided or found.
|
88
|
+
"""
|
89
|
+
# If no run is provided, use the current global run if available.
|
90
|
+
if run is None:
|
91
|
+
if wandb.run is None:
|
92
|
+
raise wandb.Error(
|
93
|
+
"You must call `wandb.init()` before instantiating WandbDSPyCallback()."
|
94
|
+
)
|
95
|
+
run = wandb.run
|
96
|
+
|
97
|
+
self.log_results = log_results
|
98
|
+
|
99
|
+
with wandb.wandb_lib.telemetry.context(run=run) as tel:
|
100
|
+
tel.feature.dspy_callback = True
|
101
|
+
|
102
|
+
self._run = run
|
103
|
+
self._did_log_config: bool = False
|
104
|
+
self._program_info: dict[str, Any] = {}
|
105
|
+
self._program_table: wandb.Table | None = None
|
106
|
+
self._row_idx: int = 0
|
107
|
+
|
108
|
+
def _flatten_dict(
|
109
|
+
self, nested: Any, parent_key: str = "", sep: str = "."
|
110
|
+
) -> dict[str, Any]:
|
111
|
+
"""Recursively flatten arbitrarily nested mappings and sequences.
|
112
|
+
|
113
|
+
Args:
|
114
|
+
nested (Any): Nested structure of mappings/lists to flatten.
|
115
|
+
parent_key (str): Prefix to prepend to keys in the flattened output.
|
116
|
+
sep (str): Key separator for nested fields.
|
117
|
+
|
118
|
+
Returns:
|
119
|
+
dict[str, Any]: Flattened dictionary representation.
|
120
|
+
"""
|
121
|
+
flat: dict[str, Any] = {}
|
122
|
+
|
123
|
+
def _walk(obj: Any, base: str) -> None:
|
124
|
+
if isinstance(obj, Mapping):
|
125
|
+
for k, v in obj.items():
|
126
|
+
new_key = f"{base}{sep}{k}" if base else str(k)
|
127
|
+
_walk(v, new_key)
|
128
|
+
elif isinstance(obj, Sequence) and not isinstance(
|
129
|
+
obj, (str, bytes, bytearray)
|
130
|
+
):
|
131
|
+
for idx, v in enumerate(obj):
|
132
|
+
new_key = f"{base}{sep}{idx}" if base else str(idx)
|
133
|
+
_walk(v, new_key)
|
134
|
+
else:
|
135
|
+
# Base can be empty only if the top-level is a scalar; guard against that.
|
136
|
+
key = base if base else ""
|
137
|
+
if key:
|
138
|
+
flat[key] = obj
|
139
|
+
|
140
|
+
_walk(nested, parent_key)
|
141
|
+
return flat
|
142
|
+
|
143
|
+
def _extract_fields(self, fields: list[dict[str, Any]]) -> dict[str, str]:
|
144
|
+
"""Convert signature fields to a flat mapping of strings.
|
145
|
+
|
146
|
+
Note:
|
147
|
+
The input is expected to be a dict-like mapping from field names to
|
148
|
+
field metadata. Values are stringified for logging.
|
149
|
+
|
150
|
+
Args:
|
151
|
+
fields (list[dict[str, Any]]): Mapping of field name to metadata.
|
152
|
+
|
153
|
+
Returns:
|
154
|
+
dict[str, str]: Mapping of field name to string value.
|
155
|
+
"""
|
156
|
+
return {k: str(v) for k, v in fields.items()}
|
157
|
+
|
158
|
+
def _extract_program_info(self, program_obj: Any) -> dict[str, Any]:
|
159
|
+
"""Extract signature-related info from a DSPy program.
|
160
|
+
|
161
|
+
Attempts to read the program signature, instructions, input and output
|
162
|
+
fields from a DSPy `Predict` parameter if available.
|
163
|
+
|
164
|
+
Args:
|
165
|
+
program_obj (Any): DSPy program/module instance.
|
166
|
+
|
167
|
+
Returns:
|
168
|
+
dict[str, Any]: Flattened dictionary of signature metadata.
|
169
|
+
"""
|
170
|
+
info_dict = {}
|
171
|
+
|
172
|
+
if program_obj is None:
|
173
|
+
return info_dict
|
174
|
+
|
175
|
+
try:
|
176
|
+
sig = next(
|
177
|
+
param.signature
|
178
|
+
for _, param in program_obj.named_parameters()
|
179
|
+
if isinstance(param, dspy.Predict)
|
180
|
+
)
|
181
|
+
|
182
|
+
if getattr(sig, "signature", None):
|
183
|
+
info_dict["signature"] = sig.signature
|
184
|
+
if getattr(sig, "instructions", None):
|
185
|
+
info_dict["instructions"] = sig.instructions
|
186
|
+
if getattr(sig, "input_fields", None):
|
187
|
+
input_fields = sig.input_fields
|
188
|
+
info_dict["input_fields"] = self._extract_fields(input_fields)
|
189
|
+
if getattr(sig, "output_fields", None):
|
190
|
+
output_fields = sig.output_fields
|
191
|
+
info_dict["output_fields"] = self._extract_fields(output_fields)
|
192
|
+
|
193
|
+
return self._flatten_dict(info_dict)
|
194
|
+
except Exception as e:
|
195
|
+
logger.warning(
|
196
|
+
"Failed to extract program info from Evaluate instance: %s", e
|
197
|
+
)
|
198
|
+
return info_dict
|
199
|
+
|
200
|
+
def on_evaluate_start(
|
201
|
+
self,
|
202
|
+
call_id: str,
|
203
|
+
instance: Any,
|
204
|
+
inputs: dict[str, Any],
|
205
|
+
) -> None:
|
206
|
+
"""Handle start of a DSPy evaluation call.
|
207
|
+
|
208
|
+
Logs non-private fields from the evaluator instance to W&B config and
|
209
|
+
captures program signature info for later logging.
|
210
|
+
|
211
|
+
Args:
|
212
|
+
call_id (str): Unique identifier for the evaluation call.
|
213
|
+
instance (Any): The evaluation instance (e.g., `dspy.Evaluate`).
|
214
|
+
inputs (dict[str, Any]): Inputs passed to the evaluation (may
|
215
|
+
include a `program` key with the DSPy program).
|
216
|
+
"""
|
217
|
+
if not self._did_log_config:
|
218
|
+
instance_vars = vars(instance) if hasattr(instance, "__dict__") else {}
|
219
|
+
serializable = {
|
220
|
+
k: v for k, v in instance_vars.items() if not k.startswith("_")
|
221
|
+
}
|
222
|
+
if "devset" in serializable:
|
223
|
+
# we don't want to log the devset in the config
|
224
|
+
del serializable["devset"]
|
225
|
+
|
226
|
+
self._run.config.update(serializable)
|
227
|
+
self._did_log_config = True
|
228
|
+
|
229
|
+
# 2) Build/append program signature tables from the 'program' inputs
|
230
|
+
if program_obj := inputs.get("program"):
|
231
|
+
self._program_info = self._extract_program_info(program_obj)
|
232
|
+
|
233
|
+
def on_evaluate_end(
|
234
|
+
self,
|
235
|
+
call_id: str,
|
236
|
+
outputs: Any | None,
|
237
|
+
exception: Exception | None = None,
|
238
|
+
) -> None:
|
239
|
+
"""Handle end of a DSPy evaluation call.
|
240
|
+
|
241
|
+
If available, logs a numeric `score` metric and (optionally) per-step
|
242
|
+
prediction tables. Always appends a row to the program-signature table.
|
243
|
+
|
244
|
+
Args:
|
245
|
+
call_id (str): Unique identifier for the evaluation call.
|
246
|
+
outputs (Any | None): Evaluation outputs; supports
|
247
|
+
`dspy.evaluate.evaluate.EvaluationResult`.
|
248
|
+
exception (Exception | None): Exception raised during evaluation, if any.
|
249
|
+
"""
|
250
|
+
# The `BaseCallback` does not define the interface for the `outputs` parameter,
|
251
|
+
# Currently, we know of `EvaluationResult` which is a subclass of `dspy.Prediction`.
|
252
|
+
# We currently support this type and will warn the user if a different type is passed.
|
253
|
+
score: float | None = None
|
254
|
+
if exception is None:
|
255
|
+
if isinstance(outputs, dspy.evaluate.evaluate.EvaluationResult):
|
256
|
+
# log the float score as a wandb metric
|
257
|
+
score = outputs.score
|
258
|
+
wandb.log({"score": float(score)}, step=self._row_idx)
|
259
|
+
|
260
|
+
# Log the predictions as a separate table for each eval end.
|
261
|
+
# We know that results if of type `list[tuple["dspy.Example", "dspy.Example", Any]]`
|
262
|
+
results = outputs.results
|
263
|
+
if self.log_results:
|
264
|
+
rows = self._parse_results(results)
|
265
|
+
if rows:
|
266
|
+
self._log_predictions_table(rows)
|
267
|
+
else:
|
268
|
+
wandb.termwarn(
|
269
|
+
f"on_evaluate_end received unexpected outputs type: {type(outputs)}. "
|
270
|
+
"Expected dspy.evaluate.evaluate.EvaluationResult; skipping logging score and `log_results`."
|
271
|
+
)
|
272
|
+
else:
|
273
|
+
wandb.termwarn(
|
274
|
+
f"on_evaluate_end received exception: {exception}. "
|
275
|
+
"Skipping logging score and `log_results`."
|
276
|
+
)
|
277
|
+
|
278
|
+
# Log the program signature iteratively
|
279
|
+
if self._program_table is None:
|
280
|
+
columns = ["step", *self._program_info.keys()]
|
281
|
+
if isinstance(score, float):
|
282
|
+
columns.append("score")
|
283
|
+
self._program_table = wandb.Table(columns=columns, log_mode="INCREMENTAL")
|
284
|
+
|
285
|
+
if self._program_table is not None:
|
286
|
+
values = list(self._program_info.values())
|
287
|
+
if isinstance(score, float):
|
288
|
+
values.append(score)
|
289
|
+
|
290
|
+
self._program_table.add_data(
|
291
|
+
self._row_idx,
|
292
|
+
*values,
|
293
|
+
)
|
294
|
+
self._run.log(
|
295
|
+
{"program_signature": self._program_table}, step=self._row_idx
|
296
|
+
)
|
297
|
+
|
298
|
+
self._row_idx += 1
|
299
|
+
|
300
|
+
def _parse_results(
|
301
|
+
self,
|
302
|
+
results: list[tuple[dspy.Example, dspy.Prediction | dspy.Completions, bool]],
|
303
|
+
) -> list[dict[str, Any]]:
|
304
|
+
"""Normalize evaluation results into serializable row dicts.
|
305
|
+
|
306
|
+
Args:
|
307
|
+
results (list[tuple]): Sequence of `(example, prediction, is_correct)`
|
308
|
+
tuples from DSPy evaluation.
|
309
|
+
|
310
|
+
Returns:
|
311
|
+
list[dict[str, Any]]: Rows with `example`, `prediction`, `is_correct`.
|
312
|
+
"""
|
313
|
+
_rows: list[dict[str, Any]] = []
|
314
|
+
for example, prediction, is_correct in results:
|
315
|
+
if isinstance(prediction, dspy.Prediction):
|
316
|
+
prediction_dict = prediction.toDict()
|
317
|
+
if isinstance(prediction, dspy.Completions):
|
318
|
+
prediction_dict = prediction.items()
|
319
|
+
|
320
|
+
row: dict[str, Any] = {
|
321
|
+
"example": example.toDict(),
|
322
|
+
"prediction": prediction_dict,
|
323
|
+
"is_correct": is_correct,
|
324
|
+
}
|
325
|
+
_rows.append(row)
|
326
|
+
|
327
|
+
return _rows
|
328
|
+
|
329
|
+
def _log_predictions_table(self, rows: list[dict[str, Any]]) -> None:
|
330
|
+
"""Log a W&B Table of predictions for the current evaluation step.
|
331
|
+
|
332
|
+
Args:
|
333
|
+
rows (list[dict[str, Any]]): Prediction rows to log.
|
334
|
+
"""
|
335
|
+
rows = _flatten_rows(rows)
|
336
|
+
columns = list(rows[0].keys())
|
337
|
+
|
338
|
+
data: list[list[Any]] = [list(row.values()) for row in rows]
|
339
|
+
|
340
|
+
preds_table = wandb.Table(columns=columns, data=data, log_mode="IMMUTABLE")
|
341
|
+
self._run.log({f"predictions_{self._row_idx}": preds_table}, step=self._row_idx)
|
342
|
+
|
343
|
+
def log_best_model(
|
344
|
+
self,
|
345
|
+
model: dspy.Module,
|
346
|
+
*,
|
347
|
+
save_program: bool = True,
|
348
|
+
save_dir: str | None = None,
|
349
|
+
filetype: Literal["json", "pkl"] = "json",
|
350
|
+
aliases: Sequence[str] = ("best", "latest"),
|
351
|
+
artifact_name: str = "dspy-program",
|
352
|
+
) -> None:
|
353
|
+
"""Save and log the best DSPy program as a W&B Artifact.
|
354
|
+
|
355
|
+
You can choose to save the full program (architecture + state) or only
|
356
|
+
the state to a single file (JSON or pickle).
|
357
|
+
|
358
|
+
Args:
|
359
|
+
model (dspy.Module): DSPy module to save.
|
360
|
+
save_program (bool): Save full program directory if True; otherwise
|
361
|
+
save only the state file. Defaults to `True`.
|
362
|
+
save_dir (str): Directory to store program files before logging. Defaults to a
|
363
|
+
subdirectory `dspy_program` within the active run's files directory
|
364
|
+
(i.e., `wandb.run.dir`).
|
365
|
+
filetype (Literal["json", "pkl"]): State file format when
|
366
|
+
`save_program` is False. Defaults to `json`.
|
367
|
+
aliases (Sequence[str]): Aliases for the logged Artifact version. Defaults to `("best", "latest")`.
|
368
|
+
artifact_name (str): Base name for the Artifact. Defaults to `dspy-program`.
|
369
|
+
|
370
|
+
Examples:
|
371
|
+
Save the complete program and add aliases:
|
372
|
+
|
373
|
+
```python
|
374
|
+
callback.log_best_model(
|
375
|
+
optimized_program, save_program=True, aliases=("best", "production")
|
376
|
+
)
|
377
|
+
```
|
378
|
+
|
379
|
+
Save only the state as JSON:
|
380
|
+
|
381
|
+
```python
|
382
|
+
callback.log_best_model(
|
383
|
+
optimized_program, save_program=False, filetype="json"
|
384
|
+
)
|
385
|
+
```
|
386
|
+
"""
|
387
|
+
# Derive metadata to help discoverability in the UI
|
388
|
+
info_dict = self._extract_program_info(model)
|
389
|
+
metadata = {
|
390
|
+
"dspy_version": getattr(dspy, "__version__", "unknown"),
|
391
|
+
"module_class": model.__class__.__name__,
|
392
|
+
**info_dict,
|
393
|
+
}
|
394
|
+
artifact = wandb.Artifact(
|
395
|
+
name=f"{artifact_name}-{self._run.id}",
|
396
|
+
type="model",
|
397
|
+
metadata=metadata,
|
398
|
+
)
|
399
|
+
|
400
|
+
# Resolve and normalize the save directory in a cross-platform way
|
401
|
+
if save_dir is None:
|
402
|
+
save_dir = os.path.join(self._run.dir, "dspy_program")
|
403
|
+
save_dir = os.path.normpath(save_dir)
|
404
|
+
|
405
|
+
try:
|
406
|
+
os.makedirs(save_dir, exist_ok=True)
|
407
|
+
except Exception as exc:
|
408
|
+
wandb.termwarn(
|
409
|
+
f"Could not create or access directory '{save_dir}': {exc}. Skipping artifact logging."
|
410
|
+
)
|
411
|
+
return
|
412
|
+
# Save per requested mode
|
413
|
+
if save_program:
|
414
|
+
model.save(save_dir, save_program=True)
|
415
|
+
artifact.add_dir(save_dir)
|
416
|
+
else:
|
417
|
+
filename = f"program.{filetype}"
|
418
|
+
file_path = os.path.join(save_dir, filename)
|
419
|
+
model.save(file_path, save_program=False)
|
420
|
+
artifact.add_file(file_path)
|
421
|
+
|
422
|
+
self._run.log_artifact(artifact, aliases=list(aliases))
|
wandb/integration/weave/weave.py
CHANGED
@@ -10,6 +10,7 @@ The integration can be disabled by setting the WANDB_DISABLE_WEAVE environment v
|
|
10
10
|
|
11
11
|
from __future__ import annotations
|
12
12
|
|
13
|
+
import importlib.util
|
13
14
|
import os
|
14
15
|
import sys
|
15
16
|
import threading
|
@@ -21,6 +22,31 @@ _weave_init_lock = threading.Lock()
|
|
21
22
|
_DISABLE_WEAVE = "WANDB_DISABLE_WEAVE"
|
22
23
|
_WEAVE_PACKAGE_NAME = "weave"
|
23
24
|
|
25
|
+
# This list is adapted from https://github.com/wandb/weave/blob/master/weave/integrations/__init__.py
|
26
|
+
_AVAILABLE_WEAVE_INTEGRATIONS = [
|
27
|
+
"anthropic",
|
28
|
+
"autogen",
|
29
|
+
"cohere",
|
30
|
+
"crewai",
|
31
|
+
"dspy",
|
32
|
+
"google.genai",
|
33
|
+
"groq",
|
34
|
+
"huggingface_hub.inference",
|
35
|
+
"instructor",
|
36
|
+
"langchain",
|
37
|
+
"litellm",
|
38
|
+
"llama_index",
|
39
|
+
"mcp",
|
40
|
+
"mistral",
|
41
|
+
"notdiamond",
|
42
|
+
"openai",
|
43
|
+
"agents",
|
44
|
+
"smolagents",
|
45
|
+
"verdict",
|
46
|
+
"verifiers",
|
47
|
+
"vertexai",
|
48
|
+
]
|
49
|
+
|
24
50
|
|
25
51
|
def setup(entity: str | None, project: str | None) -> None:
|
26
52
|
"""Set up automatic Weave initialization for the current W&B run.
|
@@ -43,6 +69,7 @@ def setup(entity: str | None, project: str | None) -> None:
|
|
43
69
|
# If weave is not yet imported, we can't init it from here. Instead, we'll
|
44
70
|
# rely on the weave library itself to detect a run and init itself.
|
45
71
|
if _WEAVE_PACKAGE_NAME not in sys.modules:
|
72
|
+
_maybe_suggest_weave_installation()
|
46
73
|
return
|
47
74
|
|
48
75
|
# If weave has already been imported, initialize immediately
|
@@ -61,3 +88,31 @@ def setup(entity: str | None, project: str | None) -> None:
|
|
61
88
|
weave.init(project_path)
|
62
89
|
except Exception as e:
|
63
90
|
wandb.termwarn(f"Failed to automatically initialize Weave: {e}")
|
91
|
+
|
92
|
+
|
93
|
+
def _maybe_suggest_weave_installation() -> None:
|
94
|
+
"""Suggest Weave installation or import if any target library is imported."""
|
95
|
+
imported_libs = [lib for lib in _AVAILABLE_WEAVE_INTEGRATIONS if lib in sys.modules]
|
96
|
+
if not imported_libs:
|
97
|
+
return
|
98
|
+
|
99
|
+
weave_spec = importlib.util.find_spec(_WEAVE_PACKAGE_NAME)
|
100
|
+
if weave_spec is None:
|
101
|
+
# Weave is not installed
|
102
|
+
msg = (
|
103
|
+
"Use W&B Weave for improved LLM call tracing. Install Weave with "
|
104
|
+
"`pip install weave` then add `import weave` to the top of your script."
|
105
|
+
)
|
106
|
+
else:
|
107
|
+
# Weave is installed but not imported
|
108
|
+
msg = (
|
109
|
+
"Use W&B Weave for improved LLM call tracing. Weave is installed "
|
110
|
+
"but not imported. Add `import weave` to the top of your script."
|
111
|
+
)
|
112
|
+
|
113
|
+
wandb.termlog(f"Detected [{', '.join(imported_libs)}] in use.", repeat=False)
|
114
|
+
wandb.termlog(msg, repeat=False)
|
115
|
+
wandb.termlog(
|
116
|
+
"For more information, check out the docs at: https://weave-docs.wandb.ai/",
|
117
|
+
repeat=False,
|
118
|
+
)
|
@@ -15,9 +15,10 @@ _sym_db = _symbol_database.Default()
|
|
15
15
|
from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2
|
16
16
|
from wandb.proto import wandb_internal_pb2 as wandb_dot_proto_dot_wandb__internal__pb2
|
17
17
|
from wandb.proto import wandb_settings_pb2 as wandb_dot_proto_dot_wandb__settings__pb2
|
18
|
+
from wandb.proto import wandb_sync_pb2 as wandb_dot_proto_dot_wandb__sync__pb2
|
18
19
|
|
19
20
|
|
20
|
-
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1ewandb/proto/wandb_server.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\x1a wandb/proto/wandb_internal.proto\x1a wandb/proto/wandb_settings.proto\"k\n\x19ServerAuthenticateRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12\x10\n\x08\x62\x61se_url\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"w\n\x1aServerAuthenticateResponse\x12\x16\n\x0e\x64\x65\x66\x61ult_entity\x18\x01 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"D\n\x15ServerShutdownRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x18\n\x16ServerShutdownResponse\"B\n\x13ServerStatusRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x16\n\x14ServerStatusResponse\"r\n\x17ServerInformInitRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1a\n\x18ServerInformInitResponse\"
|
21
|
+
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1ewandb/proto/wandb_server.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\x1a wandb/proto/wandb_internal.proto\x1a wandb/proto/wandb_settings.proto\x1a\x1cwandb/proto/wandb_sync.proto\"k\n\x19ServerAuthenticateRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12\x10\n\x08\x62\x61se_url\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"w\n\x1aServerAuthenticateResponse\x12\x16\n\x0e\x64\x65\x66\x61ult_entity\x18\x01 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"D\n\x15ServerShutdownRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x18\n\x16ServerShutdownResponse\"B\n\x13ServerStatusRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x16\n\x14ServerStatusResponse\"r\n\x17ServerInformInitRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1a\n\x18ServerInformInitResponse\"H\n\x19ServerInformFinishRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformFinishResponse\"H\n\x19ServerInformAttachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"u\n\x1aServerInformAttachResponse\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"H\n\x19ServerInformDetachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformDetachResponse\"]\n\x1bServerInformTeardownRequest\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1e\n\x1cServerInformTeardownResponse\"\xee\x05\n\rServerRequest\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x30\n\x0erecord_publish\x18\x01 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12\x34\n\x12record_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12>\n\x0binform_init\x18\x03 \x01(\x0b\x32\'.wandb_internal.ServerInformInitRequestH\x00\x12\x42\n\rinform_finish\x18\x04 \x01(\x0b\x32).wandb_internal.ServerInformFinishRequestH\x00\x12\x42\n\rinform_attach\x18\x05 \x01(\x0b\x32).wandb_internal.ServerInformAttachRequestH\x00\x12\x42\n\rinform_detach\x18\x06 \x01(\x0b\x32).wandb_internal.ServerInformDetachRequestH\x00\x12\x46\n\x0finform_teardown\x18\x07 \x01(\x0b\x32+.wandb_internal.ServerInformTeardownRequestH\x00\x12\x41\n\x0c\x61uthenticate\x18\t \x01(\x0b\x32).wandb_internal.ServerAuthenticateRequestH\x00\x12:\n\tinit_sync\x18\x0b \x01(\x0b\x32%.wandb_internal.ServerInitSyncRequestH\x00\x12\x31\n\x04sync\x18\x0c \x01(\x0b\x32!.wandb_internal.ServerSyncRequestH\x00\x12>\n\x0bsync_status\x18\r \x01(\x0b\x32\'.wandb_internal.ServerSyncStatusRequestH\x00\x42\x15\n\x13server_request_typeJ\x04\x08\x08\x10\t\"\x98\x06\n\x0eServerResponse\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x34\n\x12result_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.ResultH\x00\x12H\n\x14inform_init_response\x18\x03 \x01(\x0b\x32(.wandb_internal.ServerInformInitResponseH\x00\x12L\n\x16inform_finish_response\x18\x04 \x01(\x0b\x32*.wandb_internal.ServerInformFinishResponseH\x00\x12L\n\x16inform_attach_response\x18\x05 \x01(\x0b\x32*.wandb_internal.ServerInformAttachResponseH\x00\x12L\n\x16inform_detach_response\x18\x06 \x01(\x0b\x32*.wandb_internal.ServerInformDetachResponseH\x00\x12P\n\x18inform_teardown_response\x18\x07 \x01(\x0b\x32,.wandb_internal.ServerInformTeardownResponseH\x00\x12K\n\x15\x61uthenticate_response\x18\t \x01(\x0b\x32*.wandb_internal.ServerAuthenticateResponseH\x00\x12\x44\n\x12init_sync_response\x18\x0b \x01(\x0b\x32&.wandb_internal.ServerInitSyncResponseH\x00\x12;\n\rsync_response\x18\x0c \x01(\x0b\x32\".wandb_internal.ServerSyncResponseH\x00\x12H\n\x14sync_status_response\x18\r \x01(\x0b\x32(.wandb_internal.ServerSyncStatusResponseH\x00\x42\x16\n\x14server_response_typeJ\x04\x08\x08\x10\tB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3')
|
21
22
|
|
22
23
|
|
23
24
|
|
@@ -29,8 +30,6 @@ _SERVERSTATUSREQUEST = DESCRIPTOR.message_types_by_name['ServerStatusRequest']
|
|
29
30
|
_SERVERSTATUSRESPONSE = DESCRIPTOR.message_types_by_name['ServerStatusResponse']
|
30
31
|
_SERVERINFORMINITREQUEST = DESCRIPTOR.message_types_by_name['ServerInformInitRequest']
|
31
32
|
_SERVERINFORMINITRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformInitResponse']
|
32
|
-
_SERVERINFORMSTARTREQUEST = DESCRIPTOR.message_types_by_name['ServerInformStartRequest']
|
33
|
-
_SERVERINFORMSTARTRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformStartResponse']
|
34
33
|
_SERVERINFORMFINISHREQUEST = DESCRIPTOR.message_types_by_name['ServerInformFinishRequest']
|
35
34
|
_SERVERINFORMFINISHRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformFinishResponse']
|
36
35
|
_SERVERINFORMATTACHREQUEST = DESCRIPTOR.message_types_by_name['ServerInformAttachRequest']
|
@@ -97,20 +96,6 @@ ServerInformInitResponse = _reflection.GeneratedProtocolMessageType('ServerInfor
|
|
97
96
|
})
|
98
97
|
_sym_db.RegisterMessage(ServerInformInitResponse)
|
99
98
|
|
100
|
-
ServerInformStartRequest = _reflection.GeneratedProtocolMessageType('ServerInformStartRequest', (_message.Message,), {
|
101
|
-
'DESCRIPTOR' : _SERVERINFORMSTARTREQUEST,
|
102
|
-
'__module__' : 'wandb.proto.wandb_server_pb2'
|
103
|
-
# @@protoc_insertion_point(class_scope:wandb_internal.ServerInformStartRequest)
|
104
|
-
})
|
105
|
-
_sym_db.RegisterMessage(ServerInformStartRequest)
|
106
|
-
|
107
|
-
ServerInformStartResponse = _reflection.GeneratedProtocolMessageType('ServerInformStartResponse', (_message.Message,), {
|
108
|
-
'DESCRIPTOR' : _SERVERINFORMSTARTRESPONSE,
|
109
|
-
'__module__' : 'wandb.proto.wandb_server_pb2'
|
110
|
-
# @@protoc_insertion_point(class_scope:wandb_internal.ServerInformStartResponse)
|
111
|
-
})
|
112
|
-
_sym_db.RegisterMessage(ServerInformStartResponse)
|
113
|
-
|
114
99
|
ServerInformFinishRequest = _reflection.GeneratedProtocolMessageType('ServerInformFinishRequest', (_message.Message,), {
|
115
100
|
'DESCRIPTOR' : _SERVERINFORMFINISHREQUEST,
|
116
101
|
'__module__' : 'wandb.proto.wandb_server_pb2'
|
@@ -185,44 +170,40 @@ if _descriptor._USE_C_DESCRIPTORS == False:
|
|
185
170
|
|
186
171
|
DESCRIPTOR._options = None
|
187
172
|
DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto'
|
188
|
-
_SERVERAUTHENTICATEREQUEST._serialized_start=
|
189
|
-
_SERVERAUTHENTICATEREQUEST._serialized_end=
|
190
|
-
_SERVERAUTHENTICATERESPONSE._serialized_start=
|
191
|
-
_SERVERAUTHENTICATERESPONSE._serialized_end=
|
192
|
-
_SERVERSHUTDOWNREQUEST._serialized_start=
|
193
|
-
_SERVERSHUTDOWNREQUEST._serialized_end=
|
194
|
-
_SERVERSHUTDOWNRESPONSE._serialized_start=
|
195
|
-
_SERVERSHUTDOWNRESPONSE._serialized_end=
|
196
|
-
_SERVERSTATUSREQUEST._serialized_start=
|
197
|
-
_SERVERSTATUSREQUEST._serialized_end=
|
198
|
-
_SERVERSTATUSRESPONSE._serialized_start=
|
199
|
-
_SERVERSTATUSRESPONSE._serialized_end=
|
200
|
-
_SERVERINFORMINITREQUEST._serialized_start=
|
201
|
-
_SERVERINFORMINITREQUEST._serialized_end=
|
202
|
-
_SERVERINFORMINITRESPONSE._serialized_start=
|
203
|
-
_SERVERINFORMINITRESPONSE._serialized_end=
|
204
|
-
|
205
|
-
|
206
|
-
|
207
|
-
|
208
|
-
|
209
|
-
|
210
|
-
|
211
|
-
|
212
|
-
|
213
|
-
|
214
|
-
|
215
|
-
|
216
|
-
|
217
|
-
|
218
|
-
|
219
|
-
|
220
|
-
|
221
|
-
|
222
|
-
|
223
|
-
|
224
|
-
_SERVERREQUEST._serialized_start=1385
|
225
|
-
_SERVERREQUEST._serialized_end=2020
|
226
|
-
_SERVERRESPONSE._serialized_start=2023
|
227
|
-
_SERVERRESPONSE._serialized_end=2680
|
173
|
+
_SERVERAUTHENTICATEREQUEST._serialized_start=178
|
174
|
+
_SERVERAUTHENTICATEREQUEST._serialized_end=285
|
175
|
+
_SERVERAUTHENTICATERESPONSE._serialized_start=287
|
176
|
+
_SERVERAUTHENTICATERESPONSE._serialized_end=406
|
177
|
+
_SERVERSHUTDOWNREQUEST._serialized_start=408
|
178
|
+
_SERVERSHUTDOWNREQUEST._serialized_end=476
|
179
|
+
_SERVERSHUTDOWNRESPONSE._serialized_start=478
|
180
|
+
_SERVERSHUTDOWNRESPONSE._serialized_end=502
|
181
|
+
_SERVERSTATUSREQUEST._serialized_start=504
|
182
|
+
_SERVERSTATUSREQUEST._serialized_end=570
|
183
|
+
_SERVERSTATUSRESPONSE._serialized_start=572
|
184
|
+
_SERVERSTATUSRESPONSE._serialized_end=594
|
185
|
+
_SERVERINFORMINITREQUEST._serialized_start=596
|
186
|
+
_SERVERINFORMINITREQUEST._serialized_end=710
|
187
|
+
_SERVERINFORMINITRESPONSE._serialized_start=712
|
188
|
+
_SERVERINFORMINITRESPONSE._serialized_end=738
|
189
|
+
_SERVERINFORMFINISHREQUEST._serialized_start=740
|
190
|
+
_SERVERINFORMFINISHREQUEST._serialized_end=812
|
191
|
+
_SERVERINFORMFINISHRESPONSE._serialized_start=814
|
192
|
+
_SERVERINFORMFINISHRESPONSE._serialized_end=842
|
193
|
+
_SERVERINFORMATTACHREQUEST._serialized_start=844
|
194
|
+
_SERVERINFORMATTACHREQUEST._serialized_end=916
|
195
|
+
_SERVERINFORMATTACHRESPONSE._serialized_start=918
|
196
|
+
_SERVERINFORMATTACHRESPONSE._serialized_end=1035
|
197
|
+
_SERVERINFORMDETACHREQUEST._serialized_start=1037
|
198
|
+
_SERVERINFORMDETACHREQUEST._serialized_end=1109
|
199
|
+
_SERVERINFORMDETACHRESPONSE._serialized_start=1111
|
200
|
+
_SERVERINFORMDETACHRESPONSE._serialized_end=1139
|
201
|
+
_SERVERINFORMTEARDOWNREQUEST._serialized_start=1141
|
202
|
+
_SERVERINFORMTEARDOWNREQUEST._serialized_end=1234
|
203
|
+
_SERVERINFORMTEARDOWNRESPONSE._serialized_start=1236
|
204
|
+
_SERVERINFORMTEARDOWNRESPONSE._serialized_end=1266
|
205
|
+
_SERVERREQUEST._serialized_start=1269
|
206
|
+
_SERVERREQUEST._serialized_end=2019
|
207
|
+
_SERVERRESPONSE._serialized_start=2022
|
208
|
+
_SERVERRESPONSE._serialized_end=2814
|
228
209
|
# @@protoc_insertion_point(module_scope)
|