Skip to content

Commit 709d9ac

Browse files
pablogsalpablogsal
committed
Add benchmarks to the worker
1 parent eddb717 commit 709d9ac

27 files changed

+6039
-68
lines changed

‎backend/app/admin_auth.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,23 @@ async def require_admin_auth(
112112
# is_active=True
113113
# )
114114
# return fake_session
115+
if not admin_session_token:
116+
raise HTTPException(
117+
status_code=status.HTTP_401_UNAUTHORIZED,
118+
detail="Admin authentication required",
119+
headers={"WWW-Authenticate": "Bearer"},
120+
)
115121

122+
try:
123+
session = await get_admin_session(db, admin_session_token)
124+
except Exception as e:
125+
# Log the database error but don't expose internal details
126+
logger.error(f"Database error in admin auth: {e}")
127+
raise HTTPException(
128+
status_code=status.HTTP_401_UNAUTHORIZED,
129+
detail="Authentication service unavailable",
130+
headers={"WWW-Authenticate": "Bearer"},
131+
)
116132
if not session:
117133
raise HTTPException(
118134
status_code=status.HTTP_401_UNAUTHORIZED,

‎backend/app/routers/benchmarks.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,10 +73,7 @@ async def get_diff_table(
7373
db, commit_sha=commit_sha, binary_id=binary_id, environment_id=environment_id
7474
)
7575
if not runs:
76-
raise HTTPException(
77-
status_code=404,
78-
detail="No runs found for this commit, binary, and environment",
79-
)
76+
return []
8077

8178
current_run = runs[0] # Get the latest run
8279
current_results = await crud.get_benchmark_results(db, run_id=current_run.run_id)

‎frontend/src/app/admin/components/QueryConsole.tsx

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -193,16 +193,34 @@ export default function QueryConsole() {
193193
});
194194
}
195195
} else {
196-
const errorData = await response.json();
197-
setResult({
198-
success: false,
199-
error: errorData.detail || 'Query execution failed',
200-
});
201-
toast({
202-
title: 'Error',
203-
description: errorData.detail || 'Query execution failed',
204-
variant: 'destructive',
205-
});
196+
try {
197+
const errorData = await response.json();
198+
setResult({
199+
success: false,
200+
error: errorData.detail || 'Query execution failed',
201+
});
202+
toast({
203+
title: 'Error',
204+
description: errorData.detail || 'Query execution failed',
205+
variant: 'destructive',
206+
});
207+
} catch (jsonError) {
208+
// Handle non-JSON error responses (e.g., HTML "Internal Server Error")
209+
const errorText = await response.text();
210+
const errorMessage = errorText.includes('Internal Server Error')
211+
? 'Internal server error occurred. Please check the query syntax and try again.'
212+
: errorText || 'Query execution failed';
213+
214+
setResult({
215+
success: false,
216+
error: errorMessage,
217+
});
218+
toast({
219+
title: 'Error',
220+
description: errorMessage,
221+
variant: 'destructive',
222+
});
223+
}
206224
}
207225
} catch (error) {
208226
console.error('Error executing query:', error);

‎worker/src/memory_tracker_worker/benchmarks/__init__.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,8 +110,10 @@ def run_benchmarks(venv_dir: Path, output_dir: Path, commit: git.Commit) -> None
110110
# Create temporary directory for benchmark files
111111
temp_dir = Path(tempfile.mkdtemp(prefix="benchmarks_"))
112112
try:
113-
# Copy benchmark files to temporary directory
113+
# Copy benchmark files and directories to temporary directory
114114
benchmarks_dir = Path(__file__).parent
115+
116+
# Copy all benchmark files
115117
for benchmark_file in benchmarks_dir.glob("*.py"):
116118
if benchmark_file.name == "__init__.py":
117119
continue
@@ -122,10 +124,24 @@ def run_benchmarks(venv_dir: Path, output_dir: Path, commit: git.Commit) -> None
122124
logger.info(
123125
f"Copied benchmark {benchmark_file.name} to temporary directory"
124126
)
127+
128+
# Copy all data directories
129+
for item in benchmarks_dir.iterdir():
130+
if item.is_dir():
131+
dest_dir = temp_dir / item.name
132+
shutil.copytree(item, dest_dir)
133+
logger.info(
134+
f"Copied directory {item.name} to temporary directory"
135+
)
125136

126-
# Run benchmark with memray
137+
# Run benchmarks with memray
138+
for benchmark_file in benchmarks_dir.glob("*.py"):
139+
if benchmark_file.name == "__init__.py":
140+
continue
141+
127142
benchmark_name = benchmark_file.stem
128143
logger.info(f"Running benchmark: {benchmark_name}")
144+
dest_file = temp_dir / benchmark_file.name
129145

130146
# Run memray
131147
memray_output = output_dir / f"{benchmark_name}.bin"

‎worker/src/memory_tracker_worker/benchmarks/async_tree_base.py

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
"""
2+
Benchmark for async tree workload, which calls asyncio.gather() on a tree
3+
(6 levels deep, 6 branches per level) with the leaf nodes simulating some
4+
(potentially) async work (depending on the benchmark variant). Benchmark
5+
variants include:
6+
7+
1) "none": No actual async work in the async tree.
8+
2) "io": All leaf nodes simulate async IO workload (async sleep 50ms).
9+
3) "memoization": All leaf nodes simulate async IO workload with 90% of
10+
the data memoized
11+
4) "cpu_io_mixed": Half of the leaf nodes simulate CPU-bound workload and
12+
the other half simulate the same workload as the
13+
"memoization" variant.
14+
"""
15+
16+
import sys
17+
import asyncio
18+
import math
19+
import random
20+
21+
NUM_RECURSE_LEVELS = 6
22+
NUM_RECURSE_BRANCHES = 6
23+
RANDOM_SEED = 0
24+
IO_SLEEP_TIME = 0.05
25+
MEMOIZABLE_PERCENTAGE = 90
26+
CPU_PROBABILITY = 0.5
27+
FACTORIAL_N = 500
28+
29+
30+
class AsyncTree:
31+
def __init__(self):
32+
self.cache = {}
33+
# set to deterministic random, so that the results are reproducible
34+
random.seed(RANDOM_SEED)
35+
36+
async def mock_io_call(self):
37+
await asyncio.sleep(IO_SLEEP_TIME)
38+
39+
async def workload_func(self):
40+
raise NotImplementedError("To be implemented by each variant's derived class.")
41+
42+
async def recurse(self, recurse_level):
43+
if recurse_level == 0:
44+
await self.workload_func()
45+
return
46+
47+
await asyncio.gather(
48+
*[self.recurse(recurse_level - 1) for _ in range(NUM_RECURSE_BRANCHES)]
49+
)
50+
51+
async def run(self):
52+
await self.recurse(NUM_RECURSE_LEVELS)
53+
54+
55+
class NoneAsyncTree(AsyncTree):
56+
async def workload_func(self):
57+
return
58+
59+
60+
class IOAsyncTree(AsyncTree):
61+
async def workload_func(self):
62+
await self.mock_io_call()
63+
64+
65+
class MemoizationAsyncTree(AsyncTree):
66+
async def workload_func(self):
67+
# deterministic random, seed set in AsyncTree.__init__()
68+
data = random.randint(1, 100)
69+
70+
if data <= MEMOIZABLE_PERCENTAGE:
71+
if self.cache.get(data):
72+
return data
73+
74+
self.cache[data] = True
75+
76+
await self.mock_io_call()
77+
return data
78+
79+
80+
class CpuIoMixedAsyncTree(MemoizationAsyncTree):
81+
async def workload_func(self):
82+
# deterministic random, seed set in AsyncTree.__init__()
83+
if random.random() < CPU_PROBABILITY:
84+
# mock cpu-bound call
85+
return math.factorial(FACTORIAL_N)
86+
else:
87+
return await MemoizationAsyncTree.workload_func(self)
88+
89+
90+
def add_metadata(runner):
91+
runner.metadata["description"] = "Async tree workloads."
92+
runner.metadata["async_tree_recurse_levels"] = NUM_RECURSE_LEVELS
93+
runner.metadata["async_tree_recurse_branches"] = NUM_RECURSE_BRANCHES
94+
runner.metadata["async_tree_random_seed"] = RANDOM_SEED
95+
runner.metadata["async_tree_io_sleep_time"] = IO_SLEEP_TIME
96+
runner.metadata["async_tree_memoizable_percentage"] = MEMOIZABLE_PERCENTAGE
97+
runner.metadata["async_tree_cpu_probability"] = CPU_PROBABILITY
98+
runner.metadata["async_tree_factorial_n"] = FACTORIAL_N
99+
100+
101+
def add_cmdline_args(cmd, args):
102+
cmd.append(args.benchmark)
103+
104+
105+
def add_parser_args(parser):
106+
parser.add_argument(
107+
"benchmark",
108+
choices=BENCHMARKS,
109+
help="""\
110+
Determines which benchmark to run. Options:
111+
1) "none": No actual async work in the async tree.
112+
2) "io": All leaf nodes simulate async IO workload (async sleep 50ms).
113+
3) "memoization": All leaf nodes simulate async IO workload with 90% of
114+
the data memoized
115+
4) "cpu_io_mixed": Half of the leaf nodes simulate CPU-bound workload and
116+
the other half simulate the same workload as the
117+
"memoization" variant.
118+
""",
119+
)
120+
121+
122+
BENCHMARKS = {
123+
"none": NoneAsyncTree,
124+
"io": IOAsyncTree,
125+
"memoization": MemoizationAsyncTree,
126+
"cpu_io_mixed": CpuIoMixedAsyncTree,
127+
}
128+
129+
130+
def run_benchmark(benchmark):
131+
async_tree_class = BENCHMARKS[benchmark]
132+
async_tree = async_tree_class()
133+
asyncio.run(async_tree.run())
134+
135+
136+
if __name__ == "__main__":
137+
benchmark = "none"
138+
run_benchmark(benchmark)

0 commit comments

Comments
 (0)