zenml-nightly 0.57.1.dev20240523__py3-none-any.whl → 0.58.0.dev20240528__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 (116) hide show
  1. README.md +1 -1
  2. RELEASE_NOTES.md +88 -0
  3. zenml/VERSION +1 -1
  4. zenml/cli/base.py +15 -16
  5. zenml/client.py +11 -1
  6. zenml/config/__init__.py +2 -0
  7. zenml/config/compiler.py +2 -2
  8. zenml/config/pipeline_configurations.py +2 -0
  9. zenml/config/pipeline_run_configuration.py +2 -0
  10. zenml/config/retry_config.py +27 -0
  11. zenml/config/server_config.py +13 -9
  12. zenml/config/step_configurations.py +2 -0
  13. zenml/constants.py +1 -0
  14. zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py +2 -0
  15. zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +14 -0
  16. zenml/integrations/kubernetes/orchestrators/manifest_utils.py +4 -0
  17. zenml/integrations/mlflow/steps/mlflow_deployer.py +1 -1
  18. zenml/integrations/slack/alerters/slack_alerter.py +0 -2
  19. zenml/model/model.py +77 -45
  20. zenml/models/v2/core/model_version.py +1 -1
  21. zenml/models/v2/core/pipeline_run.py +12 -0
  22. zenml/models/v2/core/step_run.py +12 -0
  23. zenml/models/v2/misc/server_models.py +9 -3
  24. zenml/new/pipelines/run_utils.py +8 -2
  25. zenml/new/steps/step_context.py +6 -6
  26. zenml/new/steps/step_decorator.py +5 -0
  27. zenml/orchestrators/step_launcher.py +71 -53
  28. zenml/orchestrators/step_runner.py +28 -134
  29. zenml/orchestrators/utils.py +158 -1
  30. zenml/steps/base_step.py +16 -9
  31. zenml/utils/dashboard_utils.py +4 -8
  32. zenml/zen_server/dashboard/assets/{404-DtHkZQ2t.js → 404-DIgYOHj5.js} +1 -1
  33. zenml/zen_server/dashboard/assets/{@reactflow-BfgHyCG7.js → @reactflow-CF_wPv-n.js} +1 -1
  34. zenml/zen_server/dashboard/assets/{AwarenessChannel--3KuLcwp.js → AwarenessChannel-CqRB1HO9.js} +1 -1
  35. zenml/zen_server/dashboard/assets/{Cards-BzhQtXRm.js → Cards-DuJ8uYGM.js} +1 -1
  36. zenml/zen_server/dashboard/assets/{CodeSnippet-DpWTCS7Y.js → CodeSnippet-CBFcxTGW.js} +1 -1
  37. zenml/zen_server/dashboard/assets/Commands-BQb0_PAa.js +1 -0
  38. zenml/zen_server/dashboard/assets/CopyButton-BOFZCr5Z.js +2 -0
  39. zenml/zen_server/dashboard/assets/{CsvVizualization-B9nkywOx.js → CsvVizualization-D4QRvdDk.js} +1 -1
  40. zenml/zen_server/dashboard/assets/{Error-BbzLqrpR.js → Error-C9i5GdiH.js} +1 -1
  41. zenml/zen_server/dashboard/assets/{Helpbox-SS5TXXFE.js → Helpbox-CnuWjxcc.js} +1 -1
  42. zenml/zen_server/dashboard/assets/{Infobox-CV3EG_9t.js → Infobox-DV6YqO5N.js} +1 -1
  43. zenml/zen_server/dashboard/assets/{InlineAvatar-wz4SSxG5.js → InlineAvatar-DE_1-SDN.js} +1 -1
  44. zenml/zen_server/dashboard/assets/{PageHeader-Bz7qVuF9.js → PageHeader-D7r1Jb6H.js} +1 -1
  45. zenml/zen_server/dashboard/assets/{Pagination-BEP_9EMH.js → Pagination-0-Hv-4DJ.js} +1 -1
  46. zenml/zen_server/dashboard/assets/{SetPassword-B0yU0HIe.js → SetPassword-074iDZCm.js} +1 -1
  47. zenml/zen_server/dashboard/assets/SuccessStep-DEeKB0Wq.js +1 -0
  48. zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-yy-qBmO9.js → UpdatePasswordSchemas-4be5el9x.js} +1 -1
  49. zenml/zen_server/dashboard/assets/{cloud-only-DhayG_Yj.js → cloud-only-lInATL8-.js} +1 -1
  50. zenml/zen_server/dashboard/assets/connectors-video-C9qY4syJ.svg +21 -0
  51. zenml/zen_server/dashboard/assets/{index-Bbgc_ynV.css → index-D4zsM9tz.css} +1 -1
  52. zenml/zen_server/dashboard/assets/{index-UnEiujvd.js → index-DKgtNG83.js} +3 -3
  53. zenml/zen_server/dashboard/assets/{login-mutation-uNylFM6V.js → login-mutation-BfEeVmSl.js} +1 -1
  54. zenml/zen_server/dashboard/assets/{not-found-CvNr1Owb.js → not-found-an_3Z-Ep.js} +1 -1
  55. zenml/zen_server/dashboard/assets/page-0AGNIDYi.js +1 -0
  56. zenml/zen_server/dashboard/assets/{page-1rE3Idho.js → page-68OLHx44.js} +2 -2
  57. zenml/zen_server/dashboard/assets/page-6Gp_NCQL.js +14 -0
  58. zenml/zen_server/dashboard/assets/{page-BAuah6Do.js → page-B90yVxNb.js} +1 -1
  59. zenml/zen_server/dashboard/assets/page-Ba5l9-eg.js +1 -0
  60. zenml/zen_server/dashboard/assets/{page-BCpiaYh8.js → page-BiyORkbM.js} +1 -1
  61. zenml/zen_server/dashboard/assets/{page-DIZKf8ZO.js → page-Bjj9GHVF.js} +1 -1
  62. zenml/zen_server/dashboard/assets/{page-CM80cPrr.js → page-BnrWbXYJ.js} +1 -1
  63. zenml/zen_server/dashboard/assets/{page-CjxYF1Om.js → page-C0opS_t4.js} +1 -1
  64. zenml/zen_server/dashboard/assets/{page-BldrIvaZ.js → page-C1Jl09zF.js} +1 -1
  65. zenml/zen_server/dashboard/assets/{page-BVZ01rdG.js → page-C6Xn5o3F.js} +1 -1
  66. zenml/zen_server/dashboard/assets/page-CUqLUGbU.js +1 -0
  67. zenml/zen_server/dashboard/assets/page-CWpGAjSi.js +1 -0
  68. zenml/zen_server/dashboard/assets/{page-Btc7hKsa.js → page-Cea5XeRB.js} +1 -1
  69. zenml/zen_server/dashboard/assets/{page-BBUhx49Q.js → page-DFQa6C6e.js} +1 -1
  70. zenml/zen_server/dashboard/assets/page-DVAcOIuw.js +1 -0
  71. zenml/zen_server/dashboard/assets/{page-aQgPR2HA.js → page-Dn87PjPv.js} +1 -1
  72. zenml/zen_server/dashboard/assets/{page-sPFYi62U.js → page-GXA0bsIT.js} +1 -1
  73. zenml/zen_server/dashboard/assets/{page-PQQVONmL.js → page-NTQrxRsu.js} +2 -2
  74. zenml/zen_server/dashboard/assets/page-NlxFzoaz.js +1 -0
  75. zenml/zen_server/dashboard/assets/page-ZJ5Cy_-x.js +1 -0
  76. zenml/zen_server/dashboard/assets/{update-server-settings-mutation-B8N8nrWI.js → update-server-settings-mutation-CdcOI8-c.js} +1 -1
  77. zenml/zen_server/dashboard/assets/url-DLJUVIKc.js +1 -0
  78. zenml/zen_server/dashboard/index.html +4 -4
  79. zenml/zen_server/dashboard_legacy/asset-manifest.json +4 -4
  80. zenml/zen_server/dashboard_legacy/index.html +1 -1
  81. zenml/zen_server/dashboard_legacy/{precache-manifest.36c2742fac555eaab6995df9d81c08ab.js → precache-manifest.05e51e85bd64770b2708d408a88d7028.js} +4 -4
  82. zenml/zen_server/dashboard_legacy/service-worker.js +1 -1
  83. zenml/zen_server/dashboard_legacy/static/js/{main.1c1aaa11.chunk.js → main.41d438ec.chunk.js} +2 -2
  84. zenml/zen_server/dashboard_legacy/static/js/{main.1c1aaa11.chunk.js.map → main.41d438ec.chunk.js.map} +1 -1
  85. zenml/zen_server/deploy/helm/Chart.yaml +1 -1
  86. zenml/zen_server/deploy/helm/README.md +2 -2
  87. zenml/zen_server/deploy/helm/templates/_environment.tpl +5 -5
  88. zenml/zen_server/deploy/helm/values.yaml +13 -9
  89. zenml/zen_server/pipeline_deployment/utils.py +6 -2
  90. zenml/zen_server/routers/auth_endpoints.py +4 -4
  91. zenml/zen_server/zen_server_api.py +1 -1
  92. zenml/zen_stores/base_zen_store.py +2 -2
  93. zenml/zen_stores/migrations/versions/0.58.0_release.py +23 -0
  94. zenml/zen_stores/schemas/artifact_schemas.py +11 -2
  95. zenml/zen_stores/schemas/pipeline_run_schemas.py +12 -0
  96. zenml/zen_stores/schemas/step_run_schemas.py +17 -1
  97. zenml/zen_stores/sql_zen_store.py +9 -5
  98. {zenml_nightly-0.57.1.dev20240523.dist-info → zenml_nightly-0.58.0.dev20240528.dist-info}/METADATA +4 -4
  99. {zenml_nightly-0.57.1.dev20240523.dist-info → zenml_nightly-0.58.0.dev20240528.dist-info}/RECORD +102 -101
  100. zenml/zen_server/dashboard/assets/Commands-fpiuAaZ1.js +0 -1
  101. zenml/zen_server/dashboard/assets/CopyButton-DVKNxLl5.js +0 -2
  102. zenml/zen_server/dashboard/assets/SuccessStep-B8Qhr2kK.js +0 -1
  103. zenml/zen_server/dashboard/assets/VideoModal-qB9JWw-2.js +0 -1
  104. zenml/zen_server/dashboard/assets/page-3NFpYOpt.js +0 -1
  105. zenml/zen_server/dashboard/assets/page-BE7NJ8PH.js +0 -1
  106. zenml/zen_server/dashboard/assets/page-BTYWG9gF.js +0 -1
  107. zenml/zen_server/dashboard/assets/page-BuBHmjwQ.js +0 -1
  108. zenml/zen_server/dashboard/assets/page-C-Jw_wjv.js +0 -1
  109. zenml/zen_server/dashboard/assets/page-CeMgQAs7.js +0 -14
  110. zenml/zen_server/dashboard/assets/page-X0_g3JVr.js +0 -1
  111. zenml/zen_server/dashboard/assets/page-ijCkM2BW.js +0 -1
  112. zenml/zen_server/dashboard/assets/refresh-CdiBx6nL.js +0 -1
  113. zenml/zen_server/dashboard/assets/url-BxOqGmlk.js +0 -1
  114. {zenml_nightly-0.57.1.dev20240523.dist-info → zenml_nightly-0.58.0.dev20240528.dist-info}/LICENSE +0 -0
  115. {zenml_nightly-0.57.1.dev20240523.dist-info → zenml_nightly-0.58.0.dev20240528.dist-info}/WHEEL +0 -0
  116. {zenml_nightly-0.57.1.dev20240523.dist-info → zenml_nightly-0.58.0.dev20240528.dist-info}/entry_points.txt +0 -0
