feat(BA-3947): Add sessions resolver to strawberry AgentV2 schema#9620
feat(BA-3947): Add sessions resolver to strawberry AgentV2 schema#9620jopemachine wants to merge 5 commits intomainfrom
sessions resolver to strawberry AgentV2 schema#9620Conversation
github-actions
bot
added
size:M
jopemachine
changed the title
Co-authored-by: octodog <mu001@lablup.com>
jopemachine
changed the title
sessions resolver to strawberry AgentV2 schema
Co-authored-by: octodog <mu001@lablup.com>
There was a problem hiding this comment.
Pull request overview
Adds the missing AgentV2.sessions field to the Strawberry-based GraphQL schema so AgentV2 can expose session lists associated with a specific agent (parity step toward legacy AgentNode replacement).
Changes:
- Add
AgentV2.sessionsfield + resolver wiring tofetch_sessions()with pagination/filter/order support. - Add a scheduler query condition (
SessionConditions.by_agent_id) to filter sessions by agent membership. - Update GraphQL reference schema docs and add a changelog entry.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
src/ai/backend/manager/repositories/scheduler/options.py |
Adds SessionConditions.by_agent_id() and updates kernel agent-id condition typing. |
src/ai/backend/manager/api/gql/agent/types.py |
Adds Strawberry AgentV2.sessions field and passes agent-based base condition into session fetcher. |
docs/manager/graphql-reference/v2-schema.graphql |
Documents the new AgentV2.sessions field in the v2 schema reference. |
docs/manager/graphql-reference/supergraph.graphql |
Documents the new AgentV2.sessions field in the supergraph schema reference. |
changes/9620.feature.md |
Records the feature addition in the changelog system. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
| @strawberry.field( # type: ignore[misc] | ||
| description="Added in 26.3.0. List of sessions running on this agent with pagination support." | ||
| ) | ||
| async def sessions( | ||
| self, | ||
| info: Info[StrawberryGQLContext], | ||
| filter: Annotated[ | ||
| SessionV2FilterGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types") | ||
| ] | ||
| | None = None, | ||
| order_by: list[ | ||
| Annotated[ | ||
| SessionV2OrderByGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types") | ||
| ] | ||
| ] | ||
| | None = None, | ||
| before: str | None = None, | ||
| after: str | None = None, | ||
| first: int | None = None, | ||
| last: int | None = None, | ||
| limit: int | None = None, | ||
| offset: int | None = None, | ||
| ) -> Annotated[ | ||
| SessionV2ConnectionGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types") | ||
| ]: | ||
| """Fetch sessions associated with this agent.""" | ||
| from ai.backend.manager.api.gql.session.fetcher.session import fetch_sessions | ||
|
|
||
| return await fetch_sessions( | ||
| info=info, | ||
| filter=filter, | ||
| order_by=order_by, | ||
| before=before, | ||
| after=after, | ||
| first=first, | ||
| last=last, | ||
| limit=limit, | ||
| offset=offset, | ||
| base_conditions=[SessionConditions.by_agent_id(self._agent_id)], | ||
| ) |
There was a problem hiding this comment.
AgentV2GQL.sessions calls fetch_sessions() directly per agent instance. In queries that return a list/connection of agents and request sessions for each node, this will trigger one session-search per agent (N+1 pattern) and can become expensive. Consider introducing a batching approach (e.g., a DataLoader keyed by (agent_id, pagination/filter args) with internal grouping, or a dedicated processor/repository method that can fetch sessions for multiple agent IDs in one call when args match) to keep the number of DB/service calls bounded per GraphQL request.
| @strawberry.field( # type: ignore[misc] | ||
| description="Added in 26.3.0. List of sessions running on this agent with pagination support." | ||
| ) | ||
| async def sessions( | ||
| self, | ||
| info: Info[StrawberryGQLContext], | ||
| filter: Annotated[ | ||
| SessionV2FilterGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types") | ||
| ] | ||
| | None = None, | ||
| order_by: list[ | ||
| Annotated[ | ||
| SessionV2OrderByGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types") | ||
| ] | ||
| ] | ||
| | None = None, | ||
| before: str | None = None, | ||
| after: str | None = None, | ||
| first: int | None = None, | ||
| last: int | None = None, | ||
| limit: int | None = None, | ||
| offset: int | None = None, | ||
| ) -> Annotated[ | ||
| SessionV2ConnectionGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types") | ||
| ]: | ||
| """Fetch sessions associated with this agent.""" | ||
| from ai.backend.manager.api.gql.session.fetcher.session import fetch_sessions | ||
|
|
||
| return await fetch_sessions( | ||
| info=info, | ||
| filter=filter, | ||
| order_by=order_by, | ||
| before=before, | ||
| after=after, | ||
| first=first, | ||
| last=last, | ||
| limit=limit, | ||
| offset=offset, | ||
| base_conditions=[SessionConditions.by_agent_id(self._agent_id)], | ||
| ) |
There was a problem hiding this comment.
The new AgentV2GQL.sessions field introduces a new query surface and filtering/pagination wiring via base_conditions=[SessionConditions.by_agent_id(...)], but there are no unit tests covering this resolver behavior. Adding a resolver-level unit test (similar to other GraphQL resolver tests that patch the fetcher) would help ensure the agent-id condition is always applied and pagination args are forwarded correctly.
github-actions
bot
added
size:M
resolves #8130 (BA-3947)
Checklist: (if applicable)
ai.backend.testdocsdirectory📚 Documentation preview 📚: https://sorna--9620.org.readthedocs.build/en/9620/
📚 Documentation preview 📚: https://sorna-ko--9620.org.readthedocs.build/ko/9620/