zenml-nightly 0.58.2.dev20240618__py3-none-any.whl → 0.58.2.dev20240619__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 (300) hide show
  1. zenml/VERSION +1 -1
  2. zenml/_hub/client.py +8 -5
  3. zenml/actions/base_action.py +8 -10
  4. zenml/artifact_stores/base_artifact_store.py +20 -15
  5. zenml/artifact_stores/local_artifact_store.py +3 -2
  6. zenml/artifacts/artifact_config.py +34 -19
  7. zenml/artifacts/external_artifact.py +18 -8
  8. zenml/artifacts/external_artifact_config.py +14 -6
  9. zenml/artifacts/unmaterialized_artifact.py +2 -11
  10. zenml/cli/__init__.py +6 -0
  11. zenml/cli/artifact.py +20 -2
  12. zenml/cli/served_model.py +0 -1
  13. zenml/cli/server.py +3 -3
  14. zenml/cli/utils.py +36 -40
  15. zenml/cli/web_login.py +2 -2
  16. zenml/client.py +198 -24
  17. zenml/client_lazy_loader.py +20 -14
  18. zenml/config/base_settings.py +5 -6
  19. zenml/config/build_configuration.py +1 -1
  20. zenml/config/compiler.py +3 -3
  21. zenml/config/docker_settings.py +27 -28
  22. zenml/config/global_config.py +33 -37
  23. zenml/config/pipeline_configurations.py +8 -11
  24. zenml/config/pipeline_run_configuration.py +6 -2
  25. zenml/config/pipeline_spec.py +3 -4
  26. zenml/config/resource_settings.py +8 -9
  27. zenml/config/schedule.py +16 -20
  28. zenml/config/secret_reference_mixin.py +6 -3
  29. zenml/config/secrets_store_config.py +16 -23
  30. zenml/config/server_config.py +50 -46
  31. zenml/config/settings_resolver.py +1 -1
  32. zenml/config/source.py +45 -35
  33. zenml/config/step_configurations.py +53 -31
  34. zenml/config/store_config.py +20 -19
  35. zenml/config/strict_base_model.py +2 -6
  36. zenml/constants.py +26 -2
  37. zenml/container_registries/base_container_registry.py +3 -2
  38. zenml/container_registries/default_container_registry.py +3 -3
  39. zenml/event_hub/base_event_hub.py +1 -1
  40. zenml/event_sources/base_event_source.py +11 -16
  41. zenml/exceptions.py +4 -0
  42. zenml/integrations/airflow/__init__.py +2 -10
  43. zenml/integrations/airflow/flavors/airflow_orchestrator_flavor.py +6 -7
  44. zenml/integrations/airflow/orchestrators/airflow_orchestrator.py +13 -249
  45. zenml/integrations/airflow/orchestrators/dag_generator.py +5 -3
  46. zenml/integrations/argilla/flavors/argilla_annotator_flavor.py +5 -4
  47. zenml/integrations/aws/__init__.py +1 -1
  48. zenml/integrations/aws/flavors/aws_container_registry_flavor.py +3 -2
  49. zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py +11 -5
  50. zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py +6 -2
  51. zenml/integrations/aws/service_connectors/aws_service_connector.py +5 -4
  52. zenml/integrations/azure/flavors/azureml_step_operator_flavor.py +4 -4
  53. zenml/integrations/azure/service_connectors/azure_service_connector.py +4 -3
  54. zenml/integrations/azure/step_operators/azureml_step_operator.py +1 -1
  55. zenml/integrations/bentoml/steps/bentoml_deployer.py +1 -1
  56. zenml/integrations/bitbucket/plugins/event_sources/bitbucket_webhook_event_source.py +8 -12
  57. zenml/integrations/comet/flavors/comet_experiment_tracker_flavor.py +1 -1
  58. zenml/integrations/evidently/__init__.py +3 -4
  59. zenml/integrations/evidently/column_mapping.py +11 -3
  60. zenml/integrations/evidently/data_validators/evidently_data_validator.py +21 -3
  61. zenml/integrations/evidently/metrics.py +5 -6
  62. zenml/integrations/evidently/tests.py +5 -6
  63. zenml/integrations/facets/models.py +2 -6
  64. zenml/integrations/feast/__init__.py +3 -1
  65. zenml/integrations/feast/feature_stores/feast_feature_store.py +0 -23
  66. zenml/integrations/gcp/__init__.py +1 -1
  67. zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py +1 -1
  68. zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py +1 -1
  69. zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +234 -103
  70. zenml/integrations/gcp/service_connectors/gcp_service_connector.py +57 -42
  71. zenml/integrations/github/code_repositories/github_code_repository.py +1 -1
  72. zenml/integrations/github/plugins/event_sources/github_webhook_event_source.py +9 -13
  73. zenml/integrations/great_expectations/__init__.py +1 -1
  74. zenml/integrations/great_expectations/data_validators/ge_data_validator.py +44 -44
  75. zenml/integrations/great_expectations/flavors/great_expectations_data_validator_flavor.py +35 -2
  76. zenml/integrations/great_expectations/ge_store_backend.py +24 -11
  77. zenml/integrations/great_expectations/materializers/ge_materializer.py +3 -3
  78. zenml/integrations/great_expectations/utils.py +5 -5
  79. zenml/integrations/huggingface/__init__.py +3 -0
  80. zenml/integrations/huggingface/flavors/huggingface_model_deployer_flavor.py +1 -1
  81. zenml/integrations/huggingface/steps/__init__.py +3 -0
  82. zenml/integrations/huggingface/steps/accelerate_runner.py +149 -0
  83. zenml/integrations/huggingface/steps/huggingface_deployer.py +2 -2
  84. zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py +1 -1
  85. zenml/integrations/hyperai/service_connectors/hyperai_service_connector.py +4 -3
  86. zenml/integrations/kubeflow/__init__.py +1 -1
  87. zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py +48 -81
  88. zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py +295 -245
  89. zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py +1 -1
  90. zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py +11 -2
  91. zenml/integrations/kubernetes/pod_settings.py +17 -31
  92. zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py +8 -7
  93. zenml/integrations/label_studio/__init__.py +1 -3
  94. zenml/integrations/label_studio/annotators/label_studio_annotator.py +3 -4
  95. zenml/integrations/label_studio/flavors/label_studio_annotator_flavor.py +2 -2
  96. zenml/integrations/langchain/materializers/document_materializer.py +44 -8
  97. zenml/integrations/mlflow/__init__.py +9 -3
  98. zenml/integrations/mlflow/experiment_trackers/mlflow_experiment_tracker.py +1 -1
  99. zenml/integrations/mlflow/flavors/mlflow_experiment_tracker_flavor.py +29 -37
  100. zenml/integrations/mlflow/model_registries/mlflow_model_registry.py +4 -4
  101. zenml/integrations/mlflow/steps/mlflow_deployer.py +1 -1
  102. zenml/integrations/neptune/flavors/neptune_experiment_tracker_flavor.py +1 -1
  103. zenml/integrations/pigeon/flavors/pigeon_annotator_flavor.py +1 -1
  104. zenml/integrations/s3/flavors/s3_artifact_store_flavor.py +9 -8
  105. zenml/integrations/seldon/seldon_client.py +52 -67
  106. zenml/integrations/seldon/services/seldon_deployment.py +3 -3
  107. zenml/integrations/seldon/steps/seldon_deployer.py +4 -4
  108. zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py +15 -5
  109. zenml/integrations/skypilot_aws/__init__.py +1 -1
  110. zenml/integrations/skypilot_aws/flavors/skypilot_orchestrator_aws_vm_flavor.py +1 -1
  111. zenml/integrations/skypilot_azure/__init__.py +1 -1
  112. zenml/integrations/skypilot_azure/flavors/skypilot_orchestrator_azure_vm_flavor.py +1 -1
  113. zenml/integrations/skypilot_gcp/__init__.py +2 -1
  114. zenml/integrations/skypilot_gcp/flavors/skypilot_orchestrator_gcp_vm_flavor.py +1 -1
  115. zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py +2 -2
  116. zenml/integrations/spark/flavors/spark_step_operator_flavor.py +1 -1
  117. zenml/integrations/tekton/__init__.py +1 -1
  118. zenml/integrations/tekton/flavors/tekton_orchestrator_flavor.py +66 -23
  119. zenml/integrations/tekton/orchestrators/tekton_orchestrator.py +547 -233
  120. zenml/integrations/tensorboard/__init__.py +1 -12
  121. zenml/integrations/tensorboard/services/tensorboard_service.py +3 -5
  122. zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py +6 -6
  123. zenml/integrations/tensorflow/__init__.py +2 -10
  124. zenml/integrations/tensorflow/materializers/keras_materializer.py +17 -9
  125. zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py +9 -14
  126. zenml/integrations/whylogs/flavors/whylogs_data_validator_flavor.py +1 -1
  127. zenml/lineage_graph/lineage_graph.py +1 -1
  128. zenml/materializers/built_in_materializer.py +3 -3
  129. zenml/materializers/pydantic_materializer.py +2 -2
  130. zenml/metadata/lazy_load.py +4 -4
  131. zenml/metadata/metadata_types.py +64 -4
  132. zenml/model/model.py +79 -54
  133. zenml/model_deployers/base_model_deployer.py +14 -12
  134. zenml/model_registries/base_model_registry.py +17 -15
  135. zenml/models/__init__.py +79 -206
  136. zenml/models/v2/base/base.py +54 -41
  137. zenml/models/v2/base/base_plugin_flavor.py +2 -6
  138. zenml/models/v2/base/filter.py +91 -76
  139. zenml/models/v2/base/page.py +2 -12
  140. zenml/models/v2/base/scoped.py +4 -7
  141. zenml/models/v2/core/api_key.py +22 -8
  142. zenml/models/v2/core/artifact.py +2 -2
  143. zenml/models/v2/core/artifact_version.py +74 -40
  144. zenml/models/v2/core/code_repository.py +37 -10
  145. zenml/models/v2/core/component.py +65 -16
  146. zenml/models/v2/core/device.py +14 -4
  147. zenml/models/v2/core/event_source.py +1 -2
  148. zenml/models/v2/core/flavor.py +74 -8
  149. zenml/models/v2/core/logs.py +68 -8
  150. zenml/models/v2/core/model.py +8 -4
  151. zenml/models/v2/core/model_version.py +25 -6
  152. zenml/models/v2/core/model_version_artifact.py +51 -21
  153. zenml/models/v2/core/model_version_pipeline_run.py +45 -13
  154. zenml/models/v2/core/pipeline.py +37 -72
  155. zenml/models/v2/core/pipeline_build.py +29 -17
  156. zenml/models/v2/core/pipeline_deployment.py +18 -6
  157. zenml/models/v2/core/pipeline_namespace.py +113 -0
  158. zenml/models/v2/core/pipeline_run.py +50 -22
  159. zenml/models/v2/core/run_metadata.py +59 -36
  160. zenml/models/v2/core/schedule.py +37 -24
  161. zenml/models/v2/core/secret.py +31 -12
  162. zenml/models/v2/core/service.py +64 -36
  163. zenml/models/v2/core/service_account.py +24 -11
  164. zenml/models/v2/core/service_connector.py +219 -44
  165. zenml/models/v2/core/stack.py +45 -17
  166. zenml/models/v2/core/step_run.py +28 -8
  167. zenml/models/v2/core/tag.py +8 -4
  168. zenml/models/v2/core/trigger.py +2 -2
  169. zenml/models/v2/core/trigger_execution.py +1 -0
  170. zenml/models/v2/core/user.py +18 -21
  171. zenml/models/v2/core/workspace.py +13 -3
  172. zenml/models/v2/misc/build_item.py +3 -3
  173. zenml/models/v2/misc/external_user.py +2 -6
  174. zenml/models/v2/misc/hub_plugin_models.py +9 -9
  175. zenml/models/v2/misc/loaded_visualization.py +2 -2
  176. zenml/models/v2/misc/service_connector_type.py +8 -17
  177. zenml/models/v2/misc/user_auth.py +7 -2
  178. zenml/new/pipelines/build_utils.py +3 -3
  179. zenml/new/pipelines/pipeline.py +17 -13
  180. zenml/new/pipelines/run_utils.py +103 -1
  181. zenml/orchestrators/base_orchestrator.py +10 -7
  182. zenml/orchestrators/local_docker/local_docker_orchestrator.py +1 -1
  183. zenml/orchestrators/step_runner.py +3 -6
  184. zenml/orchestrators/utils.py +1 -1
  185. zenml/plugins/base_plugin_flavor.py +6 -10
  186. zenml/plugins/plugin_flavor_registry.py +3 -7
  187. zenml/secret/base_secret.py +7 -8
  188. zenml/service_connectors/docker_service_connector.py +4 -3
  189. zenml/service_connectors/service_connector.py +5 -12
  190. zenml/service_connectors/service_connector_registry.py +2 -4
  191. zenml/services/container/container_service.py +1 -1
  192. zenml/services/container/container_service_endpoint.py +1 -1
  193. zenml/services/local/local_service.py +1 -1
  194. zenml/services/local/local_service_endpoint.py +1 -1
  195. zenml/services/service.py +16 -10
  196. zenml/services/service_type.py +4 -5
  197. zenml/services/terraform/terraform_service.py +1 -1
  198. zenml/stack/flavor.py +1 -5
  199. zenml/stack/flavor_registry.py +4 -4
  200. zenml/stack/stack.py +4 -1
  201. zenml/stack/stack_component.py +55 -31
  202. zenml/steps/base_step.py +34 -28
  203. zenml/steps/entrypoint_function_utils.py +3 -5
  204. zenml/steps/utils.py +12 -14
  205. zenml/utils/cuda_utils.py +50 -0
  206. zenml/utils/deprecation_utils.py +18 -20
  207. zenml/utils/dict_utils.py +1 -1
  208. zenml/utils/filesync_model.py +65 -28
  209. zenml/utils/function_utils.py +260 -0
  210. zenml/utils/json_utils.py +131 -0
  211. zenml/utils/mlstacks_utils.py +2 -2
  212. zenml/utils/pydantic_utils.py +270 -62
  213. zenml/utils/secret_utils.py +65 -12
  214. zenml/utils/source_utils.py +2 -2
  215. zenml/utils/typed_model.py +5 -3
  216. zenml/utils/typing_utils.py +243 -0
  217. zenml/utils/yaml_utils.py +1 -1
  218. zenml/zen_server/auth.py +2 -2
  219. zenml/zen_server/cloud_utils.py +6 -6
  220. zenml/zen_server/deploy/base_provider.py +1 -1
  221. zenml/zen_server/deploy/deployment.py +6 -8
  222. zenml/zen_server/deploy/docker/docker_zen_server.py +3 -4
  223. zenml/zen_server/deploy/local/local_provider.py +0 -1
  224. zenml/zen_server/deploy/local/local_zen_server.py +6 -6
  225. zenml/zen_server/deploy/terraform/terraform_zen_server.py +4 -6
  226. zenml/zen_server/exceptions.py +4 -1
  227. zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py +1 -1
  228. zenml/zen_server/pipeline_deployment/utils.py +48 -68
  229. zenml/zen_server/rbac/models.py +2 -5
  230. zenml/zen_server/rbac/utils.py +11 -14
  231. zenml/zen_server/routers/auth_endpoints.py +2 -2
  232. zenml/zen_server/routers/pipeline_builds_endpoints.py +1 -1
  233. zenml/zen_server/routers/runs_endpoints.py +1 -1
  234. zenml/zen_server/routers/secrets_endpoints.py +3 -2
  235. zenml/zen_server/routers/server_endpoints.py +1 -1
  236. zenml/zen_server/routers/steps_endpoints.py +1 -1
  237. zenml/zen_server/routers/workspaces_endpoints.py +1 -1
  238. zenml/zen_stores/base_zen_store.py +46 -9
  239. zenml/zen_stores/migrations/utils.py +42 -46
  240. zenml/zen_stores/migrations/versions/0701da9951a0_added_service_table.py +1 -1
  241. zenml/zen_stores/migrations/versions/1041bc644e0d_remove_secrets_manager.py +5 -3
  242. zenml/zen_stores/migrations/versions/10a907dad202_delete_mlmd_tables.py +1 -1
  243. zenml/zen_stores/migrations/versions/26b776ad583e_redesign_artifacts.py +8 -10
  244. zenml/zen_stores/migrations/versions/37835ce041d2_optimizing_database.py +3 -3
  245. zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py +10 -12
  246. zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py +3 -2
  247. zenml/zen_stores/migrations/versions/6917bce75069_add_pipeline_run_unique_constraint.py +4 -4
  248. zenml/zen_stores/migrations/versions/728c6369cfaa_add_name_column_to_input_artifact_pk.py +3 -2
  249. zenml/zen_stores/migrations/versions/743ec82b1b3c_update_size_of_build_images.py +2 -2
  250. zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py +3 -2
  251. zenml/zen_stores/migrations/versions/7834208cc3f6_artifact_project_scoping.py +8 -7
  252. zenml/zen_stores/migrations/versions/7b651bf6822e_track_secrets_in_db.py +6 -4
  253. zenml/zen_stores/migrations/versions/7e4a481d17f7_add_identity_table.py +2 -2
  254. zenml/zen_stores/migrations/versions/7f603e583dd7_fixed_migration.py +1 -1
  255. zenml/zen_stores/migrations/versions/a39c4184c8ce_remove_secrets_manager_flavors.py +2 -2
  256. zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py +4 -4
  257. zenml/zen_stores/migrations/versions/alembic_start.py +1 -1
  258. zenml/zen_stores/migrations/versions/fbd7f18ced1e_increase_step_run_field_lengths.py +4 -4
  259. zenml/zen_stores/rest_zen_store.py +109 -49
  260. zenml/zen_stores/schemas/api_key_schemas.py +1 -1
  261. zenml/zen_stores/schemas/artifact_schemas.py +8 -8
  262. zenml/zen_stores/schemas/artifact_visualization_schemas.py +3 -3
  263. zenml/zen_stores/schemas/code_repository_schemas.py +1 -1
  264. zenml/zen_stores/schemas/component_schemas.py +8 -3
  265. zenml/zen_stores/schemas/device_schemas.py +8 -6
  266. zenml/zen_stores/schemas/event_source_schemas.py +3 -4
  267. zenml/zen_stores/schemas/flavor_schemas.py +5 -3
  268. zenml/zen_stores/schemas/model_schemas.py +26 -1
  269. zenml/zen_stores/schemas/pipeline_build_schemas.py +1 -1
  270. zenml/zen_stores/schemas/pipeline_deployment_schemas.py +4 -4
  271. zenml/zen_stores/schemas/pipeline_run_schemas.py +6 -6
  272. zenml/zen_stores/schemas/pipeline_schemas.py +5 -2
  273. zenml/zen_stores/schemas/run_metadata_schemas.py +2 -2
  274. zenml/zen_stores/schemas/secret_schemas.py +8 -5
  275. zenml/zen_stores/schemas/server_settings_schemas.py +3 -1
  276. zenml/zen_stores/schemas/service_connector_schemas.py +1 -1
  277. zenml/zen_stores/schemas/service_schemas.py +11 -2
  278. zenml/zen_stores/schemas/stack_schemas.py +1 -1
  279. zenml/zen_stores/schemas/step_run_schemas.py +11 -11
  280. zenml/zen_stores/schemas/tag_schemas.py +6 -2
  281. zenml/zen_stores/schemas/trigger_schemas.py +2 -2
  282. zenml/zen_stores/schemas/user_schemas.py +2 -2
  283. zenml/zen_stores/schemas/workspace_schemas.py +3 -1
  284. zenml/zen_stores/secrets_stores/aws_secrets_store.py +19 -20
  285. zenml/zen_stores/secrets_stores/azure_secrets_store.py +17 -20
  286. zenml/zen_stores/secrets_stores/base_secrets_store.py +79 -12
  287. zenml/zen_stores/secrets_stores/gcp_secrets_store.py +17 -20
  288. zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py +4 -8
  289. zenml/zen_stores/secrets_stores/service_connector_secrets_store.py +10 -7
  290. zenml/zen_stores/secrets_stores/sql_secrets_store.py +5 -6
  291. zenml/zen_stores/sql_zen_store.py +196 -120
  292. zenml/zen_stores/zen_store_interface.py +33 -0
  293. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240619.dist-info}/METADATA +8 -7
  294. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240619.dist-info}/RECORD +297 -294
  295. zenml/integrations/kubeflow/utils.py +0 -95
  296. zenml/models/v2/base/internal.py +0 -37
  297. zenml/models/v2/base/update.py +0 -44
  298. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240619.dist-info}/LICENSE +0 -0
  299. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240619.dist-info}/WHEEL +0 -0
  300. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240619.dist-info}/entry_points.txt +0 -0
