diff --git a/src/ai/backend/manager/models/gql_models/kernel.py b/src/ai/backend/manager/models/gql_models/kernel.py index cbb52db685..141ef285ba 100644 --- a/src/ai/backend/manager/models/gql_models/kernel.py +++ b/src/ai/backend/manager/models/gql_models/kernel.py @@ -5,6 +5,7 @@ TYPE_CHECKING, Any, Self, + cast, ) import graphene @@ -14,15 +15,17 @@ from ai.backend.common import msgpack, redis_helper from ai.backend.common.types import AgentId, KernelId, SessionId -from ai.backend.manager.models.base import ( + +from ..base import ( batch_multiresult_in_scalar_stream, batch_multiresult_in_session, ) - from ..gql_relay import AsyncNode, Connection -from ..kernel import KernelRow, KernelStatus +from ..kernel import KernelRow from ..user import UserRole +from ..utils import get_lastest_timestamp_for_status from .image import ImageNode +from .session import SessionStatus if TYPE_CHECKING: from ..gql import GraphQueryContext @@ -113,7 +116,12 @@ def from_row(cls, ctx: GraphQueryContext, row: KernelRow) -> Self: hide_agents = False else: hide_agents = ctx.local_config["manager"]["hide-agents"] - status_history = row.status_history or {} + + timestamp = get_lastest_timestamp_for_status( + cast(list[dict[str, str]], row.status_history), SessionStatus.SCHEDULED + ) + scheduled_at = str(timestamp) if timestamp is not None else None + return KernelNode( id=row.id, # auto-converted to Relay global ID row_id=row.id, @@ -129,7 +137,7 @@ def from_row(cls, ctx: GraphQueryContext, row: KernelRow) -> Self: 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, occupied_slots=row.occupied_slots.to_json(), agent_id=row.agent if not hide_agents else None, agent_addr=row.agent_addr if not hide_agents else None, diff --git a/src/ai/backend/manager/models/gql_models/session.py b/src/ai/backend/manager/models/gql_models/session.py index 95edc71938..d0a8df13b4 100644 --- a/src/ai/backend/manager/models/gql_models/session.py +++ b/src/ai/backend/manager/models/gql_models/session.py @@ -56,7 +56,7 @@ get_permission_ctx, ) from ..user import UserRole -from ..utils import execute_with_txn_retry +from ..utils import execute_with_txn_retry, get_lastest_timestamp_for_status from .kernel import KernelConnection, KernelNode if TYPE_CHECKING: @@ -174,7 +174,7 @@ class Meta: # status_changed = GQLDateTime() # FIXME: generated attribute status_info = graphene.String() status_data = graphene.JSONString() - status_history = graphene.JSONString() + status_history = graphene.JSONString(description="Added in 24.12.0.") created_at = GQLDateTime() terminated_at = GQLDateTime() starts_at = GQLDateTime() @@ -225,7 +225,11 @@ def from_row( permissions: Optional[Iterable[ComputeSessionPermission]] = None, ) -> Self: status_history = row.status_history or {} - raw_scheduled_at = status_history.get(SessionStatus.SCHEDULED.name) + timestamp = get_lastest_timestamp_for_status( + cast(list[dict[str, str]], status_history), SessionStatus.SCHEDULED + ) + scheduled_at = str(timestamp) if timestamp is not None else None + result = cls( # identity id=row.id, # auto-converted to Relay global ID @@ -251,9 +255,7 @@ def from_row( created_at=row.created_at, starts_at=row.starts_at, terminated_at=row.terminated_at, - scheduled_at=datetime.fromisoformat(raw_scheduled_at) - if raw_scheduled_at is not None - else None, + scheduled_at=scheduled_at, startup_command=row.startup_command, result=row.result.name, # resources diff --git a/src/ai/backend/manager/models/kernel.py b/src/ai/backend/manager/models/kernel.py index 6acd91c674..3db283f899 100644 --- a/src/ai/backend/manager/models/kernel.py +++ b/src/ai/backend/manager/models/kernel.py @@ -880,6 +880,7 @@ class Meta: # status status = graphene.String() status_changed = GQLDateTime() + status_history = graphene.JSONString(description="Added in 24.12.0.") status_info = graphene.String() status_data = graphene.JSONString() created_at = GQLDateTime() @@ -931,6 +932,7 @@ def parse_row(cls, ctx: GraphQueryContext, row: KernelRow) -> Mapping[str, Any]: # status "status": row.status.name, "status_changed": row.status_changed, + "status_history": row.status_history, "status_info": row.status_info, "status_data": row.status_data, "created_at": row.created_at,