From 222dd6d13813ebc322a993283239d411895106f7 Mon Sep 17 00:00:00 2001
From: Gyubong Lee <jopemachine@naver.com>
Date: Tue, 7 May 2024 04:33:11 +0000
Subject: [PATCH] fix: Change `get_first_timestamp_for_status` return type to
 datetime

---
 src/ai/backend/client/cli/session/lifecycle.py |  5 ++---
 src/ai/backend/common/utils.py                 |  7 ++++---
 .../backend/manager/models/resource_usage.py   |  6 ++++--
 src/ai/backend/manager/models/session.py       | 18 +++++-------------
 4 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/src/ai/backend/client/cli/session/lifecycle.py b/src/ai/backend/client/cli/session/lifecycle.py
index 307d25a6ae6..4123c65d889 100644
--- a/src/ai/backend/client/cli/session/lifecycle.py
+++ b/src/ai/backend/client/cli/session/lifecycle.py
@@ -16,7 +16,6 @@
 import inquirer
 import treelib
 from async_timeout import timeout
-from dateutil.parser import isoparse
 from dateutil.tz import tzutc
 from faker import Faker
 from humanize import naturalsize
@@ -802,9 +801,9 @@ def status_history(ctx: CLIContext, session_id):
             elif (
                 terminated := get_first_timestamp_for_status(status_history, "TERMINATED")
             ) is None:
-                elapsed = datetime.now(tzutc()) - isoparse(preparing)
+                elapsed = datetime.now(tzutc()) - preparing
             else:
-                elapsed = isoparse(terminated) - isoparse(preparing)
+                elapsed = terminated - preparing
 
             print_done(f"Actual Resource Allocation Time: {elapsed.total_seconds()}")
         except Exception as e:
diff --git a/src/ai/backend/common/utils.py b/src/ai/backend/common/utils.py
index 84dd95fe69b..acadeb48d79 100644
--- a/src/ai/backend/common/utils.py
+++ b/src/ai/backend/common/utils.py
@@ -8,7 +8,7 @@
 import sys
 import uuid
 from collections import OrderedDict
-from datetime import timedelta
+from datetime import datetime, timedelta
 from itertools import chain
 from pathlib import Path
 from typing import (
@@ -25,6 +25,7 @@
 
 import aiofiles
 from async_timeout import timeout
+from dateutil.parser import parse as dtparse
 
 if TYPE_CHECKING:
     from decimal import Decimal
@@ -408,12 +409,12 @@ async def umount(
 
 def get_first_timestamp_for_status(
     status_history_records: list[dict[str, str]], status: str
-) -> str | None:
+) -> datetime | None:
     """
     Get the first occurrence time of the given status from the status history records.
     """
 
     for status_history in status_history_records:
         if status_history["status"] == status:
-            return status_history["timestamp"]
+            return dtparse(status_history["timestamp"])
     return None
diff --git a/src/ai/backend/manager/models/resource_usage.py b/src/ai/backend/manager/models/resource_usage.py
index d8997867d7a..76b33016fe9 100644
--- a/src/ai/backend/manager/models/resource_usage.py
+++ b/src/ai/backend/manager/models/resource_usage.py
@@ -15,7 +15,7 @@
 
 from ai.backend.common import redis_helper
 from ai.backend.common.types import RedisConnectionInfo
-from ai.backend.common.utils import nmget
+from ai.backend.common.utils import get_first_timestamp_for_status, nmget
 
 from .group import GroupRow
 from .kernel import LIVE_STATUS, RESOURCE_USAGE_KERNEL_STATUSES, KernelRow, KernelStatus
@@ -516,7 +516,9 @@ async def _pipe_builder(r: Redis) -> RedisPipeline:
             session_row=kern.session,
             created_at=kern.created_at,
             terminated_at=kern.terminated_at,
-            scheduled_at=kern.status_history.get(KernelStatus.SCHEDULED.name),
+            scheduled_at=str(
+                get_first_timestamp_for_status(kern.status_history, KernelStatus.SCHEDULED.name)
+            ),
             used_time=kern.used_time,
             used_days=kern.get_used_days(local_tz),
             last_stat=stat_map[kern.id],
diff --git a/src/ai/backend/manager/models/session.py b/src/ai/backend/manager/models/session.py
index 5f882980d13..3653b221a29 100644
--- a/src/ai/backend/manager/models/session.py
+++ b/src/ai/backend/manager/models/session.py
@@ -73,7 +73,7 @@
 )
 from .group import GroupRow
 from .kernel import ComputeContainer, KernelRow, KernelStatus
-from .minilang import ArrayFieldItem, JSONFieldItem
+from .minilang import ArrayFieldItem
 from .minilang.ordering import ColumnMapType, QueryOrderParser
 from .minilang.queryfilter import FieldSpecType, QueryFilterParser, enum_field_getter
 from .user import UserRow
@@ -705,10 +705,8 @@ def main_kernel(self) -> KernelRow:
         return kerns[0]
 
     @property
-    def status_changed(self) -> Optional[datetime]:
-        if scheduled_at := get_first_timestamp_for_status(self.status_history, self.status.name):
-            return datetime.fromisoformat(scheduled_at)
-        return None
+    def status_changed(self) -> datetime | None:
+        return get_first_timestamp_for_status(self.status_history, self.status.name)
 
     @property
     def resource_opts(self) -> dict[str, Any]:
@@ -1395,10 +1393,7 @@ async def resolve_status_history(self, _info: graphene.ResolveInfo) -> Mapping[s
         "created_at": ("sessions_created_at", dtparse),
         "terminated_at": ("sessions_terminated_at", dtparse),
         "starts_at": ("sessions_starts_at", dtparse),
-        "scheduled_at": (
-            JSONFieldItem("sessions_status_history", SessionStatus.SCHEDULED.name),
-            dtparse,
-        ),
+        "scheduled_at": ("scheduled_at", None),
         "startup_command": ("sessions_startup_command", None),
     }
 
@@ -1426,10 +1421,7 @@ async def resolve_status_history(self, _info: graphene.ResolveInfo) -> Mapping[s
         "created_at": ("sessions_created_at", None),
         "terminated_at": ("sessions_terminated_at", None),
         "starts_at": ("sessions_starts_at", None),
-        "scheduled_at": (
-            JSONFieldItem("sessions_status_history", SessionStatus.SCHEDULED.name),
-            None,
-        ),
+        "scheduled_at": ("scheduled_at", None),
     }
 
     @classmethod