Add repeatable test+benchmark workflow, CI artifact upload, and scheduler benchmarks#15
Conversation
Summary of ChangesHello Igor Holt (@igor-holt), I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request lays a robust foundation for continuous quality and performance monitoring within the AIOS Layer project. It introduces a standardized, repeatable workflow for executing both unit tests and Go benchmarks, ensuring consistent evaluation of code changes. This workflow is seamlessly integrated into the project's CI pipeline, automatically running checks and archiving detailed performance reports as artifacts. Furthermore, specific benchmarks have been added for the core scheduler component, enabling precise measurement and tracking of its efficiency over time. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on Gemini (@gemini-code-assist) comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant amount of new functionality, including a repeatable test and benchmark workflow, CI integration, and scheduler benchmarks. The overall structure is well-organized, and the separation of concerns into different packages is good. I've identified several areas for improvement, particularly around error handling in the CLI, benchmark correctness, and strengthening security aspects like lease ID generation. My review includes specific suggestions to address these points.
| for i := 0; i < b.N; i++ { | ||
| sched := New(gpus) | ||
| _, _ = sched.CreateLease("lease", "user", time.Second) | ||
| } |
There was a problem hiding this comment.
The benchmark setup (creating a new scheduler with New(gpus)) is inside the for loop. This means the benchmark is measuring the setup cost in every iteration, not just the CreateLease call. For an accurate benchmark of CreateLease, the scheduler should be initialized once outside the loop, and b.ResetTimer() should be called after the setup. You will also need to ensure the scheduler's state is reset for each iteration to make the benchmark valid for b.N runs.
| for i := 0; i < b.N; i++ { | ||
| sched := New(gpus) | ||
| _, _ = sched.CreateLease("lease", "user", time.Nanosecond) | ||
| sched.ReapExpired() | ||
| } |
There was a problem hiding this comment.
This benchmark has two issues. First, like BenchmarkCreateLease, it re-initializes the scheduler in every iteration, which skews the measurement. Second, creating a lease with time.Nanosecond duration doesn't guarantee it will be expired when ReapExpired is called. A better approach is to set up the scheduler outside the loop, use b.ResetTimer(), and inside the loop, use b.StopTimer()/b.StartTimer() to exclude the setup of creating an already expired lease for each iteration.
| writeError(w, http.StatusForbidden, err.Error()) | ||
| return | ||
| } | ||
| leaseID := fmt.Sprintf("lease-%d", rand.Intn(1_000_000)) |
There was a problem hiding this comment.
Using rand.Intn for lease IDs is not robust. It can lead to collisions and is not cryptographically secure, which is a concern for a system aiming to be 'secure-by-default'. Consider using a standard UUID library (like github.com/google/uuid) to generate unique and non-sequential IDs. You will need to add the dependency and import the package.
| leaseID := fmt.Sprintf("lease-%d", rand.Intn(1_000_000)) | |
| leaseID := uuid.NewString() |
| body, _ := json.Marshal(payload) | ||
| resp, err := http.Post(agent+"/v1/leases", "application/json", bytes.NewReader(body)) | ||
| if err != nil { | ||
| fmt.Printf("error: %v\n", err) | ||
| return | ||
| } | ||
| defer resp.Body.Close() | ||
| data, _ := io.ReadAll(resp.Body) | ||
| fmt.Println(string(data)) |
There was a problem hiding this comment.
Errors from json.Marshal and io.ReadAll are being ignored by using the blank identifier _. This is a bug and can hide issues like marshalling failures or problems reading the HTTP response body. These errors must be checked and handled properly.
body, err := json.Marshal(payload)
if err != nil {
fmt.Printf("error marshalling request: %v\n", err)
return
}
resp, err := http.Post(agent+"/v1/leases", "application/json", bytes.NewReader(body))
if err != nil {
fmt.Printf("error: %v\n", err)
return
}
defer resp.Body.Close()
data, err := io.ReadAll(resp.Body)
if err != nil {
fmt.Printf("error reading response: %v\n", err)
return
}
fmt.Println(string(data))| body, _ := json.Marshal(payload) | ||
| client := http.Client{Timeout: time.Second * 60} | ||
| resp, err := client.Post(endpoint, "application/json", bytes.NewReader(body)) | ||
| if err != nil { | ||
| fmt.Printf("error: %v\n", err) | ||
| return | ||
| } | ||
| defer resp.Body.Close() | ||
| data, _ := io.ReadAll(resp.Body) | ||
| fmt.Println(string(data)) |
There was a problem hiding this comment.
Similar to the requestLease function, errors from json.Marshal and io.ReadAll are being ignored. This can lead to silent failures and makes debugging difficult. Always check and handle errors returned by these functions.
body, err := json.Marshal(payload)
if err != nil {
fmt.Printf("error marshalling request: %v\n", err)
return
}
client := http.Client{Timeout: time.Second * 60}
resp, err := client.Post(endpoint, "application/json", bytes.NewReader(body))
if err != nil {
fmt.Printf("error: %v\n", err)
return
}
defer resp.Body.Close()
data, err := io.ReadAll(resp.Body)
if err != nil {
fmt.Printf("error reading response: %v\n", err)
return
}
fmt.Println(string(data))| - uses: actions/setup-go@v5 | ||
| with: | ||
| go-version: "1.22" |
There was a problem hiding this comment.
To improve the performance of the CI workflow, consider enabling caching for Go modules. The actions/setup-go action provides a simple way to do this by adding the cache: true option. This will speed up subsequent runs by caching dependencies.
- uses: actions/setup-go@v5
with:
go-version: "1.22"
cache: true| if len(parts) < 3 { | ||
| continue | ||
| } | ||
| index, err := strconv.Atoi(strings.TrimSpace(parts[0])) | ||
| if err != nil { | ||
| continue | ||
| } | ||
| name := strings.TrimSpace(parts[1]) | ||
| mem, err := strconv.Atoi(strings.TrimSpace(parts[2])) | ||
| if err != nil { | ||
| continue |
There was a problem hiding this comment.
The current implementation silently skips malformed lines from nvidia-smi. This could hide potential issues with the output format or the command itself. It would be beneficial to log these parsing errors to aid in debugging. You will need to import the log package.
| if len(parts) < 3 { | |
| continue | |
| } | |
| index, err := strconv.Atoi(strings.TrimSpace(parts[0])) | |
| if err != nil { | |
| continue | |
| } | |
| name := strings.TrimSpace(parts[1]) | |
| mem, err := strconv.Atoi(strings.TrimSpace(parts[2])) | |
| if err != nil { | |
| continue | |
| if len(parts) < 3 { | |
| log.Printf("skipping malformed line from nvidia-smi: %q", line) | |
| continue | |
| } | |
| index, err := strconv.Atoi(strings.TrimSpace(parts[0])) | |
| if err != nil { | |
| log.Printf("failed to parse GPU index from %q: %v", parts[0], err) | |
| continue | |
| } | |
| name := strings.TrimSpace(parts[1]) | |
| mem, err := strconv.Atoi(strings.TrimSpace(parts[2])) | |
| if err != nil { | |
| log.Printf("failed to parse GPU memory from %q: %v", parts[2], err) | |
| continue | |
| } |
| type GPU struct { | ||
| Index int `json:"index"` | ||
| Name string `json:"name"` | ||
| MemoryTotal int `json:"memory_total_mb"` | ||
| } |
There was a problem hiding this comment.
The GPU struct is defined here and also in the aios-layer/agent/internal/gpu package. This duplication can lead to maintenance issues if the struct needs to be changed in the future. Consider defining this struct in a shared internal package or having the scheduler package import the gpu package to reuse the definition.
| metrics := "# HELP aios_leases Active leases\n# TYPE aios_leases gauge\n" | ||
| metrics += fmt.Sprintf("aios_leases %d\n", len(sched.ListLeases())) | ||
| w.Write([]byte(metrics)) |
There was a problem hiding this comment.
Manually constructing the Prometheus metrics exposition format as a string is brittle and hard to maintain or extend. It's recommended to use the official Go client library for Prometheus (prometheus/client_golang). It provides types for metrics (like Gauges, Counters) and handles the formatting correctly, making your metrics implementation more robust.
| type openAIRequest struct { | ||
| Model string `json:"model"` | ||
| Messages []string `json:"messages"` | ||
| } |
There was a problem hiding this comment.
The openAIRequest struct is defined but never used in the code. It should be removed to improve code clarity and maintainability.
1 participant
Motivation
Description
./scripts/run_checks.shto run tests (go test -count=2) and Go benchmarks (-bench . -benchmem -count=2) and write output toreports/.agent/internal/scheduler/scheduler_bench_test.goand updateMakefileto include abenchtarget that writesreports/bench.log../scripts/run_checks.shand uploadaios-layer/reportsas an artifact via the workflow in.github/workflows/ci.yml.README.mdand.gitignoreto document the new workflow and ignorereports/; addgo.mod/go.sumentries pulled in bygo mod tidyso tests run cleanly.Testing
go mod tidyto populatego.sumand satisfy dependencies, which completed successfully../scripts/run_checks.shlocally, which rango test -count=2and the benchmarks; unit tests (scheduler) passed and benchmarks completed, with logs written toreports/test.logandreports/bench.log.BenchmarkCreateLease~666–676 ns/op,824 B/op,5 allocs/op;BenchmarkReapExpired~837–851 ns/op,808 B/op,4 allocs/op.Codex Task