Skip to content

Commit 743f439

Browse files
committed
Add benchmarks to the worker and some fixes to the frontend
1 parent eddb717 commit 743f439

31 files changed

+6098
-87
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/models.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from sqlalchemy import (
22
Column,
33
Integer,
4+
BigInteger,
45
String,
56
DateTime,
67
Text,
@@ -118,9 +119,9 @@ class BenchmarkResult(Base):
118119
id = Column(String(200), primary_key=True)
119120
run_id = Column(String(100), ForeignKey("runs.run_id"), nullable=False)
120121
benchmark_name = Column(String(100), nullable=False)
121-
high_watermark_bytes = Column(Integer, nullable=False)
122+
high_watermark_bytes = Column(BigInteger, nullable=False)
122123
allocation_histogram = Column(JSON, nullable=False) # Array of [size, count] tuples
123-
total_allocated_bytes = Column(Integer, nullable=False)
124+
total_allocated_bytes = Column(BigInteger, nullable=False)
124125
top_allocating_functions = Column(JSON, nullable=False) # Array of function objects
125126
flamegraph_html = Column(Text, nullable=True) # HTML content of the flamegraph
126127

backend/app/routers/admin.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1080,6 +1080,9 @@ async def execute_query(
10801080

10811081
start_time = time.time()
10821082

1083+
# Get admin username now while in async context
1084+
admin_username = admin_session.github_username
1085+
10831086
try:
10841087
# Execute the query
10851088
result = await db.execute(text(query))
@@ -1121,7 +1124,7 @@ async def execute_query(
11211124

11221125
logger.error(f"Query execution failed: {e}", extra={
11231126
"query": query,
1124-
"admin_user": admin_session.github_username,
1127+
"admin_user": admin_username,
11251128
})
11261129

11271130
return QueryResult(

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: 75 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -193,16 +193,36 @@ 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+
// Get response text first, then try to parse as JSON
197+
const responseText = await response.text();
198+
199+
try {
200+
const errorData = JSON.parse(responseText);
201+
setResult({
202+
success: false,
203+
error: errorData.detail || 'Query execution failed',
204+
});
205+
toast({
206+
title: 'Error',
207+
description: errorData.detail || 'Query execution failed',
208+
variant: 'destructive',
209+
});
210+
} catch (jsonError) {
211+
// Handle non-JSON error responses (e.g., HTML "Internal Server Error")
212+
const errorMessage = responseText.includes('Internal Server Error')
213+
? 'Internal server error occurred. Please check the query syntax and try again.'
214+
: responseText || 'Query execution failed';
215+
216+
setResult({
217+
success: false,
218+
error: errorMessage,
219+
});
220+
toast({
221+
title: 'Error',
222+
description: errorMessage,
223+
variant: 'destructive',
224+
});
225+
}
206226
}
207227
} catch (error) {
208228
console.error('Error executing query:', error);
@@ -537,8 +557,9 @@ WHERE expires_at < NOW() OR is_active = false;`
537557
WHERE high_watermark_bytes > 1000000000; -- 1GB threshold`
538558
},
539559
{
540-
name: '🗑️ Remove old Python version data',
541-
query: `DELETE FROM runs
560+
name: '🗑️ Remove old Python version data (Step 1: benchmark_results)',
561+
query: `-- First delete benchmark_results for old Python versions
562+
DELETE FROM benchmark_results
542563
WHERE run_id IN (
543564
SELECT r.run_id
544565
FROM runs r
@@ -547,33 +568,66 @@ WHERE run_id IN (
547568
);`
548569
},
549570
{
550-
name: '🗑️ Delete all data for specific binary',
551-
query: `-- Delete all runs and results for a specific binary
571+
name: '🗑️ Remove old Python version data (Step 2: runs)',
572+
query: `-- Then delete runs for old Python versions
573+
DELETE FROM runs
574+
WHERE run_id IN (
575+
SELECT r.run_id
576+
FROM runs r
577+
JOIN commits c ON r.commit_sha = c.sha
578+
WHERE c.python_major < 3 OR (c.python_major = 3 AND c.python_minor < 10)
579+
);`
580+
},
581+
{
582+
name: '🗑️ Delete specific Python version data',
583+
query: `-- Template: Replace X with Python minor version (e.g., 15)
584+
-- Run this query twice: first for benchmark_results, then for runs
552585
DELETE FROM benchmark_results
553586
WHERE run_id IN (
554-
SELECT run_id FROM runs WHERE binary_id = 'REPLACE_WITH_BINARY_ID'
587+
SELECT run_id FROM runs WHERE python_minor = X
555588
);
589+
-- DELETE FROM runs WHERE python_minor = X;`
590+
},
591+
{
592+
name: '🗑️ Delete binary data (Step 1: benchmark_results)',
593+
query: `-- First delete benchmark_results for specific binary
594+
DELETE FROM benchmark_results
595+
WHERE run_id IN (
596+
SELECT run_id FROM runs WHERE binary_id = 'REPLACE_WITH_BINARY_ID'
597+
);`
598+
},
599+
{
600+
name: '🗑️ Delete binary data (Step 2: runs)',
601+
query: `-- Then delete runs for specific binary
556602
DELETE FROM runs WHERE binary_id = 'REPLACE_WITH_BINARY_ID';`
557603
},
558604
{
559-
name: '🗑️ Delete all data for specific environment',
560-
query: `-- Delete all runs and results for a specific environment
605+
name: '🗑️ Delete environment data (Step 1: benchmark_results)',
606+
query: `-- First delete benchmark_results for specific environment
561607
DELETE FROM benchmark_results
562608
WHERE run_id IN (
563609
SELECT run_id FROM runs WHERE environment_id = 'REPLACE_WITH_ENVIRONMENT_ID'
564-
);
610+
);`
611+
},
612+
{
613+
name: '🗑️ Delete environment data (Step 2: runs)',
614+
query: `-- Then delete runs for specific environment
565615
DELETE FROM runs WHERE environment_id = 'REPLACE_WITH_ENVIRONMENT_ID';`
566616
},
567617
{
568-
name: '🗑️ Delete all data older than X months',
569-
query: `-- Delete all benchmark data older than 6 months
618+
name: '🗑️ Delete old data (Step 1: benchmark_results)',
619+
query: `-- First delete benchmark_results older than 6 months
570620
DELETE FROM benchmark_results
571621
WHERE run_id IN (
572622
SELECT r.run_id
573623
FROM runs r
574624
JOIN commits c ON r.commit_sha = c.sha
575625
WHERE c.timestamp < NOW() - INTERVAL '6 months'
576-
);
626+
);`
627+
},
628+
{
629+
name: '🗑️ Delete old data (Step 2: runs)',
630+
query: `-- Then delete runs older than 6 months
577631
DELETE FROM runs
578632
WHERE run_id IN (
579633
SELECT r.run_id

frontend/src/components/layout/MemrayStatusIndicator.tsx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,8 @@ export default function MemrayStatusIndicator() {
105105
))}
106106
</div>
107107
) : (
108-
<div className="p-4 bg-green-50 border border-green-200 rounded-lg">
109-
<div className="text-green-800 text-center">
108+
<div className="p-4 bg-green-50 dark:bg-green-950 border border-green-200 dark:border-green-800 rounded-lg">
109+
<div className="text-green-800 dark:text-green-200 text-center">
110110
🎉 All memray builds are currently successful across all environments!
111111
</div>
112112
</div>

frontend/src/components/memray/FailureCard.tsx

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,13 @@ export default function FailureCard({
5454
if (compact) {
5555
return (
5656
<>
57-
<div className="p-3 border rounded-lg bg-amber-50 border-amber-200">
57+
<div className="p-3 border rounded-lg bg-amber-50 dark:bg-amber-950 border-amber-200 dark:border-amber-800">
5858
<div className="flex items-center justify-between">
5959
<div className="flex-1">
60-
<div className="font-medium text-amber-900">
60+
<div className="font-medium text-amber-900 dark:text-amber-200">
6161
{binary_name} on {environment_name}
6262
</div>
63-
<div className="text-sm text-amber-700 mt-1">
63+
<div className="text-sm text-amber-700 dark:text-amber-300 mt-1">
6464
Commit: {commit_sha.substring(0, 8)}{formatTimestamp(commit_timestamp || failure_timestamp)}
6565
</div>
6666
</div>

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"

0 commit comments

Comments
 (0)