@@ -39,6 +39,7 @@ from zenml.models.v2.base.scoped import (
39
39
  WorkspaceScopedResponseMetadata,
40
40
  WorkspaceScopedResponseResources,
41
41
  )
42
+ from zenml.models.v2.core.model_version import ModelVersionResponse
42
43
 
43
44
  if TYPE_CHECKING:
44
45
  from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
@@ -197,6 +198,8 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
197
198
  class PipelineRunResponseResources(WorkspaceScopedResponseResources):
198
199
  """Class for all resource models associated with the pipeline run entity."""
199
200
 
201
+ model_version: Optional[ModelVersionResponse]
202
+
200
203
 
201
204
  class PipelineRunResponse(
202
205
  WorkspaceScopedResponse[
@@ -394,6 +397,15 @@ class PipelineRunResponse(
394
397
  """
395
398
  return self.get_metadata().orchestrator_run_id
396
399
 
400
+ @property
401
+ def model_version(self) -> Optional[ModelVersionResponse]:
402
+ """The `model_version` property.
403
+
404
+ Returns:
405
+ the value of the property.
406
+ """
407
+ return self.get_resources().model_version
408
+
397
409
 
398
410
  # ------------------ Filter Model ------------------
399
411
 
@@ -30,6 +30,7 @@ from zenml.models.v2.base.scoped import (
30
30
  WorkspaceScopedResponseMetadata,
31
31
  WorkspaceScopedResponseResources,
32
32
  )
33
+ from zenml.models.v2.core.model_version import ModelVersionResponse
33
34
 
34
35
  if TYPE_CHECKING:
35
36
  from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
@@ -219,6 +220,8 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
219
220
  class StepRunResponseResources(WorkspaceScopedResponseResources):
220
221
  """Class for all resource models associated with the step run entity."""
221
222
 
223
+ model_version: Optional[ModelVersionResponse]
224
+
222
225
 
223
226
  class StepRunResponse(
224
227
  WorkspaceScopedResponse[
@@ -435,6 +438,15 @@ class StepRunResponse(
435
438
  """
436
439
  return self.get_metadata().run_metadata
437
440
 
441
+ @property
442
+ def model_version(self) -> Optional[ModelVersionResponse]:
443
+ """The `model_version` property.
444
+
445
+ Returns:
446
+ the value of the property.
447
+ """
448
+ return self.get_resources().model_version
449
+
438
450
 
439
451
  # ------------------ Filter Model ------------------
440
452
 
@@ -78,11 +78,17 @@ class ServerModel(BaseModel):
78
78
  auth_scheme: AuthScheme = Field(
79
79
  title="The authentication scheme that the server is using.",
80
80
  )
81
- base_url: str = Field(
81
+ server_url: str = Field(
82
82
  "",
83
- title="The Base URL of the server.",
83
+ title="The URL where the ZenML server API is reachable. If not "
84
+ "specified, the clients will use the same URL used to connect them to "
85
+ "the ZenML server.",
86
+ )
87
+ dashboard_url: str = Field(
88
+ "",
89
+ title="The URL where the ZenML dashboard is reachable. If "
90
+ "not specified, the `server_url` value will be used instead.",
84
91
  )
85
-
86
92
  analytics_enabled: bool = Field(
87
93
  default=True, # We set a default for migrations from < 0.57.0
88
94
  title="Enable server-side analytics.",
@@ -17,7 +17,7 @@ from uuid import UUID
17
17
  from zenml import constants
18
18
  from zenml.client import Client
19
19
  from zenml.config.step_configurations import StepConfigurationUpdate
20
- from zenml.enums import ExecutionStatus
20
+ from zenml.enums import ExecutionStatus, ModelStages
21
21
  from zenml.logger import get_logger
22
22
  from zenml.models import (
23
23
  PipelineDeploymentBase,
@@ -166,7 +166,13 @@ def _update_new_requesters(
166
166
  try:
167
167
  model._get_model_version()
168
168
  version_existed = key not in new_versions_requested
169
- except KeyError:
169
+ except KeyError as e:
170
+ if model.version in ModelStages.values():
171
+ raise KeyError(
172
+ f"Unable to get model `{model.name}` using stage "
173
+ f"`{model.version}`, please check that the model "
174
+ "version in given stage exists before running a pipeline."
175
+ ) from e
170
176
  version_existed = False
171
177
  if not version_existed:
172
178
  model.was_created_in_this_run = True
@@ -132,7 +132,7 @@ class StepContext(metaclass=SingletonMetaClass):
132
132
  if output_materializers.keys() != output_artifact_uris.keys():
133
133
  raise StepContextError(
134
134
  f"Mismatched keys in output materializers and output artifact "
135
- f"URIs for step '{self.step_name}'. Output materializer "
135
+ f"URIs for step `{self.step_name}`. Output materializer "
136
136
  f"keys: {set(output_materializers)}, output artifact URI "
137
137
  f"keys: {set(output_artifact_uris)}"
138
138
  )
@@ -158,7 +158,7 @@ class StepContext(metaclass=SingletonMetaClass):
158
158
  if self.pipeline_run.pipeline:
159
159
  return self.pipeline_run.pipeline
160
160
  raise StepContextError(
161
- f"Unable to get pipeline in step '{self.step_name}' of pipeline "
161
+ f"Unable to get pipeline in step `{self.step_name}` of pipeline "
162
162
  f"run '{self.pipeline_run.id}': This pipeline run does not have "
163
163
  f"a pipeline associated with it."
164
164
  )
@@ -183,7 +183,7 @@ class StepContext(metaclass=SingletonMetaClass):
183
183
  model = self.pipeline_run.config.model
184
184
  else:
185
185
  raise StepContextError(
186
- f"Unable to get Model in step '{self.step_name}' of pipeline "
186
+ f"Unable to get Model in step `{self.step_name}` of pipeline "
187
187
  f"run '{self.pipeline_run.id}': it was not set in `@step` or `@pipeline`."
188
188
  )
189
189
 
@@ -235,13 +235,13 @@ class StepContext(metaclass=SingletonMetaClass):
235
235
  output_count = len(self._outputs)
236
236
  if output_count == 0:
237
237
  raise StepContextError(
238
- f"Unable to get step output for step '{self.step_name}': "
238
+ f"Unable to get step output for step `{self.step_name}`: "
239
239
  f"This step does not have any outputs."
240
240
  )
241
241
 
242
242
  if not output_name and output_count > 1:
243
243
  raise StepContextError(
244
- f"Unable to get step output for step '{self.step_name}': "
244
+ f"Unable to get step output for step `{self.step_name}`: "
245
245
  f"This step has multiple outputs ({set(self._outputs)}), "
246
246
  f"please specify which output to return."
247
247
  )
@@ -250,7 +250,7 @@ class StepContext(metaclass=SingletonMetaClass):
250
250
  if output_name not in self._outputs:
251
251
  raise StepContextError(
252
252
  f"Unable to get step output '{output_name}' for "
253
- f"step '{self.step_name}'. This step does not have an "
253
+ f"step `{self.step_name}`. This step does not have an "
254
254
  f"output with the given name, please specify one of the "
255
255
  f"available outputs: {set(self._outputs)}."
256
256
  )
@@ -33,6 +33,7 @@ if TYPE_CHECKING:
33
33
  from types import FunctionType
34
34
 
35
35
  from zenml.config.base_settings import SettingsOrDict
36
+ from zenml.config.retry_config import StepRetryConfig
36
37
  from zenml.config.source import Source
37
38
  from zenml.materializers.base_materializer import BaseMaterializer
38
39
  from zenml.model.model import Model
@@ -72,6 +73,7 @@ def step(
72
73
  on_failure: Optional["HookSpecification"] = None,
73
74
  on_success: Optional["HookSpecification"] = None,
74
75
  model: Optional["Model"] = None,
76
+ retry: Optional["StepRetryConfig"] = None,
75
77
  model_version: Optional["Model"] = None, # TODO: deprecate me
76
78
  ) -> Callable[["F"], "BaseStep"]: ...
77
79
 
@@ -92,6 +94,7 @@ def step(
92
94
  on_failure: Optional["HookSpecification"] = None,
93
95
  on_success: Optional["HookSpecification"] = None,
94
96
  model: Optional["Model"] = None,
97
+ retry: Optional["StepRetryConfig"] = None,
95
98
  model_version: Optional["Model"] = None, # TODO: deprecate me
96
99
  ) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
97
100
  """Decorator to create a ZenML step.
@@ -123,6 +126,7 @@ def step(
123
126
  function with no arguments, or a source path to such a function
124
127
  (e.g. `module.my_function`).
125
128
  model: configuration of the model in the Model Control Plane.
129
+ retry: configuration of step retry in case of step failure.
126
130
  model_version: DEPRECATED, please use `model` instead.
127
131
 
128
132
  Returns:
@@ -162,6 +166,7 @@ def step(
162
166
  on_failure=on_failure,
163
167
  on_success=on_success,
164
168
  model=model or model_version,
169
+ retry=retry,
165
170
  )
166
171
 
167
172
  return step_instance
@@ -13,6 +13,7 @@
13
13
  # permissions and limitations under the License.
14
14
  """Class to launch (run directly or using a step operator) steps."""
15
15
 
16
+ import os
16
17
  import time
17
18
  from contextlib import nullcontext
18
19
  from datetime import datetime
@@ -23,6 +24,7 @@ from zenml.config.step_configurations import Step
23
24
  from zenml.config.step_run_info import StepRunInfo
24
25
  from zenml.constants import (
25
26
  ENV_ZENML_DISABLE_STEP_LOGS_STORAGE,
27
+ ENV_ZENML_IGNORE_FAILURE_HOOK,
26
28
  STEP_SOURCE_PARAMETER_NAME,
27
29
  TEXT_FIELD_MAX_LENGTH,
28
30
  handle_bool_env_var,
@@ -31,7 +33,6 @@ from zenml.enums import ExecutionStatus
31
33
  from zenml.environment import get_run_environment_dict
32
34
  from zenml.logger import get_logger
33
35
  from zenml.logging import step_logging
34
- from zenml.model.utils import link_artifact_config_to_model
35
36
  from zenml.models import (
36
37
  ArtifactVersionResponse,
37
38
  LogsRequest,
@@ -53,7 +54,6 @@ from zenml.stack import Stack
53
54
  from zenml.utils import string_utils
54
55
 
55
56
  if TYPE_CHECKING:
56
- from zenml.model.model import Model
57
57
  from zenml.step_operators import BaseStepOperator
58
58
 
59
59
  logger = get_logger(__name__)
@@ -226,20 +226,56 @@ class StepLauncher:
226
226
 
227
227
  logger.info(f"Step `{self._step_name}` has started.")
228
228
  if execution_needed:
229
- try:
230
- self._run_step(
231
- pipeline_run=pipeline_run,
232
- step_run=step_run_response,
233
- )
234
- except BaseException as e: # noqa: E722
235
- logger.error(
236
- f"Failed to run step `{self._step_name}`."
237
- )
238
- logger.exception(e)
239
- publish_utils.publish_failed_step_run(
240
- step_run_response.id
241
- )
242
- raise
229
+ retries = 0
230
+ last_retry = True
231
+ max_retries = (
232
+ step_run_response.config.retry.max_retries
233
+ if step_run_response.config.retry
234
+ else 1
235
+ )
236
+ delay = (
237
+ step_run_response.config.retry.delay
238
+ if step_run_response.config.retry
239
+ else 0
240
+ )
241
+ backoff = (
242
+ step_run_response.config.retry.backoff
243
+ if step_run_response.config.retry
244
+ else 1
245
+ )
246
+
247
+ while retries < max_retries:
248
+ last_retry = retries == max_retries - 1
249
+ try:
250
+ self._run_step(
251
+ pipeline_run=pipeline_run,
252
+ step_run=step_run_response,
253
+ last_retry=last_retry,
254
+ )
255
+ logger.info(
256
+ f"Step `{self._step_name}` completed successfully."
257
+ )
258
+ break
259
+ except BaseException as e: # noqa: E722
260
+ retries += 1
261
+ if retries < max_retries:
262
+ logger.error(
263
+ f"Failed to run step `{self._step_name}`. Retrying..."
264
+ )
265
+ logger.exception(e)
266
+ logger.info(
267
+ f"Sleeping for {delay} seconds before retrying."
268
+ )
269
+ time.sleep(delay)
270
+ delay *= backoff
271
+ else:
272
+ logger.error(
273
+ f"Failed to run step `{self._step_name}` after {max_retries} retries. Exiting."
274
+ )
275
+ publish_utils.publish_failed_step_run(
276
+ step_run_response.id
277
+ )
278
+ raise
243
279
 
244
280
  except: # noqa: E722
245
281
  logger.error(f"Pipeline run `{pipeline_run.name}` failed.")
@@ -360,61 +396,33 @@ class StepLauncher:
360
396
  output_name: artifact.id
361
397
  for output_name, artifact in cached_outputs.items()
362
398
  }
363
- self._link_cached_artifacts_to_model(
399
+ orchestrator_utils._link_cached_artifacts_to_model(
364
400
  model_from_context=model,
365
401
  step_run=step_run,
402
+ step_source=self._step.spec.source,
366
403
  )
404
+ if self._step.config.model:
405
+ orchestrator_utils._link_pipeline_run_to_model_from_context(
406
+ pipeline_run_id=step_run.pipeline_run_id,
407
+ model=self._step.config.model,
408
+ )
367
409
  step_run.status = ExecutionStatus.CACHED
368
410
  step_run.end_time = step_run.start_time
369
411
 
370
412
  return execution_needed, step_run
371
413
 
372
- def _link_cached_artifacts_to_model(
373
- self,
374
- model_from_context: Optional["Model"],
375
- step_run: StepRunRequest,
376
- ) -> None:
377
- """Links the output artifacts of the cached step to the model version in Control Plane.
378
-
379
- Args:
380
- model_from_context: The model version of the current step.
381
- step_run: The step to run.
382
- """
383
- from zenml.artifacts.artifact_config import ArtifactConfig
384
- from zenml.steps.base_step import BaseStep
385
- from zenml.steps.utils import parse_return_type_annotations
386
-
387
- step_instance = BaseStep.load_from_source(self._step.spec.source)
388
- output_annotations = parse_return_type_annotations(
389
- step_instance.entrypoint
390
- )
391
- for output_name_, output_id in step_run.outputs.items():
392
- artifact_config_ = None
393
- if output_name_ in output_annotations:
394
- annotation = output_annotations.get(output_name_, None)
395
- if annotation and annotation.artifact_config is not None:
396
- artifact_config_ = annotation.artifact_config.copy()
397
- # no artifact config found or artifact was produced by `save_artifact`
398
- # inside the step body, so was never in annotations
399
- if artifact_config_ is None:
400
- artifact_config_ = ArtifactConfig(name=output_name_)
401
-
402
- link_artifact_config_to_model(
403
- artifact_config=artifact_config_,
404
- model=model_from_context,
405
- artifact_version_id=output_id,
406
- )
407
-
408
414
  def _run_step(
409
415
  self,
410
416
  pipeline_run: PipelineRunResponse,
411
417
  step_run: StepRunResponse,
418
+ last_retry: bool = True,
412
419
  ) -> None:
413
420
  """Runs the current step.
414
421
 
415
422
  Args:
416
423
  pipeline_run: The model of the current pipeline run.
417
424
  step_run: The model of the current step run.
425
+ last_retry: Whether this is the last retry of the step.
418
426
  """
419
427
  # Prepare step run information.
420
428
  step_run_info = StepRunInfo(
@@ -437,6 +445,7 @@ class StepLauncher:
437
445
  self._run_step_with_step_operator(
438
446
  step_operator_name=self._step.config.step_operator,
439
447
  step_run_info=step_run_info,
448
+ last_retry=last_retry,
440
449
  )
441
450
  else:
442
451
  self._run_step_without_step_operator(
@@ -445,6 +454,7 @@ class StepLauncher:
445
454
  step_run_info=step_run_info,
446
455
  input_artifacts=step_run.inputs,
447
456
  output_artifact_uris=output_artifact_uris,
457
+ last_retry=last_retry,
448
458
  )
449
459
  except: # noqa: E722
450
460
  output_utils.remove_artifact_dirs(
@@ -462,12 +472,14 @@ class StepLauncher:
462
472
  self,
463
473
  step_operator_name: str,
464
474
  step_run_info: StepRunInfo,
475
+ last_retry: bool,
465
476
  ) -> None:
466
477
  """Runs the current step with a step operator.
467
478
 
468
479
  Args:
469
480
  step_operator_name: The name of the step operator to use.
470
481
  step_run_info: Additional information needed to run the step.
482
+ last_retry: Whether this is the last retry of the step.
471
483
  """
472
484
  step_operator = _get_step_operator(
473
485
  stack=self._stack,
@@ -485,6 +497,8 @@ class StepLauncher:
485
497
  environment = orchestrator_utils.get_config_environment_vars(
486
498
  deployment=self._deployment
487
499
  )
500
+ if last_retry:
501
+ environment[ENV_ZENML_IGNORE_FAILURE_HOOK] = str(False)
488
502
  logger.info(
489
503
  "Using step operator `%s` to run step `%s`.",
490
504
  step_operator.name,
@@ -503,6 +517,7 @@ class StepLauncher:
503
517
  step_run_info: StepRunInfo,
504
518
  input_artifacts: Dict[str, ArtifactVersionResponse],
505
519
  output_artifact_uris: Dict[str, str],
520
+ last_retry: bool,
506
521
  ) -> None:
507
522
  """Runs the current step without a step operator.
508
523
 
@@ -512,7 +527,10 @@ class StepLauncher:
512
527
  step_run_info: Additional information needed to run the step.
513
528
  input_artifacts: The input artifact versions of the current step.
514
529
  output_artifact_uris: The output artifact URIs of the current step.
530
+ last_retry: Whether this is the last retry of the step.
515
531
  """
532
+ if last_retry:
533
+ os.environ[ENV_ZENML_IGNORE_FAILURE_HOOK] = "false"
516
534
  runner = StepRunner(step=self._step, stack=self._stack)
517
535
  runner.run(
518
536
  pipeline_run=pipeline_run,
@@ -23,7 +23,6 @@ from typing import (
23
23
  Dict,
24
24
  List,
25
25
  Optional,
26
- Set,
27
26
  Tuple,
28
27
  Type,
29
28
  )
@@ -33,11 +32,11 @@ from pydantic.typing import get_origin, is_union
33
32
 
34
33
  from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
35
34
  from zenml.artifacts.utils import save_artifact
36
- from zenml.client import Client
37
35
  from zenml.config.step_configurations import StepConfiguration
38
36
  from zenml.config.step_run_info import StepRunInfo
39
37
  from zenml.constants import (
40
38
  ENV_ZENML_DISABLE_STEP_LOGS_STORAGE,
39
+ ENV_ZENML_IGNORE_FAILURE_HOOK,
41
40
  handle_bool_env_var,
42
41
  )
43
42
  from zenml.exceptions import StepContextError, StepInterfaceError
@@ -52,7 +51,11 @@ from zenml.orchestrators.publish_utils import (
52
51
  publish_step_run_metadata,
53
52
  publish_successful_step_run,
54
53
  )
55
- from zenml.orchestrators.utils import is_setting_enabled
54
+ from zenml.orchestrators.utils import (
55
+ _link_pipeline_run_to_model_from_artifacts,
56
+ _link_pipeline_run_to_model_from_context,
57
+ is_setting_enabled,
58
+ )
56
59
  from zenml.steps.step_environment import StepEnvironment
57
60
  from zenml.steps.utils import (
58
61
  OutputSignature,
@@ -62,9 +65,6 @@ from zenml.steps.utils import (
62
65
  from zenml.utils import materializer_utils, source_utils
63
66
 
64
67
  if TYPE_CHECKING:
65
- from zenml.artifacts.external_artifact_config import (
66
- ExternalArtifactConfiguration,
67
- )
68
68
  from zenml.config.source import Source
69
69
  from zenml.config.step_configurations import Step
70
70
  from zenml.models import (
@@ -192,8 +192,8 @@ class StepRunner:
192
192
  input_artifacts=input_artifacts,
193
193
  )
194
194
 
195
- self._link_pipeline_run_to_model_from_context(
196
- pipeline_run=pipeline_run
195
+ _link_pipeline_run_to_model_from_context(
196
+ pipeline_run_id=pipeline_run.id
197
197
  )
198
198
 
199
199
  step_failed = False
@@ -203,15 +203,18 @@ class StepRunner:
203
203
  )
204
204
  except BaseException as step_exception: # noqa: E722
205
205
  step_failed = True
206
- failure_hook_source = (
207
- self.configuration.failure_hook_source
208
- )
209
- if failure_hook_source:
210
- logger.info("Detected failure hook. Running...")
211
- self.load_and_run_hook(
212
- failure_hook_source,
213
- step_exception=step_exception,
214
- )
206
+ if not handle_bool_env_var(
207
+ ENV_ZENML_IGNORE_FAILURE_HOOK, False
208
+ ):
209
+ if (
210
+ failure_hook_source
211
+ := self.configuration.failure_hook_source
212
+ ):
213
+ logger.info("Detected failure hook. Running...")
214
+ self.load_and_run_hook(
215
+ failure_hook_source,
216
+ step_exception=step_exception,
217
+ )
215
218
  raise
216
219
  finally:
217
220
  step_run_metadata = self._stack.get_step_run_metadata(
@@ -225,10 +228,10 @@ class StepRunner:
225
228
  info=step_run_info, step_failed=step_failed
226
229
  )
227
230
  if not step_failed:
228
- success_hook_source = (
229
- self.configuration.success_hook_source
230
- )
231
- if success_hook_source:
231
+ if (
232
+ success_hook_source
233
+ := self.configuration.success_hook_source
234
+ ):
232
235
  logger.info("Detected success hook. Running...")
233
236
  self.load_and_run_hook(
234
237
  success_hook_source,
@@ -258,8 +261,8 @@ class StepRunner:
258
261
  link_step_artifacts_to_model(
259
262
  artifact_version_ids=output_artifact_ids
260
263
  )
261
- self._link_pipeline_run_to_model_from_artifacts(
262
- pipeline_run=pipeline_run,
264
+ _link_pipeline_run_to_model_from_artifacts(
265
+ pipeline_run_id=pipeline_run.id,
263
266
  artifact_names=list(output_artifact_ids.keys()),
264
267
  external_artifacts=list(
265
268
  step_run.config.external_input_artifacts.values()
@@ -478,7 +481,7 @@ class StepRunner:
478
481
  if len(output_annotations) == 0:
479
482
  if return_values is not None:
480
483
  raise StepInterfaceError(
481
- f"Wrong step function output type for step '{step_name}': "
484
+ f"Wrong step function output type for step `{step_name}`: "
482
485
  f"Expected no outputs but the function returned something: "
483
486
  f"{return_values}."
484
487
  )
@@ -494,7 +497,7 @@ class StepRunner:
494
497
  # or tuple.
495
498
  if not isinstance(return_values, (list, tuple)):
496
499
  raise StepInterfaceError(
497
- f"Wrong step function output type for step '{step_name}': "
500
+ f"Wrong step function output type for step `{step_name}`: "
498
501
  f"Expected multiple outputs ({output_annotations}) but "
499
502
  f"the function did not return a list or tuple "
500
503
  f"(actual return value: {return_values})."
@@ -645,115 +648,6 @@ class StepRunner:
645
648
  except StepContextError:
646
649
  return
647
650
 
648
- def _get_model_versions_from_artifacts(
649
- self,
650
- artifact_names: List[str],
651
- ) -> Set[Tuple[UUID, UUID]]:
652
- """Gets the model versions from the artifacts.
653
-
654
- Args:
655
- artifact_names: The names of the published output artifacts.
656
-
657
- Returns:
658
- Set of tuples of (model_id, model_version_id).
659
- """
660
- models = set()
661
- for artifact_name in artifact_names:
662
- artifact_config = (
663
- get_step_context()._get_output(artifact_name).artifact_config
664
- )
665
- if artifact_config is not None:
666
- if (model := artifact_config._model) is not None:
667
- model_version_response = (
668
- model._get_or_create_model_version()
669
- )
670
- models.add(
671
- (
672
- model_version_response.model.id,
673
- model_version_response.id,
674
- )
675
- )
676
- else:
677
- break
678
- return models
679
-
680
- def _get_model_versions_from_config(self) -> Set[Tuple[UUID, UUID]]:
681
- """Gets the model versions from the step model version.
682
-
683
- Returns:
684
- Set of tuples of (model_id, model_version_id).
685
- """
686
- try:
687
- mc = get_step_context().model
688
- model_version = mc._get_or_create_model_version()
689
- return {(model_version.model.id, model_version.id)}
690
- except StepContextError:
691
- return set()
692
-
693
- def _link_pipeline_run_to_model_from_context(
694
- self,
695
- pipeline_run: "PipelineRunResponse",
696
- ) -> None:
697
- """Links the pipeline run to the model version using artifacts data.
698
-
699
- Args:
700
- pipeline_run: The response model of current pipeline run.
701
- """
702
- from zenml.models import ModelVersionPipelineRunRequest
703
-
704
- models = self._get_model_versions_from_config()
705
-
706
- client = Client()
707
- for model in models:
708
- client.zen_store.create_model_version_pipeline_run_link(
709
- ModelVersionPipelineRunRequest(
710
- user=Client().active_user.id,
711
- workspace=Client().active_workspace.id,
712
- pipeline_run=pipeline_run.id,
713
- model=model[0],
714
- model_version=model[1],
715
- )
716
- )
717
-
718
- def _link_pipeline_run_to_model_from_artifacts(
719
- self,
720
- pipeline_run: "PipelineRunResponse",
721
- artifact_names: List[str],
722
- external_artifacts: List["ExternalArtifactConfiguration"],
723
- ) -> None:
724
- """Links the pipeline run to the model version using artifacts data.
725
-
726
- Args:
727
- pipeline_run: The response model of current pipeline run.
728
- artifact_names: The name of the published output artifacts.
729
- external_artifacts: The external artifacts of the step.
730
- """
731
- from zenml.models import ModelVersionPipelineRunRequest
732
-
733
- models = self._get_model_versions_from_artifacts(artifact_names)
734
- client = Client()
735
-
736
- # Add models from external artifacts
737
- for external_artifact in external_artifacts:
738
- if external_artifact.model:
739
- models.add(
740
- (
741
- external_artifact.model.model_id,
742
- external_artifact.model.id,
743
- )
744
- )
745
-
746
- for model in models:
747
- client.zen_store.create_model_version_pipeline_run_link(
748
- ModelVersionPipelineRunRequest(
749
- user=client.active_user.id,
750
- workspace=client.active_workspace.id,
751
- pipeline_run=pipeline_run.id,
752
- model=model[0],
753
- model_version=model[1],
754
- )
755
- )
756
-
757
651
  def load_and_run_hook(
758
652
  self,
759
653
  hook_source: "Source",