zenml-nightly 0.70.0.dev20241122__py3-none-any.whl → 0.70.0.dev20241201__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 (103) hide show
  1. zenml/VERSION +1 -1
  2. zenml/artifact_stores/base_artifact_store.py +2 -2
  3. zenml/artifacts/artifact_config.py +7 -1
  4. zenml/artifacts/utils.py +56 -31
  5. zenml/cli/__init__.py +18 -0
  6. zenml/cli/base.py +4 -4
  7. zenml/cli/login.py +26 -0
  8. zenml/cli/pipeline.py +80 -0
  9. zenml/cli/server.py +1 -1
  10. zenml/cli/service_connectors.py +3 -3
  11. zenml/cli/stack.py +0 -3
  12. zenml/cli/stack_components.py +0 -1
  13. zenml/cli/utils.py +0 -5
  14. zenml/client.py +8 -18
  15. zenml/config/compiler.py +12 -3
  16. zenml/config/pipeline_configurations.py +20 -0
  17. zenml/config/pipeline_run_configuration.py +1 -0
  18. zenml/config/step_configurations.py +21 -0
  19. zenml/constants.py +1 -0
  20. zenml/enums.py +1 -0
  21. zenml/image_builders/local_image_builder.py +13 -3
  22. zenml/integrations/__init__.py +1 -0
  23. zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +14 -6
  24. zenml/integrations/constants.py +1 -0
  25. zenml/integrations/feast/__init__.py +1 -1
  26. zenml/integrations/feast/feature_stores/feast_feature_store.py +13 -9
  27. zenml/integrations/kubernetes/orchestrators/kube_utils.py +54 -9
  28. zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +65 -3
  29. zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py +14 -2
  30. zenml/integrations/kubernetes/orchestrators/manifest_utils.py +6 -5
  31. zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py +2 -1
  32. zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py +3 -1
  33. zenml/integrations/modal/__init__.py +46 -0
  34. zenml/integrations/modal/flavors/__init__.py +26 -0
  35. zenml/integrations/modal/flavors/modal_step_operator_flavor.py +125 -0
  36. zenml/integrations/modal/step_operators/__init__.py +22 -0
  37. zenml/integrations/modal/step_operators/modal_step_operator.py +242 -0
  38. zenml/io/filesystem.py +2 -2
  39. zenml/io/local_filesystem.py +3 -3
  40. zenml/materializers/built_in_materializer.py +18 -1
  41. zenml/materializers/structured_string_materializer.py +8 -3
  42. zenml/model/model.py +11 -85
  43. zenml/model/utils.py +18 -16
  44. zenml/models/__init__.py +12 -1
  45. zenml/models/v2/core/artifact_version.py +6 -3
  46. zenml/models/v2/core/component.py +0 -22
  47. zenml/models/v2/core/model_version.py +6 -3
  48. zenml/models/v2/core/pipeline_run.py +19 -3
  49. zenml/models/v2/core/run_metadata.py +30 -9
  50. zenml/models/v2/core/service_connector.py +4 -0
  51. zenml/models/v2/core/step_run.py +6 -4
  52. zenml/models/v2/misc/run_metadata.py +38 -0
  53. zenml/models/v2/misc/server_models.py +23 -0
  54. zenml/orchestrators/input_utils.py +19 -6
  55. zenml/orchestrators/publish_utils.py +12 -5
  56. zenml/orchestrators/step_launcher.py +7 -3
  57. zenml/orchestrators/step_run_utils.py +26 -9
  58. zenml/orchestrators/step_runner.py +40 -3
  59. zenml/orchestrators/utils.py +24 -23
  60. zenml/pipelines/pipeline_decorator.py +4 -0
  61. zenml/pipelines/pipeline_definition.py +26 -8
  62. zenml/pipelines/run_utils.py +9 -5
  63. zenml/steps/base_step.py +11 -1
  64. zenml/steps/entrypoint_function_utils.py +7 -3
  65. zenml/steps/step_decorator.py +4 -0
  66. zenml/steps/utils.py +23 -7
  67. zenml/types.py +4 -0
  68. zenml/utils/metadata_utils.py +186 -153
  69. zenml/utils/string_utils.py +41 -16
  70. zenml/utils/visualization_utils.py +4 -1
  71. zenml/zen_server/cloud_utils.py +3 -1
  72. zenml/zen_server/deploy/helm/templates/_environment.tpl +117 -0
  73. zenml/zen_server/deploy/helm/templates/server-db-job.yaml +3 -14
  74. zenml/zen_server/deploy/helm/templates/server-deployment.yaml +16 -4
  75. zenml/zen_server/deploy/helm/templates/server-secret.yaml +2 -17
  76. zenml/zen_server/rbac/endpoint_utils.py +6 -4
  77. zenml/zen_server/rbac/models.py +3 -2
  78. zenml/zen_server/rbac/utils.py +4 -7
  79. zenml/zen_server/routers/server_endpoints.py +47 -0
  80. zenml/zen_server/routers/users_endpoints.py +35 -37
  81. zenml/zen_server/routers/workspaces_endpoints.py +44 -55
  82. zenml/zen_server/template_execution/utils.py +1 -0
  83. zenml/zen_server/zen_server_api.py +45 -6
  84. zenml/zen_stores/migrations/utils.py +40 -24
  85. zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py +36 -0
  86. zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py +135 -0
  87. zenml/zen_stores/migrations/versions/ec6307720f92_simplify_model_version_links.py +7 -6
  88. zenml/zen_stores/rest_zen_store.py +38 -1
  89. zenml/zen_stores/schemas/__init__.py +5 -1
  90. zenml/zen_stores/schemas/artifact_schemas.py +12 -11
  91. zenml/zen_stores/schemas/component_schemas.py +0 -3
  92. zenml/zen_stores/schemas/model_schemas.py +13 -11
  93. zenml/zen_stores/schemas/pipeline_run_schemas.py +44 -16
  94. zenml/zen_stores/schemas/run_metadata_schemas.py +66 -31
  95. zenml/zen_stores/schemas/step_run_schemas.py +32 -12
  96. zenml/zen_stores/schemas/utils.py +47 -3
  97. zenml/zen_stores/sql_zen_store.py +130 -34
  98. {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/METADATA +1 -1
  99. {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/RECORD +102 -95
  100. zenml/utils/cloud_utils.py +0 -40
  101. {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/LICENSE +0 -0
  102. {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/WHEEL +0 -0
  103. {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/entry_points.txt +0 -0
@@ -13,7 +13,6 @@
13
13
  # permissions and limitations under the License.
14
14
  """SQLModel implementation of artifact table."""
15
15
 
16
- import json
17
16
  from datetime import datetime
18
17
  from typing import TYPE_CHECKING, Any, List, Optional
19
18
  from uuid import UUID
@@ -50,6 +49,7 @@ from zenml.zen_stores.schemas.step_run_schemas import (
50
49
  StepRunOutputArtifactSchema,
51
50
  )
52
51
  from zenml.zen_stores.schemas.user_schemas import UserSchema
52
+ from zenml.zen_stores.schemas.utils import RunMetadataInterface
53
53
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
54
54
 
55
55
  if TYPE_CHECKING:
@@ -59,7 +59,9 @@ if TYPE_CHECKING:
59
59
  from zenml.zen_stores.schemas.model_schemas import (
60
60
  ModelVersionArtifactSchema,
61
61
  )
62
- from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
62
+ from zenml.zen_stores.schemas.run_metadata_schemas import (
63
+ RunMetadataResourceSchema,
64
+ )
63
65
  from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
64
66
 
65
67
 
@@ -171,7 +173,7 @@ class ArtifactSchema(NamedSchema, table=True):
171
173
  return self
172
174
 
173
175
 
174
- class ArtifactVersionSchema(BaseSchema, table=True):
176
+ class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
175
177
  """SQL Model for artifact versions."""
176
178
 
177
179
  __tablename__ = "artifact_version"
@@ -242,12 +244,12 @@ class ArtifactVersionSchema(BaseSchema, table=True):
242
244
  workspace: "WorkspaceSchema" = Relationship(
243
245
  back_populates="artifact_versions"
244
246
  )
245
- run_metadata: List["RunMetadataSchema"] = Relationship(
246
- back_populates="artifact_version",
247
+ run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
248
+ back_populates="artifact_versions",
247
249
  sa_relationship_kwargs=dict(
248
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ArtifactVersionSchema.id)",
250
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ArtifactVersionSchema.id)",
249
251
  cascade="delete",
250
- overlaps="run_metadata",
252
+ overlaps="run_metadata_resources",
251
253
  ),
252
254
  )
253
255
  output_of_step_runs: List["StepRunOutputArtifactSchema"] = Relationship(
@@ -352,8 +354,9 @@ class ArtifactVersionSchema(BaseSchema, table=True):
352
354
  producer_step_run_id = step_run.original_step_run_id
353
355
 
354
356
  # Create the body of the model
357
+ artifact = self.artifact.to_model()
355
358
  body = ArtifactVersionResponseBody(
356
- artifact=self.artifact.to_model(),
359
+ artifact=artifact,
357
360
  version=self.version or str(self.version_number),
358
361
  user=self.user.to_model() if self.user else None,
359
362
  uri=self.uri,
@@ -375,9 +378,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
375
378
  workspace=self.workspace.to_model(),
376
379
  producer_step_run_id=producer_step_run_id,
377
380
  visualizations=[v.to_model() for v in self.visualizations],
378
- run_metadata={
379
- m.key: json.loads(m.value) for m in self.run_metadata
380
- },
381
+ run_metadata=self.fetch_metadata(),
381
382
  )
382
383
 
383
384
  resources = None
@@ -56,7 +56,6 @@ class StackComponentSchema(NamedSchema, table=True):
56
56
  flavor: str
57
57
  configuration: bytes
58
58
  labels: Optional[bytes]
59
- component_spec_path: Optional[str]
60
59
 
61
60
  workspace_id: UUID = build_foreign_key_field(
62
61
  source=__tablename__,
@@ -135,7 +134,6 @@ class StackComponentSchema(NamedSchema, table=True):
135
134
  name=request.name,
136
135
  workspace_id=request.workspace,
137
136
  user_id=request.user,
138
- component_spec_path=request.component_spec_path,
139
137
  type=request.type,
140
138
  flavor=request.flavor,
141
139
  configuration=base64.b64encode(
@@ -218,7 +216,6 @@ class StackComponentSchema(NamedSchema, table=True):
218
216
  labels=json.loads(base64.b64decode(self.labels).decode())
219
217
  if self.labels
220
218
  else None,
221
- component_spec_path=self.component_spec_path,
222
219
  connector_resource_id=self.connector_resource_id,
223
220
  connector=self.connector.to_model()
224
221
  if self.connector
@@ -13,7 +13,6 @@
13
13
  # permissions and limitations under the License.
14
14
  """SQLModel implementation of model tables."""
15
15
 
16
- import json
17
16
  from datetime import datetime
18
17
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
19
18
  from uuid import UUID
@@ -51,11 +50,16 @@ from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
51
50
  from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
52
51
  from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
53
52
  from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
54
- from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
53
+ from zenml.zen_stores.schemas.run_metadata_schemas import (
54
+ RunMetadataResourceSchema,
55
+ )
55
56
  from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
56
57
  from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
57
58
  from zenml.zen_stores.schemas.user_schemas import UserSchema
58
- from zenml.zen_stores.schemas.utils import get_page_from_list
59
+ from zenml.zen_stores.schemas.utils import (
60
+ RunMetadataInterface,
61
+ get_page_from_list,
62
+ )
59
63
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
60
64
 
61
65
  if TYPE_CHECKING:
@@ -219,7 +223,7 @@ class ModelSchema(NamedSchema, table=True):
219
223
  return self
220
224
 
221
225
 
222
- class ModelVersionSchema(NamedSchema, table=True):
226
+ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
223
227
  """SQL Model for model version."""
224
228
 
225
229
  __tablename__ = MODEL_VERSION_TABLENAME
@@ -299,12 +303,12 @@ class ModelVersionSchema(NamedSchema, table=True):
299
303
  description: str = Field(sa_column=Column(TEXT, nullable=True))
300
304
  stage: str = Field(sa_column=Column(TEXT, nullable=True))
301
305
 
302
- run_metadata: List["RunMetadataSchema"] = Relationship(
303
- back_populates="model_version",
306
+ run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
307
+ back_populates="model_versions",
304
308
  sa_relationship_kwargs=dict(
305
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ModelVersionSchema.id)",
309
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ModelVersionSchema.id)",
306
310
  cascade="delete",
307
- overlaps="run_metadata",
311
+ overlaps="run_metadata_resources",
308
312
  ),
309
313
  )
310
314
  pipeline_runs: List["PipelineRunSchema"] = Relationship(
@@ -402,9 +406,7 @@ class ModelVersionSchema(NamedSchema, table=True):
402
406
  metadata = ModelVersionResponseMetadata(
403
407
  workspace=self.workspace.to_model(),
404
408
  description=self.description,
405
- run_metadata={
406
- rm.key: json.loads(rm.value) for rm in self.run_metadata
407
- },
409
+ run_metadata=self.fetch_metadata(),
408
410
  )
409
411
 
410
412
  resources = None
@@ -15,7 +15,7 @@
15
15
 
16
16
  import json
17
17
  from datetime import datetime
18
- from typing import TYPE_CHECKING, Any, List, Optional
18
+ from typing import TYPE_CHECKING, Any, Dict, List, Optional
19
19
  from uuid import UUID
20
20
 
21
21
  from pydantic import ConfigDict
@@ -34,6 +34,7 @@ from zenml.models import (
34
34
  PipelineRunResponseBody,
35
35
  PipelineRunResponseMetadata,
36
36
  PipelineRunUpdate,
37
+ RunMetadataEntry,
37
38
  )
38
39
  from zenml.models.v2.core.pipeline_run import PipelineRunResponseResources
39
40
  from zenml.zen_stores.schemas.base_schemas import NamedSchema
@@ -48,6 +49,7 @@ from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
48
49
  from zenml.zen_stores.schemas.stack_schemas import StackSchema
49
50
  from zenml.zen_stores.schemas.trigger_schemas import TriggerExecutionSchema
50
51
  from zenml.zen_stores.schemas.user_schemas import UserSchema
52
+ from zenml.zen_stores.schemas.utils import RunMetadataInterface
51
53
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
52
54
 
53
55
  if TYPE_CHECKING:
@@ -56,13 +58,15 @@ if TYPE_CHECKING:
56
58
  ModelVersionPipelineRunSchema,
57
59
  ModelVersionSchema,
58
60
  )
59
- from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
61
+ from zenml.zen_stores.schemas.run_metadata_schemas import (
62
+ RunMetadataResourceSchema,
63
+ )
60
64
  from zenml.zen_stores.schemas.service_schemas import ServiceSchema
61
65
  from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
62
66
  from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
63
67
 
64
68
 
65
- class PipelineRunSchema(NamedSchema, table=True):
69
+ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
66
70
  """SQL Model for pipeline runs."""
67
71
 
68
72
  __tablename__ = "pipeline_run"
@@ -136,12 +140,12 @@ class PipelineRunSchema(NamedSchema, table=True):
136
140
  )
137
141
  workspace: "WorkspaceSchema" = Relationship(back_populates="runs")
138
142
  user: Optional["UserSchema"] = Relationship(back_populates="runs")
139
- run_metadata: List["RunMetadataSchema"] = Relationship(
140
- back_populates="pipeline_run",
143
+ run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
144
+ back_populates="pipeline_runs",
141
145
  sa_relationship_kwargs=dict(
142
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataSchema.resource_id)==PipelineRunSchema.id)",
146
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==PipelineRunSchema.id)",
143
147
  cascade="delete",
144
- overlaps="run_metadata",
148
+ overlaps="run_metadata_resources",
145
149
  ),
146
150
  )
