Skip to content

Commit d3cdb84

Browse files
NiveditJainNiveditJain
authored
Enhance run information retrieval and trigger graph functionality (#340)
* Enhance run information retrieval and trigger graph functionality - Updated `get_run_info` to calculate total run counts more efficiently by storing the result in a variable, improving readability and performance. - Modified the status assignment in `get_run_info` to return `RunStatusEnum.FAILED` when no runs are found, enhancing error handling. - Refactored `trigger_graph` to only insert new stores if there are any, preventing unnecessary database operations and improving efficiency. * Add run_id and status index to State model for improved query performance * Refactor get_runs to enhance performance and error handling - Updated the `get_runs` function to utilize MongoDB aggregation for improved efficiency in retrieving run statistics. - Introduced a lookup table to map run IDs to their corresponding run objects, streamlining data processing. - Enhanced error handling by returning an empty response when no runs are found, ensuring clarity in API responses. - Refactored the response structure to include detailed run statistics, improving the overall data returned to the client. * Refactor get_runs to improve data retrieval and testing - Updated the `get_runs` function to utilize MongoDB aggregation for enhanced performance in retrieving run statistics. - Modified the total count calculation to directly query the database, ensuring accurate results when no runs are found. - Refactored test cases to include new scenarios for handling empty results and verifying aggregation pipeline structure. - Improved the overall clarity and organization of the test suite, ensuring comprehensive coverage of run status calculations and edge cases. * uv run ruff check --fix
1 parent 65767e2 commit d3cdb84

4 files changed

Lines changed: 452 additions & 254 deletions

File tree

‎state-manager/app/controller/get_runs.py‎

Lines changed: 112 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
import asyncio
2-
from beanie.operators import In, NotIn
31

42
from ..models.run_models import RunsResponse, RunListItem, RunStatusEnum
53
from ..models.db.state import State
@@ -9,40 +7,129 @@
97

108
logger = LogsManager().get_logger()
119

12-
async def get_run_status(run_id: str) -> RunStatusEnum:
13-
if await State.find(State.run_id == run_id, In(State.status, [StateStatusEnum.ERRORED, StateStatusEnum.NEXT_CREATED_ERROR])).count() > 0:
14-
return RunStatusEnum.FAILED
15-
elif await State.find(State.run_id == run_id, NotIn(State.status, [StateStatusEnum.SUCCESS, StateStatusEnum.RETRY_CREATED, StateStatusEnum.PRUNED])).count() == 0:
16-
return RunStatusEnum.SUCCESS
17-
else:
18-
return RunStatusEnum.PENDING
19-
20-
async def get_run_info(run: Run) -> RunListItem:
21-
return RunListItem(
22-
run_id=run.run_id,
23-
graph_name=run.graph_name,
24-
success_count=await State.find(State.run_id == run.run_id, In(State.status, [StateStatusEnum.SUCCESS, StateStatusEnum.PRUNED])).count(),
25-
pending_count=await State.find(State.run_id == run.run_id, In(State.status, [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED])).count(),
26-
errored_count=await State.find(State.run_id == run.run_id, In(State.status, [StateStatusEnum.ERRORED, StateStatusEnum.NEXT_CREATED_ERROR])).count(),
27-
retried_count=await State.find(State.run_id == run.run_id, State.status == StateStatusEnum.RETRY_CREATED).count(),
28-
total_count=await State.find(State.run_id == run.run_id,).count(),
29-
status=await get_run_status(run.run_id),
30-
created_at=run.created_at
31-
)
32-
33-
3410
async def get_runs(namespace_name: str, page: int, size: int, x_exosphere_request_id: str) -> RunsResponse:
3511
try:
3612
logger.info(f"Getting runs for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
3713

3814
runs = await Run.find(Run.namespace_name == namespace_name).sort(-Run.created_at).skip((page - 1) * size).limit(size).to_list() # type: ignore
15+
16+
if len(runs) == 0:
17+
return RunsResponse(
18+
namespace=namespace_name,
19+
total=await Run.find(Run.namespace_name == namespace_name).count(),
20+
page=page,
21+
size=size,
22+
runs=[]
23+
)
24+
25+
look_up_table = {
26+
run.run_id: run for run in runs
27+
}
28+
viewed = set()
29+
30+
31+
data_cursor = await State.get_pymongo_collection().aggregate(
32+
[
33+
{
34+
"$match": {
35+
"run_id": {
36+
"$in": [run.run_id for run in runs]
37+
}
38+
}
39+
},
40+
{
41+
"$group": {
42+
"_id": "$run_id",
43+
"total_count": {
44+
"$sum": 1
45+
},
46+
"success_count": {
47+
"$sum": {
48+
"$cond": {
49+
"if": {"$in": ["$status", [StateStatusEnum.SUCCESS, StateStatusEnum.PRUNED]]},
50+
"then": 1,
51+
"else": 0
52+
}
53+
}
54+
},
55+
"pending_count": {
56+
"$sum": {
57+
"$cond": {
58+
"if": {"$in": ["$status", [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED]]},
59+
"then": 1,
60+
"else": 0
61+
}
62+
}
63+
},
64+
"errored_count": {
65+
"$sum": {
66+
"$cond": {
67+
"if": {"$in": ["$status", [StateStatusEnum.ERRORED, StateStatusEnum.NEXT_CREATED_ERROR]]},
68+
"then": 1,
69+
"else": 0
70+
}
71+
}
72+
},
73+
"retried_count": {
74+
"$sum": {
75+
"$cond": {
76+
"if": {"$eq": ["$status", StateStatusEnum.RETRY_CREATED]},
77+
"then": 1,
78+
"else": 0
79+
}
80+
}
81+
}
82+
}
83+
}
84+
]
85+
)
86+
data = await data_cursor.to_list()
3987

88+
runs = []
89+
for run in data:
90+
success_count = run["success_count"]
91+
pending_count = run["pending_count"]
92+
errored_count = run["errored_count"]
93+
retried_count = run["retried_count"]
94+
95+
runs.append(
96+
RunListItem(
97+
run_id=run["_id"],
98+
graph_name=look_up_table[run["_id"]].graph_name,
99+
success_count=success_count,
100+
pending_count=pending_count,
101+
errored_count=errored_count,
102+
retried_count=retried_count,
103+
total_count=run["total_count"],
104+
status=RunStatusEnum.PENDING if pending_count > 0 else RunStatusEnum.FAILED if errored_count > 0 else RunStatusEnum.SUCCESS,
105+
created_at=look_up_table[run["_id"]].created_at
106+
)
107+
)
108+
viewed.add(run["_id"])
109+
110+
if len(look_up_table) > 0:
111+
for run_id in look_up_table:
112+
if run_id not in viewed:
113+
runs.append(
114+
RunListItem(
115+
run_id=run_id,
116+
graph_name=look_up_table[run_id].graph_name,
117+
success_count=0,
118+
pending_count=0,
119+
errored_count=0,
120+
retried_count=0,
121+
total_count=0,
122+
status=RunStatusEnum.FAILED,
123+
created_at=look_up_table[run_id].created_at
124+
)
125+
)
126+
40127
return RunsResponse(
41128
namespace=namespace_name,
42129
total=await Run.find(Run.namespace_name == namespace_name).count(),
43130
page=page,
44131
size=size,
45-
runs=await asyncio.gather(*[get_run_info(run) for run in runs])
132+
runs=sorted(runs, key=lambda x: x.created_at, reverse=True)
46133
)
47134

48135
except Exception as e:

‎state-manager/app/controller/trigger_graph.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,8 @@ async def trigger_graph(namespace_name: str, graph_name: str, body: TriggerGraph
6161
) for key, value in body.store.items()
6262
]
6363

64-
await Store.insert_many(new_stores)
64+
if len(new_stores) > 0:
65+
await Store.insert_many(new_stores)
6566

6667
root = graph_template.get_root_node()
6768

‎state-manager/app/models/db/state.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,5 +93,12 @@ class Settings:
9393
],
9494
unique=True,
9595
name="uniq_fanout_retry"
96+
),
97+
IndexModel(
98+
[
99+
("run_id", 1),
100+
("status", 1),
101+
],
102+
name="run_id_status_index"
96103
)
97104
]

0 commit comments

Comments
 (0)