wandb 0.22.1__py3-none-win32.whl → 0.22.3__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 (166) hide show
  1. wandb/__init__.py +1 -1
  2. wandb/__init__.pyi +7 -4
  3. wandb/_pydantic/__init__.py +8 -1
  4. wandb/_pydantic/base.py +54 -18
  5. wandb/_pydantic/field_types.py +8 -3
  6. wandb/_pydantic/pagination.py +46 -0
  7. wandb/_pydantic/utils.py +2 -2
  8. wandb/apis/public/api.py +24 -19
  9. wandb/apis/public/artifacts.py +259 -270
  10. wandb/apis/public/registries/_utils.py +40 -54
  11. wandb/apis/public/registries/registries_search.py +70 -85
  12. wandb/apis/public/registries/registry.py +173 -156
  13. wandb/apis/public/runs.py +27 -6
  14. wandb/apis/public/utils.py +43 -20
  15. wandb/automations/_generated/create_automation.py +2 -2
  16. wandb/automations/_generated/create_generic_webhook_integration.py +4 -4
  17. wandb/automations/_generated/delete_automation.py +2 -2
  18. wandb/automations/_generated/fragments.py +31 -52
  19. wandb/automations/_generated/generic_webhook_integrations_by_entity.py +3 -3
  20. wandb/automations/_generated/get_automations.py +3 -3
  21. wandb/automations/_generated/get_automations_by_entity.py +3 -3
  22. wandb/automations/_generated/input_types.py +9 -9
  23. wandb/automations/_generated/integrations_by_entity.py +3 -3
  24. wandb/automations/_generated/operations.py +6 -6
  25. wandb/automations/_generated/slack_integrations_by_entity.py +3 -3
  26. wandb/automations/_generated/update_automation.py +2 -2
  27. wandb/automations/_utils.py +3 -3
  28. wandb/automations/actions.py +3 -3
  29. wandb/automations/automations.py +6 -5
  30. wandb/bin/gpu_stats.exe +0 -0
  31. wandb/bin/wandb-core +0 -0
  32. wandb/cli/beta.py +23 -3
  33. wandb/cli/beta_leet.py +75 -0
  34. wandb/cli/beta_sync.py +1 -1
  35. wandb/cli/cli.py +34 -7
  36. wandb/errors/term.py +8 -8
  37. wandb/jupyter.py +0 -51
  38. wandb/old/settings.py +6 -6
  39. wandb/proto/v3/wandb_api_pb2.py +86 -0
  40. wandb/proto/v3/wandb_server_pb2.py +38 -37
  41. wandb/proto/v3/wandb_settings_pb2.py +2 -2
  42. wandb/proto/v3/wandb_sync_pb2.py +19 -6
  43. wandb/proto/v4/wandb_api_pb2.py +37 -0
  44. wandb/proto/v4/wandb_server_pb2.py +38 -37
  45. wandb/proto/v4/wandb_settings_pb2.py +2 -2
  46. wandb/proto/v4/wandb_sync_pb2.py +10 -6
  47. wandb/proto/v5/wandb_api_pb2.py +38 -0
  48. wandb/proto/v5/wandb_server_pb2.py +38 -37
  49. wandb/proto/v5/wandb_settings_pb2.py +2 -2
  50. wandb/proto/v5/wandb_sync_pb2.py +10 -6
  51. wandb/proto/v6/wandb_api_pb2.py +48 -0
  52. wandb/proto/v6/wandb_server_pb2.py +38 -37
  53. wandb/proto/v6/wandb_settings_pb2.py +2 -2
  54. wandb/proto/v6/wandb_sync_pb2.py +10 -6
  55. wandb/proto/wandb_api_pb2.py +18 -0
  56. wandb/proto/wandb_generate_proto.py +1 -0
  57. wandb/sdk/artifacts/_generated/__init__.py +96 -40
  58. wandb/sdk/artifacts/_generated/add_aliases.py +3 -3
  59. wandb/sdk/artifacts/_generated/add_artifact_collection_tags.py +26 -0
  60. wandb/sdk/artifacts/_generated/artifact_by_id.py +2 -2
  61. wandb/sdk/artifacts/_generated/artifact_by_name.py +3 -3
  62. wandb/sdk/artifacts/_generated/artifact_collection_membership_file_urls.py +27 -8
  63. wandb/sdk/artifacts/_generated/artifact_collection_membership_files.py +27 -8
  64. wandb/sdk/artifacts/_generated/artifact_created_by.py +7 -20
  65. wandb/sdk/artifacts/_generated/artifact_file_urls.py +19 -6
  66. wandb/sdk/artifacts/_generated/artifact_membership_by_name.py +26 -0
  67. wandb/sdk/artifacts/_generated/artifact_type.py +5 -5
  68. wandb/sdk/artifacts/_generated/artifact_used_by.py +8 -17
  69. wandb/sdk/artifacts/_generated/artifact_version_files.py +19 -8
  70. wandb/sdk/artifacts/_generated/delete_aliases.py +3 -3
  71. wandb/sdk/artifacts/_generated/delete_artifact.py +4 -4
  72. wandb/sdk/artifacts/_generated/delete_artifact_collection_tags.py +23 -0
  73. wandb/sdk/artifacts/_generated/delete_artifact_portfolio.py +4 -4
  74. wandb/sdk/artifacts/_generated/delete_artifact_sequence.py +4 -4
  75. wandb/sdk/artifacts/_generated/delete_registry.py +21 -0
  76. wandb/sdk/artifacts/_generated/fetch_artifact_manifest.py +8 -20
  77. wandb/sdk/artifacts/_generated/fetch_linked_artifacts.py +13 -35
  78. wandb/sdk/artifacts/_generated/fetch_org_info_from_entity.py +28 -0
  79. wandb/sdk/artifacts/_generated/fetch_registries.py +18 -8
  80. wandb/sdk/{projects → artifacts}/_generated/fetch_registry.py +4 -4
  81. wandb/sdk/artifacts/_generated/fragments.py +183 -333
  82. wandb/sdk/artifacts/_generated/input_types.py +133 -7
  83. wandb/sdk/artifacts/_generated/link_artifact.py +5 -5
  84. wandb/sdk/artifacts/_generated/operations.py +1053 -548
  85. wandb/sdk/artifacts/_generated/project_artifact_collection.py +9 -77
  86. wandb/sdk/artifacts/_generated/project_artifact_collections.py +21 -9
  87. wandb/sdk/artifacts/_generated/project_artifact_type.py +3 -3
  88. wandb/sdk/artifacts/_generated/project_artifact_types.py +19 -6
  89. wandb/sdk/artifacts/_generated/project_artifacts.py +7 -8
  90. wandb/sdk/artifacts/_generated/registry_collections.py +21 -9
  91. wandb/sdk/artifacts/_generated/registry_versions.py +20 -9
  92. wandb/sdk/artifacts/_generated/rename_registry.py +25 -0
  93. wandb/sdk/artifacts/_generated/run_input_artifacts.py +5 -9
  94. wandb/sdk/artifacts/_generated/run_output_artifacts.py +5 -9
  95. wandb/sdk/artifacts/_generated/type_info.py +2 -2
  96. wandb/sdk/artifacts/_generated/unlink_artifact.py +3 -5
  97. wandb/sdk/artifacts/_generated/update_artifact.py +3 -3
  98. wandb/sdk/artifacts/_generated/update_artifact_collection_type.py +28 -0
  99. wandb/sdk/artifacts/_generated/update_artifact_portfolio.py +7 -16
  100. wandb/sdk/artifacts/_generated/update_artifact_sequence.py +7 -16
  101. wandb/sdk/artifacts/_generated/upsert_registry.py +25 -0
  102. wandb/sdk/artifacts/_gqlutils.py +170 -6
  103. wandb/sdk/artifacts/_models/__init__.py +9 -0
  104. wandb/sdk/artifacts/_models/artifact_collection.py +109 -0
  105. wandb/sdk/artifacts/_models/manifest.py +26 -0
  106. wandb/sdk/artifacts/_models/pagination.py +26 -0
  107. wandb/sdk/artifacts/_models/registry.py +100 -0
  108. wandb/sdk/artifacts/_validators.py +45 -27
  109. wandb/sdk/artifacts/artifact.py +249 -244
  110. wandb/sdk/artifacts/artifact_file_cache.py +1 -1
  111. wandb/sdk/artifacts/artifact_manifest.py +37 -32
  112. wandb/sdk/artifacts/artifact_manifest_entry.py +82 -133
  113. wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py +43 -61
  114. wandb/sdk/artifacts/storage_handler.py +18 -12
  115. wandb/sdk/artifacts/storage_handlers/azure_handler.py +11 -6
  116. wandb/sdk/artifacts/storage_handlers/gcs_handler.py +17 -12
  117. wandb/sdk/artifacts/storage_handlers/http_handler.py +9 -4
  118. wandb/sdk/artifacts/storage_handlers/local_file_handler.py +10 -6
  119. wandb/sdk/artifacts/storage_handlers/multi_handler.py +5 -4
  120. wandb/sdk/artifacts/storage_handlers/s3_handler.py +10 -8
  121. wandb/sdk/artifacts/storage_handlers/tracking_handler.py +6 -4
  122. wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py +24 -21
  123. wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py +4 -2
  124. wandb/sdk/artifacts/storage_policies/_multipart.py +187 -0
  125. wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py +61 -242
  126. wandb/sdk/artifacts/storage_policy.py +25 -12
  127. wandb/sdk/data_types/image.py +2 -2
  128. wandb/sdk/data_types/object_3d.py +67 -2
  129. wandb/sdk/interface/interface.py +72 -64
  130. wandb/sdk/interface/interface_queue.py +27 -18
  131. wandb/sdk/interface/interface_shared.py +61 -23
  132. wandb/sdk/interface/interface_sock.py +9 -5
  133. wandb/sdk/internal/_generated/server_features_query.py +4 -4
  134. wandb/sdk/internal/job_builder.py +27 -10
  135. wandb/sdk/internal/sender.py +4 -1
  136. wandb/sdk/launch/create_job.py +2 -1
  137. wandb/sdk/launch/inputs/schema.py +13 -10
  138. wandb/sdk/lib/apikey.py +8 -12
  139. wandb/sdk/lib/asyncio_compat.py +1 -1
  140. wandb/sdk/lib/asyncio_manager.py +5 -5
  141. wandb/sdk/lib/console_capture.py +38 -30
  142. wandb/sdk/lib/progress.py +151 -125
  143. wandb/sdk/lib/retry.py +3 -2
  144. wandb/sdk/lib/service/service_connection.py +2 -2
  145. wandb/sdk/lib/wb_logging.py +2 -1
  146. wandb/sdk/mailbox/mailbox.py +1 -1
  147. wandb/sdk/wandb_init.py +11 -14
  148. wandb/sdk/wandb_run.py +14 -48
  149. wandb/sdk/wandb_settings.py +114 -30
  150. {wandb-0.22.1.dist-info → wandb-0.22.3.dist-info}/METADATA +2 -1
  151. {wandb-0.22.1.dist-info → wandb-0.22.3.dist-info}/RECORD +154 -146
  152. wandb/sdk/artifacts/_generated/artifact_via_membership_by_name.py +0 -26
  153. wandb/sdk/artifacts/_generated/create_artifact_collection_tag_assignments.py +0 -36
  154. wandb/sdk/artifacts/_generated/delete_artifact_collection_tag_assignments.py +0 -25
  155. wandb/sdk/artifacts/_generated/move_artifact_collection.py +0 -35
  156. wandb/sdk/projects/_generated/__init__.py +0 -26
  157. wandb/sdk/projects/_generated/delete_project.py +0 -22
  158. wandb/sdk/projects/_generated/enums.py +0 -4
  159. wandb/sdk/projects/_generated/fragments.py +0 -41
  160. wandb/sdk/projects/_generated/input_types.py +0 -13
  161. wandb/sdk/projects/_generated/operations.py +0 -88
  162. wandb/sdk/projects/_generated/rename_project.py +0 -27
  163. wandb/sdk/projects/_generated/upsert_registry_project.py +0 -27
  164. {wandb-0.22.1.dist-info → wandb-0.22.3.dist-info}/WHEEL +0 -0
  165. {wandb-0.22.1.dist-info → wandb-0.22.3.dist-info}/entry_points.txt +0 -0
  166. {wandb-0.22.1.dist-info → wandb-0.22.3.dist-info}/licenses/LICENSE +0 -0
