zenml-nightly 0.70.0.dev20241122__py3-none-any.whl → 0.70.0.dev20241201__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- zenml/VERSION +1 -1
- zenml/artifact_stores/base_artifact_store.py +2 -2
- zenml/artifacts/artifact_config.py +7 -1
- zenml/artifacts/utils.py +56 -31
- zenml/cli/__init__.py +18 -0
- zenml/cli/base.py +4 -4
- zenml/cli/login.py +26 -0
- zenml/cli/pipeline.py +80 -0
- zenml/cli/server.py +1 -1
- zenml/cli/service_connectors.py +3 -3
- 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/constants.py +1 -0
- zenml/enums.py +1 -0
- zenml/image_builders/local_image_builder.py +13 -3
- zenml/integrations/__init__.py +1 -0
- zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +14 -6
- zenml/integrations/constants.py +1 -0
- zenml/integrations/feast/__init__.py +1 -1
- zenml/integrations/feast/feature_stores/feast_feature_store.py +13 -9
- zenml/integrations/kubernetes/orchestrators/kube_utils.py +54 -9
- zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +65 -3
- zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py +14 -2
- zenml/integrations/kubernetes/orchestrators/manifest_utils.py +6 -5
- zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py +2 -1
- zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py +3 -1
- zenml/integrations/modal/__init__.py +46 -0
- zenml/integrations/modal/flavors/__init__.py +26 -0
- zenml/integrations/modal/flavors/modal_step_operator_flavor.py +125 -0
- zenml/integrations/modal/step_operators/__init__.py +22 -0
- zenml/integrations/modal/step_operators/modal_step_operator.py +242 -0
- zenml/io/filesystem.py +2 -2
- zenml/io/local_filesystem.py +3 -3
- zenml/materializers/built_in_materializer.py +18 -1
- zenml/materializers/structured_string_materializer.py +8 -3
- zenml/model/model.py +11 -85
- zenml/model/utils.py +18 -16
- zenml/models/__init__.py +12 -1
- 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/service_connector.py +4 -0
- zenml/models/v2/core/step_run.py +6 -4
- zenml/models/v2/misc/run_metadata.py +38 -0
- zenml/models/v2/misc/server_models.py +23 -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 +26 -9
- zenml/orchestrators/step_runner.py +40 -3
- zenml/orchestrators/utils.py +24 -23
- zenml/pipelines/pipeline_decorator.py +4 -0
- zenml/pipelines/pipeline_definition.py +26 -8
- zenml/pipelines/run_utils.py +9 -5
- zenml/steps/base_step.py +11 -1
- zenml/steps/entrypoint_function_utils.py +7 -3
- 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/cloud_utils.py +3 -1
- zenml/zen_server/deploy/helm/templates/_environment.tpl +117 -0
- zenml/zen_server/deploy/helm/templates/server-db-job.yaml +3 -14
- zenml/zen_server/deploy/helm/templates/server-deployment.yaml +16 -4
- zenml/zen_server/deploy/helm/templates/server-secret.yaml +2 -17
- zenml/zen_server/rbac/endpoint_utils.py +6 -4
- zenml/zen_server/rbac/models.py +3 -2
- zenml/zen_server/rbac/utils.py +4 -7
- zenml/zen_server/routers/server_endpoints.py +47 -0
- zenml/zen_server/routers/users_endpoints.py +35 -37
- zenml/zen_server/routers/workspaces_endpoints.py +44 -55
- zenml/zen_server/template_execution/utils.py +1 -0
- zenml/zen_server/zen_server_api.py +45 -6
- zenml/zen_stores/migrations/utils.py +40 -24
- 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/migrations/versions/ec6307720f92_simplify_model_version_links.py +7 -6
- zenml/zen_stores/rest_zen_store.py +38 -1
- 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 +130 -34
- {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/METADATA +1 -1
- {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/RECORD +102 -95
- zenml/utils/cloud_utils.py +0 -40
- {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/WHEEL +0 -0
- {zenml_nightly-0.70.0.dev20241122.dist-info → zenml_nightly-0.70.0.dev20241201.dist-info}/entry_points.txt +0 -0
@@ -13,7 +13,6 @@
|
|
13
13
|
# permissions and limitations under the License.
|
14
14
|
"""SQLModel implementation of artifact table."""
|
15
15
|
|
16
|
-
import json
|
17
16
|
from datetime import datetime
|
18
17
|
from typing import TYPE_CHECKING, Any, List, Optional
|
19
18
|
from uuid import UUID
|
@@ -50,6 +49,7 @@ from zenml.zen_stores.schemas.step_run_schemas import (
|
|
50
49
|
StepRunOutputArtifactSchema,
|
51
50
|
)
|
52
51
|
from zenml.zen_stores.schemas.user_schemas import UserSchema
|
52
|
+
from zenml.zen_stores.schemas.utils import RunMetadataInterface
|
53
53
|
from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
|
54
54
|
|
55
55
|
if TYPE_CHECKING:
|
@@ -59,7 +59,9 @@ if TYPE_CHECKING:
|
|
59
59
|
from zenml.zen_stores.schemas.model_schemas import (
|
60
60
|
ModelVersionArtifactSchema,
|
61
61
|
)
|
62
|
-
from zenml.zen_stores.schemas.run_metadata_schemas import
|
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
|
+
)
|
@@ -23,6 +23,7 @@ from sqlalchemy import TEXT, Column, String
|
|
23
23
|
from sqlalchemy.dialects.mysql import MEDIUMTEXT
|
24
24
|
from sqlmodel import Field, Relationship, SQLModel
|
25
25
|
|
26
|
+
from zenml.config.pipeline_configurations import PipelineConfiguration
|
26
27
|
from zenml.config.step_configurations import Step
|
27
28
|
from zenml.constants import MEDIUMTEXT_MAX_LENGTH
|
28
29
|
from zenml.enums import (
|
@@ -50,16 +51,19 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
|
|
50
51
|
from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
|
51
52
|
from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
|
52
53
|
from zenml.zen_stores.schemas.user_schemas import UserSchema
|
54
|
+
from zenml.zen_stores.schemas.utils import RunMetadataInterface
|
53
55
|
from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
|
54
56
|
|
55
57
|
if TYPE_CHECKING:
|
56
58
|
from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
|
57
59
|
from zenml.zen_stores.schemas.logs_schemas import LogsSchema
|
58
60
|
from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
|
59
|
-
from zenml.zen_stores.schemas.run_metadata_schemas import
|
61
|
+
from zenml.zen_stores.schemas.run_metadata_schemas import (
|
62
|
+
RunMetadataResourceSchema,
|
63
|
+
)
|
60
64
|
|
61
65
|
|
62
|
-
class StepRunSchema(NamedSchema, table=True):
|
66
|
+
class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
|
63
67
|
"""SQL Model for steps of pipeline runs."""
|
64
68
|
|
65
69
|
__tablename__ = "step_run"
|
@@ -139,12 +143,12 @@ class StepRunSchema(NamedSchema, table=True):
|
|
139
143
|
deployment: Optional["PipelineDeploymentSchema"] = Relationship(
|
140
144
|
back_populates="step_runs"
|
141
145
|
)
|
142
|
-
|
143
|
-
back_populates="
|
146
|
+
run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
|
147
|
+
back_populates="step_runs",
|
144
148
|
sa_relationship_kwargs=dict(
|
145
|
-
primaryjoin=f"and_(
|
149
|
+
primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==StepRunSchema.id)",
|
146
150
|
cascade="delete",
|
147
|
-
overlaps="
|
151
|
+
overlaps="run_metadata_resources",
|
148
152
|
),
|
149
153
|
)
|
150
154
|
input_artifacts: List["StepRunInputArtifactSchema"] = Relationship(
|
@@ -163,9 +167,15 @@ class StepRunSchema(NamedSchema, table=True):
|
|
163
167
|
"primaryjoin": "StepRunParentsSchema.child_id == StepRunSchema.id",
|
164
168
|
},
|
165
169
|
)
|
170
|
+
pipeline_run: "PipelineRunSchema" = Relationship(
|
171
|
+
back_populates="step_runs"
|
172
|
+
)
|
166
173
|
model_version: "ModelVersionSchema" = Relationship(
|
167
174
|
back_populates="step_runs",
|
168
175
|
)
|
176
|
+
original_step_run: Optional["StepRunSchema"] = Relationship(
|
177
|
+
sa_relationship_kwargs={"remote_side": "StepRunSchema.id"}
|
178
|
+
)
|
169
179
|
|
170
180
|
model_config = ConfigDict(protected_namespaces=()) # type: ignore[assignment]
|
171
181
|
|
@@ -218,11 +228,6 @@ class StepRunSchema(NamedSchema, table=True):
|
|
218
228
|
RuntimeError: If the step run schema does not have a deployment_id
|
219
229
|
or a step_configuration.
|
220
230
|
"""
|
221
|
-
run_metadata = {
|
222
|
-
metadata_schema.key: json.loads(metadata_schema.value)
|
223
|
-
for metadata_schema in self.run_metadata
|
224
|
-
}
|
225
|
-
|
226
231
|
input_artifacts = {
|
227
232
|
artifact.name: StepRunInputResponse(
|
228
233
|
input_type=StepRunInputArtifactType(artifact.type),
|
@@ -248,6 +253,21 @@ class StepRunSchema(NamedSchema, table=True):
|
|
248
253
|
full_step_config = Step.model_validate(
|
249
254
|
step_configuration[self.name]
|
250
255
|
)
|
256
|
+
new_substitutions = (
|
257
|
+
full_step_config.config._get_full_substitutions(
|
258
|
+
PipelineConfiguration.model_validate_json(
|
259
|
+
self.deployment.pipeline_configuration
|
260
|
+
),
|
261
|
+
self.pipeline_run.start_time,
|
262
|
+
)
|
263
|
+
)
|
264
|
+
full_step_config = full_step_config.model_copy(
|
265
|
+
update={
|
266
|
+
"config": full_step_config.config.model_copy(
|
267
|
+
update={"substitutions": new_substitutions}
|
268
|
+
)
|
269
|
+
}
|
270
|
+
)
|
251
271
|
elif not self.step_configuration:
|
252
272
|
raise ValueError(
|
253
273
|
f"Unable to load the configuration for step `{self.name}` from the"
|
@@ -294,7 +314,7 @@ class StepRunSchema(NamedSchema, table=True):
|
|
294
314
|
pipeline_run_id=self.pipeline_run_id,
|
295
315
|
original_step_run_id=self.original_step_run_id,
|
296
316
|
parent_step_ids=[p.parent_id for p in self.parents],
|
297
|
-
run_metadata=
|
317
|
+
run_metadata=self.fetch_metadata(),
|
298
318
|
)
|
299
319
|
|
300
320
|
resources = None
|
@@ -13,11 +13,14 @@
|
|
13
13
|
# permissions and limitations under the License.
|
14
14
|
"""Utils for schemas."""
|
15
15
|
|
16
|
+
import json
|
16
17
|
import math
|
17
|
-
from typing import List, Type, TypeVar
|
18
|
+
from typing import Dict, List, Type, TypeVar
|
18
19
|
|
19
|
-
from
|
20
|
-
|
20
|
+
from sqlmodel import Relationship
|
21
|
+
|
22
|
+
from zenml.metadata.metadata_types import MetadataType
|
23
|
+
from zenml.models import BaseResponse, Page, RunMetadataEntry
|
21
24
|
from zenml.zen_stores.schemas.base_schemas import BaseSchema
|
22
25
|
|
23
26
|
S = TypeVar("S", bound=BaseSchema)
|
@@ -67,3 +70,44 @@ def get_page_from_list(
|
|
67
70
|
total=total,
|
68
71
|
items=page_items,
|
69
72
|
)
|
73
|
+
|
74
|
+
|
75
|
+
class RunMetadataInterface:
|
76
|
+
"""The interface for entities with run metadata."""
|
77
|
+
|
78
|
+
run_metadata_resources = Relationship()
|
79
|
+
|
80
|
+
def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
|
81
|
+
"""Fetches all the metadata entries related to the artifact version.
|
82
|
+
|
83
|
+
Returns:
|
84
|
+
a dictionary, where the key is the key of the metadata entry
|
85
|
+
and the values represent the list of entries with this key.
|
86
|
+
"""
|
87
|
+
metadata_collection: Dict[str, List[RunMetadataEntry]] = {}
|
88
|
+
|
89
|
+
# Fetch the metadata related to this step
|
90
|
+
for rm in self.run_metadata_resources:
|
91
|
+
if rm.run_metadata.key not in metadata_collection:
|
92
|
+
metadata_collection[rm.run_metadata.key] = []
|
93
|
+
metadata_collection[rm.run_metadata.key].append(
|
94
|
+
RunMetadataEntry(
|
95
|
+
value=json.loads(rm.run_metadata.value),
|
96
|
+
created=rm.run_metadata.created,
|
97
|
+
)
|
98
|
+
)
|
99
|
+
|
100
|
+
return metadata_collection
|
101
|
+
|
102
|
+
def fetch_metadata(self) -> Dict[str, MetadataType]:
|
103
|
+
"""Fetches the latest metadata entry related to the artifact version.
|
104
|
+
|
105
|
+
Returns:
|
106
|
+
a dictionary, where the key is the key of the metadata entry
|
107
|
+
and the values represent the latest entry with this key.
|
108
|
+
"""
|
109
|
+
metadata_collection = self.fetch_metadata_collection()
|
110
|
+
return {
|
111
|
+
k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
|
112
|
+
for k, v in metadata_collection.items()
|
113
|
+
}
|