Skip to content

Commit cc34fa9

Browse files
evacchilioraron
andauthored
feat: e2e dispatcher integration (llm-d#458)
* feat: e2e dispatcher integration Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * bump llm-d-async from v0.7.0-rc.6 to v0.7.2 Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * address comments Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * address comments Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * address comments Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * fix: ensure per-queue config is used Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * consolidate testing scripts Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * enable tests on CI Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * address comments Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> --------- Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> Co-authored-by: Lior Aronovich <lioraronpr@gmail.com>
1 parent 2f340d4 commit cc34fa9

32 files changed

Lines changed: 2821 additions & 208 deletions

.github/workflows/ci-integration-tests.yml

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ permissions:
1919

2020
jobs:
2121
integration-tests:
22-
name: Integration Tests (${{ matrix.file_client_type }}, ${{ matrix.db_client_type }}, ${{ matrix.exchange_client_type }}${{ matrix.enable_gie == 'true' && ', gie' || '' }})
22+
name: Integration Tests (${{ matrix.file_client_type }}, ${{ matrix.db_client_type }}, ${{ matrix.exchange_client_type }}${{ matrix.enable_gie == 'true' && ', gie' || '' }}${{ matrix.enable_dispatcher == 'true' && ', dispatcher' || '' }})
2323
runs-on: ubuntu-latest
2424
timeout-minutes: 30
2525
strategy:
@@ -45,6 +45,10 @@ jobs:
4545
db_client_type: postgresql
4646
exchange_client_type: redis
4747
enable_gie: "true"
48+
- file_client_type: s3
49+
db_client_type: postgresql
50+
exchange_client_type: redis
51+
enable_dispatcher: "true"
4852
steps:
4953
- uses: actions/checkout@v7
5054

@@ -65,6 +69,7 @@ jobs:
6569
DB_CLIENT_TYPE: ${{ matrix.db_client_type }}
6670
EXCHANGE_CLIENT_TYPE: ${{ matrix.exchange_client_type }}
6771
ENABLE_GIE: ${{ matrix.enable_gie || 'false' }}
72+
ENABLE_DISPATCHER: ${{ matrix.enable_dispatcher || 'false' }}
6873

6974
- name: Run integration tests
7075
run: make test-e2e

.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
bin/
2+
.build/
23
.vscode/
34
.claude/
45
.cursor/
@@ -8,3 +9,5 @@ __pycache__/
89
*.pyc
910
benchmarks/results/*
1011
!benchmarks/results/.gitkeep
12+
.dispatcher-port-forward.pid
13+
.dispatcher-sim-port-forward.pid

charts/batch-gateway/templates/processor-configmap.yaml

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,12 +42,23 @@ data:
4242
conn_max_idle_time: {{ .Values.global.dbClient.redis.connMaxIdleTime | quote }}
4343
conn_max_lifetime: {{ .Values.global.dbClient.redis.connMaxLifetime | quote }}
4444
45+
{{- if .Values.processor.config.dispatchMode }}
46+
dispatch_mode: {{ .Values.processor.config.dispatchMode | quote }}
47+
{{- end }}
48+
{{- if .Values.processor.config.asyncDispatch }}
49+
async_dispatch:
50+
result_poll_timeout: {{ .Values.processor.config.asyncDispatch.resultPollTimeout | quote }}
51+
{{- end }}
52+
4553
{{- with .Values.processor.config.globalInferenceGateway }}
4654
global_inference_gateway:
4755
url: {{ .url | quote }}
4856
{{- if .inferenceObjective }}
4957
inference_objective: {{ .inferenceObjective | quote }}
5058
{{- end }}
59+
{{- if .inferencePoolName }}
60+
inference_pool_name: {{ .inferencePoolName | quote }}
61+
{{- end }}
5162
request_timeout: {{ .requestTimeout | quote }}
5263
max_retries: {{ .maxRetries }}
5364
initial_backoff: {{ .initialBackoff | quote }}
@@ -74,15 +85,20 @@ data:
7485
model_gateways:
7586
{{- range $model, $cfg := .Values.processor.config.modelGateways }}
7687
{{ $model | quote }}:
77-
url: {{ $cfg.url | quote }}
78-
{{- if $cfg.inferenceObjective }}
79-
inference_objective: {{ $cfg.inferenceObjective | quote }}
88+
{{- if $cfg.inferencePoolName }}
89+
inference_pool_name: {{ $cfg.inferencePoolName | quote }}
8090
{{- end }}
91+
{{- if $cfg.url }}
92+
url: {{ $cfg.url | quote }}
8193
request_timeout: {{ $cfg.requestTimeout | quote }}
8294
max_retries: {{ $cfg.maxRetries }}
8395
initial_backoff: {{ $cfg.initialBackoff | quote }}
8496
max_backoff: {{ $cfg.maxBackoff | quote }}
8597
tls_insecure_skip_verify: {{ $cfg.tlsInsecureSkipVerify | default false }}
98+
{{- end }}
99+
{{- if $cfg.inferenceObjective }}
100+
inference_objective: {{ $cfg.inferenceObjective | quote }}
101+
{{- end }}
86102
{{- if $cfg.apiKeyName }}
87103
api_key_name: {{ $cfg.apiKeyName | quote }}
88104
{{- end }}

cmd/batch-processor/main.go

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -296,19 +296,29 @@ func buildProcessorClients(ctx context.Context, cfg *config.ProcessorConfig) (*c
296296
if len(resolved.PerModel) > 0 {
297297
opts = append(opts, clientset.WithPerModelInference(resolved.PerModel))
298298
}
299+
if resolved.Async != nil {
300+
opts = append(opts, clientset.WithAsyncInference(*resolved.Async))
301+
}
299302
clients, err := clientset.NewClientset(ctx, ucom.ComponentProcessor, opts...)
300303
if err != nil {
301304
logger.Error(err, "Failed to create clients")
302305
return nil, err
303306
}
304307

305-
// Validate() guarantees exactly one of resolved.Global or resolved.PerModel is set.
306-
if resolved.Global != nil {
308+
// ResolveModelGateways populates exactly one of Async, Global, or PerModel
309+
// based on the dispatch mode validated by Validate().
310+
switch {
311+
case resolved.Async != nil:
312+
logger.V(logging.INFO).Info("Processor clients initialized",
313+
"mode", "async",
314+
"numModels", len(resolved.Async.Models),
315+
"fileClientType", cfg.FileClientCfg.Type)
316+
case resolved.Global != nil:
307317
logger.V(logging.INFO).Info("Processor clients initialized",
308318
"mode", "global",
309319
"gatewayURL", resolved.Global.URL,
310320
"fileClientType", cfg.FileClientCfg.Type)
311-
} else {
321+
default:
312322
logger.V(logging.INFO).Info("Processor clients initialized",
313323
"mode", "per-model",
314324
"numModelGateways", len(resolved.PerModel),

go.mod

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ require (
1414
github.com/go-resty/resty/v2 v2.17.2
1515
github.com/google/uuid v1.6.0
1616
github.com/jackc/pgx/v5 v5.10.0
17+
github.com/llm-d-incubation/llm-d-async/api v0.7.2
18+
github.com/llm-d-incubation/llm-d-async/producer v0.7.2
1719
github.com/pashagolub/pgxmock/v4 v4.9.0
1820
github.com/prometheus/client_golang v1.23.2
1921
github.com/quasilyte/go-ruleguard/dsl v0.3.23

go.sum

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,10 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
8888
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
8989
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
9090
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
91+
github.com/llm-d-incubation/llm-d-async/api v0.7.2 h1:sf6iFDa5LpVoKDYiOOlsbnv+6Ykj5TA22yRqMdIbOfY=
92+
github.com/llm-d-incubation/llm-d-async/api v0.7.2/go.mod h1:m2zJUwD/AZypJv8RPws3uvCnfQoNW7iv+hH9wPk/Xw4=
93+
github.com/llm-d-incubation/llm-d-async/producer v0.7.2 h1:hhnLqi+MXq8kBjsQhnUCWYtTCG6uFTZLCH296CZMlpY=
94+
github.com/llm-d-incubation/llm-d-async/producer v0.7.2/go.mod h1:IrClO4XwMtgLRnlkAqO1LDsuqmwtQjELDabT2f+5d40=
9195
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
9296
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
9397
github.com/pashagolub/pgxmock/v4 v4.9.0 h1:itlO8nrVRnzkdMBXLs8pWUyyB2PC3Gku0WGIj/gGl7I=

internal/processor/config/config.go

Lines changed: 26 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -230,18 +230,6 @@ type BucketConfig struct {
230230
BucketCount int `yaml:"count"`
231231
}
232232

233-
const asyncTenantID = "$batch"
234-
235-
// RequestQueueName returns the Redis sorted-set name for submitting async requests to the given pool.
236-
func RequestQueueName(poolName string) string {
237-
return "llm-d-async:requests:" + poolName
238-
}
239-
240-
// ResultQueueName returns the Redis list name for collecting async results from the given pool.
241-
func ResultQueueName(poolName string) string {
242-
return "llm-d-async:results:" + poolName + ":" + asyncTenantID
243-
}
244-
245233
// IsAsync returns true when the processor is configured for async dispatch.
246234
func (c *ProcessorConfig) IsAsync() bool {
247235
return c.DispatchMode == DispatchModeAsync
@@ -386,14 +374,21 @@ func (c *ProcessorConfig) Validate() error {
386374
return fmt.Errorf("progress_ttl_seconds must be > 0")
387375
}
388376

389-
if err := c.validateDispatchMode(); err != nil {
377+
if err := c.validateGateways(); err != nil {
390378
return err
391379
}
392380

393381
return nil
394382
}
395383

396-
func (c *ProcessorConfig) validateDispatchMode() error {
384+
func (c *ProcessorConfig) validateGateways() error {
385+
if c.GlobalInferenceGateway == nil && len(c.ModelGateways) == 0 {
386+
return fmt.Errorf("either global_inference_gateway or model_gateways must be configured")
387+
}
388+
if c.GlobalInferenceGateway != nil && len(c.ModelGateways) > 0 {
389+
return fmt.Errorf("global_inference_gateway and model_gateways are mutually exclusive")
390+
}
391+
397392
switch c.DispatchMode {
398393
case DispatchModeSync, DispatchMode(""):
399394
c.DispatchMode = DispatchModeSync
@@ -405,21 +400,7 @@ func (c *ProcessorConfig) validateDispatchMode() error {
405400
}
406401
}
407402

408-
func (c *ProcessorConfig) validateGateways() error {
409-
if c.GlobalInferenceGateway == nil && len(c.ModelGateways) == 0 {
410-
return fmt.Errorf("either global_inference_gateway or model_gateways must be configured")
411-
}
412-
if c.GlobalInferenceGateway != nil && len(c.ModelGateways) > 0 {
413-
return fmt.Errorf("global_inference_gateway and model_gateways are mutually exclusive")
414-
}
415-
return nil
416-
}
417-
418403
func (c *ProcessorConfig) validateSyncDispatchConfig() error {
419-
if err := c.validateGateways(); err != nil {
420-
return err
421-
}
422-
423404
if c.GlobalInferenceGateway != nil {
424405
if err := validateGatewayConfig("global_inference_gateway", *c.GlobalInferenceGateway); err != nil {
425406
return err
@@ -435,15 +416,13 @@ func (c *ProcessorConfig) validateSyncDispatchConfig() error {
435416

436417
func (c *ProcessorConfig) validateAsyncDispatchConfig() error {
437418
if c.AsyncDispatchConfig.ResultPollTimeout <= 0 {
438-
return fmt.Errorf("async.result_poll_timeout must be > 0")
439-
}
440-
if err := c.validateGateways(); err != nil {
441-
return err
419+
return fmt.Errorf("async_dispatch.result_poll_timeout must be > 0")
442420
}
443421
if c.GlobalInferenceGateway != nil {
444-
if c.GlobalInferenceGateway.InferencePoolName == "" {
445-
return fmt.Errorf("global_inference_gateway.inference_pool_name must be set when dispatch_mode is %q", DispatchModeAsync)
446-
}
422+
return fmt.Errorf("global_inference_gateway is not supported with dispatch_mode %q; use model_gateways with inference_pool_name", DispatchModeAsync)
423+
}
424+
if len(c.ModelGateways) == 0 {
425+
return fmt.Errorf("model_gateways must be configured when dispatch_mode is %q", DispatchModeAsync)
447426
}
448427
for model, gw := range c.ModelGateways {
449428
if gw.InferencePoolName == "" {
@@ -580,6 +559,7 @@ func toGatewayClientConfig(gw ModelGatewayConfig, apiKey string) inference.Gatew
580559
type ResolvedGateways struct {
581560
Global *inference.GatewayClientConfig
582561
PerModel map[string]inference.GatewayClientConfig
562+
Async *inference.AsyncClientConfig
583563
}
584564

585565
// ResolveModelGateways resolves API keys for all configured gateways and returns
@@ -588,6 +568,17 @@ type ResolvedGateways struct {
588568
func ResolveModelGateways(cfg *ProcessorConfig) (*ResolvedGateways, error) {
589569
result := &ResolvedGateways{}
590570

571+
if cfg.IsAsync() {
572+
models := make(map[string]string, len(cfg.ModelGateways))
573+
for model, gw := range cfg.ModelGateways {
574+
models[model] = gw.InferencePoolName
575+
}
576+
result.Async = &inference.AsyncClientConfig{
577+
Models: models,
578+
}
579+
return result, nil
580+
}
581+
591582
if cfg.GlobalInferenceGateway != nil {
592583
apiKey, err := resolveGatewayAPIKey("global_inference_gateway", *cfg.GlobalInferenceGateway)
593584
if err != nil {

internal/processor/config/config_test.go

Lines changed: 57 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -715,12 +715,12 @@ func TestProcessorConfig_Validate_AsyncDispatch(t *testing.T) {
715715
wantErr: true,
716716
},
717717
{
718-
name: "async valid global gateway with inference_pool_name",
718+
name: "async global gateway rejected",
719719
mutate: func(c *ProcessorConfig) {
720720
c.ModelGateways = nil
721721
c.GlobalInferenceGateway = &ModelGatewayConfig{URL: "http://gw:8000", InferencePoolName: "default-pool"}
722722
},
723-
wantErr: false,
723+
wantErr: true,
724724
},
725725
{
726726
name: "async no gateways configured",
@@ -784,15 +784,6 @@ func TestProcessorConfig_Validate_AsyncDispatch(t *testing.T) {
784784
}
785785
}
786786

787-
func TestQueueNameHelpers(t *testing.T) {
788-
if got := RequestQueueName("pool-a"); got != "llm-d-async:requests:pool-a" {
789-
t.Fatalf("RequestQueueName(\"pool-a\") = %q, want %q", got, "llm-d-async:requests:pool-a")
790-
}
791-
if got := ResultQueueName("pool-a"); got != "llm-d-async:results:pool-a:$batch" {
792-
t.Fatalf("ResultQueueName(\"pool-a\") = %q, want %q", got, "llm-d-async:results:pool-a:$batch")
793-
}
794-
}
795-
796787
func TestValidate_NormalizesEmptyDispatchMode(t *testing.T) {
797788
c := NewConfig()
798789
c.ModelGateways = validPerModelConfig()
@@ -919,3 +910,58 @@ func TestProcessorConfig_InferenceObjectiveFor(t *testing.T) {
919910
})
920911
}
921912
}
913+
914+
func TestResolveModelGateways_Async(t *testing.T) {
915+
t.Run("populates Async field", func(t *testing.T) {
916+
cfg := NewConfig()
917+
cfg.DispatchMode = DispatchModeAsync
918+
cfg.AsyncDispatchConfig = AsyncDispatchConfig{
919+
ResultPollTimeout: 10 * time.Second,
920+
}
921+
cfg.ModelGateways = map[string]ModelGatewayConfig{
922+
"model-a": {InferencePoolName: "pool-a"},
923+
"model-b": {InferencePoolName: "pool-b"},
924+
}
925+
926+
resolved, err := ResolveModelGateways(cfg)
927+
if err != nil {
928+
t.Fatalf("ResolveModelGateways() error: %v", err)
929+
}
930+
931+
if resolved.Async == nil {
932+
t.Fatal("expected Async to be set")
933+
}
934+
if resolved.Global != nil {
935+
t.Error("expected Global to be nil in async mode")
936+
}
937+
if resolved.PerModel != nil {
938+
t.Error("expected PerModel to be nil in async mode")
939+
}
940+
if len(resolved.Async.Models) != 2 {
941+
t.Fatalf("Models count = %d, want 2", len(resolved.Async.Models))
942+
}
943+
if resolved.Async.Models["model-a"] != "pool-a" {
944+
t.Errorf("Models[model-a] = %q, want %q", resolved.Async.Models["model-a"], "pool-a")
945+
}
946+
if resolved.Async.Models["model-b"] != "pool-b" {
947+
t.Errorf("Models[model-b] = %q, want %q", resolved.Async.Models["model-b"], "pool-b")
948+
}
949+
})
950+
951+
t.Run("sync mode does not populate Async", func(t *testing.T) {
952+
cfg := NewConfig()
953+
cfg.ModelGateways = validPerModelConfig()
954+
955+
resolved, err := ResolveModelGateways(cfg)
956+
if err != nil {
957+
t.Fatalf("ResolveModelGateways() error: %v", err)
958+
}
959+
960+
if resolved.Async != nil {
961+
t.Error("expected Async to be nil in sync mode")
962+
}
963+
if len(resolved.PerModel) == 0 {
964+
t.Error("expected PerModel to be populated in sync mode")
965+
}
966+
})
967+
}

0 commit comments

Comments
 (0)