wandb 0.20.2rc20250616__py3-none-win32.whl → 0.21.1__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.
Files changed (140) hide show
  1. wandb/__init__.py +16 -14
  2. wandb/__init__.pyi +450 -472
  3. wandb/agents/pyagent.py +41 -12
  4. wandb/analytics/sentry.py +7 -2
  5. wandb/apis/importers/mlflow.py +1 -1
  6. wandb/apis/internal.py +3 -0
  7. wandb/apis/paginator.py +17 -4
  8. wandb/apis/public/__init__.py +1 -1
  9. wandb/apis/public/api.py +606 -359
  10. wandb/apis/public/artifacts.py +214 -16
  11. wandb/apis/public/automations.py +19 -3
  12. wandb/apis/public/files.py +177 -38
  13. wandb/apis/public/history.py +67 -15
  14. wandb/apis/public/integrations.py +25 -2
  15. wandb/apis/public/jobs.py +90 -2
  16. wandb/apis/public/projects.py +161 -69
  17. wandb/apis/public/query_generator.py +11 -1
  18. wandb/apis/public/registries/registries_search.py +7 -15
  19. wandb/apis/public/reports.py +147 -13
  20. wandb/apis/public/runs.py +315 -128
  21. wandb/apis/public/sweeps.py +222 -22
  22. wandb/apis/public/teams.py +41 -4
  23. wandb/apis/public/users.py +45 -4
  24. wandb/automations/__init__.py +10 -10
  25. wandb/automations/_filters/run_metrics.py +0 -2
  26. wandb/automations/_utils.py +0 -2
  27. wandb/automations/actions.py +0 -2
  28. wandb/automations/automations.py +0 -2
  29. wandb/automations/events.py +0 -2
  30. wandb/beta/workflows.py +66 -30
  31. wandb/bin/gpu_stats.exe +0 -0
  32. wandb/bin/wandb-core +0 -0
  33. wandb/cli/cli.py +80 -1
  34. wandb/env.py +8 -0
  35. wandb/errors/errors.py +4 -1
  36. wandb/integration/catboost/catboost.py +6 -2
  37. wandb/integration/kfp/kfp_patch.py +3 -1
  38. wandb/integration/lightning/fabric/logger.py +3 -4
  39. wandb/integration/metaflow/__init__.py +6 -0
  40. wandb/integration/metaflow/data_pandas.py +74 -0
  41. wandb/integration/metaflow/errors.py +13 -0
  42. wandb/integration/metaflow/metaflow.py +205 -190
  43. wandb/integration/openai/fine_tuning.py +1 -2
  44. wandb/integration/sb3/sb3.py +3 -3
  45. wandb/integration/ultralytics/callback.py +6 -2
  46. wandb/jupyter.py +5 -5
  47. wandb/plot/__init__.py +2 -0
  48. wandb/plot/bar.py +30 -29
  49. wandb/plot/confusion_matrix.py +75 -71
  50. wandb/plot/custom_chart.py +30 -7
  51. wandb/plot/histogram.py +26 -25
  52. wandb/plot/line.py +33 -32
  53. wandb/plot/line_series.py +100 -103
  54. wandb/plot/pr_curve.py +33 -32
  55. wandb/plot/roc_curve.py +38 -38
  56. wandb/plot/scatter.py +27 -27
  57. wandb/proto/v3/wandb_internal_pb2.py +366 -385
  58. wandb/proto/v3/wandb_settings_pb2.py +2 -2
  59. wandb/proto/v3/wandb_telemetry_pb2.py +4 -4
  60. wandb/proto/v4/wandb_internal_pb2.py +352 -356
  61. wandb/proto/v4/wandb_settings_pb2.py +2 -2
  62. wandb/proto/v4/wandb_telemetry_pb2.py +4 -4
  63. wandb/proto/v5/wandb_internal_pb2.py +352 -356
  64. wandb/proto/v5/wandb_settings_pb2.py +2 -2
  65. wandb/proto/v5/wandb_telemetry_pb2.py +4 -4
  66. wandb/proto/v6/wandb_internal_pb2.py +352 -356
  67. wandb/proto/v6/wandb_settings_pb2.py +2 -2
  68. wandb/proto/v6/wandb_telemetry_pb2.py +4 -4
  69. wandb/proto/wandb_deprecated.py +6 -0
  70. wandb/sdk/artifacts/_generated/__init__.py +12 -1
  71. wandb/sdk/artifacts/_generated/input_types.py +20 -2
  72. wandb/sdk/artifacts/_generated/link_artifact.py +21 -0
  73. wandb/sdk/artifacts/_generated/operations.py +9 -0
  74. wandb/sdk/artifacts/_internal_artifact.py +19 -8
  75. wandb/sdk/artifacts/_validators.py +48 -2
  76. wandb/sdk/artifacts/artifact.py +269 -96
  77. wandb/sdk/data_types/audio.py +38 -10
  78. wandb/sdk/data_types/base_types/media.py +15 -63
  79. wandb/sdk/data_types/base_types/wb_value.py +6 -6
  80. wandb/sdk/data_types/graph.py +48 -14
  81. wandb/sdk/data_types/helper_types/bounding_boxes_2d.py +1 -3
  82. wandb/sdk/data_types/helper_types/image_mask.py +1 -3
  83. wandb/sdk/data_types/histogram.py +34 -21
  84. wandb/sdk/data_types/html.py +35 -12
  85. wandb/sdk/data_types/image.py +104 -68
  86. wandb/sdk/data_types/molecule.py +32 -19
  87. wandb/sdk/data_types/object_3d.py +36 -17
  88. wandb/sdk/data_types/plotly.py +18 -5
  89. wandb/sdk/data_types/saved_model.py +7 -9
  90. wandb/sdk/data_types/table.py +99 -70
  91. wandb/sdk/data_types/trace_tree.py +12 -12
  92. wandb/sdk/data_types/video.py +53 -26
  93. wandb/sdk/integration_utils/auto_logging.py +2 -2
  94. wandb/sdk/interface/interface.py +8 -19
  95. wandb/sdk/interface/interface_shared.py +7 -16
  96. wandb/sdk/internal/datastore.py +18 -18
  97. wandb/sdk/internal/handler.py +3 -5
  98. wandb/sdk/internal/internal_api.py +60 -0
  99. wandb/sdk/internal/job_builder.py +6 -0
  100. wandb/sdk/internal/sender.py +23 -3
  101. wandb/sdk/internal/sender_config.py +9 -0
  102. wandb/sdk/launch/_project_spec.py +3 -3
  103. wandb/sdk/launch/agent/agent.py +11 -4
  104. wandb/sdk/launch/agent/job_status_tracker.py +3 -1
  105. wandb/sdk/launch/agent/run_queue_item_file_saver.py +2 -2
  106. wandb/sdk/launch/create_job.py +3 -1
  107. wandb/sdk/launch/inputs/internal.py +3 -4
  108. wandb/sdk/launch/inputs/schema.py +1 -0
  109. wandb/sdk/launch/runner/kubernetes_monitor.py +1 -0
  110. wandb/sdk/launch/runner/kubernetes_runner.py +328 -1
  111. wandb/sdk/launch/sweeps/scheduler.py +2 -3
  112. wandb/sdk/launch/utils.py +3 -3
  113. wandb/sdk/lib/asyncio_compat.py +3 -0
  114. wandb/sdk/lib/console_capture.py +66 -19
  115. wandb/sdk/lib/deprecate.py +1 -7
  116. wandb/sdk/lib/disabled.py +1 -1
  117. wandb/sdk/lib/hashutil.py +14 -1
  118. wandb/sdk/lib/module.py +7 -13
  119. wandb/sdk/lib/progress.py +0 -19
  120. wandb/sdk/lib/sock_client.py +0 -4
  121. wandb/sdk/wandb_init.py +67 -93
  122. wandb/sdk/wandb_login.py +18 -14
  123. wandb/sdk/wandb_metric.py +2 -0
  124. wandb/sdk/wandb_require.py +0 -1
  125. wandb/sdk/wandb_run.py +429 -527
  126. wandb/sdk/wandb_settings.py +364 -74
  127. wandb/sdk/wandb_setup.py +28 -28
  128. wandb/sdk/wandb_sweep.py +14 -13
  129. wandb/sdk/wandb_watch.py +4 -6
  130. wandb/sync/sync.py +10 -0
  131. wandb/util.py +57 -0
  132. wandb/wandb_run.py +1 -2
  133. {wandb-0.20.2rc20250616.dist-info → wandb-0.21.1.dist-info}/METADATA +1 -1
  134. {wandb-0.20.2rc20250616.dist-info → wandb-0.21.1.dist-info}/RECORD +137 -137
  135. wandb/sdk/wandb_metadata.py +0 -623
  136. wandb/vendor/pynvml/__init__.py +0 -0
  137. wandb/vendor/pynvml/pynvml.py +0 -4779
  138. {wandb-0.20.2rc20250616.dist-info → wandb-0.21.1.dist-info}/WHEEL +0 -0
  139. {wandb-0.20.2rc20250616.dist-info → wandb-0.21.1.dist-info}/entry_points.txt +0 -0
  140. {wandb-0.20.2rc20250616.dist-info → wandb-0.21.1.dist-info}/licenses/LICENSE +0 -0
