zenml-nightly 0.68.1.dev20241107__py3-none-any.whl → 0.68.1.dev20241108__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. zenml/VERSION +1 -1
  2. zenml/artifacts/external_artifact.py +2 -1
  3. zenml/artifacts/utils.py +13 -20
  4. zenml/cli/base.py +4 -4
  5. zenml/cli/model.py +1 -6
  6. zenml/cli/stack.py +1 -0
  7. zenml/client.py +21 -73
  8. zenml/enums.py +12 -4
  9. zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +1 -1
  10. zenml/integrations/azure/orchestrators/azureml_orchestrator.py +1 -1
  11. zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +1 -1
  12. zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py +60 -54
  13. zenml/metadata/lazy_load.py +20 -7
  14. zenml/model/model.py +1 -2
  15. zenml/models/__init__.py +0 -12
  16. zenml/models/v2/core/artifact_version.py +19 -7
  17. zenml/models/v2/core/model_version.py +3 -5
  18. zenml/models/v2/core/pipeline_run.py +3 -5
  19. zenml/models/v2/core/run_metadata.py +2 -217
  20. zenml/models/v2/core/step_run.py +40 -24
  21. zenml/orchestrators/input_utils.py +44 -19
  22. zenml/orchestrators/step_launcher.py +2 -2
  23. zenml/orchestrators/step_run_utils.py +19 -15
  24. zenml/orchestrators/step_runner.py +8 -3
  25. zenml/steps/base_step.py +1 -1
  26. zenml/steps/entrypoint_function_utils.py +3 -5
  27. zenml/steps/step_context.py +3 -2
  28. zenml/steps/utils.py +8 -2
  29. zenml/zen_server/rbac/utils.py +0 -2
  30. zenml/zen_server/routers/workspaces_endpoints.py +3 -4
  31. zenml/zen_server/zen_server_api.py +0 -2
  32. zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py +99 -0
  33. zenml/zen_stores/migrations/versions/b557b2871693_update_step_run_input_types.py +33 -0
  34. zenml/zen_stores/rest_zen_store.py +3 -54
  35. zenml/zen_stores/schemas/artifact_schemas.py +8 -1
  36. zenml/zen_stores/schemas/model_schemas.py +2 -2
  37. zenml/zen_stores/schemas/pipeline_run_schemas.py +1 -1
  38. zenml/zen_stores/schemas/run_metadata_schemas.py +1 -48
  39. zenml/zen_stores/schemas/step_run_schemas.py +18 -10
  40. zenml/zen_stores/sql_zen_store.py +52 -98
  41. zenml/zen_stores/zen_store_interface.py +2 -42
  42. {zenml_nightly-0.68.1.dev20241107.dist-info → zenml_nightly-0.68.1.dev20241108.dist-info}/METADATA +1 -1
  43. {zenml_nightly-0.68.1.dev20241107.dist-info → zenml_nightly-0.68.1.dev20241108.dist-info}/RECORD +46 -45
  44. zenml/zen_server/routers/run_metadata_endpoints.py +0 -96
  45. {zenml_nightly-0.68.1.dev20241107.dist-info → zenml_nightly-0.68.1.dev20241108.dist-info}/LICENSE +0 -0
  46. {zenml_nightly-0.68.1.dev20241107.dist-info → zenml_nightly-0.68.1.dev20241108.dist-info}/WHEEL +0 -0
  47. {zenml_nightly-0.68.1.dev20241107.dist-info → zenml_nightly-0.68.1.dev20241108.dist-info}/entry_points.txt +0 -0
@@ -33,13 +33,13 @@ from zenml.environment import get_run_environment_dict
33
33
  from zenml.logger import get_logger
34
34
  from zenml.logging import step_logging
35
35
  from zenml.models import (
36
- ArtifactVersionResponse,
37
36
  LogsRequest,
38
37
  PipelineDeploymentResponse,
39
38
  PipelineRunRequest,
40
39
  PipelineRunResponse,
41
40
  StepRunResponse,
42
41
  )