147
151
  logs: Optional["LogsSchema"] = Relationship(
@@ -249,6 +253,24 @@ class PipelineRunSchema(NamedSchema, table=True):
249
253
  model_version_id=request.model_version_id,
250
254
  )
251
255
 
256
+ def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
257
+ """Fetches all the metadata entries related to the pipeline run.
258
+
259
+ Returns:
260
+ a dictionary, where the key is the key of the metadata entry
261
+ and the values represent the list of entries with this key.
262
+ """
263
+ # Fetch the metadata related to this run
264
+ metadata_collection = super().fetch_metadata_collection()
265
+
266
+ # Fetch the metadata related to the steps of this run
267
+ for s in self.step_runs:
268
+ step_metadata = s.fetch_metadata_collection()
269
+ for k, v in step_metadata.items():
270
+ metadata_collection[f"{s.name}::{k}"] = v
271
+
272
+ return metadata_collection
273
+
252
274
  def to_model(
253
275
  self,
254
276
  include_metadata: bool = False,
@@ -275,15 +297,14 @@ class PipelineRunSchema(NamedSchema, table=True):
275
297
  else {}
276
298
  )
277
299
 
278
- run_metadata = {
279
- metadata_schema.key: json.loads(metadata_schema.value)
280
- for metadata_schema in self.run_metadata
281
- }
282
-
283
300
  if self.deployment is not None:
284
301
  deployment = self.deployment.to_model()
285
302
 
286
303
  config = deployment.pipeline_configuration
304
+ new_substitutions = config._get_full_substitutions(self.start_time)
305
+ config = config.model_copy(
306
+ update={"substitutions": new_substitutions}
307
+ )
287
308
  client_environment = deployment.client_environment
288
309
 
289
310
  stack = deployment.stack
@@ -323,9 +344,11 @@ class PipelineRunSchema(NamedSchema, table=True):
323
344
  build=build,
324
345
  schedule=schedule,
325
346
  code_reference=code_reference,
326
- trigger_execution=self.trigger_execution.to_model()
327
- if self.trigger_execution
328
- else None,
347
+ trigger_execution=(
348
+ self.trigger_execution.to_model()
349
+ if self.trigger_execution
350
+ else None
351
+ ),
329
352
  created=self.created,
330
353
  updated=self.updated,
331
354
  deployment_id=self.deployment_id,
@@ -344,9 +367,13 @@ class PipelineRunSchema(NamedSchema, table=True):
344
367
 
345
368
  steps = {step.name: step.to_model() for step in self.step_runs}
346
369
 
370
+ step_substitutions = {
371
+ step_name: step.config.substitutions
372
+ for step_name, step in steps.items()
373
+ }
347
374
  metadata = PipelineRunResponseMetadata(
348
375
  workspace=self.workspace.to_model(),
349
- run_metadata=run_metadata,
376
+ run_metadata=self.fetch_metadata(),
350
377
  config=config,
351
378
  steps=steps,
352
379
  start_time=self.start_time,
@@ -361,6 +388,7 @@ class PipelineRunSchema(NamedSchema, table=True):
361
388
  if self.deployment
362
389
  else None,
363
390
  is_templatable=is_templatable,
391
+ step_substitutions=step_substitutions,
364
392
  )
