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

changes/1662.fix.md

@@ -0,0 +1 @@
+Replace `status_history`'s type `map` with `list`

src/ai/backend/client/cli/session/lifecycle.py

@@ -25,8 +25,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_occurrence_time
 
 from ...compat import asyncio_run
 from ...exceptions import BackendAPIError
@@ -763,8 +766,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 +780,31 @@ 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_occurrence_time(status_history, "PREPARING")) is None:
+                elapsed = timedelta()
+            elif (terminated := get_first_occurrence_time(status_history, "TERMINATED")) is None:
+                elapsed = datetime.now(tzutc()) - isoparse(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 = isoparse(terminated) - isoparse(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

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

@@ -546,6 +546,9 @@ type ComputeSession implements Item {
   status_info: String
   status_data: JSONString
   status_history: JSONString
+
+  """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: 8b2ec7e3d22a
+Create Date: 2024-01-26 11:19:23.075014
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "8c8e90aebacd"
+down_revision = "8b2ec7e3d22a"
+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;
+    """
+    )

src/ai/backend/manager/models/kernel.py

@@ -45,6 +45,7 @@
     SessionTypes,
     VFolderMount,
 )
+from ai.backend.common.utils import get_first_occurrence_time
 
 from ..api.exceptions import (
     BackendError,
@@ -78,7 +79,11 @@
 from .minilang.ordering import ColumnMapType, QueryOrderParser
 from .minilang.queryfilter import FieldSpecType, QueryFilterParser, enum_field_getter
 from .user import users
-from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge
+from .utils import (
+    ExtendedAsyncSAEngine,
+    execute_with_retry,
+    sql_append_dict_to_list,
+)
 
 if TYPE_CHECKING:
     from .gql import GraphQueryContext
@@ -512,7 +517,14 @@ async def handle_kernel_exception(
     #         // used to prevent duplication of SessionTerminatedEvent
     #   }
     # }
-    sa.Column("status_history", pgsql.JSONB(), nullable=True, default=sa.null()),
+    sa.Column("status_history", pgsql.JSONB(), nullable=False, default=[]),
+    # status_history records all status changes
+    # e.g)
+    # [
+    #   {"status: "PENDING", "timestamp": "2022-10-22T10:22:30"},
+    #   {"status: "SCHEDULED", "timestamp": "2022-10-22T11:40:30"},
+    #   {"status: "PREPARING", "timestamp": "2022-10-25T10:22:30"}
+    # ]
     sa.Column("callback_url", URLColumn, nullable=True, default=sa.null()),
     sa.Column("startup_command", sa.Text, nullable=True),
     sa.Column(
@@ -641,12 +653,9 @@ async def set_kernel_status(
         data = {
             "status": status,
             "status_changed": now,
-            "status_history": sql_json_merge(
-                kernels.c.status_history,
-                (),
-                {
-                    status.name: now.isoformat(),  # ["PULLING", "PREPARING"]
-                },
+            "status_history": sql_append_dict_to_list(
+                KernelRow.status_history,
+                {"status": status.name, "timestamp": now.isoformat()},
             ),
         }
         if status_data is not None:
@@ -692,12 +701,9 @@ async def _update() -> bool:
                 if update_data is None:
                     update_values = {
                         "status": new_status,
-                        "status_history": sql_json_merge(
+                        "status_history": sql_append_dict_to_list(
                             KernelRow.status_history,
-                            (),
-                            {
-                                new_status.name: now.isoformat(),
-                            },
+                            {"status": new_status.name, "timestamp": now.isoformat()},
                         ),
                     }
                 else:
@@ -836,7 +842,9 @@ def parse_row(cls, ctx: GraphQueryContext, row: KernelRow) -> Mapping[str, Any]:
             hide_agents = False
         else:
             hide_agents = ctx.local_config["manager"]["hide-agents"]
-        status_history = row.status_history or {}
+        status_history = row.status_history
+        scheduled_at = get_first_occurrence_time(status_history, KernelStatus.SCHEDULED.name)
+
         return {
             # identity
             "id": row.id,
@@ -862,7 +870,7 @@ def parse_row(cls, ctx: GraphQueryContext, row: KernelRow) -> Mapping[str, Any]:
             "created_at": row.created_at,
             "terminated_at": row.terminated_at,
             "starts_at": row.starts_at,
-            "scheduled_at": status_history.get(KernelStatus.SCHEDULED.name),
+            "scheduled_at": scheduled_at[1] if scheduled_at else None,
             "occupied_slots": row.occupied_slots.to_json(),
             # resources
             "agent": row.agent if not hide_agents else None,