@@ -132,6 +132,7 @@ def get_min_supported_for_source_dict(
132
132
 
133
133
  class JobBuilder:
134
134
  _settings: SettingsStatic
135
+ _files_dir: str
135
136
  _metadatafile_path: Optional[str]
136
137
  _requirements_path: Optional[str]
137
138
  _config: Optional[Dict[str, Any]]
@@ -146,8 +147,26 @@ class JobBuilder:
146
147
  _verbose: bool
147
148
  _services: Dict[str, str]
148
149
 
149
- def __init__(self, settings: SettingsStatic, verbose: bool = False):
150
+ def __init__(
151
+ self,
152
+ settings: SettingsStatic,
153
+ verbose: bool = False,
154
+ *,
155
+ files_dir: str,
156
+ ):
157
+ """Instantiate a JobBuilder.
158
+
159
+ Args:
160
+ settings: Parameters for the job builder.
161
+ In a run, this is the run's settings.
162
+ Otherwise, this is a set of undocumented parameters,
163
+ all of which should be made explicit like files_dir.
164
+ files_dir: The directory where to write files.
165
+ In a run, this should be the run's files directory.
166
+ """
150
167
  self._settings = settings
168
+ self._files_dir = files_dir
169
+
151
170
  self._metadatafile_path = None
152
171
  self._requirements_path = None
153
172
  self._config = None
@@ -460,9 +479,7 @@ class JobBuilder:
460
479
  )