365
393
 
366
394
  resources = None
@@ -14,15 +14,16 @@
14
14
  """SQLModel implementation of pipeline run metadata tables."""
15
15
 
16
16
  from typing import TYPE_CHECKING, List, Optional
17
- from uuid import UUID
17
+ from uuid import UUID, uuid4
18
18
 
19
19
  from sqlalchemy import TEXT, VARCHAR, Column
20
- from sqlmodel import Field, Relationship
20
+ from sqlmodel import Field, Relationship, SQLModel
21
21
 
22
22
  from zenml.enums import MetadataResourceTypes
23
23
  from zenml.zen_stores.schemas.base_schemas import BaseSchema
24
24
  from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
25
25
  from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
26
+ from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
26
27
  from zenml.zen_stores.schemas.user_schemas import UserSchema
27
28
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
28
29
 
@@ -30,7 +31,6 @@ if TYPE_CHECKING:
30
31
  from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
31
32
  from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
32
33
  from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
33
- from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
34
34
 
35
35
 
36
36
  class RunMetadataSchema(BaseSchema, table=True):
@@ -38,35 +38,10 @@ class RunMetadataSchema(BaseSchema, table=True):
38
38
 
39
39
  __tablename__ = "run_metadata"
40
40
 
41
- resource_id: UUID
42
- resource_type: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
43
- pipeline_run: List["PipelineRunSchema"] = Relationship(
44
- back_populates="run_metadata",
45
- sa_relationship_kwargs=dict(
46
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataSchema.resource_id)==PipelineRunSchema.id)",
47
- overlaps="run_metadata,step_run,artifact_version,model_version",
48
- ),
49
- )
50
- step_run: List["StepRunSchema"] = Relationship(
51
- back_populates="run_metadata",
52
- sa_relationship_kwargs=dict(
53
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataSchema.resource_id)==StepRunSchema.id)",
54
- overlaps="run_metadata,pipeline_run,artifact_version,model_version",
55
- ),
56
- )
57
- artifact_version: List["ArtifactVersionSchema"] = Relationship(
58
- back_populates="run_metadata",
59
- sa_relationship_kwargs=dict(
60
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ArtifactVersionSchema.id)",
61
- overlaps="run_metadata,pipeline_run,step_run,model_version",
62
- ),
63
- )
64
- model_version: List["ModelVersionSchema"] = Relationship(
41
+ # Relationship to link to resources
42
+ resources: List["RunMetadataResourceSchema"] = Relationship(
65
43
  back_populates="run_metadata",
66
- sa_relationship_kwargs=dict(
67
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ModelVersionSchema.id)",
68
- overlaps="run_metadata,pipeline_run,step_run,artifact_version",
69
- ),
44
+ sa_relationship_kwargs={"cascade": "delete"},
70
45
  )