wandb/beta/workflows.py CHANGED
@@ -1,21 +1,31 @@
1
+ from __future__ import annotations
2
+
1
3
  import json
2
4
  import os
3
- from typing import Any, Dict, List, Optional, Union
5
+ import warnings
6
+ from typing import Any
7
+
8
+ from typing_extensions import deprecated
4
9
 
5
10
  import wandb
6
- import wandb.data_types as data_types
7
- from wandb.data_types import _SavedModel
11
+ from wandb.data_types import WBValue, _SavedModel
12
+ from wandb.proto.wandb_deprecated import Deprecated
8
13
  from wandb.sdk import wandb_setup
9
14
  from wandb.sdk.artifacts.artifact import Artifact
10
15
  from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
16
+ from wandb.sdk.lib.deprecate import deprecate as wandb_deprecate
17
+
18
+ warnings.warn(
19
+ message=f"The {__name__!r} module is deprecated and will be removed in a future version. Please use the equivalent 'wandb.Run' methods instead.",
20
+ category=DeprecationWarning,
21
+ stacklevel=2,
22
+ )
11
23
 
12
24
 
13
25
  def _add_any(
14
26
  artifact: Artifact,
15
- path_or_obj: Union[
16
- str, ArtifactManifestEntry, data_types.WBValue
17
- ], # todo: add dataframe
18
- name: Optional[str],
27
+ path_or_obj: str | ArtifactManifestEntry | WBValue, # todo: add dataframe
28
+ name: str | None,
19
29
  ) -> Any:
20
30
  """Add an object to an artifact.
21
31
 
@@ -24,11 +34,10 @@ def _add_any(
24
34
  be moved to the Artifact class in the future.
25
35
 
26
36
  Args:
27
- artifact: `Artifact` - artifact created with `wandb.Artifact(...)`
28
- path_or_obj: `Union[str, ArtifactManifestEntry, data_types.WBValue]` - either a
29
- str or valid object which indicates what to add to an artifact.
30
-
31
- name: `str` - the name of the object which is added to an artifact.
37
+ artifact: artifact created with `wandb.Artifact(...)`
38
+ path_or_obj: either a str or valid object which indicates what to add
39
+ to an artifact.
40
+ name: the name of the object which is added to an artifact.
32
41
 
33
42
  Returns:
34
43
  Type[Any] - Union[None, ArtifactManifestEntry, etc]
@@ -36,7 +45,7 @@ def _add_any(
36
45
  """
37
46
  if isinstance(path_or_obj, ArtifactManifestEntry):
38
47
  return artifact.add_reference(path_or_obj, name)
39
- elif isinstance(path_or_obj, data_types.WBValue):
48
+ elif isinstance(path_or_obj, WBValue):
40
49
  return artifact.add(path_or_obj, name)
41
50
  elif isinstance(path_or_obj, str):
42
51
  if os.path.isdir(path_or_obj):