461
480
  return None
462
481
 
463
- if not os.path.exists(
464
- os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME)
465
- ):
482
+ if not os.path.exists(os.path.join(self._files_dir, REQUIREMENTS_FNAME)):
466
483
  self._log_if_verbose(
467
484
  "No requirements.txt found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
468
485
  "warn",
@@ -471,7 +488,7 @@ class JobBuilder:
471
488
  metadata = self._handle_metadata_file()
472
489
  if metadata is None:
473
490
  self._log_if_verbose(
474
- f"Ensure read and write access to run files dir: {self._settings.files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables",
491
+ f"Ensure read and write access to run files dir: {self._files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables",
475
492
  "warn",
476
493
  )
477
494
  return None
@@ -560,15 +577,15 @@ class JobBuilder:
560
577
  f.write(json.dumps(source_info, indent=4))
561
578
 
562
579
  artifact.add_file(
563
- os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME),
580
+ os.path.join(self._files_dir, REQUIREMENTS_FNAME),
564
581
  name=FROZEN_REQUIREMENTS_FNAME,
565
582
  )
566
583
 
567
584
  if source_type == "repo":
568
585
  # add diff
569
- if os.path.exists(os.path.join(self._settings.files_dir, DIFF_FNAME)):
586
+ if os.path.exists(os.path.join(self._files_dir, DIFF_FNAME)):
570
587
  artifact.add_file(
571
- os.path.join(self._settings.files_dir, DIFF_FNAME),
588
+ os.path.join(self._files_dir, DIFF_FNAME),
572
589
  name=DIFF_FNAME,
573
590
  )