71
46
  stack_component_id: Optional[UUID] = build_foreign_key_field(
72
47
  source=__tablename__,
@@ -103,3 +78,63 @@ class RunMetadataSchema(BaseSchema, table=True):
103
78
  key: str
104
79
  value: str = Field(sa_column=Column(TEXT, nullable=False))
105
80
  type: str
81
+
82
+ publisher_step_id: Optional[UUID] = build_foreign_key_field(
83
+ source=__tablename__,
84
+ target=StepRunSchema.__tablename__,
85
+ source_column="publisher_step_id",
86
+ target_column="id",
87
+ ondelete="SET NULL",
88
+ nullable=True,
89
+ )
90
+
91
+
92
+ class RunMetadataResourceSchema(SQLModel, table=True):
93
+ """Table for linking resources to run metadata entries."""
94
+
95
+ __tablename__ = "run_metadata_resource"
96
+
97
+ id: UUID = Field(default_factory=uuid4, primary_key=True)
98
+ resource_id: UUID
99
+ resource_type: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
100
+ run_metadata_id: UUID = build_foreign_key_field(
101
+ source=__tablename__,
102
+ target=RunMetadataSchema.__tablename__,
103
+ source_column="run_metadata_id",
104
+ target_column="id",
105
+ ondelete="CASCADE",
106
+ nullable=False,
107
+ )
108
+
109
+ # Relationship back to the base metadata table
110
+ run_metadata: RunMetadataSchema = Relationship(back_populates="resources")
111
+
112
+ # Relationship to link specific resource types
113
+ pipeline_runs: List["PipelineRunSchema"] = Relationship(
114
+ back_populates="run_metadata_resources",
115
+ sa_relationship_kwargs=dict(
116
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==PipelineRunSchema.id)",
117
+ overlaps="run_metadata_resources,step_runs,artifact_versions,model_versions",
118
+ ),
119
+ )
120
+ step_runs: List["StepRunSchema"] = Relationship(
121
+ back_populates="run_metadata_resources",
122
+ sa_relationship_kwargs=dict(
123
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==StepRunSchema.id)",
124
+ overlaps="run_metadata_resources,pipeline_runs,artifact_versions,model_versions",
125
+ ),
126
+ )
127
+ artifact_versions: List["ArtifactVersionSchema"] = Relationship(
128
+ back_populates="run_metadata_resources",
129
+ sa_relationship_kwargs=dict(
130
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ArtifactVersionSchema.id)",
131
+ overlaps="run_metadata_resources,pipeline_runs,step_runs,model_versions",
132
+ ),
133
+ )
134
+ model_versions: List["ModelVersionSchema"] = Relationship(
135
+ back_populates="run_metadata_resources",
136
+ sa_relationship_kwargs=dict(
137
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ModelVersionSchema.id)",
138
+ overlaps="run_metadata_resources,pipeline_runs,step_runs,artifact_versions",
139
+ ),
140
+ )
@@ -23,6 +23,7 @@ from sqlalchemy import TEXT, Column, String
23
23
  from sqlalchemy.dialects.mysql import MEDIUMTEXT
