zenml-nightly 0.70.0.dev20241128__py3-none-any.whl → 0.70.0.dev20241130__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 (65) hide show
  1. zenml/VERSION +1 -1
  2. zenml/artifacts/artifact_config.py +7 -1
  3. zenml/artifacts/utils.py +55 -30
  4. zenml/cli/__init__.py +15 -0
  5. zenml/cli/base.py +4 -4
  6. zenml/cli/pipeline.py +80 -0
  7. zenml/cli/server.py +1 -1
  8. zenml/cli/stack.py +0 -3
  9. zenml/cli/stack_components.py +0 -1
  10. zenml/cli/utils.py +0 -5
  11. zenml/client.py +8 -18
  12. zenml/config/compiler.py +12 -3
  13. zenml/config/pipeline_configurations.py +20 -0
  14. zenml/config/pipeline_run_configuration.py +1 -0
  15. zenml/config/step_configurations.py +21 -0
  16. zenml/enums.py +1 -0
  17. zenml/integrations/feast/__init__.py +1 -1
  18. zenml/integrations/feast/feature_stores/feast_feature_store.py +13 -9
  19. zenml/materializers/built_in_materializer.py +18 -1
  20. zenml/materializers/structured_string_materializer.py +8 -3
  21. zenml/model/model.py +11 -3
  22. zenml/model/utils.py +18 -16
  23. zenml/models/__init__.py +6 -0
  24. zenml/models/v2/core/artifact_version.py +6 -3
  25. zenml/models/v2/core/component.py +0 -22
  26. zenml/models/v2/core/model_version.py +6 -3
  27. zenml/models/v2/core/pipeline_run.py +19 -3
  28. zenml/models/v2/core/run_metadata.py +30 -9
  29. zenml/models/v2/core/step_run.py +6 -4
  30. zenml/models/v2/misc/run_metadata.py +38 -0
  31. zenml/orchestrators/input_utils.py +19 -6
  32. zenml/orchestrators/publish_utils.py +12 -5
  33. zenml/orchestrators/step_launcher.py +7 -3
  34. zenml/orchestrators/step_run_utils.py +18 -6
  35. zenml/orchestrators/step_runner.py +39 -2
  36. zenml/orchestrators/utils.py +0 -21
  37. zenml/pipelines/pipeline_decorator.py +4 -0
  38. zenml/pipelines/pipeline_definition.py +14 -3
  39. zenml/pipelines/run_utils.py +9 -5
  40. zenml/steps/base_step.py +11 -1
  41. zenml/steps/entrypoint_function_utils.py +4 -2
  42. zenml/steps/step_decorator.py +4 -0
  43. zenml/steps/utils.py +23 -7
  44. zenml/types.py +4 -0
  45. zenml/utils/metadata_utils.py +186 -153
  46. zenml/utils/string_utils.py +41 -16
  47. zenml/utils/visualization_utils.py +4 -1
  48. zenml/zen_server/routers/workspaces_endpoints.py +19 -19
  49. zenml/zen_server/template_execution/utils.py +1 -0
  50. zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py +36 -0
  51. zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py +135 -0
  52. zenml/zen_stores/schemas/__init__.py +5 -1
  53. zenml/zen_stores/schemas/artifact_schemas.py +12 -11
  54. zenml/zen_stores/schemas/component_schemas.py +0 -3
  55. zenml/zen_stores/schemas/model_schemas.py +13 -11
  56. zenml/zen_stores/schemas/pipeline_run_schemas.py +44 -16
  57. zenml/zen_stores/schemas/run_metadata_schemas.py +66 -31
  58. zenml/zen_stores/schemas/step_run_schemas.py +32 -12
  59. zenml/zen_stores/schemas/utils.py +47 -3
  60. zenml/zen_stores/sql_zen_store.py +117 -34
  61. {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/METADATA +1 -1
  62. {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/RECORD +65 -62
  63. {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/LICENSE +0 -0
  64. {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/WHEEL +0 -0
  65. {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,36 @@
1
+ """Remove component spec path [b73bc71f1106].
2
+
3
+ Revision ID: b73bc71f1106
4
+ Revises: ec6307720f92
5
+ Create Date: 2024-11-29 09:36:33.089945
6
+
7
+ """
8
+
9
+ import sqlalchemy as sa
10
+ from alembic import op
11
+
12
+ # revision identifiers, used by Alembic.
13
+ revision = "b73bc71f1106"
14
+ down_revision = "ec6307720f92"
15
+ branch_labels = None
16
+ depends_on = None
17
+
18
+
19
+ def upgrade() -> None:
20
+ """Upgrade database schema and/or data, creating a new revision."""
21
+ # ### commands auto generated by Alembic - please adjust! ###
22
+ with op.batch_alter_table("stack_component", schema=None) as batch_op:
23
+ batch_op.drop_column("component_spec_path")
24
+
25
+ # ### end Alembic commands ###
26
+
27
+
28
+ def downgrade() -> None:
29
+ """Downgrade database schema and/or data back to the previous revision."""
30
+ # ### commands auto generated by Alembic - please adjust! ###
31
+ with op.batch_alter_table("stack_component", schema=None) as batch_op:
32
+ batch_op.add_column(
33
+ sa.Column("component_spec_path", sa.VARCHAR(), nullable=True)
34
+ )
35
+
36
+ # ### end Alembic commands ###
@@ -0,0 +1,135 @@
1
+ """Separate run metadata into resource link table with new UUIDs.
2
+
3
+ Revision ID: cc269488e5a9
4
+ Revises: b73bc71f1106
5
+ Create Date: 2024-11-12 09:46:46.587478
6
+ """
7
+
8
+ import uuid
9
+
10
+ import sqlalchemy as sa
11
+ import sqlmodel
12
+ from alembic import op
13
+
14
+ # revision identifiers, used by Alembic.
15
+ revision = "cc269488e5a9"
16
+ down_revision = "b73bc71f1106"
17
+ branch_labels = None
18
+ depends_on = None
19
+
20
+
21
+ def upgrade() -> None:
22
+ """Creates the 'run_metadata_resource' table and migrates data."""
23
+ op.create_table(
24
+ "run_metadata_resource",
25
+ sa.Column(
26
+ "id",
27
+ sqlmodel.sql.sqltypes.GUID(),
28
+ nullable=False,
29
+ primary_key=True,
30
+ ),
31
+ sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
32
+ sa.Column("resource_type", sa.String(length=255), nullable=False),
33
+ sa.Column(
34
+ "run_metadata_id",
35
+ sqlmodel.sql.sqltypes.GUID(),
36
+ sa.ForeignKey("run_metadata.id", ondelete="CASCADE"),
37
+ nullable=False,
38
+ ),
39
+ )
40
+
41
+ connection = op.get_bind()
42
+
43
+ run_metadata_data = connection.execute(
44
+ sa.text("""
45
+ SELECT id, resource_id, resource_type
46
+ FROM run_metadata
47
+ """)
48
+ ).fetchall()
49
+
50
+ # Prepare data with new UUIDs for bulk insert
51
+ resource_data = [
52
+ {
53
+ "id": str(uuid.uuid4()), # Generate a new UUID for each row
54
+ "resource_id": row.resource_id,
55
+ "resource_type": row.resource_type,
56
+ "run_metadata_id": row.id,
57
+ }
58
+ for row in run_metadata_data
59
+ ]
60
+
61
+ # Perform bulk insert into `run_metadata_resource`
62
+ if resource_data: # Only perform insert if there's data to migrate
63
+ op.bulk_insert(
64
+ sa.table(
65
+ "run_metadata_resource",
66
+ sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
67
+ sa.Column(
68
+ "resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False
69
+ ),
70
+ sa.Column(
71
+ "resource_type", sa.String(length=255), nullable=False
72
+ ),
73
+ sa.Column(
74
+ "run_metadata_id",
75
+ sqlmodel.sql.sqltypes.GUID(),
76
+ nullable=False,
77
+ ), # Changed to BIGINT
78
+ ),
79
+ resource_data,
80
+ )
81
+
82
+ op.drop_column("run_metadata", "resource_id")
83
+ op.drop_column("run_metadata", "resource_type")
84
+
85
+ op.add_column(
86
+ "run_metadata",
87
+ sa.Column(
88
+ "publisher_step_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
89
+ ),
90
+ )
91
+
92
+
93
+ def downgrade() -> None:
94
+ """Reverts the 'run_metadata_resource' table and migrates data back."""
95
+ # Recreate the `resource_id` and `resource_type` columns in `run_metadata`
96
+ op.add_column(
97
+ "run_metadata",
98
+ sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=True),
99
+ )
100
+ op.add_column(
101
+ "run_metadata",
102
+ sa.Column("resource_type", sa.String(length=255), nullable=True),
103
+ )
104
+
105
+ # Migrate data back from `run_metadata_resource` to `run_metadata`
106
+ connection = op.get_bind()
107
+
108
+ # Fetch data from `run_metadata_resource`
109
+ run_metadata_resource_data = connection.execute(
110
+ sa.text("""
111
+ SELECT resource_id, resource_type, run_metadata_id
112
+ FROM run_metadata_resource
113
+ """)
114
+ ).fetchall()
115
+
116
+ # Update `run_metadata` with the data from `run_metadata_resource`
117
+ for row in run_metadata_resource_data:
118
+ connection.execute(
119
+ sa.text("""
120
+ UPDATE run_metadata
121
+ SET resource_id = :resource_id, resource_type = :resource_type
122
+ WHERE id = :run_metadata_id
123
+ """),
124
+ {
125
+ "resource_id": row.resource_id,
126
+ "resource_type": row.resource_type,
127
+ "run_metadata_id": row.run_metadata_id,
128
+ },
129
+ )
130
+
131
+ # Drop the `run_metadata_resource` table
132
+ op.drop_table("run_metadata_resource")
133
+
134
+ # Drop the cached column
135
+ op.drop_column("run_metadata", "publisher_step_id")
@@ -39,7 +39,10 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
39
39
  from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
40
40
  from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
41
41
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
42
- from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
42
+ from zenml.zen_stores.schemas.run_metadata_schemas import (
43
+ RunMetadataResourceSchema,
44
+ RunMetadataSchema,
45
+ )
43
46
  from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
44
47
  from zenml.zen_stores.schemas.secret_schemas import SecretSchema
45
48
  from zenml.zen_stores.schemas.service_schemas import ServiceSchema
@@ -90,6 +93,7 @@ __all__ = [
90
93
  "PipelineDeploymentSchema",
91
94
  "PipelineRunSchema",
92
95
  "PipelineSchema",
96
+ "RunMetadataResourceSchema",
93
97
  "RunMetadataSchema",
94
98
  "ScheduleSchema",
95
99
  "SecretSchema",
@@ -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
+ )