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.
- zenml/VERSION +1 -1
- zenml/artifacts/artifact_config.py +0 -14
- zenml/artifacts/utils.py +50 -29
- zenml/cli/__init__.py +15 -0
- zenml/cli/base.py +4 -4
- zenml/cli/server.py +1 -1
- zenml/cli/stack.py +0 -3
- zenml/cli/stack_components.py +0 -1
- zenml/cli/utils.py +0 -5
- zenml/client.py +8 -18
- zenml/model/model.py +8 -4
- zenml/model/utils.py +18 -16
- zenml/models/__init__.py +6 -0
- zenml/models/v2/core/artifact_version.py +6 -3
- zenml/models/v2/core/component.py +0 -22
- zenml/models/v2/core/model_version.py +6 -3
- zenml/models/v2/core/pipeline_run.py +16 -4
- zenml/models/v2/core/run_metadata.py +30 -9
- zenml/models/v2/core/step_run.py +5 -3
- zenml/models/v2/misc/run_metadata.py +38 -0
- zenml/orchestrators/input_utils.py +19 -6
- zenml/orchestrators/publish_utils.py +11 -4
- zenml/orchestrators/step_launcher.py +2 -2
- zenml/orchestrators/step_run_utils.py +3 -0
- zenml/orchestrators/step_runner.py +7 -8
- zenml/orchestrators/utils.py +0 -24
- zenml/pipelines/run_utils.py +3 -4
- zenml/steps/utils.py +6 -2
- zenml/utils/metadata_utils.py +186 -153
- zenml/utils/string_utils.py +15 -8
- zenml/zen_server/routers/workspaces_endpoints.py +19 -19
- zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py +36 -0
- zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py +135 -0
- zenml/zen_stores/schemas/__init__.py +5 -1
- zenml/zen_stores/schemas/artifact_schemas.py +10 -10
- zenml/zen_stores/schemas/component_schemas.py +0 -3
- zenml/zen_stores/schemas/model_schemas.py +13 -11
- zenml/zen_stores/schemas/pipeline_run_schemas.py +32 -15
- zenml/zen_stores/schemas/run_metadata_schemas.py +66 -31
- zenml/zen_stores/schemas/step_run_schemas.py +13 -12
- zenml/zen_stores/schemas/utils.py +47 -3
- zenml/zen_stores/sql_zen_store.py +102 -23
- {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/METADATA +1 -1
- {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/RECORD +47 -44
- {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.70.0.dev20241129.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/WHEEL +0 -0
- {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
|
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
|
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
|
-
|
246
|
-
back_populates="
|
247
|
+
run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
|
248
|
+
back_populates="artifact_versions",
|
247
249
|
sa_relationship_kwargs=dict(
|
248
|
-
primaryjoin=f"and_(
|
250
|
+
primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ArtifactVersionSchema.id)",
|
249
251
|
cascade="delete",
|
250
|
-
overlaps="
|
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
|
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
|
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
|
-
|
303
|
-
back_populates="
|
306
|
+
run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
|
307
|
+
back_populates="model_versions",
|
304
308
|
sa_relationship_kwargs=dict(
|
305
|
-
primaryjoin=f"and_(
|
309
|
+
primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ModelVersionSchema.id)",
|
306
310
|
cascade="delete",
|
307
|
-
overlaps="
|
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
|
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
|
-
|
140
|
-
back_populates="
|
143
|
+
run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
|
144
|
+
back_populates="pipeline_runs",
|
141
145
|
sa_relationship_kwargs=dict(
|
142
|
-
primaryjoin=f"and_(
|
146
|
+
primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==PipelineRunSchema.id)",
|
143
147
|
cascade="delete",
|
144
|
-
overlaps="
|
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
|
-
|
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=
|
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
|
-
|
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
|
-
|
42
|
-
|
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=
|
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
|
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
|
-
|
144
|
-
back_populates="
|
146
|
+
run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
|
147
|
+
back_populates="step_runs",
|
145
148
|
sa_relationship_kwargs=dict(
|
146
|
-
primaryjoin=f"and_(
|
149
|
+
primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==StepRunSchema.id)",
|
147
150
|
cascade="delete",
|
148
|
-
overlaps="
|
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=
|
317
|
+
run_metadata=self.fetch_metadata(),
|
317
318
|
)
|
318
319
|
|
319
320
|
resources = None
|