zenml-nightly 0.70.0.dev20241129__py3-none-any.whl → 0.70.0.dev20241130__py3-none-any.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 (47) hide show
  1. zenml/VERSION +1 -1
  2. zenml/artifacts/artifact_config.py +0 -14
  3. zenml/artifacts/utils.py +50 -29
  4. zenml/cli/__init__.py +15 -0
  5. zenml/cli/base.py +4 -4
  6. zenml/cli/server.py +1 -1
  7. zenml/cli/stack.py +0 -3
  8. zenml/cli/stack_components.py +0 -1
  9. zenml/cli/utils.py +0 -5
  10. zenml/client.py +8 -18
  11. zenml/model/model.py +8 -4
  12. zenml/model/utils.py +18 -16
  13. zenml/models/__init__.py +6 -0
  14. zenml/models/v2/core/artifact_version.py +6 -3
  15. zenml/models/v2/core/component.py +0 -22
  16. zenml/models/v2/core/model_version.py +6 -3
  17. zenml/models/v2/core/pipeline_run.py +16 -4
  18. zenml/models/v2/core/run_metadata.py +30 -9
  19. zenml/models/v2/core/step_run.py +5 -3
  20. zenml/models/v2/misc/run_metadata.py +38 -0
  21. zenml/orchestrators/input_utils.py +19 -6
  22. zenml/orchestrators/publish_utils.py +11 -4
  23. zenml/orchestrators/step_launcher.py +2 -2
  24. zenml/orchestrators/step_run_utils.py +3 -0
  25. zenml/orchestrators/step_runner.py +7 -8
  26. zenml/orchestrators/utils.py +0 -24
  27. zenml/pipelines/run_utils.py +3 -4
  28. zenml/steps/utils.py +6 -2
  29. zenml/utils/metadata_utils.py +186 -153
  30. zenml/utils/string_utils.py +15 -8
  31. zenml/zen_server/routers/workspaces_endpoints.py +19 -19
  32. zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py +36 -0
  33. zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py +135 -0
  34. zenml/zen_stores/schemas/__init__.py +5 -1
  35. zenml/zen_stores/schemas/artifact_schemas.py +10 -10
  36. zenml/zen_stores/schemas/component_schemas.py +0 -3
  37. zenml/zen_stores/schemas/model_schemas.py +13 -11
  38. zenml/zen_stores/schemas/pipeline_run_schemas.py +32 -15
  39. zenml/zen_stores/schemas/run_metadata_schemas.py +66 -31
  40. zenml/zen_stores/schemas/step_run_schemas.py +13 -12
  41. zenml/zen_stores/schemas/utils.py +47 -3
  42. zenml/zen_stores/sql_zen_store.py +102 -23
  43. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/METADATA +1 -1
  44. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/RECORD +47 -44
  45. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/LICENSE +0 -0
  46. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/WHEEL +0 -0
  47. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/entry_points.txt +0 -0
@@ -13,28 +13,30 @@
13
13
  # permissions and limitations under the License.
14
14
  """Utility functions to handle metadata for ZenML entities."""
15
15
 
16
- import contextlib
17
- from typing import Dict, Optional, Union, overload
16
+ from typing import Dict, List, Optional, Union, overload
18
17
  from uuid import UUID
19
18
 
20
19
  from zenml.client import Client
21
- from zenml.enums import MetadataResourceTypes
20
+ from zenml.enums import MetadataResourceTypes, ModelStages
22
21
  from zenml.logger import get_logger
23
22
  from zenml.metadata.metadata_types import MetadataType
23
+ from zenml.models import RunMetadataResource
24
24
  from zenml.steps.step_context import get_step_context
25
25
 
26
26
  logger = get_logger(__name__)
27
27
 
28
28
 
29
29
  @overload
30
- def log_metadata(metadata: Dict[str, MetadataType]) -> None: ...
30
+ def log_metadata(
31
+ metadata: Dict[str, MetadataType],
32
+ ) -> None: ...
31
33
 
32
34
 
33
35
  @overload
34
36
  def log_metadata(
35
37
  *,
36
38
  metadata: Dict[str, MetadataType],
37
- artifact_version_id: UUID,
39
+ step_id: UUID,
38
40
  ) -> None: ...
