Move step durations to body

zenml-io · Sep 26, 2024 · abc77b5 · abc77b5
1 parent a43a43b
commit abc77b5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/src/zenml/models/v2/core/step_run.py b/src/zenml/models/v2/core/step_run.py
@@ -156,6 +156,14 @@ class StepRunResponseBody(WorkspaceScopedResponseBody):
     """Response body for step runs."""
 
     status: ExecutionStatus = Field(title="The status of the step.")
+    start_time: Optional[datetime] = Field(
+        title="The start time of the step run.",
+        default=None,
+    )
+    end_time: Optional[datetime] = Field(
+        title="The end time of the step run.",
+        default=None,
+    )
     inputs: Dict[str, "ArtifactVersionResponse"] = Field(
         title="The input artifact versions of the step run.",
         default={},
@@ -201,16 +209,6 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
 
-    # Timestamps
-    start_time: Optional[datetime] = Field(
-        title="The start time of the step run.",
-        default=None,
-    )
-    end_time: Optional[datetime] = Field(
-        title="The end time of the step run.",
-        default=None,
-    )
-
     # References
     logs: Optional["LogsResponse"] = Field(
         title="Logs associated with this step run.",
@@ -409,7 +407,7 @@ def start_time(self) -> Optional[datetime]:
         Returns:
             the value of the property.
         """
-        return self.get_metadata().start_time
+        return self.get_body().start_time
 
     @property
     def end_time(self) -> Optional[datetime]:
@@ -418,7 +416,7 @@ def end_time(self) -> Optional[datetime]:
         Returns:
             the value of the property.
         """
-        return self.get_metadata().end_time
+        return self.get_body().end_time
 
     @property
     def logs(self) -> Optional["LogsResponse"]:

diff --git a/src/zenml/zen_stores/schemas/step_run_schemas.py b/src/zenml/zen_stores/schemas/step_run_schemas.py
@@ -262,6 +262,8 @@ def to_model(
         body = StepRunResponseBody(
             user=self.user.to_model() if self.user else None,
             status=ExecutionStatus(self.status),
+            start_time=self.start_time,
+            end_time=self.end_time,
             inputs=input_artifacts,
             outputs=output_artifacts,
             created=self.created,
@@ -278,8 +280,6 @@ def to_model(
                 code_hash=self.code_hash,
                 docstring=self.docstring,
                 source_code=self.source_code,
-                start_time=self.start_time,
-                end_time=self.end_time,
                 logs=self.logs.to_model() if self.logs else None,
                 deployment_id=self.deployment_id,
                 pipeline_run_id=self.pipeline_run_id,