24
24
  from sqlmodel import Field, Relationship, SQLModel
25
25
 
26
+ from zenml.config.pipeline_configurations import PipelineConfiguration
26
27
  from zenml.config.step_configurations import Step
27
28
  from zenml.constants import MEDIUMTEXT_MAX_LENGTH
28
29
  from zenml.enums import (
@@ -50,16 +51,19 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
50
51
  from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
51
52
  from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
52
53
  from zenml.zen_stores.schemas.user_schemas import UserSchema
54
+ from zenml.zen_stores.schemas.utils import RunMetadataInterface
53
55
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
54
56
 
55
57
  if TYPE_CHECKING:
56
58
  from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
57
59
  from zenml.zen_stores.schemas.logs_schemas import LogsSchema
58
60
  from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
59
- from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
61
+ from zenml.zen_stores.schemas.run_metadata_schemas import (
62
+ RunMetadataResourceSchema,
63
+ )
60
64
 
61
65
 
62
- class StepRunSchema(NamedSchema, table=True):
66
+ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
63
67
  """SQL Model for steps of pipeline runs."""
64
68
 
65
69
  __tablename__ = "step_run"
@@ -139,12 +143,12 @@ class StepRunSchema(NamedSchema, table=True):
139
143
  deployment: Optional["PipelineDeploymentSchema"] = Relationship(
140
144
  back_populates="step_runs"
141
145
  )
142
- run_metadata: List["RunMetadataSchema"] = Relationship(
143
- back_populates="step_run",
146
+ run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
147
+ back_populates="step_runs",
144
148
  sa_relationship_kwargs=dict(
145
- primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataSchema.resource_id)==StepRunSchema.id)",
149
+ primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==StepRunSchema.id)",
146
150
  cascade="delete",