@@ -30,8 +30,12 @@ from typing import (
30
30
  )
31
31
  from uuid import UUID
32
32
 
33
- from pydantic import BaseModel, Field, root_validator, validator
34
- from pydantic.typing import get_args
33
+ from pydantic import (
34
+ BaseModel,
35
+ Field,
36
+ field_validator,
37
+ model_validator,
38
+ )
35
39
  from sqlmodel import SQLModel
36
40
 
37
41
  from zenml.constants import (
@@ -43,9 +47,11 @@ from zenml.constants import (
43
47
  from zenml.enums import GenericFilterOps, LogicalOperators, SorterOps
44
48
  from zenml.exceptions import ValidationError
45
49
  from zenml.logger import get_logger
50
+ from zenml.utils.pydantic_utils import before_validator_handler
51
+ from zenml.utils.typing_utils import get_args
46
52
 
47
53
  if TYPE_CHECKING:
48
- from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
54
+ from sqlalchemy.sql.elements import ColumnElement
49
55
 
50
56
  from zenml.zen_stores.schemas import BaseSchema
51
57
 
@@ -73,14 +79,15 @@ class Filter(BaseModel, ABC):
73
79
 
74
80
  operation: GenericFilterOps
75
81
  column: str
76
- value: Any
82
+ value: Optional[Any] = None
77
83
 
78
- @validator("operation", pre=True)
79
- def validate_operation(cls, op: str) -> str:
84
+ @field_validator("operation", mode="before")
85
+ @classmethod
86
+ def validate_operation(cls, value: Any) -> Any:
80
87
  """Validate that the operation is a valid op for the field type.
81
88
 
82
89
  Args:
83
- op: The operation of this filter.
90
+ value: The operation of this filter.
84
91
 
85
92
  Returns:
86
93
  The operation if it is valid.
@@ -88,18 +95,18 @@ class Filter(BaseModel, ABC):
88
95
  Raises:
89
96
  ValueError: If the operation is not valid for this field type.
90
97
  """
91
- if op not in cls.ALLOWED_OPS:
98
+ if value not in cls.ALLOWED_OPS:
92
99
  raise ValueError(
93
100
  f"This datatype can not be filtered using this operation: "
94
- f"'{op}'. The allowed operations are: {cls.ALLOWED_OPS}"
101
+ f"'{value}'. The allowed operations are: {cls.ALLOWED_OPS}"
95
102
  )
96
103
  else:
97
- return op
104
+ return value
98
105
 
99
106
  def generate_query_conditions(
100
107
  self,
101
108
  table: Type[SQLModel],
102
- ) -> Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]:
109
+ ) -> Union["ColumnElement[bool]"]:
103
110
  """Generate the query conditions for the database.
104
111
 
105
112
  This method converts the Filter class into an appropriate SQLModel
@@ -204,7 +211,7 @@ class UUIDFilter(StrFilter):
204
211
  class NumericFilter(Filter):
205
212
  """Filter for all numeric fields."""
206
213
 
207
- value: Union[float, datetime]
214
+ value: Union[float, datetime] = Field(union_mode="left_to_right")
208
215
 
209
216
  ALLOWED_OPS: ClassVar[List[str]] = [
210
217
  GenericFilterOps.EQUALS,
@@ -286,25 +293,28 @@ class BaseFilter(BaseModel):
286
293
  )
287
294
 
288
295
  id: Optional[Union[UUID, str]] = Field(
289
- default=None, description="Id for this resource"
296
+ default=None,
297
+ description="Id for this resource",
298
+ union_mode="left_to_right",
290
299
  )
291
300
  created: Optional[Union[datetime, str]] = Field(
292
- default=None, description="Created"
301
+ default=None, description="Created", union_mode="left_to_right"
293
302
  )
294
303
  updated: Optional[Union[datetime, str]] = Field(
295
- default=None, description="Updated"
304
+ default=None, description="Updated", union_mode="left_to_right"
296
305
  )
297
306
 
298
307
  _rbac_configuration: Optional[
299
308
  Tuple[UUID, Dict[str, Optional[Set[UUID]]]]
300
309
  ] = None
301
310
 
302
- @validator("sort_by", pre=True)
303
- def validate_sort_by(cls, v: str) -> str:
311
+ @field_validator("sort_by", mode="before")
312
+ @classmethod
313
+ def validate_sort_by(cls, value: Any) -> Any:
304
314
  """Validate that the sort_column is a valid column with a valid operand.