574
591
 
@@ -619,8 +636,8 @@ class JobBuilder:
619
636
  def _handle_metadata_file(
620
637
  self,
621
638
  ) -> Optional[Dict]:
622
- if os.path.exists(os.path.join(self._settings.files_dir, METADATA_FNAME)):
623
- with open(os.path.join(self._settings.files_dir, METADATA_FNAME)) as f:
639
+ if os.path.exists(os.path.join(self._files_dir, METADATA_FNAME)):
640
+ with open(os.path.join(self._files_dir, METADATA_FNAME)) as f:
624
641
  metadata: Dict = json.load(f)
625
642
  return metadata
626
643
 
@@ -311,7 +311,10 @@ class SendManager:
311
311
  self._output_raw_file = None
312
312
 
313
313
  # job builder
314
- self._job_builder = JobBuilder(settings)
314
+ self._job_builder = JobBuilder(
315
+ settings,
316
+ files_dir=settings.files_dir,
317
+ )
315
318
 
316
319
  time_now = time.monotonic()
317
320
  self._debounce_config_time = time_now
@@ -417,10 +417,11 @@ def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuild
417
417
  if job_source == "code":
418
418
  job_source = "artifact"
419
419
 
420
- settings = wandb.Settings(x_files_dir=tmpdir, job_source=job_source)
420
+ settings = wandb.Settings(job_source=job_source)
421
421
  job_builder = JobBuilder(
422
422
  settings=settings, # type: ignore
423
423
  verbose=True,
424
+ files_dir=tmpdir,
424
425
  )