147
- overlaps="run_metadata",
151
+ overlaps="run_metadata_resources",
148
152
  ),
149
153
  )
150
154
  input_artifacts: List["StepRunInputArtifactSchema"] = Relationship(
@@ -163,9 +167,15 @@ class StepRunSchema(NamedSchema, table=True):
163
167
  "primaryjoin": "StepRunParentsSchema.child_id == StepRunSchema.id",
164
168
  },
165
169
  )
170
+ pipeline_run: "PipelineRunSchema" = Relationship(
171
+ back_populates="step_runs"
172
+ )
166
173
  model_version: "ModelVersionSchema" = Relationship(
167
174
  back_populates="step_runs",
168
175
  )
176
+ original_step_run: Optional["StepRunSchema"] = Relationship(
177
+ sa_relationship_kwargs={"remote_side": "StepRunSchema.id"}
178
+ )
169
179
 
170
180
  model_config = ConfigDict(protected_namespaces=()) # type: ignore[assignment]
171
181
 
@@ -218,11 +228,6 @@ class StepRunSchema(NamedSchema, table=True):
218
228
  RuntimeError: If the step run schema does not have a deployment_id
219
229
  or a step_configuration.
220
230
  """
221
- run_metadata = {
222
- metadata_schema.key: json.loads(metadata_schema.value)
223
- for metadata_schema in self.run_metadata
224
- }
225
-
226
231
  input_artifacts = {
227
232
  artifact.name: StepRunInputResponse(
228
233
  input_type=StepRunInputArtifactType(artifact.type),
@@ -248,6 +253,21 @@ class StepRunSchema(NamedSchema, table=True):
248
253
  full_step_config = Step.model_validate(
249
254
  step_configuration[self.name]
250
255
  )
256
+ new_substitutions = (
257
+ full_step_config.config._get_full_substitutions(
258
+ PipelineConfiguration.model_validate_json(
259
+ self.deployment.pipeline_configuration
260
+ ),
261
+ self.pipeline_run.start_time,
262
+ )
263
+ )
264
+ full_step_config = full_step_config.model_copy(
265
+ update={
266
+ "config": full_step_config.config.model_copy(
267
+ update={"substitutions": new_substitutions}
268
+ )
269
+ }
270
+ )
251
271
  elif not self.step_configuration:
252
272
  raise ValueError(
253
273
  f"Unable to load the configuration for step `{self.name}` from the"
@@ -294,7 +314,7 @@ class StepRunSchema(NamedSchema, table=True):
294
314
  pipeline_run_id=self.pipeline_run_id,
295
315
  original_step_run_id=self.original_step_run_id,
296
316
  parent_step_ids=[p.parent_id for p in self.parents],
297
- run_metadata=run_metadata,
317
+ run_metadata=self.fetch_metadata(),
298
318
  )
299
319
 
300
320
  resources = None
@@ -13,11 +13,14 @@
13
13
  # permissions and limitations under the License.
14
14
  """Utils for schemas."""
15
15
 
16
+ import json
16
17
  import math
17
- from typing import List, Type, TypeVar
18
+ from typing import Dict, List, Type, TypeVar
18
19
 
19
- from zenml.models.v2.base.base import BaseResponse
20
- from zenml.models.v2.base.page import Page
20
+ from sqlmodel import Relationship
21
+
22
+ from zenml.metadata.metadata_types import MetadataType
23
+ from zenml.models import BaseResponse, Page, RunMetadataEntry
21
24
  from zenml.zen_stores.schemas.base_schemas import BaseSchema
22
25
 
23
26
  S = TypeVar("S", bound=BaseSchema)
@@ -67,3 +70,44 @@ def get_page_from_list(
67
70
  total=total,
68
71
  items=page_items,
69
72
  )
73
+
74
+
75
+ class RunMetadataInterface:
76
+ """The interface for entities with run metadata."""
77
+
78
+ run_metadata_resources = Relationship()
79
+
80
+ def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
81
+ """Fetches all the metadata entries related to the artifact version.
82
+
83
+ Returns:
84
+ a dictionary, where the key is the key of the metadata entry
85
+ and the values represent the list of entries with this key.
86
+ """
87
+ metadata_collection: Dict[str, List[RunMetadataEntry]] = {}
88
+
89
+ # Fetch the metadata related to this step
90
+ for rm in self.run_metadata_resources:
91
+ if rm.run_metadata.key not in metadata_collection:
92
+ metadata_collection[rm.run_metadata.key] = []
93
+ metadata_collection[rm.run_metadata.key].append(
94
+ RunMetadataEntry(
95
+ value=json.loads(rm.run_metadata.value),
96
+ created=rm.run_metadata.created,
97
+ )
98
+ )
99
+
100
+ return metadata_collection
101
+
102
+ def fetch_metadata(self) -> Dict[str, MetadataType]:
103
+ """Fetches the latest metadata entry related to the artifact version.
104
+
105
+ Returns:
106
+ a dictionary, where the key is the key of the metadata entry
107
+ and the values represent the latest entry with this key.
108
+ """
109
+ metadata_collection = self.fetch_metadata_collection()
110
+ return {
111
+ k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
112
+ for k, v in metadata_collection.items()
113
+ }