39
41
 
40
42
 
@@ -42,8 +44,8 @@ def log_metadata(
42
44
  def log_metadata(
43
45
  *,
44
46
  metadata: Dict[str, MetadataType],
45
- artifact_name: str,
46
- artifact_version: Optional[str] = None,
47
+ step_name: str,
48
+ run_id_name_or_prefix: Union[UUID, str],
47
49
  ) -> None: ...
48
50
 
49
51
 
@@ -51,7 +53,7 @@ def log_metadata(
51
53
  def log_metadata(
52
54
  *,
53
55
  metadata: Dict[str, MetadataType],
54
- model_version_id: UUID,
56
+ run_id_name_or_prefix: Union[UUID, str],
55
57
  ) -> None: ...
56
58
 
57
59
 
@@ -59,8 +61,7 @@ def log_metadata(
59
61
  def log_metadata(
60
62
  *,
61
63
  metadata: Dict[str, MetadataType],
62
- model_name: str,
63
- model_version: str,
64
+ artifact_version_id: UUID,
64
65
  ) -> None: ...
65
66
 
66
67
 
@@ -68,7 +69,8 @@ def log_metadata(
68
69
  def log_metadata(
69
70
  *,
70
71
  metadata: Dict[str, MetadataType],
71
- step_id: UUID,
72
+ artifact_name: str,
73
+ artifact_version: Optional[str] = None,
72
74
  ) -> None: ...
73
75
 
74
76
 
@@ -76,33 +78,53 @@ def log_metadata(
76
78
  def log_metadata(
77
79
  *,
78
80
  metadata: Dict[str, MetadataType],
79
- run_id_name_or_prefix: Union[UUID, str],
81
+ infer_artifact: bool = False,
82
+ artifact_name: Optional[str] = None,
80
83
  ) -> None: ...
81
84
 
82
85
 
86
+ # Model Metadata
83
87
  @overload
84
88
  def log_metadata(
85
89
  *,
86
90
  metadata: Dict[str, MetadataType],
87
- step_name: str,
88
- run_id_name_or_prefix: Union[UUID, str],
91
+ model_version_id: UUID,
89
92
  ) -> None: ...
90
93
 
91
94
 
95
+ @overload
92
96
  def log_metadata(
97
+ *,
93
98
  metadata: Dict[str, MetadataType],
94
- # Parameters to manually log metadata for steps and runs
99
+ model_name: str,
100
+ model_version: Union[ModelStages, int, str],
101
+ ) -> None: ...
102
+
103
+
104
+ @overload
105
+ def log_metadata(
106
+ *,
107
+ metadata: Dict[str, MetadataType],
108
+ infer_model: bool = False,
109
+ ) -> None: ...
110
+
111
+
112
+ def log_metadata(
113
+ metadata: Dict[str, MetadataType],
114
+ # Steps and runs
95
115
  step_id: Optional[UUID] = None,
96
116
  step_name: Optional[str] = None,
97
117
  run_id_name_or_prefix: Optional[Union[UUID, str]] = None,
98
- # Parameters to manually log metadata for artifacts
118
+ # Artifacts
99
119
  artifact_version_id: Optional[UUID] = None,
100
120
  artifact_name: Optional[str] = None,
101
121
  artifact_version: Optional[str] = None,
102
- # Parameters to manually log metadata for models
122
+ infer_artifact: bool = False,
123
+ # Models
103
124
  model_version_id: Optional[UUID] = None,
104
125
  model_name: Optional[str] = None,
105
- model_version: Optional[str] = None,
126
+ model_version: Optional[Union[ModelStages, int, str]] = None,
127
+ infer_model: bool = False,
106
128
  ) -> None:
107
129
  """Logs metadata for various resource types in a generalized way.
108
130
 
@@ -114,9 +136,13 @@ def log_metadata(
114
136
  artifact_version_id: The ID of the artifact version
115
137
  artifact_name: The name of the artifact.
116
138
  artifact_version: The version of the artifact.
139
+ infer_artifact: Flag deciding whether the artifact version should be
140
+ inferred from the step context.
117
141
  model_version_id: The ID of the model version.
118
142
  model_name: The name of the model.
119
- model_version: The version of the model
143
+ model_version: The version of the model.
144
+ infer_model: Flag deciding whether the model version should be
145
+ inferred from the step context.
120
146
 
121
147
  Raises:
122
148
  ValueError: If no identifiers are provided and the function is not
@@ -124,141 +150,147 @@ def log_metadata(
124
150
  """
125
151
  client = Client()
126
152
 
127
- # If a step name is provided, we need a run_id_name_or_prefix and will log
128
- # metadata for the steps pipeline and model accordingly.
129
- if step_name is not None and run_id_name_or_prefix is not None:
130
- run_model = client.get_pipeline_run(
131
- name_id_or_prefix=run_id_name_or_prefix
132
- )
133
- step_model = run_model.steps[step_name]
153
+ resources: List[RunMetadataResource] = []
154
+ publisher_step_id = None
134
155
 
135
- client.create_run_metadata(
136
- metadata=metadata,
137
- resource_id=run_model.id,
138
- resource_type=MetadataResourceTypes.PIPELINE_RUN,
139
- )
140
- client.create_run_metadata(
141
- metadata=metadata,
142
- resource_id=step_model.id,
143
- resource_type=MetadataResourceTypes.STEP_RUN,
144
- )
145
- if step_model.model_version:
146
- client.create_run_metadata(
147
- metadata=metadata,
148
- resource_id=step_model.model_version.id,
149
- resource_type=MetadataResourceTypes.MODEL_VERSION,
156
+ # Log metadata to a step by ID
157
+ if step_id is not None:
158
+ resources = [
159
+ RunMetadataResource(
160
+ id=step_id, type=MetadataResourceTypes.STEP_RUN
150
161
  )
162
+ ]
151
163
 
152
- # If a step is identified by id, fetch it directly through the client,
153
- # follow a similar procedure and log metadata for its pipeline and model
154
- # as well.
155
- elif step_id is not None:
156
- step_model = client.get_run_step(step_run_id=step_id)
157
- client.create_run_metadata(
158
- metadata=metadata,
159
- resource_id=step_model.pipeline_run_id,
160
- resource_type=MetadataResourceTypes.PIPELINE_RUN,
161
- )
162
- client.create_run_metadata(
163
- metadata=metadata,
164
- resource_id=step_model.id,
165
- resource_type=MetadataResourceTypes.STEP_RUN,
164
+ # Log metadata to a step by name and run ID
165
+ elif step_name is not None and run_id_name_or_prefix is not None:
166
+ step_model_id = (
167
+ client.get_pipeline_run(name_id_or_prefix=run_id_name_or_prefix)
168
+ .steps[step_name]
169
+ .id
166
170
  )
167
- if step_model.model_version:
168
- client.create_run_metadata(
169
- metadata=metadata,
170
- resource_id=step_model.model_version.id,
171
- resource_type=MetadataResourceTypes.MODEL_VERSION,
171
+ resources = [
172
+ RunMetadataResource(
173
+ id=step_model_id, type=MetadataResourceTypes.STEP_RUN
172
174
  )
175
+ ]
173
176
 
174
- # If a pipeline run id is identified, we need to log metadata to it and its
175
- # model as well.
177
+ # Log metadata to a run by ID
176
178
  elif run_id_name_or_prefix is not None:
177
179
  run_model = client.get_pipeline_run(
178
180
  name_id_or_prefix=run_id_name_or_prefix
179
181
  )
180
- client.create_run_metadata(
181
- metadata=metadata,
182
- resource_id=run_model.id,
183
- resource_type=MetadataResourceTypes.PIPELINE_RUN,
184
- )
185
- if run_model.model_version:
186
- client.create_run_metadata(
187
- metadata=metadata,
188
- resource_id=run_model.model_version.id,
189
- resource_type=MetadataResourceTypes.MODEL_VERSION,
182
+ resources = [
183
+ RunMetadataResource(
184
+ id=run_model.id, type=MetadataResourceTypes.PIPELINE_RUN
190
185
  )
186
+ ]
191
187
 
192
- # If the user provides a model name and version, we use to model abstraction
193
- # to fetch the model version and attach the corresponding metadata to it.
188
+ # Log metadata to a model version by name and version
194
189
  elif model_name is not None and model_version is not None:
195
- from zenml import Model
196
-
197
- mv = Model(name=model_name, version=model_version)
198
- client.create_run_metadata(
199
- metadata=metadata,
200
- resource_id=mv.id,
201
- resource_type=MetadataResourceTypes.MODEL_VERSION,
190
+ model_version_model = client.get_model_version(
191
+ model_name_or_id=model_name,
192
+ model_version_name_or_number_or_id=model_version,
202
193
  )
194
+ resources = [
195
+ RunMetadataResource(
196
+ id=model_version_model.id,
197
+ type=MetadataResourceTypes.MODEL_VERSION,
198
+ )
199
+ ]
203
200
 
204
- # If the user provides a model version id, we use the client to fetch it and
205
- # attach the metadata to it.
201
+ # Log metadata to a model version by id
206
202
  elif model_version_id is not None:
207
- model_version_id = client.get_model_version(
208
- model_version_name_or_number_or_id=model_version_id
209
- ).id
210
- client.create_run_metadata(
211
- metadata=metadata,
212
- resource_id=model_version_id,
213
- resource_type=MetadataResourceTypes.MODEL_VERSION,
203
+ resources = [
204
+ RunMetadataResource(
205
+ id=model_version_id,
206
+ type=MetadataResourceTypes.MODEL_VERSION,
207
+ )
208
+ ]
209
+
210
+ # Log metadata to a model through the step context
211
+ elif infer_model is True:
212
+ try:
213
+ step_context = get_step_context()
214
+ except RuntimeError:
215
+ raise ValueError(
216
+ "If you are using the `infer_model` option, the function must "
217
+ "be called inside a step with configured `model` in decorator."
218
+ "Otherwise, you can provide a `model_version_id` or a "
219
+ "combination of `model_name` and `model_version`."
220
+ )
221
+
222
+ if step_context.model_version is None:
223
+ raise ValueError(
224
+ "The step context does not feature any model versions."
225
+ )
226
+
227
+ resources = [
228
+ RunMetadataResource(
229
+ id=step_context.model_version.id,
230
+ type=MetadataResourceTypes.MODEL_VERSION,
231
+ )
232
+ ]
233
+
234
+ # Log metadata to an artifact version by its name and version
235
+ elif artifact_name is not None and artifact_version is not None:
236
+ artifact_version_model = client.get_artifact_version(
237
+ name_id_or_prefix=artifact_name, version=artifact_version
214
238
  )
239
+ resources = [
240
+ RunMetadataResource(
241
+ id=artifact_version_model.id,
242
+ type=MetadataResourceTypes.ARTIFACT_VERSION,
243
+ )
244
+ ]
215
245
 
216
- # If the user provides an artifact name, there are three possibilities. If
217
- # an artifact version is also provided with the name, we use both to fetch
218
- # the artifact version and use it to log the metadata. If no version is
219
- # provided, if the function is called within a step we search the artifacts
220
- # of the step if not we fetch the latest version and attach the metadata
221
- # to the latest version.
222
- elif artifact_name is not None:
223
- if artifact_version:
224
- artifact_version_model = client.get_artifact_version(
225
- name_id_or_prefix=artifact_name, version=artifact_version
246
+ # Log metadata to an artifact version by its ID
247
+ elif artifact_version_id is not None:
248
+ resources = [
249
+ RunMetadataResource(
250
+ id=artifact_version_id,
251
+ type=MetadataResourceTypes.ARTIFACT_VERSION,
226
252
  )
227
- client.create_run_metadata(
228
- metadata=metadata,
229
- resource_id=artifact_version_model.id,
230
- resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
253
+ ]
254
+
255
+ # Log metadata to an artifact version through the step context
256
+ elif infer_artifact is True:
257
+ try:
258
+ step_context = get_step_context()
259
+ except RuntimeError:
260
+ raise ValueError(
261
+ "When you are using the `infer_artifact` option when you call "
262
+ "`log_metadata`, it must be called inside a step with outputs."
263
+ "Otherwise, you can provide a `artifact_version_id` or a "
264
+ "combination of `artifact_name` and `artifact_version`."
231
265
  )
232
- else:
233
- step_context = None
234
- with contextlib.suppress(RuntimeError):
235
- step_context = get_step_context()
236
266
 
237
- if step_context:
238
- step_context.add_output_metadata(
239
- metadata=metadata, output_name=artifact_name
240
- )
241
- else:
242
- artifact_version_model = client.get_artifact_version(
243
- name_id_or_prefix=artifact_name
267
+ step_output_names = list(step_context._outputs.keys())
268
+
269
+ if artifact_name is not None:
270
+ # If a name provided, ensure it is in the outputs
271
+ if artifact_name not in step_output_names:
272
+ raise ValueError(
273
+ f"The provided artifact name`{artifact_name}` does not "
274
+ f"exist in the step outputs: {step_output_names}."
244
275
  )
245
- client.create_run_metadata(
246
- metadata=metadata,
247
- resource_id=artifact_version_model.id,
248
- resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
276
+ else:
277
+ # If no name provided, ensure there is only one output
278
+ if len(step_output_names) > 1:
279
+ raise ValueError(
280
+ "There is more than one output. If you would like to use "
281
+ "the `infer_artifact` option, you need to define an "
282
+ "`artifact_name`."
249
283
  )
250
284
 
251
- # If the user directly provides an artifact_version_id, we use the client to
252
- # fetch is and attach the metadata accordingly.
253
- elif artifact_version_id is not None:
254
- artifact_version_model = client.get_artifact_version(
255
- name_id_or_prefix=artifact_version_id,
256
- )
257
- client.create_run_metadata(
258
- metadata=metadata,
259
- resource_id=artifact_version_model.id,
260
- resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
285
+ if len(step_output_names) == 0:
286
+ raise ValueError("The step does not have any outputs.")
287
+
288
+ artifact_name = step_output_names[0]
289
+
290
+ step_context.add_output_metadata(
291
+ metadata=metadata, output_name=artifact_name
261
292
  )
293
+ return
262
294
 
263
295
  # If every additional value is None, that means we are calling it bare bones
264
296
  # and this call needs to happen during a step execution. We will use the
@@ -287,22 +319,14 @@ def log_metadata(
287
319
  "of the step execution, please provide the required "
288
320
  "identifiers."
289
321
  )
290
- client.create_run_metadata(
291
- metadata=metadata,
292
- resource_id=step_context.pipeline_run.id,
293
- resource_type=MetadataResourceTypes.PIPELINE_RUN,
294
- )
295
- client.create_run_metadata(
296
- metadata=metadata,
297
- resource_id=step_context.step_run.id,
298
- resource_type=MetadataResourceTypes.STEP_RUN,
299
- )
300
- if step_context.model_version:
301
- client.create_run_metadata(
302
- metadata=metadata,
303
- resource_id=step_context.model_version.id,
304
- resource_type=MetadataResourceTypes.MODEL_VERSION,
322
+
323
+ resources = [
324
+ RunMetadataResource(
325
+ id=step_context.step_run.id,
326
+ type=MetadataResourceTypes.STEP_RUN,
305
327
  )
328
+ ]
329
+ publisher_step_id = step_context.step_run.id
306
330
 
307
331
  else:
308
332
  raise ValueError(
@@ -310,26 +334,35 @@ def log_metadata(
310
334
  Unsupported way to call the `log_metadata`. Possible combinations "
311
335
  include:
312
336
 
313
- # Inside a step
314
- # Logs the metadata to the step, its run and possibly its model
337
+ # Automatic logging to a step (within a step)
315
338
  log_metadata(metadata={})
316
339
 
317
- # Manually logging for a step
318
- # Logs the metadata to the step, its run and possibly its model
340
+ # Manual logging to a step
319
341
  log_metadata(metadata={}, step_name=..., run_id_name_or_prefix=...)
320
342
  log_metadata(metadata={}, step_id=...)
321
343
 
322
- # Manually logging for a run
323
- # Logs the metadata to the run, possibly its model
344
+ # Manual logging to a run
324
345
  log_metadata(metadata={}, run_id_name_or_prefix=...)
325
346
 
326
- # Manually logging for a model
347
+ # Automatic logging to a model (within a step)
348
+ log_metadata(metadata={}, infer_model=True)
349
+
350
+ # Manual logging to a model
327
351
  log_metadata(metadata={}, model_name=..., model_version=...)
328
352
  log_metadata(metadata={}, model_version_id=...)
329
353
 
330
- # Manually logging for an artifact
331
- log_metadata(metadata={}, artifact_name=...) # inside a step
354
+ # Automatic logging to an artifact (within a step)
355
+ log_metadata(metadata={}, infer_artifact=True) # step with single output
356
+ log_metadata(metadata={}, artifact_name=..., infer_artifact=True) # specific output of a step
357
+
358
+ # Manual logging to an artifact
332
359
  log_metadata(metadata={}, artifact_name=..., artifact_version=...)
333
360
  log_metadata(metadata={}, artifact_version_id=...)
334
361
  """
335
362
  )
363
+
364
+ client.create_run_metadata(
365
+ metadata=metadata,
366
+ resources=resources,
367
+ publisher_step_id=publisher_step_id,
368
+ )
@@ -18,7 +18,7 @@ import functools
18
18
  import random
19
19
  import string
20
20
  from datetime import datetime
21
- from typing import Any, Callable, Dict, TypeVar, cast
21
+ from typing import Any, Callable, Dict, Optional, TypeVar, cast
22
22
 
23
23
  from pydantic import BaseModel
24
24
 
@@ -147,24 +147,26 @@ def validate_name(model: BaseModel) -> None:
147
147
 
148
148
  def format_name_template(
149
149
  name_template: str,
150
- substitutions: Dict[str, str],
150
+ substitutions: Optional[Dict[str, str]] = None,
151
151
  ) -> str:
152
152
  """Formats a name template with the given arguments.
153
153
 
154
- By default, ZenML support Date and Time placeholders.
155
- E.g. `my_run_{date}_{time}` will be formatted as `my_run_1970_01_01_00_00_00`.
156
- Extra placeholders need to be explicitly passed in as kwargs.
154
+ Default substitutions for `{date}` and `{time}` placeholders will be used if
155
+ not included in the provided substitutions.
157
156
 
158
157
  Args:
159
158
  name_template: The name template to format.
160
- substitutions: A dictionary of substitutions to use in the template.
159
+ substitutions: Substitutions to use in the template.
161
160
 
162
161
  Returns:
163
162
  The formatted name template.
164
163
 
165
164
  Raises:
166
- KeyError: If a key in template is missing in the kwargs.
165
+ KeyError: If a key in template is missing in the substitutions.
166
+ ValueError: If the formatted name is empty.
167
167
  """
168
+ substitutions = substitutions or {}
169
+
168
170
  if ("date" not in substitutions and "{date}" in name_template) or (
169
171
  "time" not in substitutions and "{time}" in name_template
170
172
  ):
@@ -183,13 +185,18 @@ def format_name_template(
183
185
  substitutions.setdefault("time", start_time.strftime("%H_%M_%S_%f"))
184
186
 
185
187
  try:
186
- return name_template.format(**substitutions)
188
+ formatted_name = name_template.format(**substitutions)
187
189
  except KeyError as e:
188
190
  raise KeyError(
189
191
  f"Could not format the name template `{name_template}`. "
190
192
  f"Missing key: {e}"
191
193
  )
192
194
 
195
+ if not formatted_name:
196
+ raise ValueError("Empty names are not allowed.")
197
+
198
+ return formatted_name
199
+
193
200
 
194
201
  def substitute_string(value: V, substitution_func: Callable[[str], str]) -> V:
195
202
  """Recursively substitute strings in objects.
@@ -13,7 +13,7 @@
13
13
  # permissions and limitations under the License.
14
14
  """Endpoint definitions for workspaces."""
15
15
 
16
- from typing import Dict, List, Optional, Tuple, Union
16
+ from typing import Any, Dict, List, Optional, Tuple, Union
17
17
  from uuid import UUID
18
18
 
19
19
  from fastapi import APIRouter, Depends, Security
@@ -102,6 +102,7 @@ from zenml.zen_server.rbac.endpoint_utils import (
102
102
  )
103
103
  from zenml.zen_server.rbac.models import Action, ResourceType
104
104
  from zenml.zen_server.rbac.utils import (
105
+ batch_verify_permissions_for_models,
105
106
  dehydrate_page,
106
107
  dehydrate_response_model,
107
108
  get_allowed_resource_ids,
@@ -998,24 +999,23 @@ def create_run_metadata(
998
999
  "is not supported."
999
1000
  )
1000
1001
 
1001
- if run_metadata.resource_type == MetadataResourceTypes.PIPELINE_RUN:
1002
- run = zen_store().get_run(run_metadata.resource_id)
1003
- verify_permission_for_model(run, action=Action.UPDATE)
1004
- elif run_metadata.resource_type == MetadataResourceTypes.STEP_RUN:
1005
- step = zen_store().get_run_step(run_metadata.resource_id)
1006
- verify_permission_for_model(step, action=Action.UPDATE)
1007
- elif run_metadata.resource_type == MetadataResourceTypes.ARTIFACT_VERSION:
1008
- artifact_version = zen_store().get_artifact_version(
1009
- run_metadata.resource_id
1010
- )
1011
- verify_permission_for_model(artifact_version, action=Action.UPDATE)
1012
- elif run_metadata.resource_type == MetadataResourceTypes.MODEL_VERSION:
1013
- model_version = zen_store().get_model_version(run_metadata.resource_id)
1014
- verify_permission_for_model(model_version, action=Action.UPDATE)
1015
- else:
1016
- raise RuntimeError(
1017
- f"Unknown resource type: {run_metadata.resource_type}"
1018
- )
1002
+ verify_models: List[Any] = []
1003
+ for resource in run_metadata.resources:
1004
+ if resource.type == MetadataResourceTypes.PIPELINE_RUN:
1005
+ verify_models.append(zen_store().get_run(resource.id))
1006
+ elif resource.type == MetadataResourceTypes.STEP_RUN:
1007
+ verify_models.append(zen_store().get_run_step(resource.id))
1008
+ elif resource.type == MetadataResourceTypes.ARTIFACT_VERSION:
1009
+ verify_models.append(zen_store().get_artifact_version(resource.id))
1010
+ elif resource.type == MetadataResourceTypes.MODEL_VERSION:
1011
+ verify_models.append(zen_store().get_model_version(resource.id))
1012
+ else:
1013
+ raise RuntimeError(f"Unknown resource type: {resource.type}")
1014
+
1015
+ batch_verify_permissions_for_models(
1016
+ models=verify_models,
1017
+ action=Action.UPDATE,
1018
+ )
1019
1019
 
1020
1020
  verify_permission(
1021
1021
  resource_type=ResourceType.RUN_METADATA, action=Action.CREATE
@@ -0,0 +1,36 @@
1
+ """Remove component spec path [b73bc71f1106].
2
+
3
+ Revision ID: b73bc71f1106
4
+ Revises: ec6307720f92
5
+ Create Date: 2024-11-29 09:36:33.089945
6
+
7
+ """
8
+
9
+ import sqlalchemy as sa
10
+ from alembic import op
11
+
12
+ # revision identifiers, used by Alembic.
13
+ revision = "b73bc71f1106"
14
+ down_revision = "ec6307720f92"
15
+ branch_labels = None
16
+ depends_on = None
17
+
18
+
19
+ def upgrade() -> None:
20
+ """Upgrade database schema and/or data, creating a new revision."""
21
+ # ### commands auto generated by Alembic - please adjust! ###
22
+ with op.batch_alter_table("stack_component", schema=None) as batch_op:
23
+ batch_op.drop_column("component_spec_path")
24
+
25
+ # ### end Alembic commands ###
26
+
27
+
28
+ def downgrade() -> None:
29
+ """Downgrade database schema and/or data back to the previous revision."""
30
+ # ### commands auto generated by Alembic - please adjust! ###
31
+ with op.batch_alter_table("stack_component", schema=None) as batch_op:
32
+ batch_op.add_column(
33
+ sa.Column("component_spec_path", sa.VARCHAR(), nullable=True)
34
+ )
35
+
36
+ # ### end Alembic commands ###