42
+ from zenml.models.v2.core.step_run import StepRunInputResponse
43
43
  from zenml.orchestrators import output_utils, publish_utils, step_run_utils
44
44
  from zenml.orchestrators import utils as orchestrator_utils
45
45
  from zenml.orchestrators.step_runner import StepRunner
@@ -442,7 +442,7 @@ class StepLauncher:
442
442
  pipeline_run: PipelineRunResponse,
443
443
  step_run: StepRunResponse,
444
444
  step_run_info: StepRunInfo,
445
- input_artifacts: Dict[str, ArtifactVersionResponse],
445
+ input_artifacts: Dict[str, StepRunInputResponse],
446
446
  output_artifact_uris: Dict[str, str],
447
447
  last_retry: bool,
448
448
  ) -> None:
@@ -14,12 +14,12 @@
14
14
  """Utilities for creating step runs."""
15
15
 
16
16
  from datetime import datetime
17
- from typing import TYPE_CHECKING, Dict, Mapping, Optional, Set, Tuple
17
+ from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Set, Tuple
18
18
 
19
19
  from zenml.client import Client
20
20
  from zenml.config.step_configurations import ArtifactConfiguration, Step
21
21
  from zenml.constants import CODE_HASH_PARAMETER_NAME, TEXT_FIELD_MAX_LENGTH
22
- from zenml.enums import ExecutionStatus
22
+ from zenml.enums import ArtifactSaveType, ExecutionStatus
23
23
  from zenml.logger import get_logger
24
24
  from zenml.model.utils import link_artifact_version_to_model_version
25
25
  from zenml.models import (
@@ -104,6 +104,7 @@ class StepRunRequestFactory:
104
104
  input_name: artifact.id
105
105
  for input_name, artifact in input_artifacts.items()
106
106
  }
107
+
107
108
  request.inputs = input_artifact_ids
108
109
  request.parent_step_ids = parent_step_ids
109
110
 
@@ -142,8 +143,8 @@ class StepRunRequestFactory:
142
143
 
143
144
  request.original_step_run_id = cached_step_run.id
144
145
  request.outputs = {
145
- output_name: artifact.id
146
- for output_name, artifact in cached_step_run.outputs.items()
146
+ output_name: [artifact.id for artifact in artifacts]
147
+ for output_name, artifacts in cached_step_run.outputs.items()
147
148
  }
148
149
 
149
150
  request.status = ExecutionStatus.CACHED
