wandb 0.22.2__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 (147) hide show
  1. wandb/__init__.py +1 -1
  2. wandb/__init__.pyi +2 -2
  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 +8 -2
  33. wandb/cli/beta_leet.py +2 -1
  34. wandb/cli/beta_sync.py +1 -1
  35. wandb/errors/term.py +8 -8
  36. wandb/jupyter.py +0 -51
  37. wandb/old/settings.py +6 -6
  38. wandb/proto/v3/wandb_internal_pb2.py +351 -352
  39. wandb/proto/v3/wandb_server_pb2.py +38 -37
  40. wandb/proto/v3/wandb_settings_pb2.py +2 -2
  41. wandb/proto/v3/wandb_sync_pb2.py +19 -6
  42. wandb/proto/v4/wandb_internal_pb2.py +351 -352
  43. wandb/proto/v4/wandb_server_pb2.py +38 -37
  44. wandb/proto/v4/wandb_settings_pb2.py +2 -2
  45. wandb/proto/v4/wandb_sync_pb2.py +10 -6
  46. wandb/proto/v5/wandb_internal_pb2.py +351 -352
  47. wandb/proto/v5/wandb_server_pb2.py +38 -37
  48. wandb/proto/v5/wandb_settings_pb2.py +2 -2
  49. wandb/proto/v5/wandb_sync_pb2.py +10 -6
  50. wandb/proto/v6/wandb_internal_pb2.py +351 -352
  51. wandb/proto/v6/wandb_server_pb2.py +38 -37
  52. wandb/proto/v6/wandb_settings_pb2.py +2 -2
  53. wandb/proto/v6/wandb_sync_pb2.py +10 -6
  54. wandb/sdk/artifacts/_generated/__init__.py +96 -40
  55. wandb/sdk/artifacts/_generated/add_aliases.py +3 -3
  56. wandb/sdk/artifacts/_generated/add_artifact_collection_tags.py +26 -0
  57. wandb/sdk/artifacts/_generated/artifact_by_id.py +2 -2
  58. wandb/sdk/artifacts/_generated/artifact_by_name.py +3 -3
  59. wandb/sdk/artifacts/_generated/artifact_collection_membership_file_urls.py +27 -8
  60. wandb/sdk/artifacts/_generated/artifact_collection_membership_files.py +27 -8
  61. wandb/sdk/artifacts/_generated/artifact_created_by.py +7 -20
  62. wandb/sdk/artifacts/_generated/artifact_file_urls.py +19 -6
  63. wandb/sdk/artifacts/_generated/artifact_membership_by_name.py +26 -0
  64. wandb/sdk/artifacts/_generated/artifact_type.py +5 -5
  65. wandb/sdk/artifacts/_generated/artifact_used_by.py +8 -17
  66. wandb/sdk/artifacts/_generated/artifact_version_files.py +19 -8
  67. wandb/sdk/artifacts/_generated/delete_aliases.py +3 -3
  68. wandb/sdk/artifacts/_generated/delete_artifact.py +4 -4
  69. wandb/sdk/artifacts/_generated/delete_artifact_collection_tags.py +23 -0
  70. wandb/sdk/artifacts/_generated/delete_artifact_portfolio.py +4 -4
  71. wandb/sdk/artifacts/_generated/delete_artifact_sequence.py +4 -4
  72. wandb/sdk/artifacts/_generated/delete_registry.py +21 -0
  73. wandb/sdk/artifacts/_generated/fetch_artifact_manifest.py +8 -20
  74. wandb/sdk/artifacts/_generated/fetch_linked_artifacts.py +13 -35
  75. wandb/sdk/artifacts/_generated/fetch_org_info_from_entity.py +28 -0
  76. wandb/sdk/artifacts/_generated/fetch_registries.py +18 -8
  77. wandb/sdk/{projects → artifacts}/_generated/fetch_registry.py +4 -4
  78. wandb/sdk/artifacts/_generated/fragments.py +183 -333
  79. wandb/sdk/artifacts/_generated/input_types.py +133 -7
  80. wandb/sdk/artifacts/_generated/link_artifact.py +5 -5
  81. wandb/sdk/artifacts/_generated/operations.py +1053 -548
  82. wandb/sdk/artifacts/_generated/project_artifact_collection.py +9 -77
  83. wandb/sdk/artifacts/_generated/project_artifact_collections.py +21 -9
  84. wandb/sdk/artifacts/_generated/project_artifact_type.py +3 -3
  85. wandb/sdk/artifacts/_generated/project_artifact_types.py +19 -6
  86. wandb/sdk/artifacts/_generated/project_artifacts.py +7 -8
  87. wandb/sdk/artifacts/_generated/registry_collections.py +21 -9
  88. wandb/sdk/artifacts/_generated/registry_versions.py +20 -9
  89. wandb/sdk/artifacts/_generated/rename_registry.py +25 -0
  90. wandb/sdk/artifacts/_generated/run_input_artifacts.py +5 -9
  91. wandb/sdk/artifacts/_generated/run_output_artifacts.py +5 -9
  92. wandb/sdk/artifacts/_generated/type_info.py +2 -2
  93. wandb/sdk/artifacts/_generated/unlink_artifact.py +3 -5
  94. wandb/sdk/artifacts/_generated/update_artifact.py +3 -3
  95. wandb/sdk/artifacts/_generated/update_artifact_collection_type.py +28 -0
  96. wandb/sdk/artifacts/_generated/update_artifact_portfolio.py +7 -16
  97. wandb/sdk/artifacts/_generated/update_artifact_sequence.py +7 -16
  98. wandb/sdk/artifacts/_generated/upsert_registry.py +25 -0
  99. wandb/sdk/artifacts/_gqlutils.py +170 -6
  100. wandb/sdk/artifacts/_models/__init__.py +9 -0
  101. wandb/sdk/artifacts/_models/artifact_collection.py +109 -0
  102. wandb/sdk/artifacts/_models/manifest.py +26 -0
  103. wandb/sdk/artifacts/_models/pagination.py +26 -0
  104. wandb/sdk/artifacts/_models/registry.py +100 -0
  105. wandb/sdk/artifacts/_validators.py +45 -27
  106. wandb/sdk/artifacts/artifact.py +220 -215
  107. wandb/sdk/artifacts/artifact_file_cache.py +1 -1
  108. wandb/sdk/artifacts/artifact_manifest.py +37 -32
  109. wandb/sdk/artifacts/artifact_manifest_entry.py +80 -125
  110. wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py +43 -61
  111. wandb/sdk/artifacts/storage_handlers/gcs_handler.py +8 -6
  112. wandb/sdk/data_types/image.py +2 -2
  113. wandb/sdk/interface/interface.py +72 -64
  114. wandb/sdk/interface/interface_queue.py +27 -18
  115. wandb/sdk/interface/interface_shared.py +61 -23
  116. wandb/sdk/interface/interface_sock.py +9 -5
  117. wandb/sdk/internal/_generated/server_features_query.py +4 -4
  118. wandb/sdk/launch/inputs/schema.py +13 -10
  119. wandb/sdk/lib/apikey.py +8 -12
  120. wandb/sdk/lib/asyncio_compat.py +1 -1
  121. wandb/sdk/lib/asyncio_manager.py +5 -5
  122. wandb/sdk/lib/console_capture.py +38 -30
  123. wandb/sdk/lib/progress.py +159 -64
  124. wandb/sdk/lib/retry.py +3 -2
  125. wandb/sdk/lib/service/service_connection.py +2 -2
  126. wandb/sdk/lib/wb_logging.py +2 -1
  127. wandb/sdk/mailbox/mailbox.py +1 -1
  128. wandb/sdk/wandb_init.py +10 -13
  129. wandb/sdk/wandb_run.py +9 -46
  130. wandb/sdk/wandb_settings.py +102 -19
  131. {wandb-0.22.2.dist-info → wandb-0.22.3.dist-info}/METADATA +2 -1
  132. {wandb-0.22.2.dist-info → wandb-0.22.3.dist-info}/RECORD +135 -134
  133. wandb/sdk/artifacts/_generated/artifact_via_membership_by_name.py +0 -26
  134. wandb/sdk/artifacts/_generated/create_artifact_collection_tag_assignments.py +0 -36
  135. wandb/sdk/artifacts/_generated/delete_artifact_collection_tag_assignments.py +0 -25
  136. wandb/sdk/artifacts/_generated/move_artifact_collection.py +0 -35
  137. wandb/sdk/projects/_generated/__init__.py +0 -26
  138. wandb/sdk/projects/_generated/delete_project.py +0 -22
  139. wandb/sdk/projects/_generated/enums.py +0 -4
  140. wandb/sdk/projects/_generated/fragments.py +0 -41
  141. wandb/sdk/projects/_generated/input_types.py +0 -13
  142. wandb/sdk/projects/_generated/operations.py +0 -88
  143. wandb/sdk/projects/_generated/rename_project.py +0 -27
  144. wandb/sdk/projects/_generated/upsert_registry_project.py +0 -27
  145. {wandb-0.22.2.dist-info → wandb-0.22.3.dist-info}/WHEEL +0 -0
  146. {wandb-0.22.2.dist-info → wandb-0.22.3.dist-info}/entry_points.txt +0 -0
  147. {wandb-0.22.2.dist-info → wandb-0.22.3.dist-info}/licenses/LICENSE +0 -0
