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.
Files changed (27) hide show
  1. zenml/VERSION +1 -1
  2. zenml/artifacts/artifact_config.py +32 -4
  3. zenml/artifacts/utils.py +12 -24
  4. zenml/cli/base.py +1 -1
  5. zenml/client.py +4 -19
  6. zenml/model/utils.py +0 -24
  7. zenml/models/v2/core/artifact_version.py +25 -2
  8. zenml/models/v2/core/model_version.py +0 -4
  9. zenml/models/v2/core/model_version_artifact.py +19 -76
  10. zenml/models/v2/core/model_version_pipeline_run.py +6 -39
  11. zenml/orchestrators/step_launcher.py +0 -1
  12. zenml/orchestrators/step_run_utils.py +4 -17
  13. zenml/orchestrators/step_runner.py +3 -1
  14. zenml/service_connectors/service_connector_registry.py +68 -57
  15. zenml/zen_server/routers/model_versions_endpoints.py +59 -0
  16. zenml/zen_server/routers/workspaces_endpoints.py +0 -130
  17. zenml/zen_stores/base_zen_store.py +2 -1
  18. zenml/zen_stores/migrations/versions/ec6307720f92_simplify_model_version_links.py +118 -0
  19. zenml/zen_stores/rest_zen_store.py +4 -4
  20. zenml/zen_stores/schemas/model_schemas.py +10 -94
  21. zenml/zen_stores/schemas/user_schemas.py +0 -8
  22. zenml/zen_stores/schemas/workspace_schemas.py +0 -14
  23. {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/METADATA +1 -1
  24. {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/RECORD +27 -26
  25. {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/LICENSE +0 -0
  26. {zenml_nightly-0.70.0.dev20241119.dist-info → zenml_nightly-0.70.0.dev20241121.dist-info}/WHEEL +0 -0
  27. {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.dev20241119
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
- is_model_artifact: Whether the artifact is a model artifact.
51
- is_deployment_artifact: Whether the artifact is a deployment artifact.
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
- is_model_artifact: bool = False
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.10.30", # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
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 id of the user to filter by.
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.model_id == ModelSchema.id,
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, model_validator
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(WorkspaceScopedRequest):
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(WorkspaceScopedFilter):
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
- *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
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
- *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
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.is_model_artifact.is_(False), # type: ignore[attr-defined]
266
- ModelVersionArtifactSchema.is_deployment_artifact.is_(False), # type: ignore[attr-defined]
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.is_model_artifact.is_(True), # type: ignore[attr-defined]
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.is_deployment_artifact.is_(True), # type: ignore[attr-defined]
219
+ ModelVersionArtifactSchema.artifact_version_id
220
+ == ArtifactVersionSchema.id,
221
+ ArtifactVersionSchema.type == "ServiceArtifact",
279
222
  )
280
223
  custom_filters.append(deployment_artifact_filter)
281
224