zenml-nightly 0.58.2.dev20240618__py3-none-any.whl → 0.58.2.dev20240620__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.dev20240620.dist-info}/METADATA +8 -7
  294. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240620.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.dev20240620.dist-info}/LICENSE +0 -0
  299. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240620.dist-info}/WHEEL +0 -0
  300. {zenml_nightly-0.58.2.dev20240618.dist-info → zenml_nightly-0.58.2.dev20240620.dist-info}/entry_points.txt +0 -0
@@ -61,14 +61,14 @@ def huggingface_model_deployer_step(
61
61
  step_name = context.step_run.name
62
62
 
63
63
  # update the step configuration with the real pipeline runtime information
64
- service_config = service_config.copy()
64
+ service_config = service_config.model_copy()
65
65
  service_config.pipeline_name = pipeline_name
66
66
  service_config.pipeline_step_name = step_name
67
67
 
68
68
  # fetch existing services with same pipeline name, step name and
69
69
  # model name
70
70
  existing_services = model_deployer.find_model_server(
71
- config=service_config.dict()
71
+ config=service_config.model_dump()
72
72
  )
73
73
 
74
74
  # even when the deploy decision is negative, if an existing model server
@@ -43,7 +43,7 @@ class HyperAIOrchestratorSettings(BaseSettings):
43
43
  mounts_from_to: Dict[str, str] = {}
44
44
 
45
45
 
46
- class HyperAIOrchestratorConfig( # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
46
+ class HyperAIOrchestratorConfig(
47
47
  BaseOrchestratorConfig, HyperAIOrchestratorSettings
48
48
  ):
49
49
  """Configuration for the HyperAI orchestrator.
@@ -22,7 +22,7 @@ import io
22
22
  from typing import Any, List, Optional, Type
23
23
 
24
24
  import paramiko
25
- from pydantic import Field, SecretStr
25
+ from pydantic import Field
26
26
 
27
27
  from zenml.exceptions import AuthorizationException
28
28
  from zenml.integrations.hyperai import (
@@ -40,6 +40,7 @@ from zenml.service_connectors.service_connector import (
40
40
  ServiceConnector,
41
41
  )
42
42
  from zenml.utils.enum_utils import StrEnum
43
+ from zenml.utils.secret_utils import PlainSerializedSecretStr
43
44
 
44
45
  logger = get_logger(__name__)
45
46
 
@@ -47,10 +48,10 @@ logger = get_logger(__name__)
47
48
  class HyperAICredentials(AuthenticationConfig):
48
49
  """HyperAI client authentication credentials."""
49
50
 
50
- base64_ssh_key: SecretStr = Field(
51
+ base64_ssh_key: PlainSerializedSecretStr = Field(
51
52
  title="SSH key (base64)",
52
53
  )
53
- ssh_passphrase: Optional[SecretStr] = Field(
54
+ ssh_passphrase: Optional[PlainSerializedSecretStr] = Field(
54
55
  default=None,
55
56
  title="SSH key passphrase",
56
57
  )
@@ -30,7 +30,7 @@ class KubeflowIntegration(Integration):
30
30
  """Definition of Kubeflow Integration for ZenML."""
31
31
 
32
32
  NAME = KUBEFLOW
33
- REQUIREMENTS = ["kfp==1.8.22"] # Only 1.x version that supports pyyaml 6
33
+ REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"] # Only 1.x version that supports pyyaml 6
34
34
 
35
35
  @classmethod
36
36
  def flavors(cls) -> List[Type[Flavor]]:
@@ -15,7 +15,7 @@
15
15
 
16
16
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast
17
17
 
18
- from pydantic import root_validator
18
+ from pydantic import model_validator
19
19
 
20
20
  from zenml.config.base_settings import BaseSettings
21
21
  from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
@@ -24,6 +24,7 @@ from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
24
24
  from zenml.logger import get_logger
25
25
  from zenml.models import ServiceConnectorRequirements
26
26
  from zenml.orchestrators import BaseOrchestratorConfig, BaseOrchestratorFlavor
27
+ from zenml.utils.pydantic_utils import before_validator_handler
27
28
  from zenml.utils.secret_utils import SecretField
28
29
 
29
30
  if TYPE_CHECKING:
@@ -52,8 +53,6 @@ class KubeflowOrchestratorSettings(BaseSettings):
52
53
  and `client_password` need to be set together.
53
54
  user_namespace: The user namespace to use when creating experiments
54
55
  and runs.
55
- node_selectors: Deprecated: Node selectors to apply to KFP pods.
56
- node_affinity: Deprecated: Node affinities to apply to KFP pods.
57
56
  pod_settings: Pod settings to apply.
58
57
  """
59
58
 
@@ -61,99 +60,75 @@ class KubeflowOrchestratorSettings(BaseSettings):
61
60
  timeout: int = 1200
62
61
 
63
62
  client_args: Dict[str, Any] = {}
64
- client_username: Optional[str] = SecretField()
65
- client_password: Optional[str] = SecretField()
63
+ client_username: Optional[str] = SecretField(default=None)
64
+ client_password: Optional[str] = SecretField(default=None)
66
65
  user_namespace: Optional[str] = None
67
- node_selectors: Dict[str, str] = {}
68
- node_affinity: Dict[str, List[str]] = {}
69
66
  pod_settings: Optional[KubernetesPodSettings] = None
70
67
 
71
- @root_validator
68
+ @model_validator(mode="before")
69
+ @classmethod
70
+ @before_validator_handler
72
71
  def _validate_and_migrate_pod_settings(
73
- cls, values: Dict[str, Any]
72
+ cls, data: Dict[str, Any]
74
73
  ) -> Dict[str, Any]:
75
74
  """Validates settings and migrates pod settings from older version.
76
75
 
77
76
  Args:
78
- values: Dict representing user-specified runtime settings.
77
+ data: Dict representing user-specified runtime settings.
79
78
 
80
79
  Returns:
81
80
  Validated settings.
82
81
 
83
82
  Raises:
84
- AssertionError: If old and new settings are used together.
85
83
  ValueError: If username and password are not specified together.
86
84
  """
87
- has_pod_settings = bool(values.get("pod_settings"))
88
-
89
- node_selectors = cast(
90
- Dict[str, str], values.get("node_selectors") or {}
91
- )
85
+ node_selectors = cast(Dict[str, str], data.get("node_selectors") or {})
92
86
  node_affinity = cast(
93
- Dict[str, List[str]], values.get("node_affinity") or {}
87
+ Dict[str, List[str]], data.get("node_affinity") or {}
94
88
  )
95
89
 
96
- has_old_settings = any([node_selectors, node_affinity])
97
-
98
- if has_old_settings:
99
- logger.warning(
100
- "The attributes `node_selectors` and `node_affinity` of the "
101
- "Kubeflow settings will be deprecated soon. Use the "
102
- "attribute `pod_settings` instead.",
103
- )
104
-
105
- if has_pod_settings and has_old_settings:
106
- raise AssertionError(
107
- "Got Kubeflow pod settings using both the deprecated "
108
- "attributes `node_selectors` and `node_affinity` as well as "
109
- "the new attribute `pod_settings`. Please specify Kubeflow "
110
- "pod settings only using the new `pod_settings` attribute."
111
- )
112
- elif has_old_settings:
90
+ affinity = {}
91
+ if node_affinity:
113
92
  from kubernetes import client as k8s_client
114
93
 
115
- affinity = {}
116
- if node_affinity:
117
- match_expressions = [
118
- k8s_client.V1NodeSelectorRequirement(
119
- key=key,
120
- operator="In",
121
- values=values,
122
- )
123
- for key, values in node_affinity.items()
124
- ]
125
-
126
- affinity = k8s_client.V1Affinity(
127
- node_affinity=k8s_client.V1NodeAffinity(
128
- required_during_scheduling_ignored_during_execution=k8s_client.V1NodeSelector(
129
- node_selector_terms=[
130
- k8s_client.V1NodeSelectorTerm(
131
- match_expressions=match_expressions
132
- )
133
- ]
134
- )
94
+ match_expressions = [
95
+ k8s_client.V1NodeSelectorRequirement(
96
+ key=key,
97
+ operator="In",
98
+ values=values,
99
+ )
100
+ for key, values in node_affinity.items()
101
+ ]
102
+
103
+ affinity = k8s_client.V1Affinity(
104
+ node_affinity=k8s_client.V1NodeAffinity(
105
+ required_during_scheduling_ignored_during_execution=k8s_client.V1NodeSelector(
106
+ node_selector_terms=[
107
+ k8s_client.V1NodeSelectorTerm(
108
+ match_expressions=match_expressions
109
+ )
110
+ ]
135
111
  )
136
112
  )
137
- pod_settings = KubernetesPodSettings(
138
- node_selectors=node_selectors, affinity=affinity
139
113
  )
140
- values["pod_settings"] = pod_settings
141
- values["node_affinity"] = {}
142
- values["node_selectors"] = {}
114
+ pod_settings = KubernetesPodSettings(
115
+ node_selectors=node_selectors, affinity=affinity
116
+ )
117
+ data["pod_settings"] = pod_settings
143
118
 
144
119
  # Validate username and password for auth cookie logic
145
- username = values.get("client_username")
146
- password = values.get("client_password")
120
+ username = data.get("client_username")
121
+ password = data.get("client_password")
147
122
  client_creds_error = "`client_username` and `client_password` both need to be set together."
148
123
  if username and password is None:
149
124
  raise ValueError(client_creds_error)
150
125
  if password and username is None:
151
126
  raise ValueError(client_creds_error)
152
127
 
153
- return values
128
+ return data
154
129
 
155
130
 
156
- class KubeflowOrchestratorConfig( # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
131
+ class KubeflowOrchestratorConfig(
157
132
  BaseOrchestratorConfig, KubeflowOrchestratorSettings
158
133
  ):
159
134
  """Configuration for the Kubeflow orchestrator.
@@ -170,25 +145,17 @@ class KubeflowOrchestratorConfig( # type: ignore[misc] # https://github.com/pyd
170
145
  Kubeflow Pipelines deployment (i.e. when `kubeflow_hostname` is
171
146
  set) or if the stack component is linked to a Kubernetes service
172
147
  connector.
173
- local: If `True`, the orchestrator will assume it is connected to a
174
- local kubernetes cluster and will perform additional validations and
175
- operations to allow using the orchestrator in combination with other
176
- local stack components that store data in the local filesystem
177
- (i.e. it will mount the local stores directory into the pipeline
178
- containers).
179
- skip_local_validations: If `True`, the local validations will be
180
- skipped.
181
148
  """
182
149
 
183
150
  kubeflow_hostname: Optional[str] = None
184
151
  kubeflow_namespace: str = "kubeflow"
185
152
  kubernetes_context: Optional[str] # TODO: Potential setting
186
- local: bool = False
187
- skip_local_validations: bool = False
188
153
 
189
- @root_validator(pre=True)
154
+ @model_validator(mode="before")
155
+ @classmethod
156
+ @before_validator_handler
190
157
  def _validate_deprecated_attrs(
191
- cls, values: Dict[str, Any]
158
+ cls, data: Dict[str, Any]
192
159
  ) -> Dict[str, Any]:
193
160
  """Pydantic root_validator for deprecated attributes.
194
161
 
@@ -197,7 +164,7 @@ class KubeflowOrchestratorConfig( # type: ignore[misc] # https://github.com/pyd
197
164
  mandatory in the meantime.
198
165
 
199
166
  Args:
200
- values: Values passed to the object constructor
167
+ data: Values passed to the object constructor
201
168
 
202
169
  Returns:
203
170
  Values passed to the object constructor
@@ -210,10 +177,10 @@ class KubeflowOrchestratorConfig( # type: ignore[misc] # https://github.com/pyd
210
177
 
211
178
  # remove deprecated attributes from values dict
212
179
  for attr in provisioning_attrs:
213
- if attr in values:
214
- del values[attr]
180
+ if attr in data:
181
+ del data[attr]
215
182
 
216
- return values
183
+ return data
217
184
 
218
185
  @property
219
186
  def is_remote(self) -> bool:
@@ -226,7 +193,7 @@ class KubeflowOrchestratorConfig( # type: ignore[misc] # https://github.com/pyd
226
193
  Returns:
227
194
  True if this config is for a remote component, False otherwise.
228
195
  """
229
- return not self.local
196
+ return True
230
197
 
231
198
  @property
232
199
  def is_local(self) -> bool:
@@ -235,7 +202,7 @@ class KubeflowOrchestratorConfig( # type: ignore[misc] # https://github.com/pyd
235
202
  Returns:
236
203
  True if this config is for a local component, False otherwise.
237
204
  """
238
- return self.local
205
+ return False
239
206
 
240
207
  @property
241
208
  def is_synchronous(self) -> bool: