zenml-nightly 0.70.0.dev20241129__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 (47) hide show
  1. zenml/VERSION +1 -1
  2. zenml/artifacts/artifact_config.py +0 -14
  3. zenml/artifacts/utils.py +50 -29
  4. zenml/cli/__init__.py +15 -0
  5. zenml/cli/base.py +4 -4
  6. zenml/cli/server.py +1 -1
  7. zenml/cli/stack.py +0 -3
  8. zenml/cli/stack_components.py +0 -1
  9. zenml/cli/utils.py +0 -5
  10. zenml/client.py +8 -18
  11. zenml/model/model.py +8 -4
  12. zenml/model/utils.py +18 -16
  13. zenml/models/__init__.py +6 -0
  14. zenml/models/v2/core/artifact_version.py +6 -3
  15. zenml/models/v2/core/component.py +0 -22
  16. zenml/models/v2/core/model_version.py +6 -3
  17. zenml/models/v2/core/pipeline_run.py +16 -4
  18. zenml/models/v2/core/run_metadata.py +30 -9
  19. zenml/models/v2/core/step_run.py +5 -3
  20. zenml/models/v2/misc/run_metadata.py +38 -0
  21. zenml/orchestrators/input_utils.py +19 -6
  22. zenml/orchestrators/publish_utils.py +11 -4
  23. zenml/orchestrators/step_launcher.py +2 -2
  24. zenml/orchestrators/step_run_utils.py +3 -0
  25. zenml/orchestrators/step_runner.py +7 -8
  26. zenml/orchestrators/utils.py +0 -24
  27. zenml/pipelines/run_utils.py +3 -4
  28. zenml/steps/utils.py +6 -2
  29. zenml/utils/metadata_utils.py +186 -153
  30. zenml/utils/string_utils.py +15 -8
  31. zenml/zen_server/routers/workspaces_endpoints.py +19 -19
  32. zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py +36 -0
  33. zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py +135 -0
  34. zenml/zen_stores/schemas/__init__.py +5 -1
  35. zenml/zen_stores/schemas/artifact_schemas.py +10 -10
  36. zenml/zen_stores/schemas/component_schemas.py +0 -3
  37. zenml/zen_stores/schemas/model_schemas.py +13 -11
  38. zenml/zen_stores/schemas/pipeline_run_schemas.py +32 -15
  39. zenml/zen_stores/schemas/run_metadata_schemas.py +66 -31
  40. zenml/zen_stores/schemas/step_run_schemas.py +13 -12
  41. zenml/zen_stores/schemas/utils.py +47 -3
  42. zenml/zen_stores/sql_zen_store.py +102 -23
  43. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/METADATA +1 -1
  44. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/RECORD +47 -44
  45. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/LICENSE +0 -0
  46. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/WHEEL +0 -0
  47. {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/entry_points.txt +0 -0
@@ -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(
@@ -376,9 +378,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
376
378
  workspace=self.workspace.to_model(),
377
379
  producer_step_run_id=producer_step_run_id,
378
380
  visualizations=[v.to_model() for v in self.visualizations],
379
- run_metadata={
380
- m.key: json.loads(m.value) for m in self.run_metadata
381
- },
381
+ run_metadata=self.fetch_metadata(),
382
382
  )
383
383
 
384
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,11 +297,6 @@ 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
 
@@ -350,13 +367,13 @@ class PipelineRunSchema(NamedSchema, table=True):
350
367
 
351
368
  steps = {step.name: step.to_model() for step in self.step_runs}
352
369
 
353
- steps_substitutions = {
370
+ step_substitutions = {
354
371
  step_name: step.config.substitutions
355
372
  for step_name, step in steps.items()
356
373
  }
357
374
  metadata = PipelineRunResponseMetadata(
358
375
  workspace=self.workspace.to_model(),
359
- run_metadata=run_metadata,
376
+ run_metadata=self.fetch_metadata(),
360
377
  config=config,
361
378
  steps=steps,
362
379
  start_time=self.start_time,
@@ -371,7 +388,7 @@ class PipelineRunSchema(NamedSchema, table=True):
371
388
  if self.deployment
372
389
  else None,
373
390
  is_templatable=is_templatable,
374
- steps_substitutions=steps_substitutions,
391
+ step_substitutions=step_substitutions,
375
392
  )
376
393
 
377
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
+ )
@@ -51,16 +51,19 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
51
51
  from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
52
52
  from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
53
53
  from zenml.zen_stores.schemas.user_schemas import UserSchema
54
+ from zenml.zen_stores.schemas.utils import RunMetadataInterface
54
55
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
55
56
 
56
57
  if TYPE_CHECKING:
57
58
  from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
58
59
  from zenml.zen_stores.schemas.logs_schemas import LogsSchema
59
60
  from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
60
- from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
61
+ from zenml.zen_stores.schemas.run_metadata_schemas import (
62
+ RunMetadataResourceSchema,
63
+ )
61
64
 
62
65
 
63
- class StepRunSchema(NamedSchema, table=True):
66
+ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
64
67
  """SQL Model for steps of pipeline runs."""
65
68
 
66
69
  __tablename__ = "step_run"
@@ -140,12 +143,12 @@ class StepRunSchema(NamedSchema, table=True):
140
143
  deployment: Optional["PipelineDeploymentSchema"] = Relationship(
141
144
  back_populates="step_runs"
142
145
  )
143
- run_metadata: List["RunMetadataSchema"] = Relationship(
144
- back_populates="step_run",
146
+ run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
147
+ back_populates="step_runs",
145
148
  sa_relationship_kwargs=dict(
146
- 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)",
147
150
  cascade="delete",
148
- overlaps="run_metadata",
151
+ overlaps="run_metadata_resources",
149
152
  ),
150
153
  )
151
154
  input_artifacts: List["StepRunInputArtifactSchema"] = Relationship(
@@ -170,6 +173,9 @@ class StepRunSchema(NamedSchema, table=True):
170
173
  model_version: "ModelVersionSchema" = Relationship(
171
174
  back_populates="step_runs",
172
175
  )
176
+ original_step_run: Optional["StepRunSchema"] = Relationship(
177
+ sa_relationship_kwargs={"remote_side": "StepRunSchema.id"}
178
+ )
173
179
 
174
180
  model_config = ConfigDict(protected_namespaces=()) # type: ignore[assignment]
175
181
 
@@ -222,11 +228,6 @@ class StepRunSchema(NamedSchema, table=True):
222
228
  RuntimeError: If the step run schema does not have a deployment_id
223
229
  or a step_configuration.
224
230
  """
225
- run_metadata = {
226
- metadata_schema.key: json.loads(metadata_schema.value)
227
- for metadata_schema in self.run_metadata
228
- }
229
-
230
231
  input_artifacts = {
231
232
  artifact.name: StepRunInputResponse(
232
233
  input_type=StepRunInputArtifactType(artifact.type),
@@ -313,7 +314,7 @@ class StepRunSchema(NamedSchema, table=True):
313
314
  pipeline_run_id=self.pipeline_run_id,
314
315
  original_step_run_id=self.original_step_run_id,
315
316
  parent_step_ids=[p.parent_id for p in self.parents],
316
- run_metadata=run_metadata,
317
+ run_metadata=self.fetch_metadata(),
317
318
  )
318
319
 
319
320
  resources = None