305
315
 
306
316
  Args:
307
- v: The sort_by field value.
317
+ value: The sort_by field value.
308
318
 
309
319
  Returns:
310
320
  The validated sort_by field value.
@@ -316,13 +326,13 @@ class BaseFilter(BaseModel):
316
326
  # Somehow pydantic allows you to pass in int values, which will be
317
327
  # interpreted as string, however within the validator they are still
318
328
  # integers, which don't have a .split() method
319
- if not isinstance(v, str):
329
+ if not isinstance(value, str):
320
330
  raise ValidationError(
321
331
  f"str type expected for the sort_by field. "
322
- f"Received a {type(v)}"
332
+ f"Received a {type(value)}"
323
333
  )
324
- column = v
325
- split_value = v.split(":", 1)
334
+ column = value
335
+ split_value = value.split(":", 1)
326
336
  if len(split_value) == 2:
327
337
  column = split_value[1]
328
338
 
@@ -334,31 +344,33 @@ class BaseFilter(BaseModel):
334
344
  SorterOps.values(),
335
345
  column,
336
346
  )
337
- v = column
347
+ value = column
338
348
 
339
349
  if column in cls.FILTER_EXCLUDE_FIELDS:
340
350
  raise ValueError(
341
- f"This resource can not be sorted by this field: '{v}'"
351
+ f"This resource can not be sorted by this field: '{value}'"
342
352
  )