@@ -551,7 +552,7 @@ def link_pipeline_run_to_model_version(
551
552
 
552
553
 
553
554
  def link_output_artifacts_to_model_version(
554
- artifacts: Dict[str, ArtifactVersionResponse],
555
+ artifacts: Dict[str, List[ArtifactVersionResponse]],
555
556
  output_configurations: Mapping[str, ArtifactConfiguration],
556
557
  model_version: ModelVersionResponse,
557
558
  ) -> None:
@@ -562,13 +563,16 @@ def link_output_artifacts_to_model_version(
562
563
  output_configurations: The output configurations for the step.
563
564
  model_version: The model version to link.
564
565
  """
565
- for output_name, output_artifact in artifacts.items():
566
- artifact_config = None
567
- if output_config := output_configurations.get(output_name, None):
568
- artifact_config = output_config.artifact_config
569
-
570
- link_artifact_version_to_model_version(
571
- artifact_version=output_artifact,
572
- model_version=model_version,
573
- artifact_config=artifact_config,
574
- )
566
+ for output_name, output_artifacts in artifacts.items():
567
+ for output_artifact in output_artifacts:
568
+ artifact_config = None
569
+ if output_artifact.save_type == ArtifactSaveType.STEP_OUTPUT and (
570
+ output_config := output_configurations.get(output_name, None)
571
+ ):
572
+ artifact_config = output_config.artifact_config
573
+
574
+ link_artifact_version_to_model_version(
575
+ artifact_version=output_artifact,
576
+ model_version=model_version,
577
+ artifact_config=artifact_config,
578
+ )
@@ -37,10 +37,12 @@ from zenml.constants import (
37
37
  ENV_ZENML_IGNORE_FAILURE_HOOK,
38
38
  handle_bool_env_var,
39
39
  )
40
+ from zenml.enums import ArtifactSaveType
40
41
  from zenml.exceptions import StepInterfaceError
41
42
  from zenml.logger import get_logger
42
43
  from zenml.logging.step_logging import StepLogsStorageContext, redirected
43
44
  from zenml.materializers.base_materializer import BaseMaterializer
45
+ from zenml.models.v2.core.step_run import StepRunInputResponse
44
46
  from zenml.orchestrators.publish_utils import (
45
47
  publish_step_run_metadata,
46
48
  publish_successful_step_run,
@@ -99,7 +101,7 @@ class StepRunner:
99
101
  self,
100
102
  pipeline_run: "PipelineRunResponse",
101
103
  step_run: "StepRunResponse",
102
- input_artifacts: Dict[str, "ArtifactVersionResponse"],
104
+ input_artifacts: Dict[str, StepRunInputResponse],
103
105
  output_artifact_uris: Dict[str, str],
104
106
  step_run_info: StepRunInfo,
105
107
  ) -> None:
@@ -242,7 +244,9 @@ class StepRunner:
242
244
  from zenml.orchestrators import step_run_utils
243
245
 
244
246
  step_run_utils.link_output_artifacts_to_model_version(
245
- artifacts=output_artifacts,
247
+ artifacts={
248
+ k: [v] for k, v in output_artifacts.items()
249
+ },
246
250
  output_configurations=step_run.config.outputs,
247
251
  model_version=model_version,
248
252
  )
@@ -303,7 +307,7 @@ class StepRunner:
303
307
  self,
304
308
  args: List[str],
305
309
  annotations: Dict[str, Any],
306
- input_artifacts: Dict[str, "ArtifactVersionResponse"],
310
+ input_artifacts: Dict[str, StepRunInputResponse],
307
311
  ) -> Dict[str, Any]:
308
312
  """Parses the inputs for a step entrypoint function.
309
313
 
@@ -606,6 +610,7 @@ class StepRunner:
606
610
  has_custom_name=has_custom_name,
607
611
  version=version,
608
612
  tags=tags,
613
+ save_type=ArtifactSaveType.STEP_OUTPUT,
609
614
  metadata=user_metadata,
610
615
  )
611
616
  artifact_requests.append(artifact_request)
zenml/steps/base_step.py CHANGED
@@ -327,12 +327,12 @@ class BaseStep:
327
327
  The artifacts, external artifacts, model version artifacts/metadata and parameters for the step.
328
328
  """
329
329
  from zenml.artifacts.external_artifact import ExternalArtifact
330
+ from zenml.metadata.lazy_load import LazyRunMetadataResponse
330
331
  from zenml.model.lazy_load import ModelVersionDataLazyLoader
331
332
  from zenml.models.v2.core.artifact_version import (
332
333
  ArtifactVersionResponse,
333
334
  LazyArtifactVersionResponse,
334
335
  )
335
- from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
336
336
 
337
337
  signature = inspect.signature(self.entrypoint, follow_wrapped=True)
338
338
 
@@ -32,6 +32,7 @@ from zenml.constants import ENFORCE_TYPE_ANNOTATIONS
32
32
  from zenml.exceptions import StepInterfaceError
33
33
  from zenml.logger import get_logger
34
34
  from zenml.materializers.base_materializer import BaseMaterializer
35
+ from zenml.metadata.lazy_load import LazyRunMetadataResponse
35
36
  from zenml.steps.utils import (
36
37
  OutputSignature,
37
38
  parse_return_type_annotations,
@@ -136,10 +137,7 @@ class EntrypointFunctionDefinition(NamedTuple):
136
137
  UnmaterializedArtifact,
137
138
  )
138
139
  from zenml.client_lazy_loader import ClientLazyLoader
139
- from zenml.models import (
140
- ArtifactVersionResponse,
141
- RunMetadataResponse,
142
- )
140
+ from zenml.models import ArtifactVersionResponse
143
141
 
144
142
  if key not in self.inputs:
145
143
  raise KeyError(
@@ -154,8 +152,8 @@ class EntrypointFunctionDefinition(NamedTuple):
154
152
  StepArtifact,
155
153
  ExternalArtifact,
156
154
  ArtifactVersionResponse,
157
- RunMetadataResponse,
158
155
  ClientLazyLoader,
156
+ LazyRunMetadataResponse,
159
157
  ),
160
158
  ):
161
159
  # If we were to do any type validation for artifacts here, we
@@ -35,11 +35,12 @@ if TYPE_CHECKING:
35
35
  from zenml.metadata.metadata_types import MetadataType
36
36
  from zenml.model.model import Model
37
37
  from zenml.models import (
38
- ArtifactVersionResponse,
39
38
  PipelineResponse,
40
39
  PipelineRunResponse,
41
40
  StepRunResponse,
42
41
  )
42
+ from zenml.models.v2.core.step_run import StepRunInputResponse
43
+
43
44
 
44
45
  logger = get_logger(__name__)
45
46
 
@@ -191,7 +192,7 @@ class StepContext(metaclass=SingletonMetaClass):
191
192
  return self.model_version.to_model_class()
192
193
 
193
194
  @property
194
- def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
195
+ def inputs(self) -> Dict[str, "StepRunInputResponse"]:
195
196
  """Returns the input artifacts of the current step.
196
197
 
197
198
  Returns:
zenml/steps/utils.py CHANGED
@@ -26,7 +26,11 @@ from typing_extensions import Annotated
26
26
 
27
27
  from zenml.artifacts.artifact_config import ArtifactConfig
28
28
  from zenml.client import Client
29
- from zenml.enums import ExecutionStatus, MetadataResourceTypes
29
+ from zenml.enums import (
30
+ ArtifactSaveType,
31
+ ExecutionStatus,
32
+ MetadataResourceTypes,
33
+ )
30
34
  from zenml.exceptions import StepInterfaceError
31
35
  from zenml.logger import get_logger
32
36
  from zenml.metadata.metadata_types import MetadataType
@@ -547,8 +551,10 @@ def run_as_single_step_pipeline(
547
551
  # 4. Load output artifacts
548
552
  step_run = next(iter(run.steps.values()))
549
553
  outputs = [
550
- step_run.outputs[output_name].load()
554
+ artifact_version.load()
551
555
  for output_name in step_run.config.outputs.keys()
556
+ for artifact_version in step_run.outputs[output_name]
557
+ if artifact_version.save_type == ArtifactSaveType.STEP_OUTPUT
552
558
  ]
553
559
 
554
560
  if len(outputs) == 0:
@@ -404,7 +404,6 @@ def get_resource_type_for_model(
404
404
  PipelineDeploymentResponse,
405
405
  PipelineResponse,
406
406
  PipelineRunResponse,
407
- RunMetadataResponse,
408
407
  RunTemplateResponse,
409
408
  SecretResponse,
410
409
  ServiceAccountResponse,
@@ -437,7 +436,6 @@ def get_resource_type_for_model(
437
436
  ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
438
437
  WorkspaceResponse: ResourceType.WORKSPACE,
439
438
  UserResponse: ResourceType.USER,
440
- RunMetadataResponse: ResourceType.RUN_METADATA,
441
439
  PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
442
440
  PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
443
441
  PipelineRunResponse: ResourceType.PIPELINE_RUN,
@@ -74,7 +74,6 @@ from zenml.models import (
74
74
  PipelineRunRequest,
75
75
  PipelineRunResponse,
76
76
  RunMetadataRequest,
77
- RunMetadataResponse,
78
77
  RunTemplateFilter,
79
78
  RunTemplateRequest,
80
79
  RunTemplateResponse,
@@ -977,7 +976,6 @@ def get_or_create_pipeline_run(
977
976
 
978
977
  @router.post(
979
978
  WORKSPACES + "/{workspace_name_or_id}" + RUN_METADATA,
980
- response_model=List[RunMetadataResponse],
981
979
  responses={401: error_response, 409: error_response, 422: error_response},
982
980
  )
983
981
  @handle_exceptions
@@ -985,7 +983,7 @@ def create_run_metadata(
985
983
  workspace_name_or_id: Union[str, UUID],
986
984
  run_metadata: RunMetadataRequest,
987
985
  auth_context: AuthContext = Security(authorize),
988
- ) -> List[RunMetadataResponse]:
986
+ ) -> None:
989
987
  """Creates run metadata.
990
988
 
991
989
  Args:
@@ -1039,7 +1037,8 @@ def create_run_metadata(
1039
1037
  resource_type=ResourceType.RUN_METADATA, action=Action.CREATE
1040
1038
  )
1041
1039
 
1042
- return zen_store().create_run_metadata(run_metadata)
1040
+ zen_store().create_run_metadata(run_metadata)
1041
+ return None
1043
1042
 
1044
1043
 
1045
1044
  @router.post(
@@ -70,7 +70,6 @@ from zenml.zen_server.routers import (
70
70
  pipeline_deployments_endpoints,
71
71
  pipelines_endpoints,
72
72
  plugin_endpoints,
73
- run_metadata_endpoints,
74
73
  run_templates_endpoints,
75
74
  runs_endpoints,
76
75
  schedule_endpoints,
@@ -423,7 +422,6 @@ app.include_router(pipelines_endpoints.router)
423
422
  app.include_router(pipeline_builds_endpoints.router)
424
423
  app.include_router(pipeline_deployments_endpoints.router)
425
424
  app.include_router(runs_endpoints.router)
426
- app.include_router(run_metadata_endpoints.router)
427
425
  app.include_router(run_templates_endpoints.router)
428
426
  app.include_router(schedule_endpoints.router)
429
427
  app.include_router(secrets_endpoints.router)
@@ -0,0 +1,99 @@
1
+ """move artifact save type [1cb6477f72d6].
2
+
3
+ Revision ID: 1cb6477f72d6
4
+ Revises: c22561cbb3a9
5
+ Create Date: 2024-10-10 15:44:09.465210
6
+
7
+ """
8
+
9
+ import sqlalchemy as sa
10
+ from alembic import op
11
+
12
+ # revision identifiers, used by Alembic.
13
+ revision = "1cb6477f72d6"
14
+ down_revision = "c22561cbb3a9"
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
+ # Step 1: Add nullable save_type column to artifact_version
22
+ with op.batch_alter_table("artifact_version", schema=None) as batch_op:
23
+ batch_op.add_column(sa.Column("save_type", sa.TEXT(), nullable=True))
24
+
25
+ # Step 2: Move data from step_run_output_artifact.type to artifact_version.save_type
26
+ op.execute("""
27
+ UPDATE artifact_version
28
+ SET save_type = (
29
+ SELECT max(step_run_output_artifact.type)
30
+ FROM step_run_output_artifact
31
+ WHERE step_run_output_artifact.artifact_id = artifact_version.id
32
+ GROUP BY artifact_id
33
+ )
34
+ """)
35
+ op.execute("""
36
+ UPDATE artifact_version
37
+ SET save_type = 'step_output'
38
+ WHERE artifact_version.save_type = 'default'
39
+ """)
40
+ op.execute("""
41
+ UPDATE artifact_version
42
+ SET save_type = 'external'
43
+ WHERE save_type is NULL
44
+ """)
45
+
46
+ # # Step 3: Set save_type to non-nullable
47
+ with op.batch_alter_table("artifact_version", schema=None) as batch_op:
48
+ batch_op.alter_column(
49
+ "save_type",
50
+ existing_type=sa.TEXT(),
51
+ nullable=False,
52
+ )
53
+
54
+ # Step 4: Remove type column from step_run_output_artifact
55
+ with op.batch_alter_table(
56
+ "step_run_output_artifact", schema=None
57
+ ) as batch_op:
58
+ batch_op.drop_column("type")
59
+
60
+
61
+ def downgrade() -> None:
62
+ """Downgrade database schema and/or data back to the previous revision."""
63
+ # Add type column back to step_run_output_artifact
64
+ with op.batch_alter_table(
65
+ "step_run_output_artifact", schema=None
66
+ ) as batch_op:
67
+ batch_op.add_column(
68
+ sa.Column("type", sa.TEXT(), nullable=True),
69
+ )
70
+
71
+ # Move data back from artifact_version.save_type to step_run_output_artifact.type
72
+ op.execute("""
73
+ UPDATE step_run_output_artifact
74
+ SET type = (
75
+ SELECT max(artifact_version.save_type)
76
+ FROM artifact_version
77
+ WHERE step_run_output_artifact.artifact_id = artifact_version.id
78
+ GROUP BY artifact_id
79
+ )
80
+ """)
81
+ op.execute("""
82
+ UPDATE step_run_output_artifact
83
+ SET type = 'default'
84
+ WHERE step_run_output_artifact.type = 'step_output'
85
+ """)
86
+
87
+ # Set type to non-nullable
88
+ with op.batch_alter_table(
89
+ "step_run_output_artifact", schema=None
90
+ ) as batch_op:
91
+ batch_op.alter_column(
92
+ "type",
93
+ existing_type=sa.TEXT(),
94
+ nullable=False,
95
+ )
96
+
97
+ # Remove save_type column from artifact_version
98
+ with op.batch_alter_table("artifact_version", schema=None) as batch_op:
99
+ batch_op.drop_column("save_type")
@@ -0,0 +1,33 @@
1
+ """Update step run input types [b557b2871693].
2
+
3
+ Revision ID: b557b2871693
4
+ Revises: 1cb6477f72d6
5
+ Create Date: 2024-10-30 13:06:55.147202
6
+
7
+ """
8
+
9
+ from alembic import op
10
+
11
+ # revision identifiers, used by Alembic.
12
+ revision = "b557b2871693"
13
+ down_revision = "1cb6477f72d6"
14
+ branch_labels = None
15
+ depends_on = None
16
+
17
+
18
+ def upgrade() -> None:
19
+ """Upgrade database schema and/or data, creating a new revision."""
20
+ op.execute("""
21
+ UPDATE step_run_input_artifact
22
+ SET type = 'step_output'
23
+ WHERE type = 'default'
24
+ """)
25
+
26
+
27
+ def downgrade() -> None:
28
+ """Downgrade database schema and/or data back to the previous revision."""
29
+ op.execute("""
30
+ UPDATE step_run_input_artifact
31
+ SET type = 'default'
32
+ WHERE type = 'step_output'
33
+ """)
@@ -201,9 +201,7 @@ from zenml.models import (
201
201
  PipelineRunResponse,
202
202
  PipelineRunUpdate,
203
203
  PipelineUpdate,
204
- RunMetadataFilter,
205
204
  RunMetadataRequest,
206
- RunMetadataResponse,
207
205
  RunTemplateFilter,
208
206
  RunTemplateRequest,
209
207
  RunTemplateResponse,
@@ -2014,9 +2012,7 @@ class RestZenStore(BaseZenStore):
2014
2012
 
2015
2013
  # ----------------------------- Run Metadata -----------------------------
2016
2014
 
2017
- def create_run_metadata(
2018
- self, run_metadata: RunMetadataRequest
2019
- ) -> List[RunMetadataResponse]:
2015
+ def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
2020
2016
  """Creates run metadata.
2021
2017
 
2022
2018
  Args:
@@ -2026,55 +2022,8 @@ class RestZenStore(BaseZenStore):
2026
2022
  The created run metadata.
2027
2023
  """
2028
2024
  route = f"{WORKSPACES}/{str(run_metadata.workspace)}{RUN_METADATA}"
2029
- response_body = self.post(f"{route}", body=run_metadata)
2030
- result: List[RunMetadataResponse] = []
2031
- if isinstance(response_body, list):
2032
- for metadata in response_body or []:
2033
- result.append(RunMetadataResponse.model_validate(metadata))
2034
- return result
2035
-
2036
- def get_run_metadata(
2037
- self, run_metadata_id: UUID, hydrate: bool = True
2038
- ) -> RunMetadataResponse:
2039
- """Gets run metadata with the given ID.
2040
-
2041
- Args:
2042
- run_metadata_id: The ID of the run metadata to get.
2043
- hydrate: Flag deciding whether to hydrate the output model(s)
2044
- by including metadata fields in the response.
2045
-
2046
- Returns:
2047
- The run metadata.
2048
- """
2049
- return self._get_resource(
2050
- resource_id=run_metadata_id,
2051
- route=RUN_METADATA,
2052
- response_model=RunMetadataResponse,
2053
- params={"hydrate": hydrate},
2054
- )
2055
-
2056
- def list_run_metadata(
2057
- self,
2058
- run_metadata_filter_model: RunMetadataFilter,
2059
- hydrate: bool = False,
2060
- ) -> Page[RunMetadataResponse]:
2061
- """List run metadata.
2062
-
2063
- Args:
2064
- run_metadata_filter_model: All filter parameters including
2065
- pagination params.
2066
- hydrate: Flag deciding whether to hydrate the output model(s)
2067
- by including metadata fields in the response.
2068
-
2069
- Returns:
2070
- The run metadata.
2071
- """
2072
- return self._list_paginated_resources(
2073
- route=RUN_METADATA,
2074
- response_model=RunMetadataResponse,
2075
- filter_model=run_metadata_filter_model,
2076
- params={"hydrate": hydrate},
2077
- )
2025
+ self.post(f"{route}", body=run_metadata)
2026
+ return None
2078
2027
 
2079
2028
  # ----------------------------- Schedules -----------------------------
2080
2029
 
@@ -13,6 +13,7 @@
13
13
  # permissions and limitations under the License.
14
14
  """SQLModel implementation of artifact table."""
15
15
 
16
+ import json
16
17
  from datetime import datetime
17
18
  from typing import TYPE_CHECKING, Any, List, Optional
18
19
  from uuid import UUID
@@ -23,6 +24,7 @@ from sqlmodel import Field, Relationship
23
24
 
24
25
  from zenml.config.source import Source
25
26
  from zenml.enums import (
27
+ ArtifactSaveType,
26
28
  ArtifactType,
27
29
  ExecutionStatus,
28
30
  MetadataResourceTypes,
@@ -196,6 +198,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
196
198
  overlaps="tags",
197
199
  ),
198
200
  )
201
+ save_type: str = Field(sa_column=Column(TEXT, nullable=False))
199
202
 
200
203
  # Foreign keys
201
204
  artifact_id: UUID = build_foreign_key_field(
@@ -300,6 +303,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
300
303
  uri=artifact_version_request.uri,
301
304
  materializer=artifact_version_request.materializer.model_dump_json(),
302
305
  data_type=artifact_version_request.data_type.model_dump_json(),
306
+ save_type=artifact_version_request.save_type.value,
303
307
  )
304
308
 
305
309
  def to_model(
@@ -360,6 +364,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
360
364
  updated=self.updated,
361
365
  tags=[t.tag.to_model() for t in self.tags],
362
366
  producer_pipeline_run_id=producer_pipeline_run_id,
367
+ save_type=ArtifactSaveType(self.save_type),
363
368
  artifact_store_id=self.artifact_store_id,
364
369
  )
365
370
 
@@ -370,7 +375,9 @@ class ArtifactVersionSchema(BaseSchema, table=True):
370
375
  workspace=self.workspace.to_model(),
371
376
  producer_step_run_id=producer_step_run_id,
372
377
  visualizations=[v.to_model() for v in self.visualizations],
373
- run_metadata={m.key: m.to_model() for m in self.run_metadata},
378
+ run_metadata={
379
+ m.key: json.loads(m.value) for m in self.run_metadata
380
+ },
374
381
  )
375
382
 
376
383
  resources = None
@@ -13,6 +13,7 @@
13
13
  # permissions and limitations under the License.
14
14
  """SQLModel implementation of model tables."""
15
15
 
16
+ import json
16
17
  from datetime import datetime
17
18
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
18
19
  from uuid import UUID
@@ -379,8 +380,7 @@ class ModelVersionSchema(NamedSchema, table=True):
379
380
  workspace=self.workspace.to_model(),
380
381
  description=self.description,
381
382
  run_metadata={
382
- rm.key: rm.to_model(include_metadata=True)
383
- for rm in self.run_metadata
383
+ rm.key: json.loads(rm.value) for rm in self.run_metadata
384
384
  },
385
385
  )
386
386
 
@@ -271,7 +271,7 @@ class PipelineRunSchema(NamedSchema, table=True):
271
271
  )
272
272
 
273
273
  run_metadata = {
274
- metadata_schema.key: metadata_schema.to_model()
274
+ metadata_schema.key: json.loads(metadata_schema.value)
275
275
  for metadata_schema in self.run_metadata
276
276
  }
277
277
 
@@ -13,20 +13,13 @@
13
13
  # permissions and limitations under the License.
14
14
  """SQLModel implementation of pipeline run metadata tables."""
15
15
 
16
- import json
17
- from typing import TYPE_CHECKING, Any, List, Optional
16
+ from typing import TYPE_CHECKING, List, Optional
18
17
  from uuid import UUID
19
18
 
20
19
  from sqlalchemy import TEXT, VARCHAR, Column
21
20
  from sqlmodel import Field, Relationship
22
21
 
23
22
  from zenml.enums import MetadataResourceTypes
24
- from zenml.metadata.metadata_types import MetadataTypeEnum
25
- from zenml.models import (
26
- RunMetadataResponse,
27
- RunMetadataResponseBody,
28
- RunMetadataResponseMetadata,
29
- )
30
23
  from zenml.zen_stores.schemas.base_schemas import BaseSchema
31
24
  from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
32
25
  from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
@@ -110,43 +103,3 @@ class RunMetadataSchema(BaseSchema, table=True):
110
103
  key: str
111
104
  value: str = Field(sa_column=Column(TEXT, nullable=False))
112
105
  type: str
113
-
114
- def to_model(
115
- self,
116
- include_metadata: bool = False,
117
- include_resources: bool = False,
118
- **kwargs: Any,
119
- ) -> "RunMetadataResponse":
120
- """Convert a `RunMetadataSchema` to a `RunMetadataResponse`.
121
-
122
- Args:
123
- include_metadata: Whether the metadata will be filled.
124
- include_resources: Whether the resources will be filled.
125
- **kwargs: Keyword arguments to allow schema specific logic
126
-
127
-
128
- Returns:
129
- The created `RunMetadataResponse`.
130
- """
131
- body = RunMetadataResponseBody(
132
- user=self.user.to_model() if self.user else None,
133
- key=self.key,
134
- created=self.created,
135
- updated=self.updated,
136
- value=json.loads(self.value),
137
- type=MetadataTypeEnum(self.type),
138
- )
139
- metadata = None
140
- if include_metadata:
141
- metadata = RunMetadataResponseMetadata(
142
- workspace=self.workspace.to_model(),
143
- resource_id=self.resource_id,
144
- resource_type=MetadataResourceTypes(self.resource_type),
145
- stack_component_id=self.stack_component_id,
146
- )
147
-
148
- return RunMetadataResponse(
149
- id=self.id,
150
- body=body,
151
- metadata=metadata,
152
- )