fetch_ml/tests/benchmarks/json_serialization_bench_test.go
Jeremie Fraeys 7305e2bc21
test: add comprehensive test coverage and command improvements
- Add logs and debug end-to-end tests
- Add test helper utilities
- Improve test fixtures and templates
- Update API server and config lint commands
- Add multi-user database initialization
2026-02-16 20:38:15 -05:00

185 lines
4.1 KiB
Go

package benchmarks
import (
"encoding/json"
"fmt"
"testing"
"time"
"github.com/jfraeys/fetch_ml/internal/queue"
)
// BenchmarkTaskJSONMarshal profiles Task serialization to JSON
// Tier 3 C++ candidate: Zero-copy binary serialization
func BenchmarkTaskJSONMarshal(b *testing.B) {
task := createTestTask()
b.ResetTimer()
b.ReportAllocs()
for i := 0; i < b.N; i++ {
_, err := json.Marshal(task)
if err != nil {
b.Fatal(err)
}
}
}
// BenchmarkTaskJSONUnmarshal profiles Task deserialization from JSON
// Tier 3 C++ candidate: Memory-mapped binary format
func BenchmarkTaskJSONUnmarshal(b *testing.B) {
task := createTestTask()
data, err := json.Marshal(task)
if err != nil {
b.Fatal(err)
}
b.ResetTimer()
b.ReportAllocs()
for i := 0; i < b.N; i++ {
var t queue.Task
err := json.Unmarshal(data, &t)
if err != nil {
b.Fatal(err)
}
}
}
// BenchmarkTaskJSONRoundTrip profiles full serialize/deserialize cycle
func BenchmarkTaskJSONRoundTrip(b *testing.B) {
task := createTestTask()
b.ResetTimer()
b.ReportAllocs()
for i := 0; i < b.N; i++ {
data, err := json.Marshal(task)
if err != nil {
b.Fatal(err)
}
var t queue.Task
err = json.Unmarshal(data, &t)
if err != nil {
b.Fatal(err)
}
}
}
// BenchmarkPrewarmStateJSONMarshal profiles PrewarmState serialization
// Used in worker prewarm coordination
func BenchmarkPrewarmStateJSONMarshal(b *testing.B) {
state := queue.PrewarmState{
WorkerID: "worker-12345",
TaskID: "task-67890",
SnapshotID: "snap-abc123",
StartedAt: time.Now().UTC().Format(time.RFC3339),
UpdatedAt: time.Now().UTC().Format(time.RFC3339),
Phase: "running",
EnvImage: "fetchml/python:3.11",
DatasetCnt: 3,
EnvHit: 5,
EnvMiss: 1,
EnvBuilt: 2,
EnvTimeNs: 1500000000,
}
b.ResetTimer()
b.ReportAllocs()
for i := 0; i < b.N; i++ {
_, err := json.Marshal(state)
if err != nil {
b.Fatal(err)
}
}
}
// BenchmarkPrewarmStateJSONUnmarshal profiles PrewarmState deserialization
func BenchmarkPrewarmStateJSONUnmarshal(b *testing.B) {
state := queue.PrewarmState{
WorkerID: "worker-12345",
TaskID: "task-67890",
SnapshotID: "snap-abc123",
StartedAt: time.Now().UTC().Format(time.RFC3339),
UpdatedAt: time.Now().UTC().Format(time.RFC3339),
Phase: "running",
EnvImage: "fetchml/python:3.11",
DatasetCnt: 3,
EnvHit: 5,
EnvMiss: 1,
EnvBuilt: 2,
EnvTimeNs: 1500000000,
}
data, err := json.Marshal(state)
if err != nil {
b.Fatal(err)
}
b.ResetTimer()
b.ReportAllocs()
for i := 0; i < b.N; i++ {
var s queue.PrewarmState
err := json.Unmarshal(data, &s)
if err != nil {
b.Fatal(err)
}
}
}
// BenchmarkTaskBatchJSONMarshal profiles batch task serialization
// Critical for queue index operations with many tasks
func BenchmarkTaskBatchJSONMarshal(b *testing.B) {
tasks := make([]queue.Task, 100)
for i := 0; i < 100; i++ {
tasks[i] = createTestTaskWithID(i)
}
b.ResetTimer()
b.ReportAllocs()
for i := 0; i < b.N; i++ {
for _, task := range tasks {
_, err := json.Marshal(task)
if err != nil {
b.Fatal(err)
}
}
}
}
// Helper to create a test task
func createTestTask() queue.Task {
return createTestTaskWithID(0)
}
func createTestTaskWithID(id int) queue.Task {
now := time.Now()
return queue.Task{
ID: fmt.Sprintf("task-550e8400-e29b-41d4-a716-44665544%03d", id),
JobName: "ml-training-job",
Args: "--epochs=100 --batch_size=32 --learning_rate=0.001",
Status: "queued",
Priority: 1000,
CreatedAt: now,
DatasetSpecs: []queue.DatasetSpec{
{Name: "training-data", Version: "v1.2.3", Checksum: "sha256:abc123", URI: "s3://bucket/data/train"},
{Name: "validation-data", Version: "v1.0.0", Checksum: "sha256:def456", URI: "s3://bucket/data/val"},
},
Datasets: []string{"dataset1", "dataset2"},
Metadata: map[string]string{
"experiment_id": "exp-12345",
"user_email": "user@example.com",
"model_type": "transformer",
},
CPU: 4,
MemoryGB: 16,
GPU: 1,
GPUMemory: "24GB",
UserID: "user-550e8400-e29b-41d4-a716",
Username: "ml_researcher",
CreatedBy: "admin",
MaxRetries: 3,
}
}