343
- elif column in cls.__fields__:
344
- return v
353
+ elif column in cls.model_fields:
354
+ return value
345
355
  else:
346
356
  raise ValueError(
347
357
  "You can only sort by valid fields of this resource"
348
358
  )
349
359
 
350
- @root_validator(pre=True)
351
- def filter_ops(cls, values: Dict[str, Any]) -> Dict[str, Any]:
360
+ @model_validator(mode="before")
361
+ @classmethod
362
+ @before_validator_handler
363
+ def filter_ops(cls, data: Dict[str, Any]) -> Dict[str, Any]:
352
364
  """Parse incoming filters to ensure all filters are legal.
353
365
 
354
366
  Args:
355
- values: The values of the class.
367
+ data: The values of the class.
356
368
 
357
369
  Returns:
358
370
  The values of the class.
359
371
  """
360
- cls._generate_filter_list(values)
361
- return values
372
+ cls._generate_filter_list(data)
373
+ return data
362
374
 
363
375
  @property
364
376
  def list_of_filters(self) -> List[Filter]:
@@ -368,7 +380,7 @@ class BaseFilter(BaseModel):
368
380
  A list of Filter models.
369
381
  """
370
382
  return self._generate_filter_list(
371
- {key: getattr(self, key) for key in self.__fields__}
383
+ {key: getattr(self, key) for key in self.model_fields}
372
384
  )
373
385
 
374
386
  @property
@@ -403,14 +415,14 @@ class BaseFilter(BaseModel):
403
415
  column_allowed_ids: Set of IDs per column to limit the query to.
404
416
  If given, the remaining filters will be applied to entities
405
417
  within this set only. If `None`, the remaining filters will
406
- applied to all entries in the table.
418
+ be applied to all entries in the table.
407
419
  """
