refactor: Replace sessions, kernels's status_history's type map with list

changes/1662.fix.md

@@ -0,0 +1 @@
+Change the type of `status_history` from a mapping of status and timestamps to a list of log entries containing status and timestamps, to preserve timestamps when revisiting session/kernel statuses (e.g., after session restarts).

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
@@ -25,8 +24,11 @@
 from ai.backend.cli.main import main
 from ai.backend.cli.params import CommaSeparatedListType, OptionalType
 from ai.backend.cli.types import ExitCode, Undefined, undefined
+from ai.backend.client.cli.extensions import pass_ctx_obj
+from ai.backend.client.cli.types import CLIContext
 from ai.backend.common.arch import DEFAULT_IMAGE_ARCH
 from ai.backend.common.types import ClusterMode
+from ai.backend.common.utils import get_first_timestamp_for_status
 
 from ...compat import asyncio_run
 from ...exceptions import BackendAPIError
@@ -763,8 +765,9 @@ def logs(session_id):
 
 
 @session.command("status-history")
+@pass_ctx_obj
 @click.argument("session_id", metavar="SESSID")
-def status_history(session_id):
+def status_history(ctx: CLIContext, session_id):
     """
     Shows the status transition history of the compute session.
 
@@ -776,31 +779,33 @@ def status_history(session_id):
         kernel = session.ComputeSession(session_id)
         try:
             status_history = kernel.get_status_history().get("result")
-            print_info(f"status_history: {status_history}")
-            if (preparing := status_history.get("preparing")) is None:
-                result = {
-                    "result": {
-                        "seconds": 0,
-                        "microseconds": 0,
-                    },
-                }
-            elif (terminated := status_history.get("terminated")) is None:
-                alloc_time_until_now: timedelta = datetime.now(tzutc()) - isoparse(preparing)
-                result = {
-                    "result": {
-                        "seconds": alloc_time_until_now.seconds,
-                        "microseconds": alloc_time_until_now.microseconds,
-                    },
-                }
+
+            prev_time = None
+
+            for status_record in status_history:
+                timestamp = datetime.fromisoformat(status_record["timestamp"])
+
+                if prev_time:
+                    time_diff = timestamp - prev_time
+                    status_record["time_elapsed"] = str(time_diff)
+
+                prev_time = timestamp
+
+            ctx.output.print_list(
+                status_history,
+                [FieldSpec("status"), FieldSpec("timestamp"), FieldSpec("time_elapsed")],
+            )
+
+            if (preparing := get_first_timestamp_for_status(status_history, "PREPARING")) is None:
+                elapsed = timedelta()
+            elif (
+                terminated := get_first_timestamp_for_status(status_history, "TERMINATED")
+            ) is None:
+                elapsed = datetime.now(tzutc()) - preparing
             else:
-                alloc_time: timedelta = isoparse(terminated) - isoparse(preparing)
-                result = {
-                    "result": {
-                        "seconds": alloc_time.seconds,
-                        "microseconds": alloc_time.microseconds,
-                    },
-                }
-            print_done(f"Actual Resource Allocation Time: {result}")
+                elapsed = terminated - preparing
+
+            print_done(f"Actual Resource Allocation Time: {elapsed.total_seconds()}")
         except Exception as e:
             print_error(e)
             sys.exit(ExitCode.FAILURE)

src/ai/backend/client/output/fields.py

@@ -181,6 +181,8 @@
     FieldSpec("created_user_id"),
     FieldSpec("status"),
     FieldSpec("status_info"),
+    FieldSpec("status_history"),
+    FieldSpec("status_history_log"),
     FieldSpec("status_data", formatter=nested_dict_formatter),
     FieldSpec("status_changed", "Last Updated"),
     FieldSpec("created_at"),

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
@@ -404,3 +405,16 @@ async def umount(
             fstab = Fstab(fp)
             await fstab.remove_by_mountpoint(str(mountpoint))
     return True
+
+
+def get_first_timestamp_for_status(
+    status_history_records: list[dict[str, str]], status: str
+) -> 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 dtparse(status_history["timestamp"])
+    return None

src/ai/backend/manager/api/resource.py

@@ -463,7 +463,7 @@ async def _pipe_builder(r: Redis) -> RedisPipeline:
             "status": row["status"].name,
             "status_info": row["status_info"],
             "status_changed": str(row["status_changed"]),
-            "status_history": row["status_history"] or {},
+            "status_history": row["status_history"],
             "cluster_mode": row["cluster_mode"],
         }
         if group_id not in objs_per_group:

src/ai/backend/manager/api/schema.graphql

@@ -545,7 +545,10 @@
   status_changed: DateTime
   status_info: String
   status_data: JSONString
-  status_history: JSONString
+  status_history: JSONString @deprecated(reason: "Deprecated since 24.09.0; use `status_history_log`")
+
+  """Added in 24.09.0"""
+  status_history_log: JSONString
   created_at: DateTime
   terminated_at: DateTime
   starts_at: DateTime

src/ai/backend/manager/api/session.py

@@ -2151,6 +2151,36 @@ async def get_container_logs(request: web.Request, params: Any) -> web.Response:
     return web.json_response(resp, status=200)
 
 
+@server_status_required(READ_ALLOWED)
+@auth_required
+@check_api_params(
+    t.Dict({
+        tx.AliasedKey(["session_name", "sessionName", "task_id", "taskId"]) >> "kernel_id": tx.UUID,
+        t.Key("owner_access_key", default=None): t.Null | t.String,
+    })
+)
+async def get_status_history(request: web.Request, params: Any) -> web.Response:
+    root_ctx: RootContext = request.app["_root.context"]
+    session_name: str = request.match_info["session_name"]
+    requester_access_key, owner_access_key = await get_access_key_scopes(request, params)
+    log.info(
+        "GET_STATUS_HISTORY (ak:{}/{}, s:{})", requester_access_key, owner_access_key, session_name
+    )
+    resp: dict[str, Mapping] = {"result": {}}
+
+    async with root_ctx.db.begin_readonly_session() as db_sess:
+        compute_session = await SessionRow.get_session(
+            db_sess,
+            session_name,
+            owner_access_key,
+            allow_stale=True,
+            kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
+        )
+        resp["result"] = compute_session.status_history
+
+    return web.json_response(resp, status=200)
+
+
 @server_status_required(READ_ALLOWED)
 @auth_required
 @check_api_params(
@@ -2286,6 +2316,7 @@ def create_app(
         app.router.add_route("GET", "/{session_name}/direct-access-info", get_direct_access_info)
     )
     cors.add(app.router.add_route("GET", "/{session_name}/logs", get_container_logs))
+    cors.add(app.router.add_route("GET", "/{session_name}/status-history", get_status_history))
     cors.add(app.router.add_route("POST", "/{session_name}/rename", rename_session))
     cors.add(app.router.add_route("POST", "/{session_name}/interrupt", interrupt))
     cors.add(app.router.add_route("POST", "/{session_name}/complete", complete))

src/ai/backend/manager/models/alembic/versions/8c8e90aebacd_replace_status_history_to_list.py

@@ -0,0 +1,93 @@
+"""Replace sessions, kernels's status_history's type map with list
+
+Revision ID: 8c8e90aebacd
+Revises: dddf9be580f5
+Create Date: 2024-01-26 11:19:23.075014
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "8c8e90aebacd"
+down_revision = "dddf9be580f5"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.execute(
+        """
+        WITH data AS (
+            SELECT id,
+                (jsonb_each(status_history)).key AS status,
+                (jsonb_each(status_history)).value AS timestamp
+            FROM kernels
+        )
+        UPDATE kernels
+        SET status_history = (
+            SELECT jsonb_agg(
+                jsonb_build_object('status', status, 'timestamp', timestamp)
+            )
+            FROM data
+            WHERE data.id = kernels.id
+        );
+    """
+    )
+
+    op.execute(
+        """
+        WITH data AS (
+            SELECT id,
+                (jsonb_each(status_history)).key AS status,
+                (jsonb_each(status_history)).value AS timestamp
+            FROM sessions
+        )
+        UPDATE sessions
+        SET status_history = (
+            SELECT jsonb_agg(
+                jsonb_build_object('status', status, 'timestamp', timestamp)
+            )
+            FROM data
+            WHERE data.id = sessions.id
+        );
+    """
+    )
+
+
+def downgrade():
+    op.execute(
+        """
+        WITH data AS (
+            SELECT id,
+                jsonb_object_agg(
+                    elem->>'status', elem->>'timestamp'
+                ) AS new_status_history
+            FROM kernels,
+            jsonb_array_elements(status_history) AS elem
+            GROUP BY id
+        )
+        UPDATE kernels
+        SET status_history = data.new_status_history
+        FROM data
+        WHERE data.id = kernels.id;
+    """
+    )
+
+    op.execute(
+        """
+        WITH data AS (
+            SELECT id,
+                jsonb_object_agg(
+                    elem->>'status', elem->>'timestamp'
+                ) AS new_status_history
+            FROM sessions,
+            jsonb_array_elements(status_history) AS elem
+            GROUP BY id
+        )
+        UPDATE sessions
+        SET status_history = data.new_status_history
+        FROM data
+        WHERE data.id = sessions.id;
+    """
+    )