425
426
  job_builder._partial = True
426
427
  # never allow notebook runs
@@ -7,6 +7,9 @@ META_SCHEMA = {
7
7
  },
8
8
  "title": {"type": "string"},
9
9
  "description": {"type": "string"},
10
+ "label": {"type": "string"},
11
+ "placeholder": {"type": "string"},
12
+ "required": {"type": "boolean"},
10
13
  "format": {"type": "string"},
11
14
  "enum": {"type": "array", "items": {"type": ["integer", "number", "string"]}},
12
15
  "properties": {"type": "object", "patternProperties": {".*": {"$ref": "#"}}},
@@ -19,24 +22,24 @@ META_SCHEMA = {
19
22
  },
20
23
  "allOf": [
21
24
  {
22
- "if": {"properties": {"type": {"const": "number"}}},
25
+ "if": {"properties": {"type": {"const": "integer"}}},
23
26
  "then": {
24
27
  "properties": {
25
- "minimum": {"type": ["integer", "number"]},
26
- "maximum": {"type": ["integer", "number"]},
27
- "exclusiveMinimum": {"type": ["integer", "number"]},
28
- "exclusiveMaximum": {"type": ["integer", "number"]},
28
+ "minimum": {"type": "integer"},
29
+ "maximum": {"type": "integer"},
30
+ "exclusiveMinimum": {"type": "integer"},
31
+ "exclusiveMaximum": {"type": "integer"},
29
32
  }
30
33
  },
31
34
  },
32
35
  {
33
- "if": {"properties": {"type": {"const": "integer"}}},
36
+ "if": {"properties": {"type": {"const": "number"}}},
34
37
  "then": {
35
38
  "properties": {
36
- "minimum": {"type": "integer"},
37
- "maximum": {"type": "integer"},
38
- "exclusiveMinimum": {"type": "integer"},
39
- "exclusiveMaximum": {"type": "integer"},
39
+ "minimum": {"type": ["integer", "number"]},
40
+ "maximum": {"type": ["integer", "number"]},
41
+ "exclusiveMinimum": {"type": ["integer", "number"]},
42
+ "exclusiveMaximum": {"type": ["integer", "number"]},
40
43
  }
41
44
  },
42
45
  },
wandb/sdk/lib/apikey.py CHANGED
@@ -136,12 +136,6 @@ def prompt_api_key( # noqa: C901
136
136
  if (jupyter and not settings.login_timeout) or no_create:
137
137
  choices.remove(LOGIN_CHOICE_NEW)
138
138
 
139
- if jupyter and "google.colab" in sys.modules:
140
- log_string = term.LOG_STRING_NOCOLOR
141
- key = wandb.jupyter.attempt_colab_login(app_url) # type: ignore
142
- if key is not None:
143
- return key # type: ignore
144
-
145
139
  if anon_mode == "must":
146
140
  result = LOGIN_CHOICE_ANON
147
141
  # If we're not in an interactive environment, default to dry-run.
@@ -236,9 +230,9 @@ def check_netrc_access(
236
230
  def write_netrc(host: str, entity: str, key: str):
237
231
  """Add our host and key to .netrc."""
238
232
  _, key_suffix = key.split("-", 1) if "-" in key else ("", key)
239
- if len(key_suffix) != 40:
233
+ if len(key_suffix) < 40:
240
234
  raise ValueError(
241
- f"API-key must be exactly 40 characters long: {key_suffix} ({len(key_suffix)} chars)"
235
+ f"API-key must be at least 40 characters long: {key_suffix} ({len(key_suffix)} chars)"
242
236
  )
243
237
 
244
238
  normalized_host = urlparse(host).netloc
@@ -305,12 +299,14 @@ def write_key(
305
299
  # TODO(jhr): api shouldn't be optional or it shouldn't be passed, clean up callers
306
300
  api = api or InternalApi()
307
301
 
308
- # Normal API keys are 40-character hex strings. On-prem API keys have a
309
- # variable-length prefix, a dash, then the 40-char string.
302
+ # API keys are strings of at least 40 characters. On-prem API keys have a
303
+ # variable-length prefix, a dash, then the string of at least 40 chars.
310
304
  _, suffix = key.split("-", 1) if "-" in key else ("", key)
311
305
 
312
- if len(suffix) != 40:
313
- raise ValueError(f"API key must be 40 characters long, yours was {len(key)}")
306
+ if len(suffix) < 40:
307
+ raise ValueError(
308
+ f"API key must be at least 40 characters long, yours was {len(key)}"
309
+ )
314
310
 
315
311
  write_netrc(settings.base_url, "user", key)
316
312
 
@@ -133,7 +133,7 @@ class TaskGroup:
133
133
  """Object that `open_task_group()` yields."""
134
134
 
135
135
  def __init__(self) -> None:
136
- self._tasks: list[asyncio.Task] = []
136
+ self._tasks: list[asyncio.Task[None]] = []
137
137
 
138
138
  def start_soon(self, coro: Coroutine[Any, Any, Any]) -> None:
139
139
  """Schedule a task in the group.
@@ -7,7 +7,7 @@ import concurrent.futures
7
7
  import contextlib
8
8
  import logging
9
9
  import threading
10
- from typing import Any, Callable, Coroutine, TypeVar
10
+ from typing import Awaitable, Callable, TypeVar
11
11
 
12
12
  from . import asyncio_compat
13
13
 
@@ -104,7 +104,7 @@ class AsyncioManager:
104
104
  # This only matters if the KeyboardInterrupt is suppressed.
105
105
  self._runner.cancel()
106
106
 
107
- def run(self, fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
107
+ def run(self, fn: Callable[[], Awaitable[_T]]) -> _T:
108
108
  """Run an async function to completion.
109
109
 
110
110
  The function is called in the asyncio thread. Blocks until start()
@@ -148,7 +148,7 @@ class AsyncioManager:
148
148
 
149
149
  def run_soon(
150
150
  self,
151
- fn: Callable[[], Coroutine[Any, Any, None]],
151
+ fn: Callable[[], Awaitable[None]],
152
152
  *,
153
153
  daemon: bool = False,
154
154
  name: str | None = None,
@@ -186,7 +186,7 @@ class AsyncioManager:
186
186
 
187
187
  def _schedule(
188
188
  self,
189
- fn: Callable[[], Coroutine[Any, Any, _T]],
189
+ fn: Callable[[], Awaitable[_T]],
190
190
  daemon: bool,
191
191
  name: str | None = None,
192
192
  ) -> concurrent.futures.Future[_T]:
@@ -207,7 +207,7 @@ class AsyncioManager:
207
207
 
208
208
  async def _wrap(
209
209
  self,
210
- fn: Callable[[], Coroutine[Any, Any, _T]],
210
+ fn: Callable[[], Awaitable[_T]],
211
211
  daemon: bool,
212
212
  name: str | None,
213
213
  ) -> _T:
@@ -75,9 +75,12 @@ class _WriteCallback(Protocol):
75
75
  """
76
76
 
77
77
 
78
- # A reentrant lock is used to catch callbacks that write to stderr/stdout.
79
- _module_rlock = threading.RLock()
78
+ _module_lock = threading.Lock()
80
79
  _is_writing = False
80
+ """Prevents infinite print-capture loops.
81
+
82
+ If a capture callback prints, that output is not captured.
83
+ """
81
84
 
82
85
  _patch_exception: CannotCaptureConsoleError | None = None
83
86
 
@@ -99,7 +102,7 @@ def capture_stdout(callback: _WriteCallback) -> Callable[[], None]:
99
102
  Raises:
100
103
  CannotCaptureConsoleError: If patching failed on import.
101
104
  """
102
- with _module_rlock:
105
+ with _module_lock:
103
106
  if _patch_exception:
104
107
  raise _patch_exception
105
108
 
@@ -121,7 +124,7 @@ def capture_stderr(callback: _WriteCallback) -> Callable[[], None]:
121
124
  Raises:
122
125
  CannotCaptureConsoleError: If patching failed on import.
123
126
  """
124
- with _module_rlock:
127
+ with _module_lock:
125
128
  if _patch_exception:
126
129
  raise _patch_exception
127
130
 
@@ -144,7 +147,7 @@ def _insert_disposably(
144
147
  def dispose() -> None:
145
148
  nonlocal disposed
146
149
 
147
- with _module_rlock:
150
+ with _module_lock:
148
151
  if disposed:
149
152
  return
150
153
 
@@ -167,38 +170,43 @@ def _patch(
167
170
  global _is_writing
168
171
  n = orig_write(s)
169
172
 
170
- # NOTE: Since _module_rlock is reentrant, this is safe. It will not
171
- # deadlock if a callback invokes write() again.
172
- with _module_rlock:
173
+ with _module_lock:
173
174
  if _is_writing:
174
175
  return n
175
-
176
176
  _is_writing = True
177
- try:
178
- for cb in callbacks.values():
179
- cb(s, n)
180
-
181
- except BaseException as e:
182
- # Clear all callbacks on any exception to avoid infinite loops:
183
- #
184
- # * If we re-raise, an exception handler is likely to print
185
- # the exception to the console and trigger callbacks again
186
- # * If we log, we can't guarantee that this doesn't print
187
- # to console.
188
- #
189
- # This is especially important for KeyboardInterrupt.
177
+
178
+ # Invoke callbacks outside of the lock to avoid deadlocks.
179
+ # 1. A callback may print, invoking this again.
180
+ # 2. A callback may block on a different thread which then prints.
181
+ callback_list = list(callbacks.values())
182
+
183
+ try:
184
+ for cb in callback_list:
185
+ cb(s, n)
186
+
187
+ except BaseException as e:
188
+ # Clear all callbacks on any exception to avoid infinite loops:
189
+ #
190
+ # * If we re-raise, an exception handler is likely to print
191
+ # the exception to the console and trigger callbacks again
192
+ # * If we log, we can't guarantee that this doesn't print
193
+ # to console.
194
+ #
195
+ # This is especially important for KeyboardInterrupt.
196
+ with _module_lock:
190
197
  _stderr_callbacks.clear()
191
198
  _stdout_callbacks.clear()
192
199
 
193
- if isinstance(e, Exception):
194
- # We suppress Exceptions so that bugs in W&B code don't
195
- # cause the user's print() statements to raise errors.
196
- _logger.exception("Error in console callback, clearing all!")
197
- else:
198
- # Re-raise errors like KeyboardInterrupt.
199
- raise
200
+ if isinstance(e, Exception):
201
+ # We suppress Exceptions so that bugs in W&B code don't
202
+ # cause the user's print() statements to raise errors.
203
+ _logger.exception("Error in console callback, clearing all!")
204
+ else:
205
+ # Re-raise errors like KeyboardInterrupt.
206
+ raise
200
207
 
201
- finally:
208
+ finally:
209
+ with _module_lock:
202
210
  _is_writing = False
203
211
 
204
212
  return n