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.
- zenml/VERSION +1 -1
- zenml/artifacts/artifact_config.py +7 -1
- zenml/artifacts/utils.py +55 -30
- zenml/cli/__init__.py +15 -0
- zenml/cli/base.py +4 -4
- zenml/cli/pipeline.py +80 -0
- 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/config/compiler.py +12 -3
- zenml/config/pipeline_configurations.py +20 -0
- zenml/config/pipeline_run_configuration.py +1 -0
- zenml/config/step_configurations.py +21 -0
- zenml/enums.py +1 -0
- zenml/integrations/feast/__init__.py +1 -1
- zenml/integrations/feast/feature_stores/feast_feature_store.py +13 -9
- zenml/materializers/built_in_materializer.py +18 -1
- zenml/materializers/structured_string_materializer.py +8 -3
- zenml/model/model.py +11 -3
- 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 +19 -3
- zenml/models/v2/core/run_metadata.py +30 -9
- zenml/models/v2/core/step_run.py +6 -4
- zenml/models/v2/misc/run_metadata.py +38 -0
- zenml/orchestrators/input_utils.py +19 -6
- zenml/orchestrators/publish_utils.py +12 -5
- zenml/orchestrators/step_launcher.py +7 -3
- zenml/orchestrators/step_run_utils.py +18 -6
- zenml/orchestrators/step_runner.py +39 -2
- zenml/orchestrators/utils.py +0 -21
- zenml/pipelines/pipeline_decorator.py +4 -0
- zenml/pipelines/pipeline_definition.py +14 -3
- zenml/pipelines/run_utils.py +9 -5
- zenml/steps/base_step.py +11 -1
- zenml/steps/entrypoint_function_utils.py +4 -2
- zenml/steps/step_decorator.py +4 -0
- zenml/steps/utils.py +23 -7
- zenml/types.py +4 -0
- zenml/utils/metadata_utils.py +186 -153
- zenml/utils/string_utils.py +41 -16
- zenml/utils/visualization_utils.py +4 -1
- zenml/zen_server/routers/workspaces_endpoints.py +19 -19
- zenml/zen_server/template_execution/utils.py +1 -0
- 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 +12 -11
- 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 +44 -16
- zenml/zen_stores/schemas/run_metadata_schemas.py +66 -31
- zenml/zen_stores/schemas/step_run_schemas.py +32 -12
- zenml/zen_stores/schemas/utils.py +47 -3
- zenml/zen_stores/sql_zen_store.py +117 -34
- {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/METADATA +1 -1
- {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/RECORD +65 -62
- {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/WHEEL +0 -0
- {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
|
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(
|
@@ -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=
|
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
|
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,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=
|
327
|
-
|
328
|
-
|
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=
|
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
|
-
|
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
|
+
)
|