408
420
  self._rbac_configuration = (authenticated_user_id, column_allowed_ids)
409
421
 
410
422
  def generate_rbac_filter(
411
423
  self,
412
424
  table: Type["AnySchema"],
413
- ) -> Optional["BooleanClauseList[Any]"]:
425
+ ) -> Optional["ColumnElement[bool]"]:
414
426
  """Generates an optional RBAC filter.
415
427
 
416
428
  Args:
@@ -563,7 +575,7 @@ class BaseFilter(BaseModel):
563
575
 
564
576
  # Handle unsupported datatypes
565
577
  logger.warning(
566
- f"The Datatype {cls.__fields__[column].type_} might not be "
578
+ f"The Datatype {cls.model_fields[column].annotation} might not be "
567
579
  "supported for filtering. Defaulting to a string filter."
568
580
  )
569
581
  return StrFilter(
@@ -572,6 +584,37 @@ class BaseFilter(BaseModel):
572
584
  value=str(value),
573
585
  )
574
586
 
587
+ @classmethod
588
+ def check_field_annotation(cls, k: str, type_: Any) -> bool:
589
+ """Checks whether a model field has a certain annotation.
590
+
591
+ Args:
592
+ k: The name of the field.
593
+ type_: The type to check.
594
+
595
+ Raises:
596
+ ValueError: if the model field within does not have an annotation.
597
+
598
+ Returns:
599
+ True if the annotation of the field matches the given type, False
600
+ otherwise.
601
+ """
602
+ try:
603
+ annotation = cls.model_fields[k].annotation
604
+
605
+ if annotation is not None:
606
+ return (
607
+ issubclass(type_, get_args(annotation))
608
+ or annotation is type_
609
+ )
610
+ else:
611
+ raise ValueError(
612
+ f"The field '{k}' inside the model {cls.__name__} "
613
+ "does not have an annotation."
614
+ )
615
+ except TypeError:
616
+ return False
617
+
575
618
  @classmethod
576
619
  def is_datetime_field(cls, k: str) -> bool:
577
620
  """Checks if it's a datetime field.