@@ -56,12 +65,12 @@ def _add_any(
56
65
  def _log_artifact_version(
57
66
  name: str,
58
67
  type: str,
59
- entries: Dict[str, Union[str, ArtifactManifestEntry, data_types.WBValue]],
60
- aliases: Optional[Union[str, List[str]]] = None,
61
- description: Optional[str] = None,
62
- metadata: Optional[dict] = None,
63
- project: Optional[str] = None,
64
- scope_project: Optional[bool] = None,
68
+ entries: dict[str, str | ArtifactManifestEntry | WBValue],
69
+ aliases: str | list[str] | None = None,
70
+ description: str | None = None,
71
+ metadata: dict | None = None,
72
+ project: str | None = None,
73
+ scope_project: bool | None = None,
65
74
  job_type: str = "auto",
66
75
  ) -> Artifact:
67
76
  """Create an artifact, populate it, and log it with a run.
@@ -113,16 +122,20 @@ def _log_artifact_version(
113
122
  return art
114
123
 
115
124
 
125
+ _LOG_MODEL_DEPRECATION_MSG = "`log_model` is deprecated and will be removed in a future version. Please use `Run.log_artifact` instead."
126
+
127
+
128
+ @deprecated(_LOG_MODEL_DEPRECATION_MSG)
116
129
  def log_model(
117
130
  model_obj: Any,
118
131
  name: str = "model",
119
- aliases: Optional[Union[str, List[str]]] = None,
120
- description: Optional[str] = None,
121
- metadata: Optional[dict] = None,
122
- project: Optional[str] = None,
123
- scope_project: Optional[bool] = None,
124
- **kwargs: Dict[str, Any],
125
- ) -> "_SavedModel":
132
+ aliases: str | list[str] | None = None,
133
+ description: str | None = None,
134
+ metadata: dict | None = None,
135
+ project: str | None = None,
136
+ scope_project: bool | None = None,
137
+ **kwargs: dict[str, Any],
138
+ ) -> _SavedModel:
126
139
  """Log a model object to enable model-centric workflows in the UI.
127
140
 
128
141
  Supported frameworks include PyTorch, Keras, Tensorflow, Scikit-learn, etc. Under
@@ -168,7 +181,12 @@ def log_model(
168
181
  ```
169
182
 
170
183
  """
171
- model = data_types._SavedModel.init(model_obj, **kwargs)
184
+ wandb_deprecate(
185
+ field_name=Deprecated.beta__workflows__log_model,
186
+ warning_message=_LOG_MODEL_DEPRECATION_MSG,
187
+ )
188
+
189
+ model = _SavedModel.init(model_obj, **kwargs)
172
190
  _ = _log_artifact_version(
173
191
  name=name,
174
192
  type="model",
@@ -186,7 +204,11 @@ def log_model(
186
204
  return model
187
205
 
188
206
 
189
- def use_model(aliased_path: str, unsafe: bool = False) -> "_SavedModel":
207
+ _USE_MODEL_DEPRECATION_MSG = "`use_model` is deprecated and will be removed in a future version. Please update your code to use `Run.use_artifact` instead."
208
+
209
+
210
+ @deprecated(_USE_MODEL_DEPRECATION_MSG)
211
+ def use_model(aliased_path: str, unsafe: bool = False) -> _SavedModel:
190
212
  """Fetch a saved model from an alias.
191
213
 
192
214
  Under the hood, we use the alias to fetch the model artifact containing the
@@ -209,6 +231,11 @@ def use_model(aliased_path: str, unsafe: bool = False) -> "_SavedModel":
209
231
  model = sm.model_obj()
210
232
  ```
211
233
  """
234
+ wandb_deprecate(
235
+ field_name=Deprecated.beta__workflows__use_model,
236
+ warning_message=_USE_MODEL_DEPRECATION_MSG,
237
+ )
238
+
212
239
  if not unsafe:
213
240
  raise ValueError("The 'unsafe' parameter must be set to True to load a model.")
214
241
 
@@ -234,10 +261,14 @@ def use_model(aliased_path: str, unsafe: bool = False) -> "_SavedModel":
234
261
  )
235
262
 
236
263
 
264
+ _LINK_MODEL_DEPRECATION_MSG = "`link_model` is deprecated and will be removed in a future version. Please use `Run.link_artifact` instead."
265
+
266
+
267
+ @deprecated(_LINK_MODEL_DEPRECATION_MSG)
237
268
  def link_model(
238
- model: "_SavedModel",
269
+ model: _SavedModel,
239
270
  target_path: str,
240
- aliases: Optional[Union[str, List[str]]] = None,
271
+ aliases: str | list[str] | None = None,
241
272
  ) -> None:
242
273
  """Link the given model to a portfolio.
243
274
 
@@ -261,6 +292,11 @@ def link_model(
261
292
  link_model(sm, "my-portfolio")
262
293
 
263
294
  """
295
+ wandb_deprecate(
296
+ field_name=Deprecated.beta__workflows__link_model,
297
+ warning_message=_LINK_MODEL_DEPRECATION_MSG,
298
+ )
299
+
264
300
  aliases = wandb.util._resolve_aliases(aliases)
265
301
 
266
302
  if run := wandb_setup.singleton().most_recent_active_run:
wandb/bin/gpu_stats.exe CHANGED
Binary file
wandb/bin/wandb-core CHANGED
Binary file
wandb/cli/cli.py CHANGED
@@ -14,7 +14,7 @@ import textwrap
14
14
  import time
15
15
  import traceback
16
16
  from functools import wraps
17
- from typing import Any, Dict, Optional
17
+ from typing import Any, Dict, Optional, Tuple
18
18
 
19
19
  import click
20
20
  import yaml
@@ -95,6 +95,39 @@ class ClickWandbException(ClickException):
95
95
  )
96
96
 
97
97
 
98
+ def parse_service_config(
99
+ ctx: Optional[click.Context],
100
+ param: Optional[click.Parameter],
101
+ value: Optional[Tuple[str, ...]],
102
+ ) -> Dict[str, str]:
103
+ """Parse service configurations in format serviceName=policy."""
104
+ if not value:
105
+ return {}
106
+
107
+ result = {}
108
+ for config in value:
109
+ if "=" not in config:
110
+ raise click.BadParameter(
111
+ f"Service must be in format 'serviceName=policy', got '{config}'"
112
+ )
113
+
114
+ service_name, policy = config.split("=", 1)
115
+ service_name = service_name.strip()
116
+ policy = policy.strip()
117
+ if not service_name:
118
+ raise click.BadParameter("Service name cannot be empty")
119
+
120
+ # Simple validation for two policies
121
+ if policy not in ["always", "never"]:
122
+ raise click.BadParameter(
123
+ f"Policy must be 'always' or 'never', got '{policy}'"
124
+ )
125
+
126
+ result[service_name] = policy
127
+
128
+ return result
129
+
130
+
98
131
  def display_error(func):
99
132
  """Function decorator for catching common errors and re-raising as wandb.Error."""
100
133
 
@@ -454,6 +487,10 @@ def init(ctx, project, entity, reset, mode):
454
487
  @click.option("--show", default=5, help="Number of runs to show")
455
488
  @click.option("--append", is_flag=True, default=False, help="Append run")
456
489
  @click.option("--skip-console", is_flag=True, default=False, help="Skip console logs")
490
+ @click.option(
491
+ "--replace-tags",
492
+ help="Replace tags in the format 'old_tag1=new_tag1,old_tag2=new_tag2'",
493
+ )
457
494
  @display_error
458
495
  def sync(
459
496
  ctx,
@@ -479,6 +516,7 @@ def sync(
479
516
  clean_force=None,
480
517
  append=None,
481
518
  skip_console=None,
519
+ replace_tags=None,
482
520
  ):
483
521
  api = _get_cling_api()
484
522
  if not api.is_authenticated:
@@ -493,6 +531,10 @@ def sync(
493
531
  if exclude_globs:
494
532
  exclude_globs = exclude_globs.split(",")
495
533
 
534
+ replace_tags_dict = _parse_sync_replace_tags(replace_tags)
535
+ if replace_tags and replace_tags_dict is None:
536
+ return # Error already printed by helper function
537
+
496
538
  def _summary():
497
539
  all_items = get_runs(
498
540
  include_online=True,
@@ -548,6 +590,7 @@ def sync(
548
590
  log_path=_wandb_log_path,
549
591
  append=append,
550
592
  skip_console=skip_console,
593
+ replace_tags=replace_tags_dict,
551
594
  )
552
595
  for p in _path:
553
596
  sm.add(p)
@@ -633,6 +676,31 @@ def sync(
633
676
  _summary()
634
677
 
635
678
 
679
+ def _parse_sync_replace_tags(replace_tags: str) -> Optional[Dict[str, str]]:
680
+ """Parse replace_tags string into a dictionary.
681
+
682
+ Args:
683
+ replace_tags: String in format 'old_tag1=new_tag1,old_tag2=new_tag2'
684
+
685
+ Returns:
686
+ Mapping of old tags to new tags, or None if format is invalid
687
+ """
688
+ if not replace_tags:
689
+ return {}
690
+
691
+ replace_tags_dict = {}
692
+ for pair in replace_tags.split(","):
693
+ if "=" not in pair:
694
+ wandb.termerror(
695
+ f"Invalid replace-tags format: {pair}. Use 'old_tag=new_tag' format."
696
+ )
697
+ return None
698
+ old_tag, new_tag = pair.split("=", 1)
699
+ replace_tags_dict[old_tag.strip()] = new_tag.strip()
700
+
701
+ return replace_tags_dict
702
+
703
+
636
704
  @cli.command(
637
705
  context_settings=CONTEXT,
638
706
  help="Initialize a hyperparameter sweep. Search for hyperparameters that optimizes a cost function of a machine learning model by testing various combinations.",
@@ -1829,6 +1897,15 @@ def describe(job):
1829
1897
  "job_type",
1830
1898
  type=click.Choice(("git", "code", "image")),
1831
1899
  )
1900
+ @click.option(
1901
+ "--service",
1902
+ "-s",
1903
+ "services",
1904
+ multiple=True,
1905
+ callback=parse_service_config,
1906
+ help="Service configurations in format serviceName=policy. Valid policies: always, never",
1907
+ hidden=True,
1908
+ )
1832
1909
  @click.argument("path")
1833
1910
  def create(
1834
1911
  path,
@@ -1844,6 +1921,7 @@ def create(
1844
1921
  build_context,
1845
1922
  base_image,
1846
1923
  dockerfile,
1924
+ services,
1847
1925
  ):
1848
1926
  """Create a job from a source, without a wandb run.
1849
1927
 
@@ -1893,6 +1971,7 @@ def create(
1893
1971
  build_context=build_context,
1894
1972
  base_image=base_image,
1895
1973
  dockerfile=dockerfile,
1974
+ services=services,
1896
1975
  )
1897
1976
  if not artifact:
1898
1977
  wandb.termerror("Job creation failed")
wandb/env.py CHANGED
@@ -159,6 +159,14 @@ def is_offline(env: MutableMapping | None = None) -> bool:
159
159
  return env.get(MODE) == "offline"
160
160
 
161
161
 
162
+ def is_quiet() -> bool:
163
+ return _env_as_bool(QUIET, default="false")
164
+
165
+
166
+ def is_silent() -> bool:
167
+ return _env_as_bool(SILENT, default="false")
168
+
169
+
162
170
  def error_reporting_enabled() -> bool:
163
171
  return _env_as_bool(ERROR_REPORTING, default="True")
164
172
 
wandb/errors/errors.py CHANGED
@@ -2,7 +2,10 @@ from typing import Optional
2
2
 
3
3
 
4
4
  class Error(Exception):
5
- """Base W&B Error."""
5
+ """Base W&B Error.
6
+
7
+ <!-- lazydoc-ignore-class: internal -->
8
+ """
6
9
 
7
10
  def __init__(self, message, context: Optional[dict] = None) -> None:
8
11
  super().__init__(message)
@@ -22,7 +22,9 @@ class WandbCallback:
22
22
 
23
23
  Example:
24
24
  ```
25
- train_pool = Pool(train[features], label=train["label"], cat_features=cat_features)
25
+ train_pool = Pool(
26
+ train[features], label=train["label"], cat_features=cat_features
27
+ )
26
28
  test_pool = Pool(test[features], label=test["label"], cat_features=cat_features)
27
29
 
28
30
  model = CatBoostRegressor(
@@ -130,7 +132,9 @@ def log_summary(
130
132
 
131
133
  Example:
132
134
  ```python
133
- train_pool = Pool(train[features], label=train["label"], cat_features=cat_features)
135
+ train_pool = Pool(
136
+ train[features], label=train["label"], cat_features=cat_features
137
+ )
134
138
  test_pool = Pool(test[features], label=test["label"], cat_features=cat_features)
135
139
 
136
140
  model = CatBoostRegressor(
@@ -270,7 +270,9 @@ def create_component_from_func(
270
270
  Example of a component function declaring file input and output::
271
271
 
272
272
  def catboost_train_classifier(
273
- training_data_path: InputPath("CSV"), # Path to input data file of type "CSV"
273
+ training_data_path: InputPath(
274
+ "CSV"
275
+ ), # Path to input data file of type "CSV"
274
276
  trained_model_path: OutputPath(
275
277
  "CatBoostModel"
276
278
  ), # Path to output data file of type "CatBoostModel"
@@ -9,7 +9,6 @@ from typing_extensions import override
9
9
  import wandb
10
10
  from wandb import Artifact
11
11
  from wandb.sdk.lib import telemetry
12
- from wandb.sdk.wandb_run import Run
13
12
 
14
13
  try:
15
14
  import lightning
@@ -295,7 +294,7 @@ class WandbLogger(Logger):
295
294
  anonymous: Optional[bool] = None,
296
295
  project: Optional[str] = None,
297
296
  log_model: Union[Literal["all"], bool] = False,
298
- experiment: Optional["Run"] = None,
297
+ experiment: Optional["wandb.Run"] = None,
299
298
  prefix: str = "",
300
299
  checkpoint_name: Optional[str] = None,
301
300
  log_checkpoint_on: Union[Literal["success"], Literal["all"]] = "success",
@@ -362,7 +361,7 @@ class WandbLogger(Logger):
362
361
 
363
362
  @property
364
363
  @rank_zero_experiment
365
- def experiment(self) -> "Run":
364
+ def experiment(self) -> "wandb.Run":
366
365
  r"""Actual wandb object.
367
366
 
368
367
  To use wandb features in your :class:`~lightning.pytorch.core.LightningModule`, do the
@@ -395,7 +394,7 @@ class WandbLogger(Logger):
395
394
  self._experiment = wandb.init(**self._wandb_init)
396
395
 
397
396
  # define default x-axis
398
- if isinstance(self._experiment, Run) and getattr(
397
+ if isinstance(self._experiment, wandb.Run) and getattr(
399
398
  self._experiment, "define_metric", None
400
399
  ):
401
400
  self._experiment.define_metric("trainer/global_step")
@@ -1,3 +1,9 @@
1
+ """W&B Integration for Metaflow.
2
+
3
+ Defines a custom step and flow decorator `wandb_log` that automatically logs
4
+ flow parameters and artifacts to W&B.
5
+ """
6
+
1
7
  from .metaflow import wandb_log, wandb_track, wandb_use
2
8
 
3
9
  __all__ = ["wandb_log", "wandb_track", "wandb_use"]
@@ -0,0 +1,74 @@
1
+ """Support for Pandas datatypes.
2
+
3
+ May raise MissingDependencyError on import.
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ from typing_extensions import Any, TypeGuard
9
+
10
+ import wandb
11
+
12
+ from . import errors
13
+
14
+ try:
15
+ import pandas as pd
16
+ except ImportError as e:
17
+ warning = (
18
+ "`pandas` not installed >>"
19
+ " @wandb_log(datasets=True) may not auto log your dataset!"
20
+ )
21
+ raise errors.MissingDependencyError(warning=warning) from e
22
+
23
+
24
+ def is_dataframe(data: Any) -> TypeGuard[pd.DataFrame]:
25
+ """Returns whether the data is a Pandas DataFrame."""
26
+ return isinstance(data, pd.DataFrame)
27
+
28
+
29
+ def use_dataframe(
30
+ name: str,
31
+ run: wandb.Run | None,
32
+ testing: bool = False,
33
+ ) -> str | None:
34
+ """Log a dependency on a DataFrame input.
35
+
36
+ Args:
37
+ name: Name of the input.
38
+ run: The run to update.
39
+ testing: True in unit tests.
40
+ """
41
+ if testing:
42
+ return "datasets"
43
+ assert run
44
+
45
+ wandb.termlog(f"Using artifact: {name} (Pandas DataFrame)")
46
+ run.use_artifact(f"{name}:latest")
47
+ return None
48
+
49
+
50
+ def track_dataframe(
51
+ name: str,
52
+ data: pd.DataFrame,
53
+ run: wandb.Run | None,
54
+ testing: bool = False,
55
+ ) -> str | None:
56
+ """Log a DataFrame output as an artifact.
57
+
58
+ Args:
59
+ name: The output's name.
60
+ data: The output's value.
61
+ run: The run to update.
62
+ testing: True in unit tests.
63
+ """
64
+ if testing:
65
+ return "pd.DataFrame"
66
+ assert run
67
+
68
+ artifact = wandb.Artifact(name, type="dataset")
69
+ with artifact.new_file(f"{name}.parquet", "wb") as f:
70
+ data.to_parquet(f, engine="pyarrow")
71
+
72
+ wandb.termlog(f"Logging artifact: {name} (Pandas DataFrame)")
73
+ run.log_artifact(artifact)
74
+ return None
@@ -0,0 +1,13 @@
1
+ import wandb
2
+
3
+
4
+ class MissingDependencyError(Exception):
5
+ """An optional dependency is missing."""
6
+
7
+ def __init__(self, *args: object, warning: str) -> None:
8
+ super().__init__(*args)
9
+ self._wb_warning = warning
10
+
11
+ def warn(self) -> None:
12
+ """Print a warning for the problem."""
13
+ wandb.termwarn(self._wb_warning)