@@ -2,29 +2,34 @@ from __future__ import annotations
2
2
 
3
3
  from typing import TYPE_CHECKING, Any, Literal
4
4
 
5
+ from pydantic import PositiveInt
5
6
  from wandb_gql import gql
6
7
 
7
8
  import wandb
8
9
  from wandb._analytics import tracked
9
- from wandb.proto.wandb_internal_pb2 import ServerFeature
10
- from wandb.sdk.artifacts._validators import REGISTRY_PREFIX, validate_project_name
11
- from wandb.sdk.internal.internal_api import Api as InternalApi
12
- from wandb.sdk.projects._generated import (
13
- DELETE_PROJECT_GQL,
10
+ from wandb.proto import wandb_internal_pb2 as pb
11
+ from wandb.sdk.artifacts._generated import (
12
+ DELETE_REGISTRY_GQL,
14
13
  FETCH_REGISTRY_GQL,
15
- RENAME_PROJECT_GQL,
16
- UPSERT_REGISTRY_PROJECT_GQL,
17
- DeleteProject,
18
- RenameProject,
19
- UpsertRegistryProject,
14
+ RENAME_REGISTRY_GQL,
15
+ UPSERT_REGISTRY_GQL,
16
+ DeleteRegistry,
17
+ FetchRegistry,
18
+ RegistryFragment,
19
+ RenameProjectInput,
20
+ RenameRegistry,
21
+ UpsertModelInput,
22
+ UpsertRegistry,
20
23
  )
24
+ from wandb.sdk.artifacts._gqlutils import server_supports
25
+ from wandb.sdk.artifacts._models import RegistryData
26
+ from wandb.sdk.artifacts._validators import REGISTRY_PREFIX, validate_project_name
21
27
 
22
28
  from ._freezable_list import AddOnlyArtifactTypesList
23
29
  from ._utils import (
30
+ Visibility,
24
31
  fetch_org_entity_from_organization,
25
- format_gql_artifact_types_input,
26
- gql_to_registry_visibility,
27
- registry_visibility_to_gql,
32
+ prepare_artifact_types_input,
28
33
  )
29
34
  from .registries_search import Collections, Versions
30
35
 
@@ -35,86 +40,86 @@ if TYPE_CHECKING:
35
40
  class Registry:
36
41
  """A single registry in the Registry."""
37
42
 
43
+ _saved: RegistryData
44
+ """The saved registry data as last fetched from the W&B server."""
45
+
46
+ _current: RegistryData
47
+ """The local, editable registry data."""
48
+
38
49
  def __init__(
39
50
  self,
40
51
  client: Client,
41
52
  organization: str,
42
53
  entity: str,
43
54
  name: str,
44
- attrs: dict[str, Any] | None = None,
55
+ attrs: RegistryFragment | None = None,
45
56
  ):
46
57
  self.client = client
47
- self._name = name
48
- self._saved_name = name
49
- self._entity = entity
50
- self._organization = organization
51
- if attrs is not None:
52
- self._update_attributes(attrs)
53
58
 
54
- def _update_attributes(self, attrs: dict[str, Any]) -> None:
55
- """Helper method to update instance attributes from a dictionary."""
56
- self._id = attrs.get("id", "")
57
- if self._id is None:
58
- raise ValueError(f"Registry {self.name}'s id is not found")
59
+ if attrs is None:
60
+ # FIXME: This is awkward and bypasses validation which seems shaky.
61
+ # Reconsider the init signature of `Registry` so this isn't necessary?
62
+ draft = RegistryData.model_construct(
63
+ organization=organization, entity=entity, name=name
64
+ )
65
+ self._saved = draft
66
+ self._current = draft.model_copy(deep=True)
67
+ else:
68
+ self._update_attributes(attrs)
59
69
 
60
- self._description = attrs.get("description", "")
61
- self._allow_all_artifact_types = attrs.get(
62
- "allowAllArtifactTypesInRegistry", False
63
- )
64
- self._artifact_types = AddOnlyArtifactTypesList(
65
- t["node"]["name"] for t in attrs.get("artifactTypes", {}).get("edges", [])
66
- )
67
- self._created_at = attrs.get("createdAt", "")
68
- self._updated_at = attrs.get("updatedAt", "")
69
- self._visibility = gql_to_registry_visibility(attrs.get("access", ""))
70
+ def _update_attributes(self, fragment: RegistryFragment) -> None:
71
+ """Internal helper method to update instance attributes from GraphQL fragment data."""
72
+ saved = RegistryData.from_fragment(fragment)
73
+ self._saved = saved
74
+ self._current = saved.model_copy(deep=True)
70
75
 
71
76
  @property
72
77
  def full_name(self) -> str:
73
78
  """Full name of the registry including the `wandb-registry-` prefix."""
74
- return f"wandb-registry-{self.name}"
79
+ return self._current.full_name
75
80
 
76
81
  @property
77
82
  def name(self) -> str:
78
83
  """Name of the registry without the `wandb-registry-` prefix."""
79
- return self._name
84
+ return self._current.name
80
85
 
81
86
  @name.setter
82
87
  def name(self, value: str):
83
- self._name = value
88
+ self._current.name = value
84
89
 
85
90
  @property
86
91
  def entity(self) -> str:
87
92
  """Organization entity of the registry."""
88
- return self._entity
93
+ return self._current.entity
89
94
 
90
95
  @property
91
96
  def organization(self) -> str:
92
97
  """Organization name of the registry."""
93
- return self._organization
98
+ return self._current.organization
94
99
 
95
100
  @property
96
101
  def description(self) -> str:
97
102
  """Description of the registry."""
98
- return self._description
103
+ return self._current.description
99
104
 
100
105
  @description.setter
101
106
  def description(self, value: str):
102
107
  """Set the description of the registry."""
103
- self._description = value
108
+ self._current.description = value
104
109
 
105
110
  @property
106
- def allow_all_artifact_types(self):
111
+ def allow_all_artifact_types(self) -> bool:
107
112
  """Returns whether all artifact types are allowed in the registry.
108
113
 
109
114
  If `True` then artifacts of any type can be added to this registry.
110
115
  If `False` then artifacts are restricted to the types in `artifact_types` for this registry.
111
116
  """
112
- return self._allow_all_artifact_types
117
+ return self._current.allow_all_artifact_types
113
118
 
114
119
  @allow_all_artifact_types.setter
115
- def allow_all_artifact_types(self, value: bool):
120
+ def allow_all_artifact_types(self, value: bool) -> None:
116
121
  """Set whether all artifact types are allowed in the registry."""
117
- self._allow_all_artifact_types = value
122
+ self._current.allow_all_artifact_types = value
118
123
 
119
124
  @property
120
125
  def artifact_types(self) -> AddOnlyArtifactTypesList:
@@ -141,20 +146,20 @@ class Registry:
141
146
  ) # Types can only be removed if it has not been saved yet