@@ -582,40 +625,31 @@ class BaseFilter(BaseModel):
582
625
  Returns:
583
626
  True if the field is a datetime field, False otherwise.
584
627
  """
585
- return (
586
- issubclass(datetime, get_args(cls.__fields__[k].type_))
587
- or cls.__fields__[k].type_ is datetime
588
- )
628
+ return cls.check_field_annotation(k=k, type_=datetime)
589
629
 
590
630
  @classmethod
591
631
  def is_uuid_field(cls, k: str) -> bool:
592
- """Checks if it's a uuid field.
632
+ """Checks if it's a UUID field.
593
633
 
594
634
  Args:
595
635
  k: The key to check.
596
636
 
597
637
  Returns:
598
- True if the field is a uuid field, False otherwise.
638
+ True if the field is a UUID field, False otherwise.
599
639
  """
600
- return (
601
- issubclass(UUID, get_args(cls.__fields__[k].type_))
602
- or cls.__fields__[k].type_ is UUID
603
- )
640
+ return cls.check_field_annotation(k=k, type_=UUID)
604
641
 
605
642
  @classmethod
606
643
  def is_int_field(cls, k: str) -> bool:
607
- """Checks if it's a int field.
644
+ """Checks if it's an int field.
608
645
 
609
646
  Args:
