Skip to content

Commit b8dd6f2

Browse files
tlrmchlsmthtlrmchlsmth
committed
Fix staged rate scheduling and completion controls
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
1 parent dadcb54 commit b8dd6f2

3 files changed

Lines changed: 231 additions & 32 deletions

File tree

‎pkg/loadgen/generator.go‎

Lines changed: 77 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,12 @@ type Generator struct {
6969
Metrics *metrics.Metrics // Optional Prometheus metrics (nil = disabled)
7070
StreamUsage bool // Request token usage stats from server (stream_options)
7171

72-
recorderPtr atomic.Pointer[recorder.Recorder] // swappable recorder for warmup→main transition
73-
recordWG sync.WaitGroup // tracks in-flight recordResult goroutines
74-
inFlight atomic.Int64
75-
evalCount atomic.Int64
76-
evalCorrect atomic.Int64
72+
recorderPtr atomic.Pointer[recorder.Recorder] // swappable recorder for warmup→main transition
73+
recordWG sync.WaitGroup // tracks in-flight recordResult goroutines
74+
inFlight atomic.Int64
75+
evalCount atomic.Int64
76+
evalCorrect atomic.Int64
77+
rateStreamID atomic.Int64 // unique across rate stages, including warmup
7778

7879
maxReqState atomic.Pointer[maxRequestsState]
7980
consecutiveErrors atomic.Int64
@@ -169,7 +170,7 @@ type Stage struct {
169170
Concurrency int
170171
ConversationPoolSize int
171172
Duration time.Duration
172-
Rampup time.Duration // stagger new stream starts over this duration
173+
Rampup time.Duration // stagger streams or ramp arrival rate over this duration
173174
MaxRequests int // stop after this many requests (0 = unlimited)
174175
Barrier bool // sync point — pool stays alive (unless BarrierDrain), onBarrier fires
175176
BarrierDrain bool // stop pool before sync, fresh pool after
@@ -191,6 +192,10 @@ func (g *Generator) RunStagesUntil(ctx context.Context, stages []Stage, onStage
191192
g.runConversationPoolStages(ctx, stages, onStage, onBarrier, onStageComplete)
192193
return
193194
}
195+
if g.Mode == ModeConstant || g.Mode == ModePoisson {
196+
g.runRateStages(ctx, stages, onStage, onBarrier, onStageComplete)
197+
return
198+
}
194199

195200
ctx, cancel := context.WithCancel(ctx)
196201
defer cancel()
@@ -317,43 +322,74 @@ func (g *Generator) runConcurrent(ctx context.Context, c *client.Client) {
317322

318323
// runRateBasedConstant sends requests at evenly-spaced intervals.
319324
func (g *Generator) runRateBasedConstant(ctx context.Context, c *client.Client, startTime time.Time) {
320-
g.runRateBased(ctx, c, startTime, false)
325+
g.runRateBased(ctx, c, startTime, false, g.Rampup, 0)
321326
}
322327

323328
// runRateBasedPoisson sends requests with exponentially-distributed inter-arrival times.
324329
func (g *Generator) runRateBasedPoisson(ctx context.Context, c *client.Client, startTime time.Time) {
325-
g.runRateBased(ctx, c, startTime, true)
330+
g.runRateBased(ctx, c, startTime, true, g.Rampup, 0)
331+
}
332+
333+
func (g *Generator) runRateStages(ctx context.Context, stages []Stage, onStage func(int, int), onBarrier func(int), onStageComplete func(int) bool) {
334+
ctx, cancel := context.WithCancel(ctx)
335+
defer cancel()
336+
g.stopFunc = cancel
337+
c := client.New(g.Target)
338+
339+
for i, stage := range stages {
340+
if ctx.Err() != nil {
341+
break
342+
}
343+
// Rate stages drain before advancing, so both barrier types are drained.
344+
if stage.Barrier {
345+
if onBarrier != nil {
346+
onBarrier(i)
347+
}
348+
continue
349+
}
350+
if onStage != nil {
351+
onStage(i, stage.Concurrency)
352+
}
353+
startTime := time.Now()
354+
stageCtx, stageCancel := context.WithTimeout(ctx, stage.Duration)
355+
g.runRateBased(stageCtx, c, startTime, g.Mode == ModePoisson, stage.Rampup, stage.MaxRequests)
356+
stageCancel()
357+
g.recordWG.Wait()
358+
if ctx.Err() == nil && onStageComplete != nil && !onStageComplete(i) {
359+
break
360+
}
361+
}
326362
}
327363

328-
func (g *Generator) runRateBased(ctx context.Context, c *client.Client, startTime time.Time, poisson bool) {
364+
func (g *Generator) runRateBased(ctx context.Context, c *client.Client, startTime time.Time, poisson bool, rampup time.Duration, maxRequests int) {
329365
var sem chan struct{}
330366
if g.MaxInFlight > 0 {
331367
sem = make(chan struct{}, g.MaxInFlight)
332368
}
333369

334370
var wg sync.WaitGroup
335-
streamID := 0
371+
dispatched := 0
336372

337373
for {
338-
if ctx.Err() != nil {
374+
if ctx.Err() != nil || (maxRequests > 0 && dispatched >= maxRequests) {
339375
break
340376
}
341-
342-
// Compute next arrival time
343-
elapsed := time.Since(startTime).Seconds()
344-
rate := g.rate(elapsed)
345-
if rate <= 0 {
346-
time.Sleep(10 * time.Millisecond)
377+
if g.Rate <= 0 {
378+
select {
379+
case <-ctx.Done():
380+
case <-time.After(10 * time.Millisecond):
381+
}
347382
continue
348383
}
349384

350-
var gap time.Duration
385+
// Compute next arrival time
386+
elapsed := time.Since(startTime).Seconds()
387+
arrival := 1.0
351388
if poisson {
352389
// Exponential inter-arrival time
353-
gap = time.Duration(float64(time.Second) * (-math.Log(1-mathrand.Float64()) / rate))
354-
} else {
355-
gap = time.Duration(float64(time.Second) / rate)
390+
arrival = -math.Log(1 - mathrand.Float64())
356391
}
392+
gap := g.arrivalGap(elapsed, rampup, arrival)
357393

358394
select {
359395
case <-ctx.Done():
@@ -378,9 +414,9 @@ func (g *Generator) runRateBased(ctx context.Context, c *client.Client, startTim
378414
}
379415

380416
conversation := g.Dataset.NextConversation()
381-
convID := fmt.Sprintf("w%d-c%d", streamID, streamID)
382-
sid := streamID
383-
streamID++
417+
sid := int(g.rateStreamID.Add(1) - 1)
418+
convID := fmt.Sprintf("w%d-c%d", sid, sid)
419+
dispatched++
384420

385421
wg.Add(1)
386422
go func() {
@@ -395,14 +431,24 @@ func (g *Generator) runRateBased(ctx context.Context, c *client.Client, startTim
395431
wg.Wait()
396432
}
397433

398-
// rate returns the effective request rate at a given elapsed time,
399-
// accounting for linear rampup.
400-
func (g *Generator) rate(elapsed float64) float64 {
401-
if g.Rampup.Seconds() <= 0 || elapsed >= g.Rampup.Seconds() {
402-
return g.Rate
434+
// arrivalGap integrates the linear ramp so the first arrival can be scheduled
435+
// even when the instantaneous rate starts at zero. arrival is one for constant
436+
// traffic or an exponential sample for Poisson traffic.
437+
func (g *Generator) arrivalGap(elapsed float64, rampup time.Duration, arrival float64) time.Duration {
438+
if g.Rate <= 0 {
439+
return 10 * time.Millisecond
440+
}
441+
rampSeconds := rampup.Seconds()
442+
if rampSeconds > 0 && elapsed < rampSeconds {
443+
slope := g.Rate / rampSeconds
444+
remaining := slope * (rampSeconds*rampSeconds - elapsed*elapsed) / 2
445+
if arrival <= remaining {
446+
gap := math.Sqrt(elapsed*elapsed+2*arrival/slope) - elapsed
447+
return time.Duration(gap * float64(time.Second))
448+
}
449+
return time.Duration((rampSeconds - elapsed + (arrival-remaining)/g.Rate) * float64(time.Second))
403450
}
404-
// Linear ramp from 0 to target rate
405-
return g.Rate * (elapsed / g.Rampup.Seconds())
451+
return time.Duration(arrival / g.Rate * float64(time.Second))
406452
}
407453

408454
func (g *Generator) runStream(ctx context.Context, c *client.Client, streamID int, delay time.Duration) {

‎pkg/loadgen/generator_test.go‎

Lines changed: 124 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -558,13 +558,14 @@ func TestRunStagesPoolResize(t *testing.T) {
558558
}
559559

560560
func TestRunStagesUntilStopsBeforeNextStage(t *testing.T) {
561-
for _, mode := range []loadgen.Mode{loadgen.ModeConcurrent, loadgen.ModeConversationPool} {
561+
for _, mode := range []loadgen.Mode{loadgen.ModeConcurrent, loadgen.ModeConversationPool, loadgen.ModeConstant, loadgen.ModePoisson} {
562562
t.Run(string(mode), func(t *testing.T) {
563563
addr := startMockServer(t)
564564
gen := &loadgen.Generator{
565565
Target: "http://" + addr + "/v1",
566566
Model: "test-model",
567567
Mode: mode,
568+
Rate: 100,
568569
Dataset: dataset.NewSynthetic(32, 10, 1, 4.0),
569570
Recorder: recorder.NewMemory(),
570571
}
@@ -586,6 +587,128 @@ func TestRunStagesUntilStopsBeforeNextStage(t *testing.T) {
586587
}
587588
}
588589

590+
func TestRateStagesLimitsDrainAndUniqueIDs(t *testing.T) {
591+
for _, mode := range []loadgen.Mode{loadgen.ModeConstant, loadgen.ModePoisson} {
592+
t.Run(string(mode), func(t *testing.T) {
593+
addr := startMockServer(t)
594+
rec := recorder.NewMemory()
595+
defer rec.Close()
596+
gen := &loadgen.Generator{
597+
Target: "http://" + addr + "/v1", Model: "test-model",
598+
Mode: mode, Rate: 1000, MaxInFlight: 1,
599+
// Each stage's zero rampup must override this generator default.
600+
Rampup: time.Hour,
601+
Dataset: dataset.NewSynthetic(32, 10, 2, 4.0), Recorder: rec,
602+
}
603+
stages := []loadgen.Stage{
604+
{Duration: 5 * time.Second, MaxRequests: 3},
605+
{Barrier: true, BarrierDrain: true},
606+
{Duration: 5 * time.Second, MaxRequests: 4},
607+
}
608+
var started, completed, barriers []int
609+
gen.RunStagesUntil(context.Background(), stages, func(i, _ int) {
610+
started = append(started, i)
611+
}, func(i int) {
612+
barriers = append(barriers, i)
613+
}, func(i int) bool {
614+
completed = append(completed, i)
615+
if gen.InFlight() != 0 {
616+
t.Error("rate stage did not drain before completion callback")
617+
}
618+
return true
619+
})
620+
if !reflect.DeepEqual(started, []int{0, 2}) || !reflect.DeepEqual(completed, []int{0, 2}) || !reflect.DeepEqual(barriers, []int{1}) {
621+
t.Fatalf("started=%v completed=%v barriers=%v", started, completed, barriers)
622+
}
623+
rec.Close()
624+
records := rec.Records()
625+
// Like concurrent mode, the arrival limit counts conversations;
626+
// every admitted conversation must finish all of its turns.
627+
if len(records) != 14 {
628+
t.Fatalf("got %d records, want 14 from seven two-turn conversations", len(records))
629+
}
630+
conversations := make(map[string]int)
631+
requests := make(map[string]bool)
632+
for _, r := range records {
633+
if r.Status != "ok" {
634+
t.Errorf("request %s failed: %s", r.RequestID, r.Error)
635+
}
636+
if requests[r.RequestID] {
637+
t.Errorf("duplicate request ID %s", r.RequestID)
638+
}
639+
requests[r.RequestID] = true
640+
conversations[r.ConversationID]++
641+
}
642+
if len(conversations) != 7 {
643+
t.Fatalf("got %d conversations, want 7", len(conversations))
644+
}
645+
for id, turns := range conversations {
646+
if turns != 2 {
647+
t.Errorf("conversation %s has %d turns, want 2", id, turns)
648+
}
649+
}
650+
})
651+
}
652+
}
653+
654+
func TestRateStageRampup(t *testing.T) {
655+
addr := startMockServer(t)
656+
rec := recorder.NewMemory()
657+
defer rec.Close()
658+
gen := &loadgen.Generator{
659+
Target: "http://" + addr + "/v1", Model: "test-model",
660+
Mode: loadgen.ModeConstant, Rate: 100,
661+
Dataset: dataset.NewSynthetic(32, 10, 1, 4.0), Recorder: rec,
662+
}
663+
var starts []time.Time
664+
gen.RunStages(context.Background(), []loadgen.Stage{
665+
{Duration: time.Second, Rampup: 500 * time.Millisecond, MaxRequests: 1},
666+
{Duration: time.Second, Rampup: 500 * time.Millisecond, MaxRequests: 1},
667+
}, func(_, _ int) { starts = append(starts, time.Now()) }, nil)
668+
rec.Close()
669+
records := rec.Records()
670+
if len(records) != 2 {
671+
t.Fatalf("got %d records, want one per stage", len(records))
672+
}
673+
sort.Slice(records, func(i, j int) bool { return records[i].StartTime < records[j].StartTime })
674+
for i, r := range records {
675+
delay := r.StartTime - recorder.TimeToFloat(starts[i])
676+
// Integrated rate gives a first arrival at 100ms; full rate gives 10ms.
677+
if delay < 0.075 {
678+
t.Errorf("stage %d first arrival after %.3fs, rampup ignored", i, delay)
679+
}
680+
}
681+
}
682+
683+
func TestRateStageDuration(t *testing.T) {
684+
for _, mode := range []loadgen.Mode{loadgen.ModeConstant, loadgen.ModePoisson} {
685+
for _, limit := range []int{0, 10000} {
686+
t.Run(fmt.Sprintf("%s/limit=%d", mode, limit), func(t *testing.T) {
687+
addr := startMockServer(t)
688+
rec := recorder.NewMemory()
689+
defer rec.Close()
690+
gen := &loadgen.Generator{
691+
Target: "http://" + addr + "/v1", Model: "test-model",
692+
Mode: mode, Rate: 100,
693+
Dataset: dataset.NewSynthetic(32, 10, 1, 4.0), Recorder: rec,
694+
}
695+
start := time.Now()
696+
gen.RunStages(context.Background(), []loadgen.Stage{{
697+
Duration: 200 * time.Millisecond, MaxRequests: limit,
698+
}}, nil, nil)
699+
elapsed := time.Since(start)
700+
rec.Close()
701+
if len(rec.Records()) == 0 {
702+
t.Fatal("expected rate arrivals before the stage deadline")
703+
}
704+
if elapsed < 180*time.Millisecond || elapsed > 2*time.Second {
705+
t.Errorf("stage lasted %v, want approximately 200ms", elapsed)
706+
}
707+
})
708+
}
709+
}
710+
}
711+
589712
func TestCacheSaltAppearsInRequest(t *testing.T) {
590713
// Start a tiny HTTP server that captures the request body
591714
var bodies []string

‎pkg/loadgen/rate_test.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package loadgen
2+
3+
import (
4+
"testing"
5+
"time"
6+
)
7+
8+
func TestArrivalGapLinearRamp(t *testing.T) {
9+
g := &Generator{Rate: 100}
10+
for _, tc := range []struct {
11+
name string
12+
elapsed float64
13+
rampup time.Duration
14+
arrival float64
15+
want time.Duration
16+
}{
17+
{"no ramp", 0, 0, 1, 10 * time.Millisecond},
18+
{"ramp start", 0, time.Second, 0.5, 100 * time.Millisecond},
19+
{"within ramp", 0.3, time.Second, 3.5, 100 * time.Millisecond},
20+
{"across ramp end", 0.9, time.Second, 19.5, 200 * time.Millisecond},
21+
{"after ramp", 2, time.Second, 2, 20 * time.Millisecond},
22+
} {
23+
t.Run(tc.name, func(t *testing.T) {
24+
got := g.arrivalGap(tc.elapsed, tc.rampup, tc.arrival)
25+
if delta := got - tc.want; delta < -time.Nanosecond || delta > time.Nanosecond {
26+
t.Errorf("arrival gap = %v, want %v", got, tc.want)
27+
}
28+
})
29+
}
30+
}

0 commit comments

Comments
 (0)