zenml-nightly 0.70.0.dev20241119__py3-none-any.whl → 0.70.0.dev20241121__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.
- zenml/VERSION +1 -1
- zenml/artifacts/artifact_config.py +32 -4
- zenml/artifacts/utils.py +12 -24
- zenml/cli/base.py +1 -1
- zenml/client.py +4 -19
- zenml/model/utils.py +0 -24
- zenml/models/v2/core/artifact_version.py +25 -2
- zenml/models/v2/core/model_version.py +0 -4
- zenml/models/v2/core/model_version_artifact.py +19 -76
- zenml/models/v2/core/model_version_pipeline_run.py +6 -39
- zenml/orchestrators/step_launcher.py +0 -1
- zenml/orchestrators/step_run_utils.py +4 -17
- zenml/orchestrators/step_runner.py +3 -1
- zenml/service_connectors/service_connector_registry.py +68 -57
- zenml/zen_server/routers/model_versions_endpoints.py +59 -0
- zenml/zen_server/routers/workspaces_endpoints.py +0 -130
- zenml/zen_stores/base_zen_store.py +2 -1
- zenml/zen_stores/migrations/versions/ec6307720f92_simplify_model_version_links.py +118 -0
- zenml/zen_stores/rest_zen_store.py +4 -4
- zenml/zen_stores/schemas/model_schemas.py +10 -94
- zenml/zen_stores/schemas/user_schemas.py +0 -8
- zenml/zen_stores/schemas/workspace_schemas.py +0 -14
- {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/METADATA +1 -1
- {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/RECORD +27 -26
- {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/WHEEL +0 -0
- {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/entry_points.txt +0 -0
zenml/VERSION
CHANGED
@@ -1 +1 @@
|
|
1
|
-
0.70.0.
|
1
|
+
0.70.0.dev20241121
|
@@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Union
|
|
17
17
|
|
18
18
|
from pydantic import BaseModel, Field, model_validator
|
19
19
|
|
20
|
+
from zenml.enums import ArtifactType
|
20
21
|
from zenml.logger import get_logger
|
21
22
|
from zenml.metadata.metadata_types import MetadataType
|
22
23
|
from zenml.utils.pydantic_utils import before_validator_handler
|
@@ -36,6 +37,7 @@ class ArtifactConfig(BaseModel):
|
|
36
37
|
int, ArtifactConfig(
|
37
38
|
name="my_artifact", # override the default artifact name
|
38
39
|
version=42, # set a custom version
|
40
|
+
artifact_type=ArtifactType.MODEL, # Specify the artifact type
|
39
41
|
tags=["tag1", "tag2"], # set custom tags
|
40
42
|
)
|
41
43
|
]:
|
@@ -47,8 +49,9 @@ class ArtifactConfig(BaseModel):
|
|
47
49
|
version: The version of the artifact.
|
48
50
|
tags: The tags of the artifact.
|
49
51
|
run_metadata: Metadata to add to the artifact.
|
50
|
-
|
51
|
-
|
52
|
+
artifact_type: Optional type of the artifact. If not given, the type
|
53
|
+
specified by the materializer that is used to save this artifact
|
54
|
+
is used.
|
52
55
|
"""
|
53
56
|
|
54
57
|
name: Optional[str] = None
|
@@ -58,8 +61,7 @@ class ArtifactConfig(BaseModel):
|
|
58
61
|
tags: Optional[List[str]] = None
|
59
62
|
run_metadata: Optional[Dict[str, MetadataType]] = None
|
60
63
|
|
61
|
-
|
62
|
-
is_deployment_artifact: bool = False
|
64
|
+
artifact_type: Optional[ArtifactType] = None
|
63
65
|
|
64
66
|
@model_validator(mode="before")
|
65
67
|
@classmethod
|
@@ -70,6 +72,10 @@ class ArtifactConfig(BaseModel):
|
|
70
72
|
Args:
|
71
73
|
data: The model data.
|
72
74
|
|
75
|
+
Raises:
|
76
|
+
ValueError: If the artifact is configured to be
|
77
|
+
both a model and a deployment artifact.
|
78
|
+
|
73
79
|
Returns:
|
74
80
|
Model data without the removed attributes.
|
75
81
|
"""
|
@@ -82,4 +88,26 @@ class ArtifactConfig(BaseModel):
|
|
82
88
|
"artifact is not supported anymore."
|
83
89
|
)
|
84
90
|
|
91
|
+
is_model_artifact = data.pop("is_model_artifact", None)
|
92
|
+
is_deployment_artifact = data.pop("is_deployment_artifact", None)
|
93
|
+
|
94
|
+
if is_model_artifact and is_deployment_artifact:
|
95
|
+
raise ValueError(
|
96
|
+
"An artifact can only be a model artifact or deployment "
|
97
|
+
"artifact."
|
98
|
+
)
|
99
|
+
elif is_model_artifact:
|
100
|
+
logger.warning(
|
101
|
+
"`ArtifactConfig.is_model_artifact` is deprecated and will be "
|
102
|
+
"removed soon. Use `ArtifactConfig.artifact_type` instead."
|
103
|
+
)
|
104
|
+
data.setdefault("artifact_type", ArtifactType.MODEL)
|
105
|
+
elif is_deployment_artifact:
|
106
|
+
logger.warning(
|
107
|
+
"`ArtifactConfig.is_deployment_artifact` is deprecated and "
|
108
|
+
"will be removed soon. Use `ArtifactConfig.artifact_type` "
|
109
|
+
"instead."
|
110
|
+
)
|
111
|
+
data.setdefault("artifact_type", ArtifactType.SERVICE)
|
112
|
+
|
85
113
|
return data
|
zenml/artifacts/utils.py
CHANGED
@@ -30,7 +30,6 @@ from typing import (
|
|
30
30
|
)
|
31
31
|
from uuid import UUID, uuid4
|
32
32
|
|
33
|
-
from zenml.artifacts.artifact_config import ArtifactConfig
|
34
33
|
from zenml.artifacts.preexisting_data_materializer import (
|
35
34
|
PreexistingDataMaterializer,
|
36
35
|
)
|
@@ -118,6 +117,7 @@ def _store_artifact_data_and_prepare_request(
|
|
118
117
|
materializer_class: Type["BaseMaterializer"],
|
119
118
|
save_type: ArtifactSaveType,
|
120
119
|
version: Optional[Union[int, str]] = None,
|
120
|
+
artifact_type: Optional[ArtifactType] = None,
|
121
121
|
tags: Optional[List[str]] = None,
|
122
122
|
store_metadata: bool = True,
|
123
123
|
store_visualizations: bool = True,
|
@@ -134,6 +134,8 @@ def _store_artifact_data_and_prepare_request(
|
|
134
134
|
artifact data.
|
135
135
|
save_type: Save type of the artifact version.
|
136
136
|
version: The artifact version.
|
137
|
+
artifact_type: The artifact type. If not given, the type will be defined
|
138
|
+
by the materializer that is used to save the artifact.
|
137
139
|
tags: Tags for the artifact version.
|
138
140
|
store_metadata: Whether to store metadata for the artifact version.
|
139
141
|
store_visualizations: Whether to store visualizations for the artifact
|
@@ -176,7 +178,7 @@ def _store_artifact_data_and_prepare_request(
|
|
176
178
|
artifact_name=name,
|
177
179
|
version=version,
|
178
180
|
tags=tags,
|
179
|
-
type=materializer.ASSOCIATED_ARTIFACT_TYPE,
|
181
|
+
type=artifact_type or materializer.ASSOCIATED_ARTIFACT_TYPE,
|
180
182
|
uri=materializer.uri,
|
181
183
|
materializer=source_utils.resolve(materializer.__class__),
|
182
184
|
data_type=source_utils.resolve(data_type),
|
@@ -198,14 +200,13 @@ def save_artifact(
|
|
198
200
|
data: Any,
|
199
201
|
name: str,
|
200
202
|
version: Optional[Union[int, str]] = None,
|
203
|
+
artifact_type: Optional[ArtifactType] = None,
|
201
204
|
tags: Optional[List[str]] = None,
|
202
205
|
extract_metadata: bool = True,
|
203
206
|
include_visualizations: bool = True,
|
204
207
|
user_metadata: Optional[Dict[str, "MetadataType"]] = None,
|
205
208
|
materializer: Optional["MaterializerClassOrSource"] = None,
|
206
209
|
uri: Optional[str] = None,
|
207
|
-
is_model_artifact: bool = False,
|
208
|
-
is_deployment_artifact: bool = False,
|
209
210
|
# TODO: remove these once external artifact does not use this function anymore
|
210
211
|
save_type: ArtifactSaveType = ArtifactSaveType.MANUAL,
|
211
212
|
has_custom_name: bool = True,
|
@@ -218,6 +219,8 @@ def save_artifact(
|
|
218
219
|
version: The version of the artifact. If not provided, a new
|
219
220
|
auto-incremented version will be used.
|
220
221
|
tags: Tags to associate with the artifact.
|
222
|
+
artifact_type: The artifact type. If not given, the type will be defined
|
223
|
+
by the materializer that is used to save the artifact.
|
221
224
|
extract_metadata: If artifact metadata should be extracted and returned.
|
222
225
|
include_visualizations: If artifact visualizations should be generated.
|
223
226
|
user_metadata: User-provided metadata to store with the artifact.
|
@@ -226,8 +229,6 @@ def save_artifact(
|
|
226
229
|
uri: The URI within the artifact store to upload the artifact
|
227
230
|
to. If not provided, the artifact will be uploaded to
|
228
231
|
`custom_artifacts/{name}/{version}`.
|
229
|
-
is_model_artifact: If the artifact is a model artifact.
|
230
|
-
is_deployment_artifact: If the artifact is a deployment artifact.
|
231
232
|
save_type: The type of save operation that created the artifact version.
|
232
233
|
has_custom_name: If the artifact name is custom and should be listed in
|
233
234
|
the dashboard "Artifacts" tab.
|
@@ -273,6 +274,7 @@ def save_artifact(
|
|
273
274
|
materializer_class=materializer_class,
|
274
275
|
save_type=save_type,
|
275
276
|
version=version,
|
277
|
+
artifact_type=artifact_type,
|
276
278
|
tags=tags,
|
277
279
|
store_metadata=extract_metadata,
|
278
280
|
store_visualizations=include_visualizations,
|
@@ -286,8 +288,6 @@ def save_artifact(
|
|
286
288
|
if save_type == ArtifactSaveType.MANUAL:
|
287
289
|
_link_artifact_version_to_the_step_and_model(
|
288
290
|
artifact_version=artifact_version,
|
289
|
-
is_model_artifact=is_model_artifact,
|
290
|
-
is_deployment_artifact=is_deployment_artifact,
|
291
291
|
)
|
292
292
|
|
293
293
|
return artifact_version
|
@@ -297,10 +297,9 @@ def register_artifact(
|
|
297
297
|
folder_or_file_uri: str,
|
298
298
|
name: str,
|
299
299
|
version: Optional[Union[int, str]] = None,
|
300
|
+
artifact_type: Optional[ArtifactType] = None,
|
300
301
|
tags: Optional[List[str]] = None,
|
301
302
|
has_custom_name: bool = True,
|
302
|
-
is_model_artifact: bool = False,
|
303
|
-
is_deployment_artifact: bool = False,
|
304
303
|
artifact_metadata: Dict[str, "MetadataType"] = {},
|
305
304
|
) -> "ArtifactVersionResponse":
|
306
305
|
"""Register existing data stored in the artifact store as a ZenML Artifact.
|
@@ -311,11 +310,11 @@ def register_artifact(
|
|
311
310
|
name: The name of the artifact.
|
312
311
|
version: The version of the artifact. If not provided, a new
|
313
312
|
auto-incremented version will be used.
|
313
|
+
artifact_type: The artifact type. If not given, the type will default
|
314
|
+
to `data`.
|
314
315
|
tags: Tags to associate with the artifact.
|
315
316
|
has_custom_name: If the artifact name is custom and should be listed in
|
316
317
|
the dashboard "Artifacts" tab.
|
317
|
-
is_model_artifact: If the artifact is a model artifact.
|
318
|
-
is_deployment_artifact: If the artifact is a deployment artifact.
|
319
318
|
artifact_metadata: Metadata dictionary to attach to the artifact version.
|
320
319
|
|
321
320
|
Returns:
|
@@ -346,7 +345,7 @@ def register_artifact(
|
|
346
345
|
artifact_name=name,
|
347
346
|
version=version,
|
348
347
|
tags=tags,
|
349
|
-
type=ArtifactType.DATA,
|
348
|
+
type=artifact_type or ArtifactType.DATA,
|
350
349
|
save_type=ArtifactSaveType.PREEXISTING,
|
351
350
|
uri=folder_or_file_uri,
|
352
351
|
materializer=source_utils.resolve(PreexistingDataMaterializer),
|
@@ -365,8 +364,6 @@ def register_artifact(
|
|
365
364
|
|
366
365
|
_link_artifact_version_to_the_step_and_model(
|
367
366
|
artifact_version=artifact_version,
|
368
|
-
is_model_artifact=is_model_artifact,
|
369
|
-
is_deployment_artifact=is_deployment_artifact,
|
370
367
|
)
|
371
368
|
|
372
369
|
return artifact_version
|
@@ -674,8 +671,6 @@ def _check_if_artifact_with_given_uri_already_registered(
|
|
674
671
|
|
675
672
|
def _link_artifact_version_to_the_step_and_model(
|
676
673
|
artifact_version: ArtifactVersionResponse,
|
677
|
-
is_model_artifact: bool,
|
678
|
-
is_deployment_artifact: bool,
|
679
674
|
) -> None:
|
680
675
|
"""Link an artifact version to the step and its' context model.
|
681
676
|
|
@@ -685,8 +680,6 @@ def _link_artifact_version_to_the_step_and_model(
|
|
685
680
|
|
686
681
|
Args:
|
687
682
|
artifact_version: The artifact version to link.
|
688
|
-
is_model_artifact: Whether the artifact is a model artifact.
|
689
|
-
is_deployment_artifact: Whether the artifact is a deployment artifact.
|
690
683
|
"""
|
691
684
|
client = Client()
|
692
685
|
try:
|
@@ -706,14 +699,9 @@ def _link_artifact_version_to_the_step_and_model(
|
|
706
699
|
link_artifact_version_to_model_version,
|
707
700
|
)
|
708
701
|
|
709
|
-
artifact_config = ArtifactConfig(
|
710
|
-
is_model_artifact=is_model_artifact,
|
711
|
-
is_deployment_artifact=is_deployment_artifact,
|
712
|
-
)
|
713
702
|
link_artifact_version_to_model_version(
|
714
703
|
artifact_version=artifact_version,
|
715
704
|
model_version=step_context.model_version,
|
716
|
-
artifact_config=artifact_config,
|
717
705
|
)
|
718
706
|
except (RuntimeError, StepContextError):
|
719
707
|
logger.debug(f"Unable to link saved artifact to {error_message}.")
|
zenml/cli/base.py
CHANGED
@@ -79,7 +79,7 @@ class ZenMLProjectTemplateLocation(BaseModel):
|
|
79
79
|
ZENML_PROJECT_TEMPLATES = dict(
|
80
80
|
e2e_batch=ZenMLProjectTemplateLocation(
|
81
81
|
github_url="zenml-io/template-e2e-batch",
|
82
|
-
github_tag="2024.
|
82
|
+
github_tag="2024.11.20", # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
|
83
83
|
),
|
84
84
|
starter=ZenMLProjectTemplateLocation(
|
85
85
|
github_url="zenml-io/template-starter",
|
zenml/client.py
CHANGED
@@ -4220,6 +4220,7 @@ class Client(metaclass=ClientMetaClass):
|
|
4220
4220
|
materializer: Optional[str] = None,
|
4221
4221
|
workspace_id: Optional[Union[str, UUID]] = None,
|
4222
4222
|
user_id: Optional[Union[str, UUID]] = None,
|
4223
|
+
model_version_id: Optional[Union[str, UUID]] = None,
|
4223
4224
|
only_unused: Optional[bool] = False,
|
4224
4225
|
has_custom_name: Optional[bool] = None,
|
4225
4226
|
user: Optional[Union[UUID, str]] = None,
|
@@ -4249,7 +4250,8 @@ class Client(metaclass=ClientMetaClass):
|
|
4249
4250
|
uri: The uri of the artifact to filter by.
|
4250
4251
|
materializer: The materializer of the artifact to filter by.
|
4251
4252
|
workspace_id: The id of the workspace to filter by.
|
4252
|
-
user_id: The
|
4253
|
+
user_id: The id of the user to filter by.
|
4254
|
+
model_version_id: Filter by model version ID.
|
4253
4255
|
only_unused: Only return artifact versions that are not used in
|
4254
4256
|
any pipeline runs.
|
4255
4257
|
has_custom_name: Filter artifacts with/without custom names.
|
@@ -4283,6 +4285,7 @@ class Client(metaclass=ClientMetaClass):
|
|
4283
4285
|
materializer=materializer,
|
4284
4286
|
workspace_id=workspace_id,
|
4285
4287
|
user_id=user_id,
|
4288
|
+
model_version_id=model_version_id,
|
4286
4289
|
only_unused=only_unused,
|
4287
4290
|
has_custom_name=has_custom_name,
|
4288
4291
|
tag=tag,
|
@@ -6435,9 +6438,6 @@ class Client(metaclass=ClientMetaClass):
|
|
6435
6438
|
logical_operator: LogicalOperators = LogicalOperators.AND,
|
6436
6439
|
created: Optional[Union[datetime, str]] = None,
|
6437
6440
|
updated: Optional[Union[datetime, str]] = None,
|
6438
|
-
workspace_id: Optional[Union[UUID, str]] = None,
|
6439
|
-
user_id: Optional[Union[UUID, str]] = None,
|
6440
|
-
model_id: Optional[Union[UUID, str]] = None,
|
6441
6441
|
model_version_id: Optional[Union[UUID, str]] = None,
|
6442
6442
|
artifact_version_id: Optional[Union[UUID, str]] = None,
|
6443
6443
|
artifact_name: Optional[str] = None,
|
@@ -6457,9 +6457,6 @@ class Client(metaclass=ClientMetaClass):
|
|
6457
6457
|
logical_operator: Which logical operator to use [and, or]
|
6458
6458
|
created: Use to filter by time of creation
|
6459
6459
|
updated: Use the last updated date for filtering
|
6460
|
-
workspace_id: Use the workspace id for filtering
|
6461
|
-
user_id: Use the user id for filtering
|
6462
|
-
model_id: Use the model id for filtering
|
6463
6460
|
model_version_id: Use the model version id for filtering
|
6464
6461
|
artifact_version_id: Use the artifact id for filtering
|
6465
6462
|
artifact_name: Use the artifact name for filtering
|
@@ -6482,9 +6479,6 @@ class Client(metaclass=ClientMetaClass):
|
|
6482
6479
|
size=size,
|
6483
6480
|
created=created,
|
6484
6481
|
updated=updated,
|
6485
|
-
workspace_id=workspace_id,
|
6486
|
-
user_id=user_id,
|
6487
|
-
model_id=model_id,
|
6488
6482
|
model_version_id=model_version_id,
|
6489
6483
|
artifact_version_id=artifact_version_id,
|
6490
6484
|
artifact_name=artifact_name,
|
@@ -6556,9 +6550,6 @@ class Client(metaclass=ClientMetaClass):
|
|
6556
6550
|
logical_operator: LogicalOperators = LogicalOperators.AND,
|
6557
6551
|
created: Optional[Union[datetime, str]] = None,
|
6558
6552
|
updated: Optional[Union[datetime, str]] = None,
|
6559
|
-
workspace_id: Optional[Union[UUID, str]] = None,
|
6560
|
-
user_id: Optional[Union[UUID, str]] = None,
|
6561
|
-
model_id: Optional[Union[UUID, str]] = None,
|
6562
6553
|
model_version_id: Optional[Union[UUID, str]] = None,
|
6563
6554
|
pipeline_run_id: Optional[Union[UUID, str]] = None,
|
6564
6555
|
pipeline_run_name: Optional[str] = None,
|
@@ -6574,9 +6565,6 @@ class Client(metaclass=ClientMetaClass):
|
|
6574
6565
|
logical_operator: Which logical operator to use [and, or]
|
6575
6566
|
created: Use to filter by time of creation
|
6576
6567
|
updated: Use the last updated date for filtering
|
6577
|
-
workspace_id: Use the workspace id for filtering
|
6578
|
-
user_id: Use the user id for filtering
|
6579
|
-
model_id: Use the model id for filtering
|
6580
6568
|
model_version_id: Use the model version id for filtering
|
6581
6569
|
pipeline_run_id: Use the pipeline run id for filtering
|
6582
6570
|
pipeline_run_name: Use the pipeline run name for filtering
|
@@ -6595,9 +6583,6 @@ class Client(metaclass=ClientMetaClass):
|
|
6595
6583
|
size=size,
|
6596
6584
|
created=created,
|
6597
6585
|
updated=updated,
|
6598
|
-
workspace_id=workspace_id,
|
6599
|
-
user_id=user_id,
|
6600
|
-
model_id=model_id,
|
6601
6586
|
model_version_id=model_version_id,
|
6602
6587
|
pipeline_run_id=pipeline_run_id,
|
6603
6588
|
pipeline_run_name=pipeline_run_name,
|
zenml/model/utils.py
CHANGED
@@ -16,7 +16,6 @@
|
|
16
16
|
from typing import Dict, Optional, Union
|
17
17
|
from uuid import UUID
|
18
18
|
|
19
|
-
from zenml.artifacts.artifact_config import ArtifactConfig
|
20
19
|
from zenml.client import Client
|
21
20
|
from zenml.enums import ModelStages
|
22
21
|
from zenml.exceptions import StepContextError
|
@@ -81,32 +80,18 @@ def log_model_metadata(
|
|
81
80
|
def link_artifact_version_to_model_version(
|
82
81
|
artifact_version: ArtifactVersionResponse,
|
83
82
|
model_version: ModelVersionResponse,
|
84
|
-
artifact_config: Optional[ArtifactConfig] = None,
|
85
83
|
) -> None:
|
86
84
|
"""Link an artifact version to a model version.
|
87
85
|
|
88
86
|
Args:
|
89
87
|
artifact_version: The artifact version to link.
|
90
88
|
model_version: The model version to link.
|
91
|
-
artifact_config: Output artifact configuration.
|
92
89
|
"""
|
93
|
-
if artifact_config:
|
94
|
-
is_model_artifact = artifact_config.is_model_artifact
|
95
|
-
is_deployment_artifact = artifact_config.is_deployment_artifact
|
96
|
-
else:
|
97
|
-
is_model_artifact = False
|
98
|
-
is_deployment_artifact = False
|
99
|
-
|
100
90
|
client = Client()
|
101
91
|
client.zen_store.create_model_version_artifact_link(
|
102
92
|
ModelVersionArtifactRequest(
|
103
|
-
user=client.active_user.id,
|
104
|
-
workspace=client.active_workspace.id,
|
105
93
|
artifact_version=artifact_version.id,
|
106
|
-
model=model_version.model.id,
|
107
94
|
model_version=model_version.id,
|
108
|
-
is_model_artifact=is_model_artifact,
|
109
|
-
is_deployment_artifact=is_deployment_artifact,
|
110
95
|
)
|
111
96
|
)
|
112
97
|
|
@@ -114,16 +99,12 @@ def link_artifact_version_to_model_version(
|
|
114
99
|
def link_artifact_to_model(
|
115
100
|
artifact_version: ArtifactVersionResponse,
|
116
101
|
model: Optional["Model"] = None,
|
117
|
-
is_model_artifact: bool = False,
|
118
|
-
is_deployment_artifact: bool = False,
|
119
102
|
) -> None:
|
120
103
|
"""Link the artifact to the model.
|
121
104
|
|
122
105
|
Args:
|
123
106
|
artifact_version: The artifact version to link.
|
124
107
|
model: The model to link to.
|
125
|
-
is_model_artifact: Whether the artifact is a model artifact.
|
126
|
-
is_deployment_artifact: Whether the artifact is a deployment artifact.
|
127
108
|
|
128
109
|
Raises:
|
129
110
|
RuntimeError: If called outside of a step.
|
@@ -145,14 +126,9 @@ def link_artifact_to_model(
|
|
145
126
|
)
|
146
127
|
|
147
128
|
model_version = model._get_or_create_model_version()
|
148
|
-
artifact_config = ArtifactConfig(
|
149
|
-
is_model_artifact=is_model_artifact,
|
150
|
-
is_deployment_artifact=is_deployment_artifact,
|
151
|
-
)
|
152
129
|
link_artifact_version_to_model_version(
|
153
130
|
artifact_version=artifact_version,
|
154
131
|
model_version=model_version,
|
155
|
-
artifact_config=artifact_config,
|
156
132
|
)
|
157
133
|
|
158
134
|
|
@@ -37,7 +37,7 @@ from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
|
|
37
37
|
from zenml.enums import ArtifactSaveType, ArtifactType, GenericFilterOps
|
38
38
|
from zenml.logger import get_logger
|
39
39
|
from zenml.metadata.metadata_types import MetadataType
|
40
|
-
from zenml.models.v2.base.filter import StrFilter
|
40
|
+
from zenml.models.v2.base.filter import FilterGenerator, StrFilter
|
41
41
|
from zenml.models.v2.base.scoped import (
|
42
42
|
WorkspaceScopedRequest,
|
43
43
|
WorkspaceScopedResponse,
|
@@ -474,6 +474,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
|
|
474
474
|
"user",
|
475
475
|
"model",
|
476
476
|
"pipeline_run",
|
477
|
+
"model_version_id",
|
477
478
|
"run_metadata",
|
478
479
|
]
|
479
480
|
artifact_id: Optional[Union[UUID, str]] = Field(
|
@@ -525,6 +526,11 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
|
|
525
526
|
description="User that produced this artifact",
|
526
527
|
union_mode="left_to_right",
|
527
528
|
)
|
529
|
+
model_version_id: Optional[Union[UUID, str]] = Field(
|
530
|
+
default=None,
|
531
|
+
description="ID of the model version that is associated with this artifact version.",
|
532
|
+
union_mode="left_to_right",
|
533
|
+
)
|
528
534
|
only_unused: Optional[bool] = Field(
|
529
535
|
default=False, description="Filter only for unused artifacts"
|
530
536
|
)
|
@@ -568,6 +574,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
|
|
568
574
|
ArtifactVersionSchema,
|
569
575
|
ModelSchema,
|
570
576
|
ModelVersionArtifactSchema,
|
577
|
+
ModelVersionSchema,
|
571
578
|
PipelineRunSchema,
|
572
579
|
RunMetadataSchema,
|
573
580
|
StepRunInputArtifactSchema,
|
@@ -600,6 +607,20 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
|
|
600
607
|
)
|
601
608
|
custom_filters.append(unused_filter)
|
602
609
|
|
610
|
+
if self.model_version_id:
|
611
|
+
value, operator = self._resolve_operator(self.model_version_id)
|
612
|
+
|
613
|
+
model_version_filter = and_(
|
614
|
+
ArtifactVersionSchema.id
|
615
|
+
== ModelVersionArtifactSchema.artifact_version_id,
|
616
|
+
ModelVersionArtifactSchema.model_version_id
|
617
|
+
== ModelVersionSchema.id,
|
618
|
+
FilterGenerator(ModelVersionSchema)
|
619
|
+
.define_filter(column="id", value=value, operator=operator)
|
620
|
+
.generate_query_conditions(ModelVersionSchema),
|
621
|
+
)
|
622
|
+
custom_filters.append(model_version_filter)
|
623
|
+
|
603
624
|
if self.has_custom_name is not None:
|
604
625
|
custom_name_filter = and_(
|
605
626
|
ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
|
@@ -622,7 +643,9 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
|
|
622
643
|
model_filter = and_(
|
623
644
|
ArtifactVersionSchema.id
|
624
645
|
== ModelVersionArtifactSchema.artifact_version_id,
|
625
|
-
ModelVersionArtifactSchema.
|
646
|
+
ModelVersionArtifactSchema.model_version_id
|
647
|
+
== ModelVersionSchema.id,
|
648
|
+
ModelVersionSchema.model_id == ModelSchema.id,
|
626
649
|
self.generate_name_or_id_query_conditions(
|
627
650
|
value=self.model, table=ModelSchema
|
628
651
|
),
|
@@ -576,10 +576,6 @@ class ModelVersionResponse(
|
|
576
576
|
force=force,
|
577
577
|
)
|
578
578
|
|
579
|
-
# TODO in https://zenml.atlassian.net/browse/OSS-2433
|
580
|
-
# def generate_model_card(self, template_name: str) -> str:
|
581
|
-
# """Return HTML/PDF based on input template"""
|
582
|
-
|
583
579
|
|
584
580
|
# ------------------ Filter Model ------------------
|
585
581
|
|
@@ -16,20 +16,17 @@
|
|
16
16
|
from typing import TYPE_CHECKING, List, Optional, Union
|
17
17
|
from uuid import UUID
|
18
18
|
|
19
|
-
from pydantic import ConfigDict, Field
|
19
|
+
from pydantic import ConfigDict, Field
|
20
20
|
|
21
21
|
from zenml.enums import GenericFilterOps
|
22
22
|
from zenml.models.v2.base.base import (
|
23
23
|
BaseDatedResponseBody,
|
24
24
|
BaseIdentifiedResponse,
|
25
|
+
BaseRequest,
|
25
26
|
BaseResponseMetadata,
|
26
27
|
BaseResponseResources,
|
27
28
|
)
|
28
|
-
from zenml.models.v2.base.filter import StrFilter
|
29
|
-
from zenml.models.v2.base.scoped import (
|
30
|
-
WorkspaceScopedFilter,
|
31
|
-
WorkspaceScopedRequest,
|
32
|
-
)
|
29
|
+
from zenml.models.v2.base.filter import BaseFilter, StrFilter
|
33
30
|
|
34
31
|
if TYPE_CHECKING:
|
35
32
|
from sqlalchemy.sql.elements import ColumnElement
|
@@ -40,14 +37,11 @@ if TYPE_CHECKING:
|
|
40
37
|
# ------------------ Request Model ------------------
|
41
38
|
|
42
39
|
|
43
|
-
class ModelVersionArtifactRequest(
|
40
|
+
class ModelVersionArtifactRequest(BaseRequest):
|
44
41
|
"""Request model for links between model versions and artifacts."""
|
45
42
|
|
46
|
-
model: UUID
|
47
43
|
model_version: UUID
|
48
44
|
artifact_version: UUID
|
49
|
-
is_model_artifact: bool = False
|
50
|
-
is_deployment_artifact: bool = False
|
51
45
|
|
52
46
|
# TODO: In Pydantic v2, the `model_` is a protected namespaces for all
|
53
47
|
# fields defined under base models. If not handled, this raises a warning.
|
@@ -57,15 +51,6 @@ class ModelVersionArtifactRequest(WorkspaceScopedRequest):
|
|
57
51
|
# careful we might overwrite some fields protected by pydantic.
|
58
52
|
model_config = ConfigDict(protected_namespaces=())
|
59
53
|
|
60
|
-
@model_validator(mode="after")
|
61
|
-
def _validate_is_endpoint_artifact(self) -> "ModelVersionArtifactRequest":
|
62
|
-
if self.is_model_artifact and self.is_deployment_artifact:
|
63
|
-
raise ValueError(
|
64
|
-
"Artifact cannot be a model artifact and deployment artifact "
|
65
|
-
"at the same time."
|
66
|
-
)
|
67
|
-
return self
|
68
|
-
|
69
54
|
|
70
55
|
# ------------------ Update Model ------------------
|
71
56
|
|
@@ -77,11 +62,8 @@ class ModelVersionArtifactRequest(WorkspaceScopedRequest):
|
|
77
62
|
class ModelVersionArtifactResponseBody(BaseDatedResponseBody):
|
78
63
|
"""Response body for links between model versions and artifacts."""
|
79
64
|
|
80
|
-
model: UUID
|
81
65
|
model_version: UUID
|
82
66
|
artifact_version: "ArtifactVersionResponse"
|
83
|
-
is_model_artifact: bool = False
|
84
|
-
is_deployment_artifact: bool = False
|
85
67
|
|
86
68
|
# TODO: In Pydantic v2, the `model_` is a protected namespaces for all
|
87
69
|
# fields defined under base models. If not handled, this raises a warning.
|
@@ -105,16 +87,6 @@ class ModelVersionArtifactResponse(
|
|
105
87
|
):
|
106
88
|
"""Response model for links between model versions and artifacts."""
|
107
89
|
|
108
|
-
# Body and metadata properties
|
109
|
-
@property
|
110
|
-
def model(self) -> UUID:
|
111
|
-
"""The `model` property.
|
112
|
-
|
113
|
-
Returns:
|
114
|
-
the value of the property.
|
115
|
-
"""
|
116
|
-
return self.get_body().model
|
117
|
-
|
118
90
|
@property
|
119
91
|
def model_version(self) -> UUID:
|
120
92
|
"""The `model_version` property.
|
@@ -133,34 +105,16 @@ class ModelVersionArtifactResponse(
|
|
133
105
|
"""
|
134
106
|
return self.get_body().artifact_version
|
135
107
|
|
136
|
-
@property
|
137
|
-
def is_model_artifact(self) -> bool:
|
138
|
-
"""The `is_model_artifact` property.
|
139
|
-
|
140
|
-
Returns:
|
141
|
-
the value of the property.
|
142
|
-
"""
|
143
|
-
return self.get_body().is_model_artifact
|
144
|
-
|
145
|
-
@property
|
146
|
-
def is_deployment_artifact(self) -> bool:
|
147
|
-
"""The `is_deployment_artifact` property.
|
148
|
-
|
149
|
-
Returns:
|
150
|
-
the value of the property.
|
151
|
-
"""
|
152
|
-
return self.get_body().is_deployment_artifact
|
153
|
-
|
154
108
|
|
155
109
|
# ------------------ Filter Model ------------------
|
156
110
|
|
157
111
|
|
158
|
-
class ModelVersionArtifactFilter(
|
112
|
+
class ModelVersionArtifactFilter(BaseFilter):
|
159
113
|
"""Model version pipeline run links filter model."""
|
160
114
|
|
161
115
|
# Artifact name and type are not DB fields and need to be handled separately
|
162
116
|
FILTER_EXCLUDE_FIELDS = [
|
163
|
-
*
|
117
|
+
*BaseFilter.FILTER_EXCLUDE_FIELDS,
|
164
118
|
"artifact_name",
|
165
119
|
"only_data_artifacts",
|
166
120
|
"only_model_artifacts",
|
@@ -169,34 +123,16 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
|
|
169
123
|
"user",
|
170
124
|
]
|
171
125
|
CLI_EXCLUDE_FIELDS = [
|
172
|
-
*
|
126
|
+
*BaseFilter.CLI_EXCLUDE_FIELDS,
|
173
127
|
"only_data_artifacts",
|
174
128
|
"only_model_artifacts",
|
175
129
|
"only_deployment_artifacts",
|
176
130
|
"has_custom_name",
|
177
|
-
"model_id",
|
178
131
|
"model_version_id",
|
179
|
-
"user_id",
|
180
|
-
"workspace_id",
|
181
132
|
"updated",
|
182
133
|
"id",
|
183
134
|
]
|
184
135
|
|
185
|
-
workspace_id: Optional[Union[UUID, str]] = Field(
|
186
|
-
default=None,
|
187
|
-
description="The workspace of the Model Version",
|
188
|
-
union_mode="left_to_right",
|
189
|
-
)
|
190
|
-
user_id: Optional[Union[UUID, str]] = Field(
|
191
|
-
default=None,
|
192
|
-
description="The user of the Model Version",
|
193
|
-
union_mode="left_to_right",
|
194
|
-
)
|
195
|
-
model_id: Optional[Union[UUID, str]] = Field(
|
196
|
-
default=None,
|
197
|
-
description="Filter by model ID",
|
198
|
-
union_mode="left_to_right",
|
199
|
-
)
|
200
136
|
model_version_id: Optional[Union[UUID, str]] = Field(
|
201
137
|
default=None,
|
202
138
|
description="Filter by model version ID",
|
@@ -236,7 +172,7 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
|
|
236
172
|
"""
|
237
173
|
custom_filters = super().get_custom_filters()
|
238
174
|
|
239
|
-
from sqlmodel import and_
|
175
|
+
from sqlmodel import and_, col
|
240
176
|
|
241
177
|
from zenml.zen_stores.schemas import (
|
242
178
|
ArtifactSchema,
|
@@ -262,20 +198,27 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
|
|
262
198
|
|
263
199
|
if self.only_data_artifacts:
|
264
200
|
data_artifact_filter = and_(
|
265
|
-
ModelVersionArtifactSchema.
|
266
|
-
|
201
|
+
ModelVersionArtifactSchema.artifact_version_id
|
202
|
+
== ArtifactVersionSchema.id,
|
203
|
+
col(ArtifactVersionSchema.type).not_in(
|
204
|
+
["ServiceArtifact", "ModelArtifact"]
|
205
|
+
),
|
267
206
|
)
|
268
207
|
custom_filters.append(data_artifact_filter)
|
269
208
|
|
270
209
|
if self.only_model_artifacts:
|
271
210
|
model_artifact_filter = and_(
|
272
|
-
ModelVersionArtifactSchema.
|
211
|
+
ModelVersionArtifactSchema.artifact_version_id
|
212
|
+
== ArtifactVersionSchema.id,
|
213
|
+
ArtifactVersionSchema.type == "ModelArtifact",
|
273
214
|
)
|
274
215
|
custom_filters.append(model_artifact_filter)
|
275
216
|
|
276
217
|
if self.only_deployment_artifacts:
|
277
218
|
deployment_artifact_filter = and_(
|
278
|
-
ModelVersionArtifactSchema.
|
219
|
+
ModelVersionArtifactSchema.artifact_version_id
|
220
|
+
== ArtifactVersionSchema.id,
|
221
|
+
ArtifactVersionSchema.type == "ServiceArtifact",
|
279
222
|
)
|
280
223
|
custom_filters.append(deployment_artifact_filter)
|
281
224
|
|