610
647
  k: The key to check.
611
648
 
612
649
  Returns:
613
- True if the field is a int field, False otherwise.
650
+ True if the field is an int field, False otherwise.
614
651
  """
615
- return (
616
- issubclass(int, get_args(cls.__fields__[k].type_))
617
- or cls.__fields__[k].type_ is int
618
- )
652
+ return cls.check_field_annotation(k=k, type_=int)
619
653
 
620
654
  @classmethod
621
655
  def is_bool_field(cls, k: str) -> bool:
@@ -627,10 +661,7 @@ class BaseFilter(BaseModel):
627
661
  Returns:
628
662
  True if the field is a bool field, False otherwise.
629
663
  """
630
- return (
631
- issubclass(bool, get_args(cls.__fields__[k].type_))
632
- or cls.__fields__[k].type_ is bool
633
- )
664
+ return cls.check_field_annotation(k=k, type_=bool)
634
665
 
635
666
  @classmethod
636
667
  def is_str_field(cls, k: str) -> bool:
@@ -642,10 +673,7 @@ class BaseFilter(BaseModel):
642
673
  Returns:
643
674
  True if the field is a string field, False otherwise.
644
675
  """
645
- return (
646
- issubclass(str, get_args(cls.__fields__[k].type_))
647
- or cls.__fields__[k].type_ is str
648
- )
676
+ return cls.check_field_annotation(k=k, type_=str)
649
677
 
650
678
  @classmethod
651
679
  def is_sort_by_field(cls, k: str) -> bool:
@@ -657,10 +685,7 @@ class BaseFilter(BaseModel):
657
685
  Returns:
658
686
  True if the field is a sort by field, False otherwise.
659
687
  """
660
- return (
661
- issubclass(str, get_args(cls.__fields__[k].type_))
662
- or cls.__fields__[k].type_ == str
663
- ) and k == "sort_by"
688
+ return cls.check_field_annotation(k=k, type_=str) and k == "sort_by"
664
689
 
665
690
  @staticmethod
666
691
  def _define_datetime_filter(
@@ -772,7 +797,7 @@ class BaseFilter(BaseModel):
772
797
 
773
798
  def generate_filter(
774
799
  self, table: Type[SQLModel]
775
- ) -> Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]:
800
+ ) -> Union["ColumnElement[bool]"]:
776
801
  """Generate the filter for the query.
777
802
 
778
803
  Args:
@@ -784,8 +809,7 @@ class BaseFilter(BaseModel):
784
809
  Raises:
785
810
  RuntimeError: If a valid logical operator is not supplied.
786
811
  """
787
- from sqlalchemy import and_
788
- from sqlmodel import or_
812
+ from sqlmodel import and_, or_
789
813
 
790
814
  filters = []
791
815
  for column_filter in self.list_of_filters:
@@ -801,9 +825,7 @@ class BaseFilter(BaseModel):
801
825
  else:
802
826
  raise RuntimeError("No valid logical operator was supplied.")
803
827
 