142
147
  ```
143
148
  """
144
- return self._artifact_types
149
+ return self._current.artifact_types
145
150
 
146
151
  @property
147
152
  def created_at(self) -> str:
148
153
  """Timestamp of when the registry was created."""
149
- return self._created_at
154
+ return self._current.created_at
150
155
 
151
156
  @property
152
157
  def updated_at(self) -> str:
153
158
  """Timestamp of when the registry was last updated."""
154
- return self._updated_at
159
+ return self._current.updated_at
155
160
 
156
161
  @property
157
- def path(self):
162
+ def path(self) -> list[str]:
158
163
  return [self.entity, self.full_name]
159
164
 
160
165
  @property
@@ -168,7 +173,7 @@ class Registry:
168
173
  - "restricted": Only invited members via the UI can access this registry.
169
174
  Public sharing is disabled.
170
175
  """
171
- return self._visibility
176
+ return self._current.visibility.name
172
177
 
173
178
  @visibility.setter
174
179
  def visibility(self, value: Literal["organization", "restricted"]):
@@ -181,23 +186,34 @@ class Registry:
181
186
  - "restricted": Only invited members via the UI can access this registry.
182
187
  Public sharing is disabled.
183
188
  """
184
- self._visibility = value
189
+ self._current.visibility = value
185
190
 
186
191
  @tracked
187
- def collections(self, filter: dict[str, Any] | None = None) -> Collections:
192
+ def collections(
193
+ self, filter: dict[str, Any] | None = None, per_page: PositiveInt = 100
194
+ ) -> Collections:
188
195
  """Returns the collections belonging to the registry."""
189
- registry_filter = {
190
- "name": self.full_name,
191
- }
192
- return Collections(self.client, self.organization, registry_filter, filter)
196
+ return Collections(
197
+ client=self.client,
198
+ organization=self.organization,
199
+ registry_filter={"name": self.full_name},
200
+ collection_filter=filter,
201
+ per_page=per_page,
202
+ )
193
203
 
194
204
  @tracked
195
- def versions(self, filter: dict[str, Any] | None = None) -> Versions:
205
+ def versions(
206
+ self, filter: dict[str, Any] | None = None, per_page: PositiveInt = 100
207
+ ) -> Versions:
196
208
  """Returns the versions belonging to the registry."""
197
- registry_filter = {
198
- "name": self.full_name,
199
- }
200
- return Versions(self.client, self.organization, registry_filter, None, filter)
209
+ return Versions(
210
+ client=self.client,
211
+ organization=self.organization,
212
+ registry_filter={"name": self.full_name},
213
+ collection_filter=None,
214
+ artifact_filter=filter,
215
+ per_page=per_page,
216
+ )
201
217
 
202
218
  @classmethod
203
219
  @tracked
@@ -230,141 +246,142 @@ class Registry:
230
246
  ValueError: If a registry with the same name already exists in the
231
247
  organization or if the creation fails.
232
248
  """
233
- org_entity = fetch_org_entity_from_organization(client, organization)
234
- full_name = REGISTRY_PREFIX + name
235
- validate_project_name(full_name)
236
- accepted_artifact_types = []
237
- if artifact_types:
238
- accepted_artifact_types = format_gql_artifact_types_input(artifact_types)
239
- visibility_value = registry_visibility_to_gql(visibility)
240
- registry_creation_error = (
249
+ failed_msg = (
241
250
  f"Failed to create registry {name!r} in organization {organization!r}."
242
251
  )
252
+
253
+ org_entity = fetch_org_entity_from_organization(client, organization)
254
+
255
+ gql_op = gql(UPSERT_REGISTRY_GQL)
256
+ gql_input = UpsertModelInput(
257
+ description=description,
258
+ entity_name=org_entity,
259
+ name=validate_project_name(f"{REGISTRY_PREFIX}{name}"),
260
+ access=Visibility.from_python(visibility).value,
261
+ allow_all_artifact_types_in_registry=not artifact_types,
262
+ artifact_types=prepare_artifact_types_input(artifact_types),
263
+ )
264
+ gql_vars = {"input": gql_input.model_dump()}
243
265
  try:
244
- response = client.execute(
245
- gql(UPSERT_REGISTRY_PROJECT_GQL),
246
- variable_values={
247
- "description": description,
248
- "entityName": org_entity,
249
- "name": full_name,
250
- "access": visibility_value,
251
- "allowAllArtifactTypesInRegistry": not accepted_artifact_types,
252
- "artifactTypes": accepted_artifact_types,
253
- },
254
- )
255
- except Exception:
256
- raise ValueError(registry_creation_error)
257
- if not response["upsertModel"]["inserted"]:
258
- raise ValueError(registry_creation_error)
266
+ data = client.execute(gql_op, variable_values=gql_vars)
267
+ result = UpsertRegistry.model_validate(data).upsert_model
268
+ except Exception as e:
269
+ raise ValueError(failed_msg) from e
270
+ if not (result and result.inserted and (registry_project := result.project)):
271
+ raise ValueError(failed_msg)
259
272
 
260
273
  return Registry(
261
274
  client,
262
- organization,
263
- org_entity,
264
- name,
265
- response["upsertModel"]["project"],
275
+ organization=organization,
276
+ entity=org_entity,
277
+ name=name,
278
+ attrs=registry_project,
266
279
  )
267
280
 
268
281
  @tracked
269
282
  def delete(self) -> None:
270
283
  """Delete the registry. This is irreversible."""
284
+ failed_msg = f"Failed to delete registry {self.name!r} in organization {self.organization!r}"
285
+
286
+ gql_op = gql(DELETE_REGISTRY_GQL)
287
+ gql_vars = {"id": self._saved.id}
271
288
  try:
272
- response = self.client.execute(
273
- gql(DELETE_PROJECT_GQL), variable_values={"id": self._id}
274
- )
275
- result = DeleteProject.model_validate(response)
276
- except Exception:
277
- raise ValueError(
278
- f"Failed to delete registry: {self.name!r} in organization: {self.organization!r}"
279
- )
280
- if not result.delete_model.success:
281
- raise ValueError(
282
- f"Failed to delete registry: {self.name!r} in organization: {self.organization!r}"
283
- )
289
+ data = self.client.execute(gql_op, variable_values=gql_vars)
290
+ result = DeleteRegistry.model_validate(data).delete_model
291
+ except Exception as e:
292
+ raise ValueError(failed_msg) from e
293
+ if not (result and result.success):
294
+ raise ValueError(failed_msg)
284
295
 
285
296
  @tracked
286
297
  def load(self) -> None:
287
298
  """Load the registry attributes from the backend to reflect the latest saved state."""
288
- load_failure_message = (
289
- f"Failed to load registry {self.name!r} "
290
- f"in organization {self.organization!r}."
291
- )
299
+ failed_msg = f"Failed to load registry {self.name!r} in organization {self.organization!r}."
300
+
301
+ gql_op = gql(FETCH_REGISTRY_GQL)
302
+ gql_vars = {"name": self.full_name, "entity": self.entity}
292
303
  try:
293
- response = self.client.execute(
294
- gql(FETCH_REGISTRY_GQL),
295
- variable_values={
296
- "name": self.full_name,
297
- "entityName": self.entity,
298
- },
299
- )
300
- except Exception:
301
- raise ValueError(load_failure_message)
302
- if response["entity"] is None:
303
- raise ValueError(load_failure_message)
304
- self.attrs = response["entity"]["project"]
305
- if self.attrs is None:
306
- raise ValueError(load_failure_message)
307
- self._update_attributes(self.attrs)
304
+ data = self.client.execute(gql_op, variable_values=gql_vars)
305
+ result = FetchRegistry.model_validate(data)
306
+ except Exception as e:
307
+ raise ValueError(failed_msg) from e
308
+
309
+ if not ((entity := result.entity) and (registry_project := entity.project)):
310
+ raise ValueError(failed_msg)
311
+
312
+ self._update_attributes(registry_project)
308
313
 
309
314
  @tracked
310
315
  def save(self) -> None:
311
316
  """Save registry attributes to the backend."""
312
- if not InternalApi()._server_supports(
313
- ServerFeature.INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION
317
+ if not server_supports(
318
+ self.client, pb.INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION
314
319
  ):
315
320
  raise RuntimeError(
316
- "saving the registry is not enabled on this wandb server version. "
321
+ "Saving the registry is not enabled on this wandb server version. "
317
322
  "Please upgrade your server version or contact support at support@wandb.com."
318
323
  )
319
324
 
320
- if self._no_updating_registry_types():
325
+ # If `artifact_types.draft` has items, it means the user has added artifact types that aren't saved yet.
326
+ if (
327
+ new_artifact_types := self.artifact_types.draft
328
+ ) and self.allow_all_artifact_types:
321
329
  raise ValueError(
322
330
  f"Cannot update artifact types when `allows_all_artifact_types` is {True!r}. Set it to {False!r} first."
323
331
  )
324
332
 
325
- validate_project_name(self.full_name)
326
- visibility_value = registry_visibility_to_gql(self.visibility)
327
- newly_added_types = format_gql_artifact_types_input(self.artifact_types.draft)
328
- registry_save_error = f"Failed to save and update registry: {self.name} in organization: {self.organization}"
329
- full_saved_name = f"{REGISTRY_PREFIX}{self._saved_name}"
333
+ failed_msg = f"Failed to save registry {self.name!r} in organization {self.organization!r}"
334
+
335
+ old_project_name = validate_project_name(self._saved.full_name)
336
+ new_project_name = validate_project_name(self._current.full_name)
337
+
338
+ upsert_op = gql(UPSERT_REGISTRY_GQL)
339
+ upsert_input = UpsertModelInput(
340
+ description=self.description,
341
+ entity_name=self.entity,
342
+ name=old_project_name,
343
+ access=self._current.visibility.value,
344
+ allow_all_artifact_types_in_registry=self.allow_all_artifact_types,
345
+ artifact_types=prepare_artifact_types_input(new_artifact_types),
346
+ )
347
+ upsert_vars = {"input": upsert_input.model_dump()}
330
348
  try:
331
- response = self.client.execute(
332
- gql(UPSERT_REGISTRY_PROJECT_GQL),
333
- variable_values={
334
- "description": self.description,
335
- "entityName": self.entity,
336
- "name": full_saved_name, # this makes it so we are updating the original registry in case the name has changed
337
- "access": visibility_value,
338
- "allowAllArtifactTypesInRegistry": self.allow_all_artifact_types,
339
- "artifactTypes": newly_added_types,
340
- },
341
- )
342
- result = UpsertRegistryProject.model_validate(response)
343
- except Exception:
344
- raise ValueError(registry_save_error)
345
- if result.upsert_model.inserted:
349
+ data = self.client.execute(upsert_op, variable_values=upsert_vars)
350
+ result = UpsertRegistry.model_validate(data).upsert_model
351
+ except Exception as e:
352
+ raise ValueError(failed_msg) from e
353
+
354
+ if result and result.inserted:
346
355
  # This is not suppose trigger unless the user has messed with the `_saved_name` variable
347
356
  wandb.termlog(
348
357
  f"Created registry {self.name!r} in organization {self.organization!r} on save"
349
358
  )
350
- self._update_attributes(response["upsertModel"]["project"])
359
+
360
+ if not (result and (registry_project := result.project)):
361
+ raise ValueError(failed_msg)
362
+
363
+ self._update_attributes(registry_project)
351
364
 
352
365
  # Update the name of the registry if it has changed
353
- if self._saved_name != self.name:
354
- response = self.client.execute(
355
- gql(RENAME_PROJECT_GQL),
356
- variable_values={
357
- "entityName": self.entity,
358
- "oldProjectName": full_saved_name,
359
- "newProjectName": self.full_name,
360
- },
366
+ if old_project_name != new_project_name:
367
+ rename_op = gql(RENAME_REGISTRY_GQL)
368
+ rename_input = RenameProjectInput(
369
+ entity_name=self.entity,
370
+ old_project_name=old_project_name,
371
+ new_project_name=new_project_name,
361
372
  )
362
- result = RenameProject.model_validate(response)
363
- self._saved_name = self.name
364
- if result.rename_project.inserted:
373
+ rename_vars = {"input": rename_input.model_dump()}
374
+ data = self.client.execute(rename_op, variable_values=rename_vars)
375
+ result = RenameRegistry.model_validate(data).rename_project
376
+ if not (result and (registry_project := result.project)):
377
+ raise ValueError(failed_msg)
378
+
379
+ if result.inserted:
365
380
  # This is not suppose trigger unless the user has messed with the `_saved_name` variable
366
381
  wandb.termlog(f"Created new registry {self.name!r} on save")
367
382
 
383
+ self._update_attributes(registry_project)
384
+
368
385
  def _no_updating_registry_types(self) -> bool:
369
386
  # artifact types draft means user assigned types to add that are not yet saved
370
387
  return len(self.artifact_types.draft) > 0 and self.allow_all_artifact_types
wandb/apis/public/runs.py CHANGED
@@ -745,21 +745,24 @@ class Run(Attrs):
745
745
  withRuns=False,
746
746
  )
747
747
 
748
- if not self._is_loaded:
748
+ if not self._is_loaded or force:
749
749
  # Always set _project_internal_id if projectId is available, regardless of fragment type
750
750
  if "projectId" in self._attrs:
751
751
  self._project_internal_id = int(self._attrs["projectId"])
752
752
  else:
753
753
  self._project_internal_id = None
754
754
 
755
- # Only call _load_from_attrs when using the full fragment or when the fields are actually present
755
+ # Always call _load_from_attrs when using the full fragment or when the fields are actually present
756
756
  if fragment_name == RUN_FRAGMENT_NAME or (
757
757
  "config" in self._attrs
758
758
  or "summaryMetrics" in self._attrs
759
759
  or "systemMetrics" in self._attrs
760
760
  ):
761
761
  self._load_from_attrs()
762
- self._is_loaded = True
762
+
763
+ # Only mark as loaded for lightweight fragments, not full fragments
764
+ if fragment_name == LIGHTWEIGHT_RUN_FRAGMENT_NAME:
765
+ self._is_loaded = True
763
766
 
764
767
  return self._attrs
765
768
 
@@ -1305,7 +1308,13 @@ class Run(Attrs):
1305
1308
  """Get run config. Auto-loads full data if in lazy mode."""
1306
1309
  if self._lazy and not self._full_data_loaded and "config" not in self._attrs:
1307
1310
  self.load_full_data()
1308
- return self._attrs.get("config", {})
1311
+
1312
+ # Ensure config is always converted to dict (defensive against conversion issues)
1313
+ config_value = self._attrs.get("config", {})
1314
+ # _convert_to_dict handles dict inputs (noop) and converts str/bytes/bytearray to dict
1315
+ config_value = _convert_to_dict(config_value)
1316
+ self._attrs["config"] = config_value
1317
+ return config_value
1309
1318
 
1310
1319
  @property
1311
1320
  def summary(self):
@@ -1332,7 +1341,13 @@ class Run(Attrs):
1332
1341
  and "systemMetrics" not in self._attrs
1333
1342
  ):
1334
1343
  self.load_full_data()
1335
- return self._attrs.get("systemMetrics", {})
1344
+
1345
+ # Ensure systemMetrics is always converted to dict (defensive against conversion issues)
1346
+ system_metrics_value = self._attrs.get("systemMetrics", {})
1347
+ # _convert_to_dict handles dict inputs (noop) and converts str/bytes/bytearray to dict
1348
+ system_metrics_value = _convert_to_dict(system_metrics_value)
1349
+ self._attrs["systemMetrics"] = system_metrics_value
1350
+ return system_metrics_value
1336
1351
 
1337
1352
  @property
1338
1353
  def summary_metrics(self):
@@ -1343,7 +1358,13 @@ class Run(Attrs):
1343
1358
  and "summaryMetrics" not in self._attrs
1344
1359
  ):
1345
1360
  self.load_full_data()
1346
- return self._attrs.get("summaryMetrics", {})
1361
+
1362
+ # Ensure summaryMetrics is always converted to dict (defensive against conversion issues)
1363
+ summary_metrics_value = self._attrs.get("summaryMetrics", {})
1364
+ # _convert_to_dict handles dict inputs (noop) and converts str/bytes/bytearray to dict
1365
+ summary_metrics_value = _convert_to_dict(summary_metrics_value)
1366
+ self._attrs["summaryMetrics"] = summary_metrics_value
1367
+ return summary_metrics_value
1347
1368
 
1348
1369
  @property
1349
1370
  def rawconfig(self):
@@ -6,7 +6,9 @@ from typing import Any, Iterable, Mapping
6
6
  from urllib.parse import urlparse
7
7
 
8
8
  from wandb_gql import gql
9
+ from wandb_graphql import TypeInfo
9
10
  from wandb_graphql.language import ast, visitor
11
+ from wandb_graphql.validation.validation import ValidationContext
10
12
 
11
13
  from wandb._iterutils import one
12
14
  from wandb.sdk.artifacts._validators import is_artifact_registry_project
@@ -112,11 +114,6 @@ def fetch_org_from_settings_or_entity(
112
114
  class _GQLCompatRewriter(visitor.Visitor):
113
115
  """GraphQL AST visitor to rewrite queries/mutations to be compatible with older server versions."""
114
116
 
115
- omit_variables: set[str]
116
- omit_fragments: set[str]
117
- omit_fields: set[str]
118
- rename_fields: dict[str, str]
119
-
120
117
  def __init__(
121
118
  self,
122
119
  omit_variables: Iterable[str] | None = None,
@@ -129,25 +126,52 @@ class _GQLCompatRewriter(visitor.Visitor):
129
126
  self.omit_fields = set(omit_fields or ())
130
127
  self.rename_fields = dict(rename_fields or {})
131
128
 
132
- def enter_VariableDefinition(self, node: ast.VariableDefinition, *_, **__) -> Any: # noqa: N802
133
- if node.variable.name.value in self.omit_variables:
129
+ def leave_Document(self, node: ast.Document, *_, **__) -> Any: # noqa: N802
130
+ # After rewriting the GQL document, prune "orphan" (unused) fragment definitions.
131
+ # Note: The ValidationContext doesn't require a schema here, as we only use it to check for reachable fragments.
132
+ ctx = ValidationContext(schema=None, ast=node, type_info=TypeInfo(schema=None))
133
+ operation_defns = {
134
+ dfn for dfn in node.definitions if isinstance(dfn, ast.OperationDefinition)
135
+ }
136
+ used_fragment_defns = {
137
+ frag
138
+ for op in operation_defns
139
+ for frag in ctx.get_recursively_referenced_fragments(op)
140
+ }
141
+ # Preserve original defintion order
142
+ allowed_defns = operation_defns | used_fragment_defns
143
+ node.definitions = [dfn for dfn in node.definitions if (dfn in allowed_defns)]
144
+
145
+ def enter_Variable(self, node: ast.Variable, *_, **__) -> Any: # noqa: N802
146
+ if node.name.value in self.omit_variables:
134
147
  return visitor.REMOVE
135
148
 
136
- def enter_ObjectField(self, node: ast.ObjectField, *_, **__) -> Any: # noqa: N802
137
- # For context, note that e.g.:
149
+ def leave_VariableDefinition(self, node: ast.VariableDefinition, *_, **__) -> Any: # noqa: N802
150
+ # For context, consider the `$varName: String` variable definition below:
151
+ # (..., $varName: String, ...)
152
+ #
153
+ # On ENTERING, the AST looks like:
154
+ # VariableDefinition(variable=Variable(name=Name(value='varName')), ...)
138
155
  #
139
- # {description: $description
140
- # ...}
156
+ # On LEAVING, if `$varName` was removed, the AST looks like:
157
+ # VariableDefinition(variable=REMOVE, ...)
158
+ if node.variable is visitor.REMOVE:
159
+ return visitor.REMOVE
160
+
161
+ def leave_ObjectField(self, node: ast.ObjectField, *_, **__) -> Any: # noqa: N802
162
+ # For context, consider `argName: $varName` in the input args below:
163
+ # input: {..., argName: $varName, ...}
141
164
  #
142
- # Is parsed as:
165
+ # On ENTERING, the AST for `argName: $varName` looks like:
166
+ # ObjectField(
167
+ # name=Name(value='argName'), value=Variable(name=Name(value='varName')),
168
+ # )
143
169
  #
144
- # ObjectValue(fields=[
145
- # ObjectField(name=Name(value='description'), value=Variable(name=Name(value='description'))),
146
- # ...])
147
- if (
148
- isinstance(var := node.value, ast.Variable)
149
- and var.name.value in self.omit_variables
150
- ):
170
+ # On LEAVING, if `$varName` was removed, the AST looks like:
171
+ # ObjectField(
172
+ # name=Name(value='argName'), value=REMOVE,
173
+ # )
174
+ if node.value is visitor.REMOVE:
151
175
  return visitor.REMOVE
152
176
 
153
177
  def enter_Argument(self, node: ast.Argument, *_, **__) -> Any: # noqa: N802
@@ -167,7 +191,6 @@ class _GQLCompatRewriter(visitor.Visitor):
167
191
  return visitor.REMOVE
168
192
  if new_name := self.rename_fields.get(node.name.value):
169
193
  node.name.value = new_name
170
- return node
171
194
 
172
195
  def leave_Field(self, node: ast.Field, *_, **__) -> Any: # noqa: N802
173
196
  # If the field had a selection set, but now it's empty, remove the field entirely
@@ -5,12 +5,12 @@ from __future__ import annotations
5
5
 
6
6
  from typing import Optional
7
7
 
8
- from wandb._pydantic import GQLBase
8
+ from wandb._pydantic import GQLResult
9
9
 
10
10
  from .fragments import CreateAutomationResult
11
11
 
12
12
 
13
- class CreateAutomation(GQLBase):
13
+ class CreateAutomation(GQLResult):
14
14
  result: Optional[CreateAutomationResult]
15
15
 
16
16