804
- def get_custom_filters(
805
- self,
806
- ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
828
+ def get_custom_filters(self) -> List["ColumnElement[bool]"]:
807
829
  """Get custom filters.
808
830
 
809
831
  This can be overridden by subclasses to define custom filters that are
@@ -839,10 +861,3 @@ class BaseFilter(BaseModel):
839
861
  query = query.where(filters)
840
862
 
841
863
  return query
842
-
843
- class Config:
844
- """Pydantic configuration class."""
845
-
846
- # all attributes with leading underscore are private and therefore
847
- # are mutable and not included in serialization
848
- underscore_attrs_are_private = True
@@ -15,8 +15,7 @@
15
15
 
16
16
  from typing import Generator, Generic, List, TypeVar
17
17
 
18
- from pydantic import SecretStr
19
- from pydantic.generics import GenericModel
18
+ from pydantic import BaseModel
20
19
  from pydantic.types import NonNegativeInt, PositiveInt
21
20
 
22
21
  from zenml.models.v2.base.base import BaseResponse
@@ -25,7 +24,7 @@ from zenml.models.v2.base.filter import BaseFilter
25
24
  B = TypeVar("B", bound=BaseResponse) # type: ignore[type-arg]
26
25
 
27
26
 
28
- class Page(GenericModel, Generic[B]):
27
+ class Page(BaseModel, Generic[B]):
29
28
  """Return Model for List Models to accommodate pagination."""
30
29
 
31
30
  index: PositiveInt
@@ -91,12 +90,3 @@ class Page(GenericModel, Generic[B]):
91
90
  Whether the item is in the page.
92
91
  """
93
92
  return item in self.items
94
-
95
- class Config:
96
- """Pydantic configuration class."""
97
-
98
- # This is needed to allow the REST API server to unpack SecretStr
99
- # values correctly before sending them to the client.
100
- json_encoders = {
101
- SecretStr: lambda v: v.get_secret_value() if v else None
102
- }
@@ -23,7 +23,6 @@ from typing import (
23
23
  Optional,
24
24
  Type,
25
25
  TypeVar,
26
- Union,
27
26
  )
28
27
  from uuid import UUID
29
28
 
@@ -40,7 +39,7 @@ from zenml.models.v2.base.base import (
40
39
  from zenml.models.v2.base.filter import AnyQuery, BaseFilter
41
40
 
42
41
  if TYPE_CHECKING:
43
- from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
42
+ from sqlalchemy.sql.elements import ColumnElement
44
43
 
45
44
  from zenml.models.v2.core.user import UserResponse
46
45
  from zenml.models.v2.core.workspace import WorkspaceResponse
@@ -99,7 +98,7 @@ class UserScopedResponseBody(BaseDatedResponseBody):
99
98
  """Base user-owned body."""
100
99
 
101
100
  user: Optional["UserResponse"] = Field(
102
- title="The user who created this resource."
101
+ title="The user who created this resource.", default=None
103
102
  )
104
103
 
105
104
 
@@ -332,9 +331,7 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
332
331
 
333
332
  return query
334
333
 
335
- def get_custom_filters(
336
- self,
337
- ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
334
+ def get_custom_filters(self) -> List["ColumnElement[bool]"]:
338
335
  """Get custom tag filters.
339
336
 
340
337
  Returns:
@@ -344,6 +341,6 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
344
341
 
345
342
  custom_filters = super().get_custom_filters()
346
343
  if self.tag:
347
- custom_filters.append(col(TagSchema.name) == self.tag) # type: ignore[arg-type]
344
+ custom_filters.append(col(TagSchema.name) == self.tag)
348
345
 
349
346
  return custom_filters
@@ -31,9 +31,9 @@ from zenml.models.v2.base.base import (
31
31
  BaseRequest,
32
32
  BaseResponseMetadata,
33
33
  BaseResponseResources,
34
+ BaseUpdate,
34
35
  )
35
36
  from zenml.models.v2.base.filter import AnyQuery, BaseFilter
36
- from zenml.models.v2.base.update import update_model
37
37
  from zenml.utils.string_utils import b64_decode, b64_encode
38
38
 
39
39
  if TYPE_CHECKING:
@@ -64,7 +64,7 @@ class APIKey(BaseModel):
64
64
  encoded_key = encoded_key[len(ZENML_API_KEY_PREFIX) :]
65
65
  try:
66
66
  json_key = b64_decode(encoded_key)
67
- return cls.parse_raw(json_key)
67
+ return cls.model_validate_json(json_key)
68
68
  except Exception:
69
69
  raise ValueError("Invalid API key.")
70
70
 
@@ -74,7 +74,7 @@ class APIKey(BaseModel):
74
74
  Returns:
75
75
  The encoded API key.
76
76
  """
77
- encoded_key = b64_encode(self.json())
77
+ encoded_key = b64_encode(self.model_dump_json())
78
78
  return f"{ZENML_API_KEY_PREFIX}{encoded_key}"
79
79
 
80
80
 
@@ -109,13 +109,22 @@ class APIKeyRotateRequest(BaseModel):
109
109
  # ------------------ Update Model ------------------
110
110
 
111
111
 
112
- @update_model
113
- class APIKeyUpdate(APIKeyRequest):
112
+ class APIKeyUpdate(BaseUpdate):
114
113
  """Update model for API keys."""
115
114
 
115
+ name: Optional[str] = Field(
116
+ title="The name of the API Key.",
117
+ max_length=STR_FIELD_MAX_LENGTH,
118
+ default=None,
119
+ )
120
+ description: Optional[str] = Field(
121
+ title="The description of the API Key.",
122
+ max_length=TEXT_FIELD_MAX_LENGTH,
123
+ default=None,
124
+ )
116
125
  active: Optional[bool] = Field(
117
- default=True,
118
126
  title="Whether the API key is active.",
127
+ default=None,
119
128
  )
120
129
 
121
130
 
@@ -349,12 +358,17 @@ class APIKeyFilter(BaseFilter):
349
358
  active: Optional[Union[bool, str]] = Field(
350
359
  default=None,
351
360
  title="Whether the API key is active.",
361
+ union_mode="left_to_right",
352
362
  )
353
363
  last_login: Optional[Union[datetime, str]] = Field(
354
- default=None, title="Time when the API key was last used to log in."
364
+ default=None,
365
+ title="Time when the API key was last used to log in.",
366
+ union_mode="left_to_right",
355
367
  )
356
368
  last_rotated: Optional[Union[datetime, str]] = Field(
357
- default=None, title="Time when the API key was last rotated."
369
+ default=None,
370
+ title="Time when the API key was last rotated.",
371
+ union_mode="left_to_right",
358
372
  )
359
373
 
360
374
  def set_service_account(self, service_account_id: UUID) -> None:
@@ -74,8 +74,8 @@ class ArtifactResponseBody(BaseDatedResponseBody):
74
74
  tags: List[TagResponse] = Field(
75
75
  title="Tags associated with the model",
76
76
  )
77
- latest_version_name: Optional[str]
78
- latest_version_id: Optional[UUID]
77
+ latest_version_name: Optional[str] = None
78
+ latest_version_id: Optional[UUID] = None
79
79
 
80
80
 
81
81
  class ArtifactResponseMetadata(BaseResponseMetadata):