diff --git a/.github/workflows/ci-integration-tests.yml b/.github/workflows/ci-integration-tests.yml index 3c273fe32..02ed3eefc 100644 --- a/.github/workflows/ci-integration-tests.yml +++ b/.github/workflows/ci-integration-tests.yml @@ -19,7 +19,7 @@ permissions: jobs: integration-tests: - name: Integration Tests (${{ matrix.file_client_type }}, ${{ matrix.db_client_type }}, ${{ matrix.exchange_client_type }}${{ matrix.enable_gie == 'true' && ', gie' || '' }}) + 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' || '' }}) runs-on: ubuntu-latest timeout-minutes: 30 strategy: @@ -45,6 +45,10 @@ jobs: db_client_type: postgresql exchange_client_type: redis enable_gie: "true" + - file_client_type: s3 + db_client_type: postgresql + exchange_client_type: redis + enable_dispatcher: "true" steps: - uses: actions/checkout@v7 @@ -65,6 +69,7 @@ jobs: DB_CLIENT_TYPE: ${{ matrix.db_client_type }} EXCHANGE_CLIENT_TYPE: ${{ matrix.exchange_client_type }} ENABLE_GIE: ${{ matrix.enable_gie || 'false' }} + ENABLE_DISPATCHER: ${{ matrix.enable_dispatcher || 'false' }} - name: Run integration tests run: make test-e2e diff --git a/.gitignore b/.gitignore index 45fc4ee1b..25953c7e9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ bin/ +.build/ .vscode/ .claude/ .cursor/ @@ -8,3 +9,5 @@ __pycache__/ *.pyc benchmarks/results/* !benchmarks/results/.gitkeep +.dispatcher-port-forward.pid +.dispatcher-sim-port-forward.pid diff --git a/charts/batch-gateway/templates/processor-configmap.yaml b/charts/batch-gateway/templates/processor-configmap.yaml index 59894708a..ba5ad416f 100644 --- a/charts/batch-gateway/templates/processor-configmap.yaml +++ b/charts/batch-gateway/templates/processor-configmap.yaml @@ -42,12 +42,23 @@ data: conn_max_idle_time: {{ .Values.global.dbClient.redis.connMaxIdleTime | quote }} conn_max_lifetime: {{ .Values.global.dbClient.redis.connMaxLifetime | quote }} + {{- if .Values.processor.config.dispatchMode }} + dispatch_mode: {{ .Values.processor.config.dispatchMode | quote }} + {{- end }} + {{- if .Values.processor.config.asyncDispatch }} + async_dispatch: + result_poll_timeout: {{ .Values.processor.config.asyncDispatch.resultPollTimeout | quote }} + {{- end }} + {{- with .Values.processor.config.globalInferenceGateway }} global_inference_gateway: url: {{ .url | quote }} {{- if .inferenceObjective }} inference_objective: {{ .inferenceObjective | quote }} {{- end }} + {{- if .inferencePoolName }} + inference_pool_name: {{ .inferencePoolName | quote }} + {{- end }} request_timeout: {{ .requestTimeout | quote }} max_retries: {{ .maxRetries }} initial_backoff: {{ .initialBackoff | quote }} @@ -74,15 +85,20 @@ data: model_gateways: {{- range $model, $cfg := .Values.processor.config.modelGateways }} {{ $model | quote }}: - url: {{ $cfg.url | quote }} - {{- if $cfg.inferenceObjective }} - inference_objective: {{ $cfg.inferenceObjective | quote }} + {{- if $cfg.inferencePoolName }} + inference_pool_name: {{ $cfg.inferencePoolName | quote }} {{- end }} + {{- if $cfg.url }} + url: {{ $cfg.url | quote }} request_timeout: {{ $cfg.requestTimeout | quote }} max_retries: {{ $cfg.maxRetries }} initial_backoff: {{ $cfg.initialBackoff | quote }} max_backoff: {{ $cfg.maxBackoff | quote }} tls_insecure_skip_verify: {{ $cfg.tlsInsecureSkipVerify | default false }} + {{- end }} + {{- if $cfg.inferenceObjective }} + inference_objective: {{ $cfg.inferenceObjective | quote }} + {{- end }} {{- if $cfg.apiKeyName }} api_key_name: {{ $cfg.apiKeyName | quote }} {{- end }} diff --git a/cmd/batch-processor/main.go b/cmd/batch-processor/main.go index a030ec667..a2862fdfb 100644 --- a/cmd/batch-processor/main.go +++ b/cmd/batch-processor/main.go @@ -296,19 +296,29 @@ func buildProcessorClients(ctx context.Context, cfg *config.ProcessorConfig) (*c if len(resolved.PerModel) > 0 { opts = append(opts, clientset.WithPerModelInference(resolved.PerModel)) } + if resolved.Async != nil { + opts = append(opts, clientset.WithAsyncInference(*resolved.Async)) + } clients, err := clientset.NewClientset(ctx, ucom.ComponentProcessor, opts...) if err != nil { logger.Error(err, "Failed to create clients") return nil, err } - // Validate() guarantees exactly one of resolved.Global or resolved.PerModel is set. - if resolved.Global != nil { + // ResolveModelGateways populates exactly one of Async, Global, or PerModel + // based on the dispatch mode validated by Validate(). + switch { + case resolved.Async != nil: + logger.V(logging.INFO).Info("Processor clients initialized", + "mode", "async", + "numModels", len(resolved.Async.Models), + "fileClientType", cfg.FileClientCfg.Type) + case resolved.Global != nil: logger.V(logging.INFO).Info("Processor clients initialized", "mode", "global", "gatewayURL", resolved.Global.URL, "fileClientType", cfg.FileClientCfg.Type) - } else { + default: logger.V(logging.INFO).Info("Processor clients initialized", "mode", "per-model", "numModelGateways", len(resolved.PerModel), diff --git a/go.mod b/go.mod index 06a8f9e7f..753eebaa8 100644 --- a/go.mod +++ b/go.mod @@ -14,6 +14,8 @@ require ( github.com/go-resty/resty/v2 v2.17.2 github.com/google/uuid v1.6.0 github.com/jackc/pgx/v5 v5.10.0 + github.com/llm-d-incubation/llm-d-async/api v0.7.2 + github.com/llm-d-incubation/llm-d-async/producer v0.7.2 github.com/pashagolub/pgxmock/v4 v4.9.0 github.com/prometheus/client_golang v1.23.2 github.com/quasilyte/go-ruleguard/dsl v0.3.23 diff --git a/go.sum b/go.sum index 9918f93e5..6085c0834 100644 --- a/go.sum +++ b/go.sum @@ -88,6 +88,10 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/llm-d-incubation/llm-d-async/api v0.7.2 h1:sf6iFDa5LpVoKDYiOOlsbnv+6Ykj5TA22yRqMdIbOfY= +github.com/llm-d-incubation/llm-d-async/api v0.7.2/go.mod h1:m2zJUwD/AZypJv8RPws3uvCnfQoNW7iv+hH9wPk/Xw4= +github.com/llm-d-incubation/llm-d-async/producer v0.7.2 h1:hhnLqi+MXq8kBjsQhnUCWYtTCG6uFTZLCH296CZMlpY= +github.com/llm-d-incubation/llm-d-async/producer v0.7.2/go.mod h1:IrClO4XwMtgLRnlkAqO1LDsuqmwtQjELDabT2f+5d40= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/pashagolub/pgxmock/v4 v4.9.0 h1:itlO8nrVRnzkdMBXLs8pWUyyB2PC3Gku0WGIj/gGl7I= diff --git a/internal/processor/config/config.go b/internal/processor/config/config.go index 17d46fbe5..5e795f6ec 100644 --- a/internal/processor/config/config.go +++ b/internal/processor/config/config.go @@ -230,18 +230,6 @@ type BucketConfig struct { BucketCount int `yaml:"count"` } -const asyncTenantID = "$batch" - -// RequestQueueName returns the Redis sorted-set name for submitting async requests to the given pool. -func RequestQueueName(poolName string) string { - return "llm-d-async:requests:" + poolName -} - -// ResultQueueName returns the Redis list name for collecting async results from the given pool. -func ResultQueueName(poolName string) string { - return "llm-d-async:results:" + poolName + ":" + asyncTenantID -} - // IsAsync returns true when the processor is configured for async dispatch. func (c *ProcessorConfig) IsAsync() bool { return c.DispatchMode == DispatchModeAsync @@ -386,14 +374,21 @@ func (c *ProcessorConfig) Validate() error { return fmt.Errorf("progress_ttl_seconds must be > 0") } - if err := c.validateDispatchMode(); err != nil { + if err := c.validateGateways(); err != nil { return err } return nil } -func (c *ProcessorConfig) validateDispatchMode() error { +func (c *ProcessorConfig) validateGateways() error { + if c.GlobalInferenceGateway == nil && len(c.ModelGateways) == 0 { + return fmt.Errorf("either global_inference_gateway or model_gateways must be configured") + } + if c.GlobalInferenceGateway != nil && len(c.ModelGateways) > 0 { + return fmt.Errorf("global_inference_gateway and model_gateways are mutually exclusive") + } + switch c.DispatchMode { case DispatchModeSync, DispatchMode(""): c.DispatchMode = DispatchModeSync @@ -405,21 +400,7 @@ func (c *ProcessorConfig) validateDispatchMode() error { } } -func (c *ProcessorConfig) validateGateways() error { - if c.GlobalInferenceGateway == nil && len(c.ModelGateways) == 0 { - return fmt.Errorf("either global_inference_gateway or model_gateways must be configured") - } - if c.GlobalInferenceGateway != nil && len(c.ModelGateways) > 0 { - return fmt.Errorf("global_inference_gateway and model_gateways are mutually exclusive") - } - return nil -} - func (c *ProcessorConfig) validateSyncDispatchConfig() error { - if err := c.validateGateways(); err != nil { - return err - } - if c.GlobalInferenceGateway != nil { if err := validateGatewayConfig("global_inference_gateway", *c.GlobalInferenceGateway); err != nil { return err @@ -435,15 +416,13 @@ func (c *ProcessorConfig) validateSyncDispatchConfig() error { func (c *ProcessorConfig) validateAsyncDispatchConfig() error { if c.AsyncDispatchConfig.ResultPollTimeout <= 0 { - return fmt.Errorf("async.result_poll_timeout must be > 0") - } - if err := c.validateGateways(); err != nil { - return err + return fmt.Errorf("async_dispatch.result_poll_timeout must be > 0") } if c.GlobalInferenceGateway != nil { - if c.GlobalInferenceGateway.InferencePoolName == "" { - return fmt.Errorf("global_inference_gateway.inference_pool_name must be set when dispatch_mode is %q", DispatchModeAsync) - } + return fmt.Errorf("global_inference_gateway is not supported with dispatch_mode %q; use model_gateways with inference_pool_name", DispatchModeAsync) + } + if len(c.ModelGateways) == 0 { + return fmt.Errorf("model_gateways must be configured when dispatch_mode is %q", DispatchModeAsync) } for model, gw := range c.ModelGateways { if gw.InferencePoolName == "" { @@ -580,6 +559,7 @@ func toGatewayClientConfig(gw ModelGatewayConfig, apiKey string) inference.Gatew type ResolvedGateways struct { Global *inference.GatewayClientConfig PerModel map[string]inference.GatewayClientConfig + Async *inference.AsyncClientConfig } // ResolveModelGateways resolves API keys for all configured gateways and returns @@ -588,6 +568,17 @@ type ResolvedGateways struct { func ResolveModelGateways(cfg *ProcessorConfig) (*ResolvedGateways, error) { result := &ResolvedGateways{} + if cfg.IsAsync() { + models := make(map[string]string, len(cfg.ModelGateways)) + for model, gw := range cfg.ModelGateways { + models[model] = gw.InferencePoolName + } + result.Async = &inference.AsyncClientConfig{ + Models: models, + } + return result, nil + } + if cfg.GlobalInferenceGateway != nil { apiKey, err := resolveGatewayAPIKey("global_inference_gateway", *cfg.GlobalInferenceGateway) if err != nil { diff --git a/internal/processor/config/config_test.go b/internal/processor/config/config_test.go index c0d5ab136..c4df0b815 100644 --- a/internal/processor/config/config_test.go +++ b/internal/processor/config/config_test.go @@ -715,12 +715,12 @@ func TestProcessorConfig_Validate_AsyncDispatch(t *testing.T) { wantErr: true, }, { - name: "async valid global gateway with inference_pool_name", + name: "async global gateway rejected", mutate: func(c *ProcessorConfig) { c.ModelGateways = nil c.GlobalInferenceGateway = &ModelGatewayConfig{URL: "http://gw:8000", InferencePoolName: "default-pool"} }, - wantErr: false, + wantErr: true, }, { name: "async no gateways configured", @@ -784,15 +784,6 @@ func TestProcessorConfig_Validate_AsyncDispatch(t *testing.T) { } } -func TestQueueNameHelpers(t *testing.T) { - if got := RequestQueueName("pool-a"); got != "llm-d-async:requests:pool-a" { - t.Fatalf("RequestQueueName(\"pool-a\") = %q, want %q", got, "llm-d-async:requests:pool-a") - } - if got := ResultQueueName("pool-a"); got != "llm-d-async:results:pool-a:$batch" { - t.Fatalf("ResultQueueName(\"pool-a\") = %q, want %q", got, "llm-d-async:results:pool-a:$batch") - } -} - func TestValidate_NormalizesEmptyDispatchMode(t *testing.T) { c := NewConfig() c.ModelGateways = validPerModelConfig() @@ -919,3 +910,58 @@ func TestProcessorConfig_InferenceObjectiveFor(t *testing.T) { }) } } + +func TestResolveModelGateways_Async(t *testing.T) { + t.Run("populates Async field", func(t *testing.T) { + cfg := NewConfig() + cfg.DispatchMode = DispatchModeAsync + cfg.AsyncDispatchConfig = AsyncDispatchConfig{ + ResultPollTimeout: 10 * time.Second, + } + cfg.ModelGateways = map[string]ModelGatewayConfig{ + "model-a": {InferencePoolName: "pool-a"}, + "model-b": {InferencePoolName: "pool-b"}, + } + + resolved, err := ResolveModelGateways(cfg) + if err != nil { + t.Fatalf("ResolveModelGateways() error: %v", err) + } + + if resolved.Async == nil { + t.Fatal("expected Async to be set") + } + if resolved.Global != nil { + t.Error("expected Global to be nil in async mode") + } + if resolved.PerModel != nil { + t.Error("expected PerModel to be nil in async mode") + } + if len(resolved.Async.Models) != 2 { + t.Fatalf("Models count = %d, want 2", len(resolved.Async.Models)) + } + if resolved.Async.Models["model-a"] != "pool-a" { + t.Errorf("Models[model-a] = %q, want %q", resolved.Async.Models["model-a"], "pool-a") + } + if resolved.Async.Models["model-b"] != "pool-b" { + t.Errorf("Models[model-b] = %q, want %q", resolved.Async.Models["model-b"], "pool-b") + } + }) + + t.Run("sync mode does not populate Async", func(t *testing.T) { + cfg := NewConfig() + cfg.ModelGateways = validPerModelConfig() + + resolved, err := ResolveModelGateways(cfg) + if err != nil { + t.Fatalf("ResolveModelGateways() error: %v", err) + } + + if resolved.Async != nil { + t.Error("expected Async to be nil in sync mode") + } + if len(resolved.PerModel) == 0 { + t.Error("expected PerModel to be populated in sync mode") + } + }) +} diff --git a/internal/processor/worker/executor.go b/internal/processor/worker/executor.go index 9e4eb3b56..9e2dd9f15 100644 --- a/internal/processor/worker/executor.go +++ b/internal/processor/worker/executor.go @@ -279,18 +279,34 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte // This ensures the first real error reaches errCh before any context.Canceled // from other models whose contexts were cancelled by requestAbortFn. go func(safeModelID, modelID string) { - err := p.processModel( - requestAbortCtx, - ctx, - sloCtx, - userCancelCtx, - inputFile, - plansDir, safeModelID, modelID, - writers, - progress, - passThroughHeaders, - tenantID, - ) + var err error + if p.asyncInference != nil { + err = p.processModelAsync( + requestAbortCtx, + ctx, + sloCtx, + userCancelCtx, + inputFile, + plansDir, safeModelID, modelID, + writers, + progress, + passThroughHeaders, + tenantID, + ) + } else { + err = p.processModel( + requestAbortCtx, + ctx, + sloCtx, + userCancelCtx, + inputFile, + plansDir, safeModelID, modelID, + writers, + progress, + passThroughHeaders, + tenantID, + ) + } // Abort all sibling models when any model hits a fatal I/O error // (e.g. output file write failure). modelErr is only set for local // I/O failures — not inference errors, which are recorded normally @@ -507,67 +523,198 @@ dispatch: return } - // If user-initiated cancel arrived while this request was in-flight, - // overwrite the result as batch_cancelled and write to the error file - // so that output lines + error lines == total requests. - // SLO expiry does not overwrite in-flight results — only user cancel does. - if sloCtx.Err() == nil && userCancelCtx.Err() != nil { - result.Response = nil - result.Error = &outputError{ - Code: string(batch_types.ErrCodeBatchCancelled), - Message: "This request was cancelled while in progress.", - } - progress.record(requestAbortCtx, false) - - lineBytes, marshalErr := json.Marshal(result) - if marshalErr != nil { - errOnce.Do(func() { - modelErr = fmt.Errorf("marshal cancelled output line at offset %d: %w", entry.Offset, marshalErr) - }) - return - } - lineBytes = append(lineBytes, '\n') - if writeErr := writers.write(lineBytes, true); writeErr != nil { - errOnce.Do(func() { modelErr = fmt.Errorf("write cancelled output line at offset %d: %w", entry.Offset, writeErr) }) - } - return - } - if result.Error != nil && mainCtx.Err() != nil { shutdownCancelled.Add(1) } - progress.record(requestAbortCtx, result.isSuccess()) + if err := writeResult(result, sloCtx, userCancelCtx, requestAbortCtx, writers, progress); err != nil { + errOnce.Do(func() { modelErr = err }) + } + }(entry) + } - lineBytes, marshalErr := json.Marshal(result) - if marshalErr != nil { - errOnce.Do(func() { modelErr = fmt.Errorf("marshal output line at offset %d: %w", entry.Offset, marshalErr) }) - return + wg.Wait() + + return p.drainAndFinalize(requestAbortCtx, mainCtx, sloCtx, userCancelCtx, + inputFile, entries[dispatchedCount:], writers, progress, modelErr, logger, len(entries), + shutdownCancelled.Load()) +} + +// processModelAsync processes all plan entries for a single model using the +// async submit/collect pattern. All requests are submitted to the queue first, +// then results are collected as they arrive on a shared channel. +func (p *Processor) processModelAsync( + requestAbortCtx context.Context, + mainCtx context.Context, + sloCtx context.Context, + userCancelCtx context.Context, + inputFile *os.File, + plansDir, safeModelID, modelID string, + writers *outputWriters, + progress *executionProgress, + passThroughHeaders map[string]string, + tenantID string, +) error { + logger := logr.FromContextOrDiscard(requestAbortCtx).WithValues("model", modelID) + requestAbortCtx = logr.NewContext(requestAbortCtx, logger) + + planPath := filepath.Join(plansDir, safeModelID+".plan") + entries, err := readPlanEntries(planPath) + if err != nil { + return fmt.Errorf("model setup failed: read plan for model %s: %w", modelID, err) + } + + logger.V(logging.INFO).Info("Processing requests for model (async)", "numEntries", len(entries)) + + asyncClient := p.asyncInference.ClientFor(modelID) + if asyncClient == nil { + logger.V(logging.INFO).Info("No async client for model, draining as model_not_found") + p.drainUnprocessedRequests( + requestAbortCtx, inputFile, entries, writers, progress, + inference.ErrCodeModelNotFound) + return nil + } + defer func() { + if err := asyncClient.Close(); err != nil { + logger.Error(err, "Failed to close async client") + } + }() + + // ── Phase 1: Submit ──────────────────────────────────────────────────── + type pendingRequest struct { + batchReqID string + customID string + } + + pending := make(map[string]*pendingRequest) + var submitCount int + + for _, entry := range entries { + if requestAbortCtx.Err() != nil { + logger.V(logging.INFO).Info("Async submit aborted", "submitted", len(pending), "total", len(entries), "reason", requestAbortCtx.Err()) + break + } + + req, batchReqID, parseErr, readErr := readRequestLine(inputFile, entry, logger) + if readErr != nil { + return readErr + } + if parseErr != nil { + lineBytes, err := json.Marshal(parseErr) + if err != nil { + return fmt.Errorf("marshal parse error line: %w", err) } lineBytes = append(lineBytes, '\n') + if err := writers.write(lineBytes, true); err != nil { + return fmt.Errorf("write parse error line: %w", err) + } + progress.record(requestAbortCtx, false) + submitCount++ + continue + } - // Write to error file only for non-HTTP errors (error field populated). - // HTTP error responses (4xx/5xx) go to output file since they carry a valid - // response object with status_code and body per the OpenAI batch spec. - isError := result.Error != nil - if writeErr := writers.write(lineBytes, isError); writeErr != nil { - kind := "output" - if isError { - kind = "error" - } - errOnce.Do(func() { modelErr = fmt.Errorf("write %s line at offset %d: %w", kind, entry.Offset, writeErr) }) + if errors.Is(sloCtx.Err(), context.DeadlineExceeded) { + break + } + + headers := maps.Clone(passThroughHeaders) + headers = mergeInferenceHeaders(headers, sloCtx, p.cfg.InferenceObjectiveFor(modelID), p.fairnessID(tenantID)) + + inferReq := &inference.GenerateRequest{ + RequestID: batchReqID, + Endpoint: req.URL, + Params: req.Body, + Headers: headers, + } + + if submitErr := asyncClient.Submit(requestAbortCtx, inferReq); submitErr != nil { + out := newErrorOutputLine(batchReqID, req.CustomID, + string(submitErr.Category), submitErr.Message) + lineBytes, err := json.Marshal(out) + if err != nil { + return fmt.Errorf("marshal submit error line: %w", err) } - }(entry) + lineBytes = append(lineBytes, '\n') + if err := writers.write(lineBytes, true); err != nil { + return fmt.Errorf("write submit error line: %w", err) + } + progress.record(requestAbortCtx, false) + submitCount++ + continue + } + + pending[batchReqID] = &pendingRequest{ + batchReqID: batchReqID, + customID: req.CustomID, + } + submitCount++ } - wg.Wait() + logger.V(logging.INFO).Info("Submit phase complete", "submitted", len(pending), "total", submitCount) - // Drain undispatched entries to the error file based on the termination reason, and return the - // appropriate sentinel so executeJob can route without re-examining context state. - // Priority: SLO expiry > user cancel > system error > pod shutdown. - // Use sloCtx.Err() rather than requestAbortCtx.Err(): requestAbortCtx may report Canceled if - // requestAbortFn() was called by another goroutine before the sloCtx deadline propagated. - undispatched := entries[dispatchedCount:] + // ── Phase 2: Collect ─────────────────────────────────────────────────── + var modelErr error + + for len(pending) > 0 { + resp, err := asyncClient.GetResult(requestAbortCtx) + if err != nil { + if requestAbortCtx.Err() == nil { + logger.Error(err, "Failed to collect async result", "pendingCount", len(pending)) + modelErr = fmt.Errorf("async result collection failed: %w", err) + } + break + } + + pr, ok := pending[resp.RequestID] + if !ok { + logger.V(logging.TRACE).Info("Ignoring result for unknown request", "requestID", resp.RequestID) + continue + } + + out := buildOutputLine(pr.batchReqID, pr.customID, modelID, resp.RequestID, resp, nil, logger) + if err := writeResult(out, sloCtx, userCancelCtx, requestAbortCtx, writers, progress); err != nil { + modelErr = err + break + } + delete(pending, resp.RequestID) + } + + // Drain submitted-but-uncollected requests as errors so that + // output_lines + error_lines == total_requests. + for _, pr := range pending { + out := newErrorOutputLine(pr.batchReqID, pr.customID, + string(batch_types.ErrCodeBatchExpired), "result not collected before deadline") + lineBytes, err := json.Marshal(out) + if err != nil { + return fmt.Errorf("marshal uncollected error line: %w", err) + } + lineBytes = append(lineBytes, '\n') + if err := writers.write(lineBytes, true); err != nil { + return fmt.Errorf("write uncollected error line: %w", err) + } + progress.record(requestAbortCtx, false) + } + + return p.drainAndFinalize(requestAbortCtx, mainCtx, sloCtx, userCancelCtx, + inputFile, entries[submitCount:], writers, progress, modelErr, logger, len(entries), 0) +} + +// drainAndFinalize drains undispatched entries based on termination reason and +// returns the appropriate sentinel error. Shared by processModel and processModelAsync. +func (p *Processor) drainAndFinalize( + requestAbortCtx context.Context, + mainCtx context.Context, + sloCtx context.Context, + userCancelCtx context.Context, + inputFile *os.File, + undispatched []planEntry, + writers *outputWriters, + progress *executionProgress, + modelErr error, + logger logr.Logger, + totalEntries int, + shutdownCancelledCount int32, +) error { var returnErr error switch { case errors.Is(sloCtx.Err(), context.DeadlineExceeded): @@ -598,7 +745,7 @@ dispatch: returnErr = modelErr default: - if mainCtx.Err() != nil && (len(undispatched) > 0 || shutdownCancelled.Load() > 0) { + if mainCtx.Err() != nil && (len(undispatched) > 0 || shutdownCancelledCount > 0) { // Pod shutdown (SIGTERM): main processor context is cancelled. // Do not drain here — the job will be left for the orphan // reconciler to transition to a terminal state. The undispatched @@ -617,7 +764,7 @@ dispatch: } siblingAbort := returnErr == nil && requestAbortCtx.Err() != nil - logger.V(logging.INFO).Info("Finished processing model", "numEntries", len(entries), "hasError", returnErr != nil, "siblingAbort", siblingAbort) + logger.V(logging.INFO).Info("Finished processing model", "numEntries", totalEntries, "hasError", returnErr != nil, "siblingAbort", siblingAbort) return returnErr } @@ -656,14 +803,7 @@ func (p *Processor) drainUnprocessedRequests( requestID := uuid.NewString() - line := &outputLine{ - ID: newBatchRequestID(requestID), - CustomID: customID, - Error: &outputError{ - Code: string(errCode), - Message: errMessage, - }, - } + line := newErrorOutputLine(newBatchRequestID(requestID), customID, string(errCode), errMessage) lineBytes, err := json.Marshal(line) if err != nil { @@ -743,6 +883,34 @@ func mergeInferenceHeaders(headers map[string]string, sloCtx context.Context, in return headers } +// readRequestLine reads a single plan entry from the input file, parses it, and +// generates a batch request ID. Returns the parsed request and batch request ID +// on success, an outputLine on parse error, or a fatal error on I/O failure. +func readRequestLine(inputFile *os.File, entry planEntry, logger logr.Logger) (*batch_types.Request, string, *outputLine, error) { + buf := make([]byte, entry.Length) + if _, err := inputFile.ReadAt(buf, entry.Offset); err != nil { + return nil, "", nil, fmt.Errorf("%w at offset %d: %w", errRequestInputRead, entry.Offset, err) + } + trimmed := bytes.TrimSuffix(buf, []byte{'\n'}) + batchReqID := newBatchRequestID(uuid.NewString()) + + var req batch_types.Request + if err := json.Unmarshal(trimmed, &req); err != nil { + logger.Error(err, "failed to parse request line, recording as error") + return nil, batchReqID, newErrorOutputLine(batchReqID, "", string(httpclient.ErrCategoryParse), + fmt.Sprintf("failed to parse request line: %v", err)), nil + } + + return &req, batchReqID, nil, nil +} + +func (p *Processor) fairnessID(tenantID string) string { + if p.cfg.SendFairnessHeader { + return tenantID + } + return "" +} + // executeOneRequest reads a single input line from the input file at the given plan entry offset, // sends it to the inference gateway, and returns the formatted output line. func (p *Processor) executeOneRequest( @@ -754,64 +922,31 @@ func (p *Processor) executeOneRequest( passThroughHeaders map[string]string, tenantID string, ) (*outputLine, error) { - // read the request line from input.jsonl at the given offset and length - buf := make([]byte, entry.Length) - if _, err := inputFile.ReadAt(buf, entry.Offset); err != nil { - return nil, fmt.Errorf("%w at offset %d: %w", errRequestInputRead, entry.Offset, err) + logger := logr.FromContextOrDiscard(ctx) + req, batchReqID, parseErr, readErr := readRequestLine(inputFile, entry, logger) + if readErr != nil { + return nil, readErr + } + if parseErr != nil { + return parseErr, nil } - // trim the newline character from the request line - trimmed := bytes.TrimSuffix(buf, []byte{'\n'}) - - // generate a new request ID - requestID := uuid.NewString() + logger = logger.WithValues("customId", req.CustomID, "requestId", batchReqID) - // parse the request line into a batch_types.Request object - var req batch_types.Request - if err := json.Unmarshal(trimmed, &req); err != nil { - logr.FromContextOrDiscard(ctx).Error(err, "failed to parse request line, recording as error") - return &outputLine{ - ID: newBatchRequestID(requestID), - Error: &outputError{ - Code: string(httpclient.ErrCategoryParse), - Message: fmt.Sprintf("failed to parse request line: %v", err), - }, - }, nil - } - - // model id, job id and tenant id are already set in the context - logger := logr.FromContextOrDiscard(ctx).WithValues("customId", req.CustomID, "requestId", requestID) - - // Per-model mode rejects unregistered models at ingestion (fast path). ClientFor can - // still return nil after gateway config changes between ingestion and execution, or - // during recovery when model_map/plan files predate the current resolver — treat as - // a request-level error so the rest of the batch can complete. inferClient := p.inference.ClientFor(modelID) if inferClient == nil { logger.V(logging.INFO).Info("ClientFor returned nil during execution (expected rejection at ingestion)", "model", modelID) - result := &outputLine{ - ID: newBatchRequestID(requestID), - CustomID: req.CustomID, - Error: &outputError{ - Code: inference.ErrCodeModelNotFound, - Message: fmt.Sprintf("model %q is not configured in any gateway", modelID), - }, - } metrics.RecordRequestError(modelID) - return result, nil - } - - fairnessID := "" - if p.cfg.SendFairnessHeader { - fairnessID = tenantID + return newErrorOutputLine(batchReqID, req.CustomID, inference.ErrCodeModelNotFound, + fmt.Sprintf("model %q is not configured in any gateway", modelID)), nil } headers := maps.Clone(passThroughHeaders) - headers = mergeInferenceHeaders(headers, sloCtx, p.cfg.InferenceObjectiveFor(modelID), fairnessID) + headers = mergeInferenceHeaders(headers, sloCtx, p.cfg.InferenceObjectiveFor(modelID), p.fairnessID(tenantID)) inferReq := &inference.GenerateRequest{ - RequestID: newBatchRequestID(requestID), + RequestID: batchReqID, Endpoint: req.URL, Params: req.Body, Headers: headers, @@ -819,14 +954,8 @@ func (p *Processor) executeOneRequest( if sloCtx.Err() == context.DeadlineExceeded { logger.V(logging.INFO).Info("SLO expired during execution, skipping request", "error", sloCtx.Err()) - result := &outputLine{ - ID: newBatchRequestID(requestID), - CustomID: req.CustomID, - Error: &outputError{ - Code: string(batch_types.ErrCodeBatchExpired), - Message: batch_types.ErrCodeBatchExpired.Message(), - }, - } + result := newErrorOutputLine(batchReqID, req.CustomID, + string(batch_types.ErrCodeBatchExpired), batch_types.ErrCodeBatchExpired.Message()) metrics.RecordRequestError(modelID) return result, nil } @@ -842,9 +971,58 @@ func (p *Processor) executeOneRequest( metrics.DecProcessorInflightRequests() metrics.RecordModelRequestExecutionDuration(time.Since(start), modelID) + result := buildOutputLine(batchReqID, req.CustomID, modelID, inferReq.RequestID, inferResp, inferErr, logger) + return result, nil +} + +func newErrorOutputLine(batchReqID, customID, code, message string) *outputLine { + return &outputLine{ + ID: batchReqID, + CustomID: customID, + Error: &outputError{Code: code, Message: message}, + } +} + +// writeResult applies user-cancel overwrite if needed, records progress, marshals +// the output line, and writes it to the appropriate file. Returns an error only +// for marshal/write failures. +func writeResult( + out *outputLine, + sloCtx, userCancelCtx, progressCtx context.Context, + writers *outputWriters, + progress *executionProgress, +) error { + if sloCtx.Err() == nil && userCancelCtx.Err() != nil { + out.Response = nil + out.Error = &outputError{ + Code: string(batch_types.ErrCodeBatchCancelled), + Message: "This request was cancelled while in progress.", + } + } + + progress.record(progressCtx, out.isSuccess()) + lineBytes, err := json.Marshal(out) + if err != nil { + return fmt.Errorf("marshal output line for %s: %w", out.ID, err) + } + lineBytes = append(lineBytes, '\n') + if writeErr := writers.write(lineBytes, out.Error != nil); writeErr != nil { + return fmt.Errorf("write output line for %s: %w", out.ID, writeErr) + } + return nil +} + +// buildOutputLine converts an inference response and/or error into an outputLine. +// Used by both executeOneRequest (sync path) and processModelAsync (async path). +func buildOutputLine( + batchReqID, customID, modelID, serverRequestID string, + inferResp *inference.GenerateResponse, + inferErr *inference.ClientError, + logger logr.Logger, +) *outputLine { result := &outputLine{ - ID: newBatchRequestID(requestID), - CustomID: req.CustomID, + ID: batchReqID, + CustomID: customID, } // Response handling by case. @@ -865,7 +1043,7 @@ func (p *Processor) executeOneRequest( Message: batch_types.ErrCodeBatchExpired.Message(), } metrics.RecordRequestError(modelID) - return result, nil + return result } // HTTP error (4xx/5xx) — populate response with status code and original body // per OpenAI spec, error field is only for non-HTTP errors @@ -885,7 +1063,7 @@ func (p *Processor) executeOneRequest( } result.Response = &batch_types.ResponseData{ StatusCode: inferErr.StatusCode, - RequestID: inferReq.RequestID, + RequestID: serverRequestID, Body: body, } } else { @@ -931,7 +1109,7 @@ func (p *Processor) executeOneRequest( if !result.isSuccess() { metrics.RecordRequestError(modelID) } - return result, nil + return result } // recordTokenUsageFromBody extracts prompt and completion token counts from the diff --git a/internal/processor/worker/executor_test.go b/internal/processor/worker/executor_test.go index faf03278a..83d736b0d 100644 --- a/internal/processor/worker/executor_test.go +++ b/internal/processor/worker/executor_test.go @@ -3771,3 +3771,308 @@ func TestProcessModel_EndpointLimitNil_DrainsAsModelNotFound(t *testing.T) { t.Fatalf("output buffer should be empty, got %d bytes", outBuf.Len()) } } + +// ===================================================================== +// Tests: processModelAsync +// ===================================================================== + +// newAsyncTestProcessorEnv creates a Processor wired for async dispatch testing. +func newAsyncTestProcessorEnv(t *testing.T, cfg *config.ProcessorConfig, asyncResolver *inference.AsyncGatewayResolver) *testProcessorEnv { + t.Helper() + + dbClient := newMockBatchDBClient() + pqClient := mockdb.NewMockBatchPriorityQueueClient() + statusClient := mockdb.NewMockBatchStatusClient() + + p, err := NewProcessor(cfg, &clientset.Clientset{ + BatchDB: dbClient, + FileDB: newMockFileDBClient(), + File: mockfiles.NewMockBatchFilesClient(t.TempDir()), + Queue: pqClient, + Status: statusClient, + Event: mockdb.NewMockBatchEventChannelClient(), + InFlight: mockdb.NewMockInFlightClient(), + AsyncInference: asyncResolver, + }, "test-processor", testLogger(t)) + if err != nil { + t.Fatalf("NewProcessor: %v", err) + } + p.tokens, err = semaphore.New(cfg.NumWorkers, nil) + if err != nil { + t.Fatalf("worker semaphore: %v", err) + } + p.globalSem, err = semaphore.New(cfg.Concurrency.Global, nil) + if err != nil { + t.Fatalf("global semaphore: %v", err) + } + p.poller = NewPoller(pqClient, dbClient) + + return &testProcessorEnv{ + p: p, + dbClient: dbClient, + pqClient: pqClient, + updater: NewStatusUpdater(dbClient, statusClient, 86400), + } +} + +// setupAsyncExecutionJob creates a complete job directory and wires an async resolver. +func setupAsyncExecutionJob( + t *testing.T, + cfg *config.ProcessorConfig, + asyncResolver *inference.AsyncGatewayResolver, + requests []batch_types.Request, + modelToSafe map[string]string, +) (*testProcessorEnv, *batch_types.JobInfo) { + t.Helper() + + env := newAsyncTestProcessorEnv(t, cfg, asyncResolver) + + jobID := "test-job" + tenantID := "tenant-1" + + jobRootDir, err := env.p.jobRootDir(jobID, tenantID) + if err != nil { + t.Fatalf("jobRootDir: %v", err) + } + if err := os.MkdirAll(jobRootDir, 0o755); err != nil { + t.Fatalf("MkdirAll: %v", err) + } + + inputPath := filepath.Join(jobRootDir, "input.jsonl") + rawInput := writeInputJSONL(t, inputPath, requests) + + allEntries := planEntriesFromLines(rawInput) + + safeToModel := make(map[string]string, len(modelToSafe)) + modelEntries := make(map[string][]planEntry) + for model, safe := range modelToSafe { + safeToModel[safe] = model + } + for i, req := range requests { + safe := modelToSafe[req.Body["model"].(string)] + modelEntries[safe] = append(modelEntries[safe], allEntries[i]) + } + + plansDir := filepath.Join(jobRootDir, "plans") + for safe, entries := range modelEntries { + writePlanFile(t, plansDir, safe, entries) + } + + writeModelMap(t, jobRootDir, modelMapFile{ + ModelToSafe: modelToSafe, + SafeToModel: safeToModel, + LineCount: int64(len(requests)), + }) + + return env, &batch_types.JobInfo{JobID: jobID, TenantID: tenantID} +} + +func TestProcessModelAsync(t *testing.T) { + t.Run("Success", func(t *testing.T) { + cfg := config.NewConfig() + cfg.WorkDir = t.TempDir() + + results := make(chan *inference.GenerateResponse, 5) + var submitted []string + + resolver := inference.NewTestAsyncResolver(map[string]func() inference.AsyncInferenceClient{ + "m1": func() inference.AsyncInferenceClient { + return &mockAsyncInferenceClient{ + submitFn: func(_ context.Context, req *inference.GenerateRequest) *inference.ClientError { + submitted = append(submitted, req.RequestID) + results <- &inference.GenerateResponse{ + RequestID: req.RequestID, + Response: []byte(`{"choices":[{"message":{"content":"ok"}}]}`), + } + return nil + }, + getResultFn: func(ctx context.Context) (*inference.GenerateResponse, error) { + select { + case r := <-results: + return r, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + } + }, + }) + + requests := []batch_types.Request{ + {CustomID: "a", Method: "POST", URL: "/v1/chat/completions", Body: map[string]interface{}{"model": "m1"}}, + {CustomID: "b", Method: "POST", URL: "/v1/chat/completions", Body: map[string]interface{}{"model": "m1"}}, + {CustomID: "c", Method: "POST", URL: "/v1/chat/completions", Body: map[string]interface{}{"model": "m1"}}, + } + env, jobInfo := setupAsyncExecutionJob(t, cfg, resolver, requests, map[string]string{"m1": "m1"}) + + inputPath, _ := env.p.jobInputFilePath(jobInfo.JobID, jobInfo.TenantID) + inputFile, _ := os.Open(inputPath) + defer inputFile.Close() + + plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) + + var outBuf, errBuf bytes.Buffer + writers := &outputWriters{output: bufio.NewWriter(&outBuf), errors: bufio.NewWriter(&errBuf)} + progress := &executionProgress{total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID} + + ctx := testLoggerCtx(t) + err := env.p.processModelAsync(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + if err != nil { + t.Fatalf("processModelAsync error: %v", err) + } + + _ = writers.output.Flush() + _ = writers.errors.Flush() + + if len(submitted) != len(requests) { + t.Fatalf("submitted = %d, want %d", len(submitted), len(requests)) + } + + counts := progress.counts() + if counts.Completed != int64(len(requests)) { + t.Fatalf("completed = %d, want %d", counts.Completed, len(requests)) + } + + lines := bytes.Split(bytes.TrimSpace(outBuf.Bytes()), []byte{'\n'}) + if len(lines) != len(requests) { + t.Fatalf("output lines = %d, want %d", len(lines), len(requests)) + } + + if errBuf.Len() != 0 { + t.Errorf("expected empty error file, got: %s", errBuf.String()) + } + }) + + t.Run("CancelDuringCollect", func(t *testing.T) { + cfg := config.NewConfig() + cfg.WorkDir = t.TempDir() + + results := make(chan *inference.GenerateResponse, 5) + submittedIDs := make(chan string, 5) + + resolver := inference.NewTestAsyncResolver(map[string]func() inference.AsyncInferenceClient{ + "m1": func() inference.AsyncInferenceClient { + return &mockAsyncInferenceClient{ + submitFn: func(_ context.Context, req *inference.GenerateRequest) *inference.ClientError { + submittedIDs <- req.RequestID + return nil + }, + getResultFn: func(ctx context.Context) (*inference.GenerateResponse, error) { + select { + case r := <-results: + return r, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + } + }, + }) + + requests := []batch_types.Request{ + {CustomID: "a", Method: "POST", URL: "/v1/chat/completions", Body: map[string]interface{}{"model": "m1"}}, + {CustomID: "b", Method: "POST", URL: "/v1/chat/completions", Body: map[string]interface{}{"model": "m1"}}, + {CustomID: "c", Method: "POST", URL: "/v1/chat/completions", Body: map[string]interface{}{"model": "m1"}}, + } + env, jobInfo := setupAsyncExecutionJob(t, cfg, resolver, requests, map[string]string{"m1": "m1"}) + + inputPath, _ := env.p.jobInputFilePath(jobInfo.JobID, jobInfo.TenantID) + inputFile, _ := os.Open(inputPath) + defer inputFile.Close() + + plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) + + var outBuf, errBuf bytes.Buffer + writers := &outputWriters{output: bufio.NewWriter(&outBuf), errors: bufio.NewWriter(&errBuf)} + progress := &executionProgress{total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID} + + ctx := testLoggerCtx(t) + abortCtx, abortFn := context.WithCancel(ctx) + + // Run processModelAsync in a goroutine so we can intercept submitted IDs. + done := make(chan error, 1) + go func() { + done <- env.p.processModelAsync(abortCtx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + }() + + // Wait for all 3 submits, deliver 1 result using the real ID, then cancel. + firstID := <-submittedIDs + <-submittedIDs + <-submittedIDs + + results <- &inference.GenerateResponse{ + RequestID: firstID, + Response: []byte(`{"choices":[{"message":{"content":"ok"}}]}`), + } + // Give the collect loop time to process the result before cancelling. + time.Sleep(50 * time.Millisecond) + abortFn() + + <-done + + _ = writers.output.Flush() + _ = writers.errors.Flush() + + counts := progress.counts() + // 1 completed + 2 expired = 3 total + if counts.Completed+counts.Failed != int64(len(requests)) { + t.Fatalf("completed+failed = %d, want %d", counts.Completed+counts.Failed, len(requests)) + } + if counts.Completed != 1 { + t.Errorf("completed = %d, want 1", counts.Completed) + } + if counts.Failed != 2 { + t.Errorf("failed = %d, want 2", counts.Failed) + } + }) + + t.Run("UnknownModel", func(t *testing.T) { + cfg := config.NewConfig() + cfg.WorkDir = t.TempDir() + + // Resolver has no model "m1" → ClientFor returns nil. + resolver := inference.NewTestAsyncResolver(map[string]func() inference.AsyncInferenceClient{}) + + requests := []batch_types.Request{ + {CustomID: "a", Method: "POST", URL: "/v1/chat/completions", Body: map[string]interface{}{"model": "m1"}}, + } + env, jobInfo := setupAsyncExecutionJob(t, cfg, resolver, requests, map[string]string{"m1": "m1"}) + + inputPath, _ := env.p.jobInputFilePath(jobInfo.JobID, jobInfo.TenantID) + inputFile, _ := os.Open(inputPath) + defer inputFile.Close() + + plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) + + var outBuf, errBuf bytes.Buffer + writers := &outputWriters{output: bufio.NewWriter(&outBuf), errors: bufio.NewWriter(&errBuf)} + progress := &executionProgress{total: 1, updater: env.updater, jobID: jobInfo.JobID} + + ctx := testLoggerCtx(t) + err := env.p.processModelAsync(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + if err != nil { + t.Fatalf("processModelAsync error: %v", err) + } + + _ = writers.output.Flush() + _ = writers.errors.Flush() + + counts := progress.counts() + if counts.Failed != 1 { + t.Fatalf("failed = %d, want 1", counts.Failed) + } + + errLines := bytes.Split(bytes.TrimSpace(errBuf.Bytes()), []byte{'\n'}) + if len(errLines) != 1 { + t.Fatalf("error lines = %d, want 1", len(errLines)) + } + + var entry outputLine + if err := json.Unmarshal(errLines[0], &entry); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if entry.Error == nil || entry.Error.Code != string(inference.ErrCodeModelNotFound) { + t.Fatalf("expected error code %s, got %+v", inference.ErrCodeModelNotFound, entry.Error) + } + }) +} diff --git a/internal/processor/worker/test_helpers_test.go b/internal/processor/worker/test_helpers_test.go index 4ee31256d..ff4a81054 100644 --- a/internal/processor/worker/test_helpers_test.go +++ b/internal/processor/worker/test_helpers_test.go @@ -102,6 +102,37 @@ func (m *mockInferenceClient) Generate(ctx context.Context, req *inference.Gener }, nil } +// --------------------------------------------------------------------------- +// Mock async inference client +// --------------------------------------------------------------------------- + +type mockAsyncInferenceClient struct { + submitFn func(ctx context.Context, req *inference.GenerateRequest) *inference.ClientError + getResultFn func(ctx context.Context) (*inference.GenerateResponse, error) + closeFn func() error +} + +func (m *mockAsyncInferenceClient) Submit(ctx context.Context, req *inference.GenerateRequest) *inference.ClientError { + if m.submitFn != nil { + return m.submitFn(ctx, req) + } + return nil +} + +func (m *mockAsyncInferenceClient) GetResult(ctx context.Context) (*inference.GenerateResponse, error) { + if m.getResultFn != nil { + return m.getResultFn(ctx) + } + return nil, ctx.Err() +} + +func (m *mockAsyncInferenceClient) Close() error { + if m.closeFn != nil { + return m.closeFn() + } + return nil +} + // --------------------------------------------------------------------------- // Mock files client (upload retry testing) // --------------------------------------------------------------------------- diff --git a/internal/processor/worker/worker.go b/internal/processor/worker/worker.go index ae628ae8e..f56cd659f 100644 --- a/internal/processor/worker/worker.go +++ b/internal/processor/worker/worker.go @@ -70,11 +70,12 @@ type Processor struct { poller *Poller updater *StatusUpdater - batchDB db.BatchDBClient // job status lookups (heartbeat DB check) - event db.BatchEventChannelClient // cancel-event subscription - inflight db.InFlightClient // in-flight job tracking for orphan recovery - inference *inference.GatewayResolver // model → gateway routing - files *fileManager + batchDB db.BatchDBClient // job status lookups (heartbeat DB check) + event db.BatchEventChannelClient // cancel-event subscription + inflight db.InFlightClient // in-flight job tracking for orphan recovery + inference *inference.GatewayResolver // model → gateway routing (sync) + asyncInference *inference.AsyncGatewayResolver // model → async client routing + files *fileManager } func NewProcessor( @@ -92,15 +93,16 @@ func NewProcessor( poller := NewPoller(clients.Queue, clients.BatchDB) updater := NewStatusUpdater(clients.BatchDB, clients.Status, cfg.ProgressTTLSeconds) return &Processor{ - cfg: cfg, - processorID: processorID, - poller: poller, - updater: updater, - batchDB: clients.BatchDB, - event: clients.Event, - inflight: clients.InFlight, - inference: clients.Inference, - files: newFileManager(clients.File, clients.FileDB), + cfg: cfg, + processorID: processorID, + poller: poller, + updater: updater, + batchDB: clients.BatchDB, + event: clients.Event, + inflight: clients.InFlight, + inference: clients.Inference, + asyncInference: clients.AsyncInference, + files: newFileManager(clients.File, clients.FileDB), }, nil } @@ -128,6 +130,17 @@ func (p *Processor) Run(ctx context.Context, onReady func()) error { logger := logr.FromContextOrDiscard(ctx) + if err := p.initConcurrencyControls(logger, stopAccepting); err != nil { + return err + } + + return p.runPollingLoop(pollingCtx, ctx) +} + +// initConcurrencyControls creates semaphores and per-endpoint AIMD controllers. +// In async mode (p.inference == nil), only the job-level worker semaphore is +// created — inference concurrency is controlled by the llm-d-async dispatcher. +func (p *Processor) initConcurrencyControls(logger logr.Logger, stopAccepting context.CancelFunc) error { // Create semaphores here (not in NewProcessor) so the double-release guard // callback can capture stopAccepting. This keeps semaphores immutable after // construction — no mutex, no OnDoubleRelease method. @@ -142,6 +155,16 @@ func (p *Processor) Run(ctx context.Context, onReady func()) error { if err != nil { return fmt.Errorf("worker semaphore (NumWorkers=%d): %w", p.cfg.NumWorkers, err) } + + if p.asyncInference != nil { + logger.V(logging.INFO).Info( + "Processor run started (async dispatch)", + "loopInterval", p.cfg.PollInterval, + "maxWorkers", p.cfg.NumWorkers, + ) + return nil + } + cc := &p.cfg.Concurrency p.globalSem, err = semaphore.New(cc.Global, makeGuard("global-concurrency")) if err != nil { @@ -191,8 +214,7 @@ func (p *Processor) Run(ctx context.Context, onReady func()) error { "concurrency.aimd.enabled", cc.AIMD.Enabled, "num_endpoints", len(clients), ) - - return p.runPollingLoop(pollingCtx, ctx) + return nil } // Stop gracefully stops the processor, waiting for all workers to finish. @@ -516,9 +538,12 @@ func (p *Processor) validate() error { if p.inflight == nil { return fmt.Errorf("in-flight client is missing") } - if p.inference == nil { + if p.inference == nil && p.asyncInference == nil { return fmt.Errorf("inference client is missing") } + if p.inference != nil && p.asyncInference != nil { + return fmt.Errorf("sync and async inference clients are mutually exclusive") + } if p.files == nil { return fmt.Errorf("file manager is missing") } diff --git a/internal/processor/worker/worker_test.go b/internal/processor/worker/worker_test.go index 298f4ec44..17a7c6c5d 100644 --- a/internal/processor/worker/worker_test.go +++ b/internal/processor/worker/worker_test.go @@ -11,6 +11,7 @@ import ( "github.com/llm-d/llm-d-batch-gateway/internal/shared/openai" "github.com/llm-d/llm-d-batch-gateway/internal/util/clientset" "github.com/llm-d/llm-d-batch-gateway/internal/util/semaphore" + "github.com/llm-d/llm-d-batch-gateway/pkg/clients/inference" ) func TestClientsetFields_Assigned(t *testing.T) { @@ -20,6 +21,40 @@ func TestClientsetFields_Assigned(t *testing.T) { } } +func TestValidate_InferenceClientRequired(t *testing.T) { + t.Run("rejects when neither sync nor async inference is set", func(t *testing.T) { + cs := validProcessorClients(t) + cs.Inference = nil + cs.AsyncInference = nil + p := mustNewProcessor(t, config.NewConfig(), cs) + + if err := p.validate(); err == nil { + t.Fatal("expected validation error when both inference clients are nil") + } + }) + + t.Run("accepts sync inference only", func(t *testing.T) { + cs := validProcessorClients(t) + cs.AsyncInference = nil + p := mustNewProcessor(t, config.NewConfig(), cs) + + if err := p.validate(); err != nil { + t.Fatalf("unexpected validation error: %v", err) + } + }) + + t.Run("accepts async inference only", func(t *testing.T) { + cs := validProcessorClients(t) + cs.Inference = nil + cs.AsyncInference = &inference.AsyncGatewayResolver{} + p := mustNewProcessor(t, config.NewConfig(), cs) + + if err := p.validate(); err != nil { + t.Fatalf("unexpected validation error: %v", err) + } + }) +} + func TestNewProcessor_InvalidNumWorkers(t *testing.T) { cfg := config.NewConfig() cfg.NumWorkers = 0 diff --git a/internal/util/clientset/clientset.go b/internal/util/clientset/clientset.go index 416be7bbf..c4329874f 100644 --- a/internal/util/clientset/clientset.go +++ b/internal/util/clientset/clientset.go @@ -23,6 +23,7 @@ import ( "context" "errors" "fmt" + "maps" "github.com/go-logr/logr" dbapi "github.com/llm-d/llm-d-batch-gateway/internal/database/api" @@ -41,14 +42,15 @@ import ( // Clientset holds all clients. type Clientset struct { - File fsapi.BatchFilesClient - BatchDB dbapi.BatchDBClient - FileDB dbapi.FileDBClient - Queue dbapi.BatchPriorityQueueClient - Event dbapi.BatchEventChannelClient - Status dbapi.BatchStatusClient - InFlight dbapi.InFlightClient - Inference *inference.GatewayResolver + File fsapi.BatchFilesClient + BatchDB dbapi.BatchDBClient + FileDB dbapi.FileDBClient + Queue dbapi.BatchPriorityQueueClient + Event dbapi.BatchEventChannelClient + Status dbapi.BatchStatusClient + InFlight dbapi.InFlightClient + Inference *inference.GatewayResolver + AsyncInference *inference.AsyncGatewayResolver } // NewFSFileClient creates a filesystem-based file storage client. @@ -150,6 +152,7 @@ type clientsetConfig struct { exchangeRedisCfg *uredis.RedisClientConfig inferenceGlobal *inference.GatewayClientConfig inferencePerModel map[string]inference.GatewayClientConfig + asyncInference *inference.AsyncClientConfig } // WithDB enables creation of batch and file database clients. @@ -183,6 +186,13 @@ func WithPerModelInference(cfgs map[string]inference.GatewayClientConfig) Option return func(c *clientsetConfig) { c.inferencePerModel = copied } } +// WithAsyncInference enables async dispatch via llm-d-async queues. +func WithAsyncInference(cfg inference.AsyncClientConfig) Option { + copied := cfg + copied.Models = maps.Clone(cfg.Models) + return func(c *clientsetConfig) { c.asyncInference = &copied } +} + // NewClientset creates the clients specified by the given options. func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) (*Clientset, error) { logger := logr.FromContextOrDiscard(ctx) @@ -266,6 +276,19 @@ func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) // build inference client(s) switch { + case cfg.asyncInference != nil: + if cfg.asyncInference.RedisURL == "" { + if cfg.exchangeRedisCfg == nil { + return nil, fmt.Errorf("async inference requires a Redis URL (set RedisURL or use WithExchange)") + } + cfg.asyncInference.RedisURL = cfg.exchangeRedisCfg.Url + } + resolver, err := inference.NewAsyncResolver(*cfg.asyncInference, logger) + if err != nil { + return nil, fmt.Errorf("failed to create async inference clients: %w", err) + } + logger.Info("Async inference clients created", "count", len(cfg.asyncInference.Models)) + cs.AsyncInference = resolver case cfg.inferenceGlobal != nil: resolver, err := inference.NewGlobalResolver(*cfg.inferenceGlobal, logger) if err != nil { @@ -322,5 +345,15 @@ func (cs *Clientset) Close() error { errs = append(errs, err) } } + if cs.Inference != nil { + if err := cs.Inference.Close(); err != nil { + errs = append(errs, err) + } + } + if cs.AsyncInference != nil { + if err := cs.AsyncInference.Close(); err != nil { + errs = append(errs, err) + } + } return errors.Join(errs...) } diff --git a/pkg/clients/inference/async_inference_client_impl.go b/pkg/clients/inference/async_inference_client_impl.go new file mode 100644 index 000000000..d91b2399e --- /dev/null +++ b/pkg/clients/inference/async_inference_client_impl.go @@ -0,0 +1,220 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package inference + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + "github.com/go-logr/logr" + asyncapi "github.com/llm-d-incubation/llm-d-async/api" + "github.com/llm-d-incubation/llm-d-async/producer" + "github.com/llm-d/llm-d-batch-gateway/internal/util/logging" + httpclient "github.com/llm-d/llm-d-batch-gateway/pkg/clients/http" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/propagation" +) + +var _ AsyncInferenceClient = (*asyncProducerClient)(nil) + +// defaultResultBufferSize is the per-job channel capacity for async results. +const defaultResultBufferSize = 100 + +// resultDispatcher reads results from the producer's shared result queue and +// routes them to the correct caller by request ID. The processor dispatches +// multiple requests per model concurrently, and results arrive in any order, +// so a single reader must demux them. +type resultDispatcher struct { + producer producer.Producer + logger logr.Logger + waiters sync.Map // requestID -> chan<- *GenerateResponse + once sync.Once + wg sync.WaitGroup + cancel context.CancelFunc +} + +func newResultDispatcher(p producer.Producer, logger logr.Logger) *resultDispatcher { + return &resultDispatcher{ + producer: p, + logger: logger, + } +} + +func (d *resultDispatcher) ensureStarted() { + d.once.Do(func() { + ctx, cancel := context.WithCancel(context.Background()) + d.cancel = cancel + d.wg.Add(1) + go d.run(ctx) + }) +} + +func (d *resultDispatcher) run(ctx context.Context) { + defer d.wg.Done() + for { + pollCtx, pollCancel := context.WithTimeout(ctx, time.Second) + result, err := d.producer.GetResult(pollCtx) + pollCancel() + if err != nil { + if ctx.Err() != nil { + return + } + if !errors.Is(err, context.DeadlineExceeded) { + d.logger.Error(err, "Failed to read from result queue") + } + continue + } + + if val, ok := d.waiters.LoadAndDelete(result.ID); ok { + ch, ok := val.(chan<- *GenerateResponse) + if !ok { + d.logger.Error(fmt.Errorf("unexpected type %T in waiters map", val), "Type assertion failed") + continue + } + resp := &GenerateResponse{ + RequestID: result.ID, + Response: []byte(result.Payload), + } + select { + case ch <- resp: + default: + d.logger.Error(fmt.Errorf("result channel full"), "Dropping result", "resultID", result.ID) + } + } else { + d.logger.Info("Dropped result with no waiter", "resultID", result.ID) + } + } +} + +func (d *resultDispatcher) register(requestID string, ch chan<- *GenerateResponse) { + d.waiters.Store(requestID, ch) + d.ensureStarted() +} + +func (d *resultDispatcher) unregister(requestID string) { + d.waiters.Delete(requestID) +} + +func (d *resultDispatcher) Close() error { + if d.cancel != nil { + d.cancel() + done := make(chan struct{}) + go func() { + d.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + return fmt.Errorf("result dispatcher did not shut down within 2s") + } + } + return nil +} + +// asyncPool holds the shared resources for one inference pool. +// Multiple per-job clients share the same pool. +type asyncPool struct { + producer producer.Producer + dispatcher *resultDispatcher + logger logr.Logger + defaultDeadline time.Duration +} + +// asyncProducerClient is a per-job client that submits requests and collects +// results via an internal channel. Each job gets its own client (and channel) +// from AsyncGatewayResolver.ClientFor, backed by a shared pool. +type asyncProducerClient struct { + pool *asyncPool + results chan *GenerateResponse + pendingIDs sync.Map // tracks submitted request IDs for cleanup + logger logr.Logger +} + +func newAsyncProducerClient(pool *asyncPool) *asyncProducerClient { + return &asyncProducerClient{ + pool: pool, + results: make(chan *GenerateResponse, defaultResultBufferSize), + logger: pool.logger, + } +} + +// Submit enqueues a request for async processing. The result will be routed +// to this client's internal channel by the shared dispatcher. +func (c *asyncProducerClient) Submit(ctx context.Context, req *GenerateRequest) *ClientError { + now := time.Now() + fallback := c.pool.defaultDeadline + if fallback == 0 { + fallback = 5 * time.Minute + } + deadline := now.Add(fallback) + if dl, ok := ctx.Deadline(); ok { + deadline = dl + } + + metadata := make(map[string]string) + otel.GetTextMapPropagator().Inject(ctx, propagation.MapCarrier(metadata)) + + reqMsg := &asyncapi.RequestMessage{ + ID: req.RequestID, + Created: now.Unix(), + Deadline: deadline.Unix(), + Payload: req.Params, + Headers: req.Headers, + Endpoint: req.Endpoint, + Metadata: metadata, + } + + c.pool.dispatcher.register(req.RequestID, c.results) + c.pendingIDs.Store(req.RequestID, struct{}{}) + + if err := c.pool.producer.SubmitRequest(ctx, reqMsg); err != nil { + c.pool.dispatcher.unregister(req.RequestID) + c.pendingIDs.Delete(req.RequestID) + return &ClientError{ + Category: httpclient.ErrCategoryServer, + Message: fmt.Sprintf("submit async request: %v", err), + RawError: err, + } + } + + c.logger.V(logging.TRACE).Info("Submitted async request", "requestID", req.RequestID) + return nil +} + +// GetResult blocks until the next result arrives or the context is cancelled. +func (c *asyncProducerClient) GetResult(ctx context.Context) (*GenerateResponse, error) { + select { + case resp := <-c.results: + c.pendingIDs.Delete(resp.RequestID) + return resp, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// Close unregisters all pending waiters from the shared dispatcher. +func (c *asyncProducerClient) Close() error { + c.pendingIDs.Range(func(key, _ any) bool { + c.pool.dispatcher.unregister(key.(string)) + return true + }) + return nil +} diff --git a/pkg/clients/inference/async_inference_client_integration_test.go b/pkg/clients/inference/async_inference_client_integration_test.go new file mode 100644 index 000000000..e9b73efd6 --- /dev/null +++ b/pkg/clients/inference/async_inference_client_integration_test.go @@ -0,0 +1,92 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package inference + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/llm-d-incubation/llm-d-async/api" + "github.com/redis/go-redis/v9" +) + +func TestAsyncProducerClient_Submit_roundtrip(t *testing.T) { + t.Run("enqueues request and returns result via GetResult", func(t *testing.T) { + mr := miniredis.RunT(t) + poolName := "test-pool" + reqQueue := asyncQueuePrefix + "requests:" + poolName + resultQueue := asyncQueuePrefix + "results:" + poolName + + pool := newTestPool(t, mr, poolName) + client := newAsyncProducerClient(pool) + defer func() { _ = client.Close() }() + + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + defer func() { _ = rdb.Close() }() + + go func() { + time.Sleep(50 * time.Millisecond) + pushResult(t, mr, resultQueue, "req-1", `{"choices":[{"text":"hello"}]}`) + }() + + if err := client.Submit(context.Background(), &GenerateRequest{ + RequestID: "req-1", + Endpoint: "/v1/completions", + Params: map[string]any{"model": "test-model", "prompt": "hello"}, + }); err != nil { + t.Fatalf("Submit() error: %s", err.Message) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + resp, err := client.GetResult(ctx) + if err != nil { + t.Fatalf("GetResult error: %v", err) + } + + if resp.RequestID != "req-1" { + t.Errorf("RequestID = %q, want %q", resp.RequestID, "req-1") + } + if resp.Response == nil { + t.Fatal("expected non-nil Response") + } + + members, zErr := rdb.ZRange(context.Background(), reqQueue, 0, -1).Result() + if zErr != nil { + t.Fatalf("ZRange: %v", zErr) + } + if len(members) != 1 { + t.Fatalf("expected 1 member in request queue, got %d", len(members)) + } + + var envelope map[string]json.RawMessage + if uErr := json.Unmarshal([]byte(members[0]), &envelope); uErr != nil { + t.Fatalf("unmarshal enqueued request: %v", uErr) + } + var data api.RequestMessage + if uErr := json.Unmarshal(envelope["data"], &data); uErr != nil { + t.Fatalf("unmarshal data field: %v", uErr) + } + if data.ID != "req-1" { + t.Errorf("enqueued request ID = %q, want %q", data.ID, "req-1") + } + }) +} diff --git a/pkg/clients/inference/async_inference_client_interface.go b/pkg/clients/inference/async_inference_client_interface.go new file mode 100644 index 000000000..808507ee9 --- /dev/null +++ b/pkg/clients/inference/async_inference_client_interface.go @@ -0,0 +1,12 @@ +package inference + +import "context" + +// AsyncInferenceClient defines the interface for non-blocking async dispatch. +// Each instance is per-job: created by AsyncGatewayResolver.ClientFor, used +// for one submit/collect cycle, then closed. +type AsyncInferenceClient interface { + Submit(ctx context.Context, req *GenerateRequest) *ClientError + GetResult(ctx context.Context) (*GenerateResponse, error) + Close() error +} diff --git a/pkg/clients/inference/async_inference_client_resolver.go b/pkg/clients/inference/async_inference_client_resolver.go new file mode 100644 index 000000000..dd3977a8f --- /dev/null +++ b/pkg/clients/inference/async_inference_client_resolver.go @@ -0,0 +1,130 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package inference + +import ( + "errors" + "fmt" + "io" + "time" + + "github.com/go-logr/logr" + "github.com/llm-d-incubation/llm-d-async/producer" + "github.com/redis/go-redis/v9" +) + +const asyncQueuePrefix = "llm-d-async:" + +// AsyncClientConfig holds the resolved configuration for async dispatch. +type AsyncClientConfig struct { + RedisURL string + Models map[string]string // model name -> pool name + DefaultDeadline time.Duration // fallback deadline when ctx has none; 0 defaults to 5m +} + +// AsyncGatewayResolver routes models to per-job AsyncInferenceClient instances. +// Each call to ClientFor creates a fresh client with its own result channel, +// backed by a shared producer and dispatcher per pool. +// Immutable after construction — safe for concurrent reads. +type AsyncGatewayResolver struct { + pools map[string]*asyncPool // model → pool + closers []io.Closer + clientFactories map[string]func() AsyncInferenceClient // test-only override +} + +// ClientFor creates a fresh per-job async client for the given model. +// Returns nil if no matching pool exists. +func (r *AsyncGatewayResolver) ClientFor(modelID string) AsyncInferenceClient { + if r.clientFactories != nil { + if factory, ok := r.clientFactories[modelID]; ok { + return factory() + } + return nil + } + pool, ok := r.pools[modelID] + if !ok { + return nil + } + return newAsyncProducerClient(pool) +} + +// NewTestAsyncResolver creates a resolver backed by factory functions instead of +// real Redis connections. Each call to ClientFor invokes the corresponding factory. +func NewTestAsyncResolver(factories map[string]func() AsyncInferenceClient) *AsyncGatewayResolver { + return &AsyncGatewayResolver{clientFactories: factories} +} + +// Close releases resources held by the resolver (dispatchers, producers, Redis). +func (r *AsyncGatewayResolver) Close() error { + var errs []error + for _, c := range r.closers { + if err := c.Close(); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +// NewAsyncResolver creates an AsyncGatewayResolver with one shared pool +// (producer + dispatcher) per model/pool pair. +func NewAsyncResolver(config AsyncClientConfig, logger logr.Logger) (*AsyncGatewayResolver, error) { + opts, err := redis.ParseURL(config.RedisURL) + if err != nil { + return nil, fmt.Errorf("failed to parse async inference Redis URL: %w", err) + } + rdb := redis.NewClient(opts) + + poolToModel := make(map[string]string, len(config.Models)) + for model, poolName := range config.Models { + if existing, ok := poolToModel[poolName]; ok { + _ = rdb.Close() + return nil, fmt.Errorf("models %q and %q both map to pool %q: each pool must have a single consumer", existing, model, poolName) + } + poolToModel[poolName] = model + } + + pools := make(map[string]*asyncPool, len(config.Models)) + var closers []io.Closer + + for model, poolName := range config.Models { + p, err := producer.NewRedisSortedSetProducer( + producer.RedisSortedSetConfig{ + RequestQueueName: asyncQueuePrefix + "requests:" + poolName, + ResultQueueName: asyncQueuePrefix + "results:" + poolName, + }, + producer.WithRedisClient(rdb), + ) + if err != nil { + for _, c := range closers { + _ = c.Close() + } + _ = rdb.Close() + return nil, fmt.Errorf("failed to create producer for model %q (pool %s): %w", model, poolName, err) + } + + poolLogger := logger.WithName("async-inference").WithValues("pool", poolName) + d := newResultDispatcher(p, poolLogger) + pool := &asyncPool{producer: p, dispatcher: d, logger: poolLogger, defaultDeadline: config.DefaultDeadline} + + pools[model] = pool + closers = append(closers, d, p) + } + + closers = append(closers, rdb) + + return &AsyncGatewayResolver{pools: pools, closers: closers}, nil +} diff --git a/pkg/clients/inference/async_inference_client_resolver_test.go b/pkg/clients/inference/async_inference_client_resolver_test.go new file mode 100644 index 000000000..6917a663e --- /dev/null +++ b/pkg/clients/inference/async_inference_client_resolver_test.go @@ -0,0 +1,127 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package inference + +import ( + "testing" + + "github.com/alicebob/miniredis/v2" +) + +func TestNewAsyncResolver(t *testing.T) { + t.Run("creates per-model clients", func(t *testing.T) { + mr := miniredis.RunT(t) + + cfg := AsyncClientConfig{ + RedisURL: "redis://" + mr.Addr(), + Models: map[string]string{ + "model-a": "pool-a", + "model-b": "pool-b", + }, + } + + r, err := NewAsyncResolver(cfg, testLogger(t)) + if err != nil { + t.Fatalf("NewAsyncResolver: %v", err) + } + + if got := r.ClientFor("model-a"); got == nil { + t.Fatal("expected non-nil client for model-a") + } + if got := r.ClientFor("model-b"); got == nil { + t.Fatal("expected non-nil client for model-b") + } + if got := r.ClientFor("unknown"); got != nil { + t.Fatalf("expected nil for unknown model, got %v", got) + } + }) + + t.Run("returns nil for unknown model", func(t *testing.T) { + mr := miniredis.RunT(t) + + r, err := NewAsyncResolver(AsyncClientConfig{ + RedisURL: "redis://" + mr.Addr(), + Models: map[string]string{"model-a": "pool-a"}, + }, testLogger(t)) + if err != nil { + t.Fatalf("NewAsyncResolver: %v", err) + } + + if got := r.ClientFor("unknown"); got != nil { + t.Fatalf("expected nil for unknown model, got %v", got) + } + }) + + t.Run("rejects duplicate pool mapping", func(t *testing.T) { + mr := miniredis.RunT(t) + + _, err := NewAsyncResolver(AsyncClientConfig{ + RedisURL: "redis://" + mr.Addr(), + Models: map[string]string{ + "model-a": "shared-pool", + "model-b": "shared-pool", + }, + }, testLogger(t)) + if err == nil { + t.Fatal("expected error for duplicate pool mapping") + } + }) + + t.Run("invalid Redis URL returns error", func(t *testing.T) { + _, err := NewAsyncResolver(AsyncClientConfig{ + RedisURL: "not-a-url", + Models: map[string]string{"model-a": "pool-a"}, + }, testLogger(t)) + if err == nil { + t.Fatal("expected error for invalid Redis URL") + } + }) + + t.Run("close releases resources", func(t *testing.T) { + mr := miniredis.RunT(t) + + r, err := NewAsyncResolver(AsyncClientConfig{ + RedisURL: "redis://" + mr.Addr(), + Models: map[string]string{"model-a": "pool-a"}, + }, testLogger(t)) + if err != nil { + t.Fatalf("NewAsyncResolver: %v", err) + } + + if err := r.Close(); err != nil { + t.Fatalf("Close() returned error: %v", err) + } + }) + + t.Run("each ClientFor call returns a fresh client", func(t *testing.T) { + mr := miniredis.RunT(t) + + r, err := NewAsyncResolver(AsyncClientConfig{ + RedisURL: "redis://" + mr.Addr(), + Models: map[string]string{"model-a": "pool-a"}, + }, testLogger(t)) + if err != nil { + t.Fatalf("NewAsyncResolver: %v", err) + } + + client1 := r.ClientFor("model-a") + client2 := r.ClientFor("model-a") + if client1 == client2 { + t.Fatal("expected fresh client per ClientFor call") + } + }) +} diff --git a/pkg/clients/inference/async_inference_client_test.go b/pkg/clients/inference/async_inference_client_test.go new file mode 100644 index 000000000..dda7e70bc --- /dev/null +++ b/pkg/clients/inference/async_inference_client_test.go @@ -0,0 +1,294 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package inference + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/llm-d-incubation/llm-d-async/api" + "github.com/llm-d-incubation/llm-d-async/producer" + "github.com/redis/go-redis/v9" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/propagation" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" +) + +func newTestPool(t *testing.T, mr *miniredis.Miniredis, poolName string) *asyncPool { + t.Helper() + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + + p, err := producer.NewRedisSortedSetProducer( + producer.RedisSortedSetConfig{ + RequestQueueName: asyncQueuePrefix + "requests:" + poolName, + ResultQueueName: asyncQueuePrefix + "results:" + poolName, + }, + producer.WithRedisClient(rdb), + ) + if err != nil { + t.Fatalf("NewRedisSortedSetProducer: %v", err) + } + + logger := testLogger(t) + d := newResultDispatcher(p, logger) + pool := &asyncPool{producer: p, dispatcher: d, logger: logger} + + t.Cleanup(func() { + _ = d.Close() + _ = p.Close() + }) + + return pool +} + +func pushResult(t *testing.T, mr *miniredis.Miniredis, queue, id, payload string) { + t.Helper() + data, err := json.Marshal(api.ResultMessage{ID: id, Payload: payload}) + if err != nil { + t.Fatalf("marshal result %s: %v", id, err) + } + if _, err := mr.Lpush(queue, string(data)); err != nil { + t.Fatalf("Lpush %s: %v", id, err) + } +} + +func TestAsyncProducerClient_Submit(t *testing.T) { + t.Run("submit and get result", func(t *testing.T) { + mr := miniredis.RunT(t) + poolName := "submit-pool" + resultQueue := asyncQueuePrefix + "results:" + poolName + + pool := newTestPool(t, mr, poolName) + client := newAsyncProducerClient(pool) + defer func() { _ = client.Close() }() + + go func() { + time.Sleep(50 * time.Millisecond) + pushResult(t, mr, resultQueue, "req-1", `{"choices":[{"text":"hello"}]}`) + }() + + if err := client.Submit(context.Background(), &GenerateRequest{ + RequestID: "req-1", + Endpoint: "/v1/completions", + Params: map[string]any{"model": "test-model"}, + }); err != nil { + t.Fatalf("Submit error: %s", err.Message) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + resp, err := client.GetResult(ctx) + if err != nil { + t.Fatalf("GetResult error: %v", err) + } + if resp.RequestID != "req-1" { + t.Errorf("RequestID = %q, want %q", resp.RequestID, "req-1") + } + }) + + t.Run("multiple submits routed correctly", func(t *testing.T) { + mr := miniredis.RunT(t) + poolName := "multi-pool" + resultQueue := asyncQueuePrefix + "results:" + poolName + + pool := newTestPool(t, mr, poolName) + client := newAsyncProducerClient(pool) + defer func() { _ = client.Close() }() + + for _, id := range []string{"s-1", "s-2", "s-3"} { + if err := client.Submit(context.Background(), &GenerateRequest{ + RequestID: id, + Endpoint: "/v1/completions", + Params: map[string]any{"model": "test-model"}, + }); err != nil { + t.Fatalf("Submit(%s) error: %s", id, err.Message) + } + } + + // Push results in reverse order + for _, id := range []string{"s-3", "s-1", "s-2"} { + pushResult(t, mr, resultQueue, id, fmt.Sprintf(`{"id":"%s"}`, id)) + } + + got := make(map[string]bool) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + for i := 0; i < 3; i++ { + resp, err := client.GetResult(ctx) + if err != nil { + t.Fatalf("GetResult %d/3 error: %v", i+1, err) + } + got[resp.RequestID] = true + } + + for _, id := range []string{"s-1", "s-2", "s-3"} { + if !got[id] { + t.Errorf("missing result for %s", id) + } + } + }) + + t.Run("close unregisters pending waiters", func(t *testing.T) { + mr := miniredis.RunT(t) + poolName := "close-pool" + resultQueue := asyncQueuePrefix + "results:" + poolName + + pool := newTestPool(t, mr, poolName) + client := newAsyncProducerClient(pool) + + if err := client.Submit(context.Background(), &GenerateRequest{ + RequestID: "c-1", + Endpoint: "/v1/completions", + Params: map[string]any{"model": "test-model"}, + }); err != nil { + t.Fatalf("Submit error: %s", err.Message) + } + + _ = client.Close() + + // Push a result — it should be dropped (no waiter) + pushResult(t, mr, resultQueue, "c-1", `{"id":"c-1"}`) + time.Sleep(200 * time.Millisecond) + + // Verify no result on the channel + select { + case <-client.results: + t.Fatal("expected no result after Close") + default: + } + }) + + t.Run("per-job isolation with shared pool", func(t *testing.T) { + mr := miniredis.RunT(t) + poolName := "isolation-pool" + resultQueue := asyncQueuePrefix + "results:" + poolName + + pool := newTestPool(t, mr, poolName) + clientA := newAsyncProducerClient(pool) + clientB := newAsyncProducerClient(pool) + defer func() { _ = clientA.Close() }() + defer func() { _ = clientB.Close() }() + + if err := clientA.Submit(context.Background(), &GenerateRequest{ + RequestID: "job-a-req", + Endpoint: "/v1/completions", + Params: map[string]any{"model": "test-model"}, + }); err != nil { + t.Fatalf("Submit A error: %s", err.Message) + } + if err := clientB.Submit(context.Background(), &GenerateRequest{ + RequestID: "job-b-req", + Endpoint: "/v1/completions", + Params: map[string]any{"model": "test-model"}, + }); err != nil { + t.Fatalf("Submit B error: %s", err.Message) + } + + // Push results + pushResult(t, mr, resultQueue, "job-b-req", `{"id":"job-b-req"}`) + pushResult(t, mr, resultQueue, "job-a-req", `{"id":"job-a-req"}`) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + respA, err := clientA.GetResult(ctx) + if err != nil { + t.Fatalf("GetResult A error: %v", err) + } + if respA.RequestID != "job-a-req" { + t.Errorf("client A got %q, want %q", respA.RequestID, "job-a-req") + } + + respB, err := clientB.GetResult(ctx) + if err != nil { + t.Fatalf("GetResult B error: %v", err) + } + if respB.RequestID != "job-b-req" { + t.Errorf("client B got %q, want %q", respB.RequestID, "job-b-req") + } + }) +} + +func TestAsyncProducerClient_SubmitPropagatesTraceContext(t *testing.T) { + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + otel.SetTracerProvider(tp) + otel.SetTextMapPropagator(propagation.TraceContext{}) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + + mr := miniredis.RunT(t) + poolName := "otel-pool" + requestQueue := asyncQueuePrefix + "requests:" + poolName + + pool := newTestPool(t, mr, poolName) + client := newAsyncProducerClient(pool) + defer func() { _ = client.Close() }() + + // Create a parent span to simulate the job runner's trace context + ctx, parentSpan := otel.Tracer("test").Start(context.Background(), "process-batch") + parentTraceID := parentSpan.SpanContext().TraceID().String() + + if err := client.Submit(ctx, &GenerateRequest{ + RequestID: "otel-req-1", + Endpoint: "/v1/completions", + Params: map[string]any{"model": "test-model", "prompt": "hello"}, + }); err != nil { + t.Fatalf("Submit error: %s", err.Message) + } + parentSpan.End() + + // Read the enqueued message from Redis and verify it carries traceparent + members, err := mr.ZMembers(requestQueue) + if err != nil { + t.Fatalf("ZMembers error: %v", err) + } + if len(members) == 0 { + t.Fatal("expected at least one message in request queue") + } + + var ir api.InternalRequest + if err := json.Unmarshal([]byte(members[0]), &ir); err != nil { + t.Fatalf("unmarshal InternalRequest: %v", err) + } + if ir.PublicRequest == nil { + t.Fatal("expected PublicRequest in InternalRequest") + } + metadata := ir.PublicRequest.ReqMetadata() + if metadata == nil { + t.Fatal("expected non-nil Metadata on enqueued request") + } + + traceparent, ok := metadata["traceparent"] + if !ok { + t.Fatal("expected 'traceparent' key in request Metadata") + } + if len(traceparent) == 0 { + t.Fatal("expected non-empty traceparent value") + } + + if !strings.Contains(traceparent, parentTraceID) { + t.Errorf("traceparent %q does not contain parent trace ID %q", traceparent, parentTraceID) + } +} diff --git a/pkg/clients/inference/inference_client_resolver.go b/pkg/clients/inference/inference_client_resolver.go index f016136f0..2b7864ec3 100644 --- a/pkg/clients/inference/inference_client_resolver.go +++ b/pkg/clients/inference/inference_client_resolver.go @@ -17,7 +17,9 @@ limitations under the License. package inference import ( + "errors" "fmt" + "io" "time" "github.com/go-logr/logr" @@ -77,6 +79,7 @@ type GatewayResolver struct { globalClient InferenceClient modelClients map[string]InferenceClient clientURLs map[InferenceClient]string + closers []io.Closer } // NewGlobalResolver creates a GatewayResolver where all models resolve to a @@ -170,6 +173,18 @@ func (r *GatewayResolver) ClientLabel(c InferenceClient) string { return "unknown" } +// Close releases resources held by the resolver (e.g. Redis connections for +// async dispatch). Safe to call on resolvers that hold no closeable resources. +func (r *GatewayResolver) Close() error { + var errs []error + for _, c := range r.closers { + if err := c.Close(); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + // NewSingleClientResolver wraps a single InferenceClient in a GatewayResolver // where all models resolve to that client. Used in tests to inject mock // inference clients into Clientset. diff --git a/scripts/dev-clean.sh b/scripts/dev-clean.sh index fb1f5efaf..6ddb6a4bf 100755 --- a/scripts/dev-clean.sh +++ b/scripts/dev-clean.sh @@ -21,6 +21,20 @@ cleanup_kubernetes_resources() { || warn "Failed to delete InferenceObjective CRDs" log "Uninstalling helm releases..." + # Clean up dispatcher port-forward + for pid_file in "${SCRIPT_DIR}/../.dispatcher-port-forward.pid" "${SCRIPT_DIR}/../.dispatcher-sim-port-forward.pid"; do + if [[ -f "${pid_file}" ]]; then + local pf_pid + pf_pid=$(cat "${pid_file}") + kill "${pf_pid}" 2>/dev/null || true + rm -f "${pid_file}" + log "Stopped port-forward (PID: ${pf_pid})" + fi + done + + helm uninstall "${DISPATCHER_RELEASE:-dispatcher}" -n "${NAMESPACE}" 2>/dev/null || warn "Failed to uninstall dispatcher (may not exist)" + helm uninstall "${DISPATCHER_SCRAPE_RELEASE:-dispatcher-scrape}" -n "${NAMESPACE}" 2>/dev/null || warn "Failed to uninstall dispatcher-scrape (may not exist)" + helm uninstall "${DISPATCHER_PROM_RELEASE:-dispatcher-prom}" -n "${NAMESPACE}" 2>/dev/null || warn "Failed to uninstall dispatcher-prom (may not exist)" helm uninstall "${HELM_RELEASE}" -n "${NAMESPACE}" 2>/dev/null || warn "Failed to uninstall ${HELM_RELEASE} (may not exist)" helm uninstall "${REDIS_RELEASE}" -n "${NAMESPACE}" 2>/dev/null || warn "Failed to uninstall ${REDIS_RELEASE} (may not exist)" helm uninstall "${POSTGRESQL_RELEASE}" -n "${NAMESPACE}" 2>/dev/null || warn "Failed to uninstall ${POSTGRESQL_RELEASE} (may not exist)" diff --git a/scripts/dev-deploy-dispatcher.sh b/scripts/dev-deploy-dispatcher.sh new file mode 100755 index 000000000..d4ae13b14 --- /dev/null +++ b/scripts/dev-deploy-dispatcher.sh @@ -0,0 +1,284 @@ +#!/bin/bash +# This script can be run standalone or sourced from dev-deploy.sh. +# When sourced, SCRIPT_DIR, REPO_ROOT, and dev-common.sh are already set. +if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then + set -euo pipefail + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + source "${SCRIPT_DIR}/dev-common.sh" +fi + +# ── Configuration ──────────────────────────────────────────────────────────── +KIND_CLUSTER_NAME="${KIND_CLUSTER_NAME:-batch-gateway-dev}" +DISPATCHER_RELEASE="${DISPATCHER_RELEASE:-dispatcher}" +DISPATCHER_VERSION="${DISPATCHER_VERSION:-v0.7.2}" +DISPATCHER_IMAGE="${DISPATCHER_IMAGE:-ghcr.io/llm-d-incubation/llm-d-async:${DISPATCHER_VERSION}}" +DISPATCHER_CHART="${DISPATCHER_CHART:-oci://ghcr.io/llm-d-incubation/charts/async-processor}" +DISPATCHER_CHART_VERSION="${DISPATCHER_CHART_VERSION:-0.7.2}" +DISPATCHER_REDIS_PORT="${DISPATCHER_REDIS_PORT:-6399}" +PID_FILE="${REPO_ROOT}/.dispatcher-port-forward.pid" +# Set DISPATCHER_SOURCE to a local llm-d-async checkout to build from source +# instead of pulling a released image. The local chart is used automatically. +# Example: DISPATCHER_SOURCE=~/src/llm-d-async ENABLE_DISPATCHER=true make dev-deploy +DISPATCHER_SOURCE="${DISPATCHER_SOURCE:-}" + +# ── Prerequisites (standalone only — dev-deploy.sh already checks these) ───── +if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then + for cmd in kubectl helm kind jq nc; do + command -v "$cmd" &>/dev/null || die "Missing required tool: $cmd" + done + + if [[ -n "${CONTAINER_TOOL:-}" ]]; then + : # caller specified + elif command -v docker &>/dev/null && docker info &>/dev/null 2>&1; then + CONTAINER_TOOL="docker" + elif command -v podman &>/dev/null; then + CONTAINER_TOOL="podman" + else + die "Neither docker (running) nor podman found. Please install one." + fi + + if ! kind get clusters 2>/dev/null | grep -qx "${KIND_CLUSTER_NAME}"; then + die "Kind cluster '${KIND_CLUSTER_NAME}' not found. Run 'make dev-deploy' first." + fi +fi + +# ── Build or pull dispatcher image ──────────────────────────────────────────── +if [[ -n "${DISPATCHER_SOURCE}" ]]; then + if [[ ! -d "${DISPATCHER_SOURCE}" ]]; then + die "DISPATCHER_SOURCE directory not found: ${DISPATCHER_SOURCE}" + fi + DISPATCHER_IMAGE="ghcr.io/llm-d-incubation/async-processor:dev-local" + DISPATCHER_CHART="${DISPATCHER_SOURCE}/charts/async-processor" + unset DISPATCHER_CHART_VERSION + step "Building async-processor image from ${DISPATCHER_SOURCE}..." + ${CONTAINER_TOOL} build -t "${DISPATCHER_IMAGE}" "${DISPATCHER_SOURCE}" +else + if ${CONTAINER_TOOL} image exists "${DISPATCHER_IMAGE}" 2>/dev/null || \ + ${CONTAINER_TOOL} inspect "${DISPATCHER_IMAGE}" &>/dev/null; then + step "Using local dispatcher image ${DISPATCHER_IMAGE}" + else + step "Pulling dispatcher image ${DISPATCHER_IMAGE}..." + ${CONTAINER_TOOL} pull "${DISPATCHER_IMAGE}" + fi +fi + +step "Loading dispatcher image into Kind cluster '${KIND_CLUSTER_NAME}'..." +if [[ "${CONTAINER_TOOL}" == "docker" ]]; then + kind load docker-image "${DISPATCHER_IMAGE}" --name "${KIND_CLUSTER_NAME}" +else + ${CONTAINER_TOOL} save "${DISPATCHER_IMAGE}" | kind load image-archive /dev/stdin --name "${KIND_CLUSTER_NAME}" +fi + +# ── Deploy async-processor via Helm ────────────────────────────────────────── +HELM_VALUES="${REPO_ROOT}/test/e2e/dispatcher/helm-values.yaml" + +if [[ ! -f "${HELM_VALUES}" ]]; then + die "Helm values file not found: ${HELM_VALUES}" +fi + +DISPATCHER_SCRAPE_RELEASE="${DISPATCHER_SCRAPE_RELEASE:-dispatcher-scrape}" +HELM_VALUES_SCRAPE="${REPO_ROOT}/test/e2e/dispatcher/helm-values-scrape.yaml" +IMAGE_REPO="$(echo "${DISPATCHER_IMAGE}" | cut -d: -f1)" +IMAGE_TAG="$(echo "${DISPATCHER_IMAGE}" | cut -d: -f2)" + +HELM_VERSION_FLAG=() +if [[ -n "${DISPATCHER_CHART_VERSION:-}" ]]; then + HELM_VERSION_FLAG=(--version "${DISPATCHER_CHART_VERSION}") +fi + +step "Deploying async-processor with redis gate (release: ${DISPATCHER_RELEASE})..." +helm upgrade --install "${DISPATCHER_RELEASE}" "${DISPATCHER_CHART}" \ + "${HELM_VERSION_FLAG[@]}" \ + --namespace "${NAMESPACE}" \ + --values "${HELM_VALUES}" \ + --set "ap.image.repository=${IMAGE_REPO}" \ + --set "ap.image.tag=${IMAGE_TAG}" \ + --wait --timeout=120s + +step "Deploying async-processor with endpoint-scrape gate (release: ${DISPATCHER_SCRAPE_RELEASE})..." +helm upgrade --install "${DISPATCHER_SCRAPE_RELEASE}" "${DISPATCHER_CHART}" \ + "${HELM_VERSION_FLAG[@]}" \ + --namespace "${NAMESPACE}" \ + --values "${HELM_VALUES_SCRAPE}" \ + --set "ap.image.repository=${IMAGE_REPO}" \ + --set "ap.image.tag=${IMAGE_TAG}" \ + --wait --timeout=120s + +DISPATCHER_PROM_RELEASE="${DISPATCHER_PROM_RELEASE:-dispatcher-prom}" +HELM_VALUES_PROM="${REPO_ROOT}/test/e2e/dispatcher/helm-values-prometheus.yaml" + +step "Deploying async-processor with prometheus-query gate (release: ${DISPATCHER_PROM_RELEASE})..." +helm upgrade --install "${DISPATCHER_PROM_RELEASE}" "${DISPATCHER_CHART}" \ + "${HELM_VERSION_FLAG[@]}" \ + --namespace "${NAMESPACE}" \ + --values "${HELM_VALUES_PROM}" \ + --set "ap.image.repository=${IMAGE_REPO}" \ + --set "ap.image.tag=${IMAGE_TAG}" \ + --wait --timeout=120s + +log "Dispatchers deployed." + +# ── Verify dispatchers ─────────────────────────────────────────────────────── +step "Waiting for dispatcher pods to be ready..." +kubectl wait --for=condition=available deployment/"${DISPATCHER_RELEASE}-async-processor" \ + --namespace "${NAMESPACE}" --timeout=60s +kubectl wait --for=condition=available deployment/"${DISPATCHER_SCRAPE_RELEASE}-async-processor" \ + --namespace "${NAMESPACE}" --timeout=60s +kubectl wait --for=condition=available deployment/"${DISPATCHER_PROM_RELEASE}-async-processor" \ + --namespace "${NAMESPACE}" --timeout=60s + +# ── Add vllm-sim to Prometheus scrape targets ──────────────────────────────── +step "Adding vllm-sim to Prometheus scrape config..." +PROM_CM="${PROMETHEUS_NAME:-prometheus}-config" +CURRENT_PROM_CONFIG=$(kubectl get configmap "${PROM_CM}" --namespace "${NAMESPACE}" -o jsonpath='{.data.prometheus\.yml}' 2>/dev/null || true) + +if echo "${CURRENT_PROM_CONFIG}" | grep -q "vllm-sim"; then + log "vllm-sim already in Prometheus scrape config, skipping" +else + # Append vllm-sim scrape job (indentation must match existing scrape_configs entries) + read -r -d '' VLLM_SIM_SCRAPE <<-SCRAPE || true +- job_name: 'vllm-sim' + metrics_path: /metrics + scrape_interval: 5s + static_configs: + - targets: ['${VLLM_SIM_NAME}.${NAMESPACE}.svc.cluster.local:8000'] + labels: + component: vllm-sim +SCRAPE + UPDATED_PROM_CONFIG="${CURRENT_PROM_CONFIG} +${VLLM_SIM_SCRAPE}" + + kubectl create configmap "${PROM_CM}" \ + --namespace "${NAMESPACE}" \ + --from-literal="prometheus.yml=${UPDATED_PROM_CONFIG}" \ + --dry-run=client -o yaml | kubectl apply -f - + + # Reload Prometheus config + kubectl rollout restart deployment/"${PROMETHEUS_NAME:-prometheus}" --namespace "${NAMESPACE}" + kubectl rollout status deployment/"${PROMETHEUS_NAME:-prometheus}" --namespace "${NAMESPACE}" --timeout=60s + log "Prometheus scrape config updated with vllm-sim target" +fi + +# ── Enable fake metrics on vllm-sim ────────────────────────────────────────── +step "Patching vllm-sim to enable --fake-metrics..." +# --fake-metrics requires a JSON argument with initial metric values. +# Replace the full args array to avoid duplicate appends on re-runs. +CURRENT_ARGS=$(kubectl get deployment "${VLLM_SIM_NAME}" --namespace "${NAMESPACE}" \ + -o jsonpath='{.spec.template.spec.containers[0].args}') + +if echo "${CURRENT_ARGS}" | grep -q "fake-metrics"; then + log "vllm-sim already has --fake-metrics, skipping patch" +else + kubectl patch deployment "${VLLM_SIM_NAME}" --namespace "${NAMESPACE}" --type=json \ + -p='[{"op":"add","path":"/spec/template/spec/containers/0/args/-","value":"--fake-metrics"},{"op":"add","path":"/spec/template/spec/containers/0/args/-","value":"{\"kv-cache-usage\": 0, \"waiting-requests\": 0, \"running-requests\": 0}"}]' + kubectl rollout status deployment/"${VLLM_SIM_NAME}" --namespace "${NAMESPACE}" --timeout=120s +fi + +# ── Reconfigure processor for async dispatch ───────────────────────────────── +PROCESSOR_ASYNC_VALUES="${REPO_ROOT}/test/e2e/dispatcher/processor-async-values.yaml" + +step "Reconfiguring batch-gateway processor for async dispatch..." + +# --reuse-values deep-merges maps, so stale sync-mode models from the initial +# deploy would persist and crash the processor (missing inferencePoolName). +# Instead, export current values, strip modelGateways, and pass as a file. +REUSED_VALUES=$(mktemp) +helm get values "${HELM_RELEASE}" -n "${NAMESPACE}" -o json | \ + jq 'del(.processor.config.modelGateways)' > "${REUSED_VALUES}" + +helm upgrade "${HELM_RELEASE}" "${REPO_ROOT}/charts/batch-gateway" \ + --namespace "${NAMESPACE}" \ + --reset-values \ + --values "${REUSED_VALUES}" \ + --values "${PROCESSOR_ASYNC_VALUES}" \ + --wait --timeout=120s +rm -f "${REUSED_VALUES}" + +step "Restarting processor to pick up new config..." +kubectl rollout restart deployment/"${HELM_RELEASE}-processor" --namespace "${NAMESPACE}" +kubectl rollout status deployment/"${HELM_RELEASE}-processor" --namespace "${NAMESPACE}" --timeout=60s + +log "Processor reconfigured for async dispatch." + +# ── Port-forward Redis to host ─────────────────────────────────────────────── +# Kind only exposes ports declared in extraPortMappings at cluster creation. +# Use kubectl port-forward to make Redis accessible from the host. +step "Setting up Redis port-forward on localhost:${DISPATCHER_REDIS_PORT}..." + +# Kill any previous port-forward +if [[ -f "${PID_FILE}" ]]; then + old_pid=$(cat "${PID_FILE}") + kill "${old_pid}" 2>/dev/null || true + rm -f "${PID_FILE}" +fi + +# Detect Redis service name +redis_svc="${REDIS_RELEASE}-master" +if [[ "${EXCHANGE_CLIENT_TYPE}" == "valkey" ]]; then + redis_svc="${REDIS_RELEASE}-valkey-primary" +fi + +kubectl port-forward "svc/${redis_svc}" "${DISPATCHER_REDIS_PORT}:6379" \ + --namespace "${NAMESPACE}" & +PORT_FORWARD_PID=$! +echo "${PORT_FORWARD_PID}" > "${PID_FILE}" + +# Wait for the port-forward to be ready +for i in $(seq 1 10); do + if nc -z localhost "${DISPATCHER_REDIS_PORT}" 2>/dev/null; then + break + fi + sleep 0.5 +done + +if ! nc -z localhost "${DISPATCHER_REDIS_PORT}" 2>/dev/null; then + die "Port-forward to Redis failed to start" +fi + +log "Redis accessible at localhost:${DISPATCHER_REDIS_PORT}" + +# ── Port-forward vLLM sim to host ──────────────────────────────────────────── +DISPATCHER_SIM_PORT="${DISPATCHER_SIM_PORT:-8099}" +SIM_PID_FILE="${REPO_ROOT}/.dispatcher-sim-port-forward.pid" + +step "Setting up vLLM sim port-forward on localhost:${DISPATCHER_SIM_PORT}..." + +if [[ -f "${SIM_PID_FILE}" ]]; then + old_pid=$(cat "${SIM_PID_FILE}") + kill "${old_pid}" 2>/dev/null || true + rm -f "${SIM_PID_FILE}" +fi + +kubectl port-forward "svc/${VLLM_SIM_NAME}" "${DISPATCHER_SIM_PORT}:8000" \ + --namespace "${NAMESPACE}" & +SIM_PF_PID=$! +echo "${SIM_PF_PID}" > "${SIM_PID_FILE}" + +for i in $(seq 1 10); do + if nc -z localhost "${DISPATCHER_SIM_PORT}" 2>/dev/null; then + break + fi + sleep 0.5 +done + +if ! nc -z localhost "${DISPATCHER_SIM_PORT}" 2>/dev/null; then + die "Port-forward to vLLM sim failed to start" +fi + +log "vLLM sim accessible at localhost:${DISPATCHER_SIM_PORT}" + +log "" +log "Dispatcher is ready." +log "" +log "Usage:" +log " ENABLE_DISPATCHER=true make test-e2e" +log " ENABLE_DISPATCHER=true TEST_REDIS_URL=redis://localhost:${DISPATCHER_REDIS_PORT} go test ./test/e2e/ -run TestDispatcher -v -count=1" +log "" +log "Jaeger UI: http://localhost:${JAEGER_PORT} (traces from both batch-gateway and async-processor)" +log "To stop port-forwards: make dev-clean" +log "" +if [[ -n "${DISPATCHER_SOURCE}" ]]; then + log "Built from local source: ${DISPATCHER_SOURCE}" + log "To rebuild after changes: DISPATCHER_SOURCE=${DISPATCHER_SOURCE} ENABLE_DISPATCHER=true make dev-deploy" +fi diff --git a/scripts/dev-deploy.sh b/scripts/dev-deploy.sh index 718772eca..fd09e0072 100755 --- a/scripts/dev-deploy.sh +++ b/scripts/dev-deploy.sh @@ -60,6 +60,12 @@ USE_KIND="${USE_KIND:-true}" # batch-sheddable (non-GIE mode) # Managed-by label: app.kubernetes.io/managed-by=batch-gateway-dev ENABLE_GIE="${ENABLE_GIE:-false}" + +# ── Async dispatcher (llm-d-async) support ────────────────────────────────── +# Set ENABLE_DISPATCHER=true to deploy llm-d-async dispatcher instances +# alongside the batch-gateway. The processor is reconfigured for async dispatch. +# Set DISPATCHER_SOURCE to a local llm-d-async checkout to build from source. +ENABLE_DISPATCHER="${ENABLE_DISPATCHER:-false}" GIE_REPO="${GIE_REPO:-}" GIE_UPSTREAM_REPO="https://github.com/kubernetes-sigs/gateway-api-inference-extension.git" GIE_VERSION="${GIE_VERSION:-v1.5.0}" @@ -136,6 +142,9 @@ nodes: - containerPort: ${MINIO_NODE_PORT} hostPort: ${MINIO_PORT} protocol: TCP + - containerPort: ${REDIS_NODE_PORT:-30479} + hostPort: ${REDIS_PORT:-6399} + protocol: TCP EOF fi @@ -1339,6 +1348,12 @@ print_usage() { echo " - Each model has its own InferencePool and InferenceObjective" echo " - InferenceObjectives: interactive-default (priority 100), ${GIE_OBJECTIVE_PREFIX} (priority -1)" fi + if [ "${ENABLE_DISPATCHER}" = "true" ]; then + echo "" + echo " Async dispatcher is enabled:" + echo " - Processor is configured for async dispatch via llm-d-async" + echo " - Run dispatcher tests: ENABLE_DISPATCHER=true make test-e2e" + fi echo "" echo " 3. Create a batch (replace FILE_ID with the id from step 2):" echo "" @@ -1448,6 +1463,9 @@ main() { if [ "${USE_KIND}" = true ]; then create_nodeport_services fi + if [ "${ENABLE_DISPATCHER}" = "true" ]; then + source "${SCRIPT_DIR}/dev-deploy-dispatcher.sh" + fi print_usage log "Deployment complete!" diff --git a/test/e2e/dispatcher/helm-values-prometheus.yaml b/test/e2e/dispatcher/helm-values-prometheus.yaml new file mode 100644 index 000000000..4243c1f91 --- /dev/null +++ b/test/e2e/dispatcher/helm-values-prometheus.yaml @@ -0,0 +1,32 @@ +ap: + imagePullPolicy: Never + messageQueueImpl: "redis-sortedset" + concurrency: 1 + prometheusURL: "http://prometheus.default.svc.cluster.local:9090" + prometheusCacheTTL: "0s" + otel: + endpoint: "http://jaeger.default.svc.cluster.local:4317" + insecure: true + sampler: "always_on" + samplerArg: "1.0" + redisTracing: false + metrics: + enabled: true + port: 9092 + secure: false + redis: + enabled: true + url: "redis://redis-master.default.svc.cluster.local:6379" + pollIntervalMs: 500 + batchSize: 10 + queuesConfig: + - queue_name: "llm-d-async:requests:sim-pool-prom" + result_queue_name: "llm-d-async:results:sim-pool-prom" + request_path_url: "/v1/completions" + igw_base_url: "http://vllm-sim.default.svc.cluster.local:8000" + gate_type: "prometheus-query" + gate_params: + query: "1 - clamp_max(vllm:num_requests_waiting / 5, 1)" + fallback: "1.0" + modelServerMonitor: + enabled: false diff --git a/test/e2e/dispatcher/helm-values-scrape.yaml b/test/e2e/dispatcher/helm-values-scrape.yaml new file mode 100644 index 000000000..fc8b566b4 --- /dev/null +++ b/test/e2e/dispatcher/helm-values-scrape.yaml @@ -0,0 +1,33 @@ +ap: + imagePullPolicy: Never + messageQueueImpl: "redis-sortedset" + concurrency: 1 + prometheusCacheTTL: "0s" + otel: + endpoint: "http://jaeger.default.svc.cluster.local:4317" + insecure: true + sampler: "always_on" + samplerArg: "1.0" + redisTracing: false + metrics: + enabled: true + port: 9091 + secure: false + redis: + enabled: true + url: "redis://redis-master.default.svc.cluster.local:6379" + pollIntervalMs: 500 + batchSize: 10 + queuesConfig: + - queue_name: "llm-d-async:requests:sim-pool-scrape" + result_queue_name: "llm-d-async:results:sim-pool-scrape" + request_path_url: "/v1/completions" + igw_base_url: "http://vllm-sim.default.svc.cluster.local:8000" + gate_type: "endpoint-scrape" + gate_params: + url: "http://vllm-sim.default.svc.cluster.local:8000/metrics" + metric: "vllm:num_requests_waiting" + max_count_per_pod: "5" + fallback: "1.0" + modelServerMonitor: + enabled: false diff --git a/test/e2e/dispatcher/helm-values.yaml b/test/e2e/dispatcher/helm-values.yaml new file mode 100644 index 000000000..70df79ade --- /dev/null +++ b/test/e2e/dispatcher/helm-values.yaml @@ -0,0 +1,37 @@ +ap: + imagePullPolicy: Never + messageQueueImpl: "redis-sortedset" + concurrency: 1 + prometheusCacheTTL: "0s" + otel: + endpoint: "http://jaeger.default.svc.cluster.local:4317" + insecure: true + sampler: "always_on" + samplerArg: "1.0" + redisTracing: false + metrics: + enabled: true + port: 9090 + secure: false + redis: + enabled: true + url: "redis://redis-master.default.svc.cluster.local:6379" + pollIntervalMs: 500 + batchSize: 10 + queuesConfig: + - queue_name: "llm-d-async:requests:sim-pool" + result_queue_name: "llm-d-async:results:sim-pool" + request_path_url: "/v1/completions" + igw_base_url: "http://vllm-sim.default.svc.cluster.local:8000" + gate_type: "redis" + gate_params: + address: "redis-master.default.svc.cluster.local:6379" + - queue_name: "llm-d-async:requests:sim-pool-gate" + result_queue_name: "llm-d-async:results:sim-pool-gate" + request_path_url: "/v1/completions" + igw_base_url: "http://vllm-sim.default.svc.cluster.local:8000" + gate_type: "redis" + gate_params: + address: "redis-master.default.svc.cluster.local:6379" + modelServerMonitor: + enabled: false diff --git a/test/e2e/dispatcher/processor-async-values.yaml b/test/e2e/dispatcher/processor-async-values.yaml new file mode 100644 index 000000000..7da77b3e4 --- /dev/null +++ b/test/e2e/dispatcher/processor-async-values.yaml @@ -0,0 +1,8 @@ +processor: + config: + dispatchMode: async + asyncDispatch: + resultPollTimeout: "30s" + modelGateways: + sim-model: + inferencePoolName: "sim-pool" diff --git a/test/e2e/dispatcher_otel_test.go b/test/e2e/dispatcher_otel_test.go new file mode 100644 index 000000000..8e73de61b --- /dev/null +++ b/test/e2e/dispatcher_otel_test.go @@ -0,0 +1,146 @@ +// Copyright 2026 The llm-d Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package e2e_test + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "testing" + "time" + + "github.com/openai/openai-go/v3" +) + +func TestDispatcherOTelTraces(t *testing.T) { + if !detectDispatcherDeployed(t) { + t.Skip("skipping: dispatcher not deployed") + } + waitForReady(t, testApiserverObsURL, 30*time.Second) + + jaegerClient := &http.Client{Timeout: 5 * time.Second} + checkResp, err := jaegerClient.Get(testJaegerURL + "/") + if err != nil { + t.Skipf("Jaeger not reachable at %s, skipping OTel trace verification: %v", testJaegerURL, err) + } + checkResp.Body.Close() + + t.Run("CrossServiceTracePropagation", func(t *testing.T) { + testCrossServiceTracePropagation(t, jaegerClient) + }) +} + +// testCrossServiceTracePropagation verifies that a batch processed through the +// async dispatcher produces a connected trace spanning both batch-gateway and +// async-processor. The batch-gateway injects trace context into +// RequestMessage.Metadata, the async processor extracts it, and both services +// export spans to the same Jaeger instance under the same trace ID. +func testCrossServiceTracePropagation(t *testing.T, jaegerClient *http.Client) { + jsonl := fmt.Sprintf( + `{"custom_id":"otel-xsvc-1","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":5,"messages":[{"role":"user","content":"trace propagation test"}]}}`, + testModel) + + fileID := mustCreateFile(t, fmt.Sprintf("otel-xsvc-%s.jsonl", testRunID), jsonl) + batchID := mustCreateBatch(t, fileID) + t.Logf("Created batch %s for cross-service trace test", batchID) + + batch, _ := waitForBatchStatus(t, batchID, 120*time.Second, openai.BatchStatusCompleted) + if batch.RequestCounts.Completed != 1 { + t.Fatalf("Expected 1 completed request, got %d", batch.RequestCounts.Completed) + } + + // Poll Jaeger for traces from batch-gateway that contain async-processor spans. + // The batch-gateway creates a "process-batch" span, and the async-processor + // creates a "process-request" child span under the same trace ID. + type jaegerSpan struct { + OperationName string `json:"operationName"` + ProcessID string `json:"processID"` + } + type jaegerTrace struct { + TraceID string `json:"traceID"` + Spans []jaegerSpan `json:"spans"` + Processes map[string]struct { + ServiceName string `json:"serviceName"` + } `json:"processes"` + } + type jaegerResponse struct { + Data []jaegerTrace `json:"data"` + } + + var found *jaegerTrace + deadline := time.After(30 * time.Second) + for found == nil { + select { + case <-deadline: + t.Fatal("Timed out waiting for cross-service trace in Jaeger") + default: + } + + resp, err := jaegerClient.Get(testJaegerURL + "/api/traces?service=async-processor&limit=20&lookback=5m") + if err != nil { + t.Logf("Jaeger query failed (retrying): %v", err) + time.Sleep(2 * time.Second) + continue + } + + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + var result jaegerResponse + if err := json.Unmarshal(body, &result); err != nil { + t.Logf("Failed to parse Jaeger response (retrying): %v", err) + time.Sleep(2 * time.Second) + continue + } + + for i := range result.Data { + trace := &result.Data[i] + services := make(map[string]bool) + for _, span := range trace.Spans { + if proc, ok := trace.Processes[span.ProcessID]; ok { + services[proc.ServiceName] = true + } + } + if services["batch-gateway"] && services["async-processor"] { + found = trace + break + } + } + + if found == nil { + time.Sleep(2 * time.Second) + } + } + + // Verify expected span operations exist in the trace + spanOps := make(map[string]bool) + for _, span := range found.Spans { + spanOps[span.OperationName] = true + } + + if !spanOps["process-request"] { + t.Error("Expected 'process-request' span from async-processor in trace") + } + + // Collect service names for logging + services := make(map[string]bool) + for _, proc := range found.Processes { + services[proc.ServiceName] = true + } + + t.Logf("Cross-service trace verified: traceID=%s, services=%v, spans=%v", + found.TraceID, services, spanOps) +} diff --git a/test/e2e/dispatcher_test.go b/test/e2e/dispatcher_test.go new file mode 100644 index 000000000..df339726b --- /dev/null +++ b/test/e2e/dispatcher_test.go @@ -0,0 +1,429 @@ +// Copyright 2026 The llm-d Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package e2e_test + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "os/exec" + "strings" + "testing" + "time" + + asyncapi "github.com/llm-d-incubation/llm-d-async/api" + "github.com/llm-d-incubation/llm-d-async/producer" + "github.com/openai/openai-go/v3" + "github.com/redis/go-redis/v9" +) + +var ( + testRedisURL = getEnvOrDefault("TEST_REDIS_URL", "redis://localhost:6399") + testSimURL = getEnvOrDefault("TEST_SIM_URL", "http://localhost:8099") + dispatcherPool = "sim-pool" + dispatcherReqQueue = "llm-d-async:requests:" + dispatcherPool + dispatcherResultQueue = "llm-d-async:results:" + dispatcherPool + + gatePool = "sim-pool-gate" + gateReqQueue = "llm-d-async:requests:" + gatePool + gateResultQueue = "llm-d-async:results:" + gatePool + dispatchGateBudgetKey = "dispatch-gate-budget" + + scrapePool = "sim-pool-scrape" + scrapeReqQueue = "llm-d-async:requests:" + scrapePool + scrapeResultQueue = "llm-d-async:results:" + scrapePool + + promPool = "sim-pool-prom" + promReqQueue = "llm-d-async:requests:" + promPool + promResultQueue = "llm-d-async:results:" + promPool +) + +func newDispatcherRedisClient(t *testing.T) *redis.Client { + t.Helper() + opts, err := redis.ParseURL(testRedisURL) + if err != nil { + t.Fatalf("Failed to parse TEST_REDIS_URL %q: %v", testRedisURL, err) + } + client := redis.NewClient(opts) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := client.Ping(ctx).Err(); err != nil { + t.Fatalf("Failed to connect to Redis at %s: %v", testRedisURL, err) + } + return client +} + +func newDispatcherProducer(t *testing.T, rdb *redis.Client, poolName string) *producer.RedisSortedSetProducer { + t.Helper() + p, err := producer.NewRedisSortedSetProducer( + producer.RedisSortedSetConfig{ + RequestQueueName: "llm-d-async:requests:" + poolName, + ResultQueueName: "llm-d-async:results:" + poolName, + }, + producer.WithRedisClient(rdb), + ) + if err != nil { + t.Fatalf("Failed to create producer for pool %s: %v", poolName, err) + } + t.Cleanup(func() { _ = p.Close() }) + return p +} + +// detectDispatcherDeployed checks whether at least one async-processor +// deployment exists in the test namespace. +func detectDispatcherDeployed(t *testing.T) bool { + t.Helper() + + out, err := exec.Command("kubectl", "get", "deployments", + "-n", testNamespace, + "-o", "name", + ).CombinedOutput() + if err != nil { + t.Logf("kubectl get deployments failed: %v", err) + return false + } + for _, line := range strings.Split(string(out), "\n") { + if strings.Contains(line, "async-processor") { + return true + } + } + return false +} + +func TestDispatcher(t *testing.T) { + if !detectDispatcherDeployed(t) { + t.Skip("skipping: dispatcher not deployed") + } + rdb := newDispatcherRedisClient(t) + defer rdb.Close() + + waitForReady(t, testApiserverObsURL, 30*time.Second) + + t.Cleanup(func() { + ctx := context.Background() + rdb.Del(ctx, dispatcherReqQueue, dispatcherResultQueue) + rdb.Del(ctx, gateReqQueue, gateResultQueue) + rdb.Del(ctx, scrapeReqQueue, scrapeResultQueue) + rdb.Del(ctx, promReqQueue, promResultQueue) + }) + + t.Run("BatchThroughDispatcher", func(t *testing.T) { + testDispatcherBatchRoundTrip(t, rdb) + }) + t.Run("MultiRequestBatch", func(t *testing.T) { + testDispatcherMultiRequestBatch(t, rdb) + }) + t.Run("DispatchGate", func(t *testing.T) { + testDispatcherRedisGate(t, rdb) + }) + t.Run("EndpointScrapeGate", func(t *testing.T) { + testDispatcherEndpointScrapeGate(t, rdb) + }) + t.Run("PrometheusGate", func(t *testing.T) { + testDispatcherPrometheusGate(t, rdb) + }) +} + +func testDispatcherBatchRoundTrip(t *testing.T, rdb *redis.Client) { + jsonl := fmt.Sprintf( + `{"custom_id":"dreq-1","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":5,"messages":[{"role":"user","content":"Hello dispatcher"}]}}`, + testModel) + + fileID := mustCreateFile(t, fmt.Sprintf("dispatcher-single-%s.jsonl", testRunID), jsonl) + batchID := mustCreateBatch(t, fileID) + t.Logf("Created batch %s with file %s", batchID, fileID) + + batch, results := waitForBatchStatus(t, batchID, 120*time.Second, openai.BatchStatusCompleted) + + if batch.RequestCounts.Total != 1 { + t.Errorf("Expected 1 total request, got %d", batch.RequestCounts.Total) + } + if batch.RequestCounts.Completed != 1 { + t.Errorf("Expected 1 completed request, got %d", batch.RequestCounts.Completed) + } + if results == nil { + t.Fatal("Expected non-nil results") + } + if results.OutputLines != 1 { + t.Errorf("Expected 1 output line, got %d", results.OutputLines) + } + + t.Logf("Batch %s completed via dispatcher", batchID) +} + +func testDispatcherMultiRequestBatch(t *testing.T, rdb *redis.Client) { + jsonl := strings.Join([]string{ + fmt.Sprintf(`{"custom_id":"dreq-m1","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":5,"messages":[{"role":"user","content":"Hello 1"}]}}`, testModel), + fmt.Sprintf(`{"custom_id":"dreq-m2","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":5,"messages":[{"role":"user","content":"Hello 2"}]}}`, testModel), + fmt.Sprintf(`{"custom_id":"dreq-m3","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":5,"messages":[{"role":"user","content":"Hello 3"}]}}`, testModel), + }, "\n") + + fileID := mustCreateFile(t, fmt.Sprintf("dispatcher-multi-%s.jsonl", testRunID), jsonl) + batchID := mustCreateBatch(t, fileID) + t.Logf("Created batch %s with 3 requests", batchID) + + batch, results := waitForBatchStatus(t, batchID, 120*time.Second, openai.BatchStatusCompleted) + + if batch.RequestCounts.Total != 3 { + t.Errorf("Expected 3 total requests, got %d", batch.RequestCounts.Total) + } + if batch.RequestCounts.Completed != 3 { + t.Errorf("Expected 3 completed requests, got %d", batch.RequestCounts.Completed) + } + if results == nil { + t.Fatal("Expected non-nil results") + } + if results.OutputLines != 3 { + t.Errorf("Expected 3 output lines, got %d", results.OutputLines) + } + + t.Logf("All 3 requests completed via dispatcher") +} + +func testDispatcherRedisGate(t *testing.T, rdb *redis.Client) { + ctx := context.Background() + + p := newDispatcherProducer(t, rdb, gatePool) + + // Close the gate before submitting any request + rdb.Set(ctx, dispatchGateBudgetKey, "0.0", 0) + defer rdb.Del(ctx, dispatchGateBudgetKey) + t.Log("Gate closed (budget=0.0)") + + // Wait for the dispatcher to see the closed gate (poll interval is 500ms) + time.Sleep(2 * time.Second) + + // Enqueue a request via the producer (bypass the processor + // so we can observe the gate independently of BRPOP timeouts) + reqID := fmt.Sprintf("gate-test-%s", testRunID) + err := p.SubmitRequest(ctx, &asyncapi.RequestMessage{ + ID: reqID, + Created: time.Now().Unix(), + Deadline: time.Now().Add(5 * time.Minute).Unix(), + Payload: map[string]any{"model": testModel, "prompt": "Hello gate", "max_tokens": 5}, + Endpoint: "/v1/completions", + }) + if err != nil { + t.Fatalf("Failed to submit request: %v", err) + } + t.Logf("Enqueued request %s while gate is closed", reqID) + + // Verify request stays in the queue (gate blocks dispatch) + time.Sleep(3 * time.Second) + queueDepth, _ := rdb.ZCard(ctx, gateReqQueue).Result() + if queueDepth == 0 { + t.Fatal("Expected request in dispatcher queue while gate is closed, but queue is empty") + } + t.Logf("Confirmed: request stuck in queue (depth=%d, gate closed)", queueDepth) + + // Open the gate + rdb.Set(ctx, dispatchGateBudgetKey, "1.0", 0) + t.Log("Gate opened (budget=1.0)") + + // Wait for the dispatcher to process the request + deadline := time.After(30 * time.Second) + for { + select { + case <-deadline: + t.Fatal("Timed out waiting for dispatcher to drain queue after gate opened") + default: + } + depth, _ := rdb.ZCard(ctx, gateReqQueue).Result() + if depth == 0 { + break + } + time.Sleep(500 * time.Millisecond) + } + + // Poll results via the producer until we find ours + pollCtx, pollCancel := context.WithTimeout(ctx, 10*time.Second) + defer pollCancel() + for { + result, err := p.GetResult(pollCtx) + if err != nil { + t.Fatalf("Failed to get result for %s: %v", reqID, err) + } + if result.ID == reqID { + t.Logf("Request %s completed after gate opened", reqID) + return + } + t.Logf("Skipped stale result %s", result.ID) + } +} + +func testDispatcherEndpointScrapeGate(t *testing.T, rdb *redis.Client) { + ctx := context.Background() + + p := newDispatcherProducer(t, rdb, scrapePool) + + // Saturate the sim — gate should close + // (endpoint-scrape gate: vllm:num_requests_waiting / max_count_per_pod >= 1 → budget 0) + setSimWaitingRequests(t, 10) + defer setSimWaitingRequests(t, 0) + t.Log("Sim saturated (waiting-requests=10, gate should close)") + + // Give the scrape gate time to poll the new metric value + time.Sleep(3 * time.Second) + + // Enqueue a request — it should stay in the queue (gate closed) + reqID := fmt.Sprintf("scrape-gate-%s", testRunID) + err := p.SubmitRequest(ctx, &asyncapi.RequestMessage{ + ID: reqID, + Created: time.Now().Unix(), + Deadline: time.Now().Add(5 * time.Minute).Unix(), + Payload: map[string]any{"model": testModel, "prompt": "Hello scrape gate", "max_tokens": 5}, + Endpoint: "/v1/completions", + }) + if err != nil { + t.Fatalf("Failed to submit request: %v", err) + } + t.Logf("Enqueued request %s while gate is closed", reqID) + + // Verify no result arrives (gate closed) + time.Sleep(3 * time.Second) + queueDepth, _ := rdb.ZCard(ctx, scrapeReqQueue).Result() + if queueDepth == 0 { + t.Fatal("Expected request in queue while gate is closed, but queue is empty") + } + t.Logf("Confirmed: request stuck in queue (depth=%d, gate closed)", queueDepth) + + // Clear saturation — gate should open + setSimWaitingRequests(t, 0) + t.Log("Sim idle (waiting-requests=0, gate should open)") + + // Wait for the scrape gate to pick up the updated metrics and dispatcher to drain + deadline := time.After(30 * time.Second) + for { + select { + case <-deadline: + t.Fatal("Timed out waiting for dispatcher to drain queue after gate opened") + default: + } + depth, _ := rdb.ZCard(ctx, scrapeReqQueue).Result() + if depth == 0 { + break + } + time.Sleep(500 * time.Millisecond) + } + + // Poll result + pollCtx, pollCancel := context.WithTimeout(ctx, 10*time.Second) + defer pollCancel() + for { + result, err := p.GetResult(pollCtx) + if err != nil { + t.Fatalf("Failed to get result for %s: %v", reqID, err) + } + if result.ID == reqID { + t.Logf("Request %s completed after endpoint-scrape gate opened", reqID) + return + } + t.Logf("Skipped stale result %s", result.ID) + } +} + +func setSimWaitingRequests(t *testing.T, count int) { + t.Helper() + body, err := json.Marshal(map[string]any{"waiting-requests": count}) + if err != nil { + t.Fatalf("Failed to marshal fake_metrics body: %v", err) + } + req, err := http.NewRequest(http.MethodPost, testSimURL+"/fake_metrics", bytes.NewReader(body)) + if err != nil { + t.Fatalf("Failed to build fake_metrics request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to set fake_metrics: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent { + t.Fatalf("fake_metrics returned %d", resp.StatusCode) + } +} + +func testDispatcherPrometheusGate(t *testing.T, rdb *redis.Client) { + ctx := context.Background() + + p := newDispatcherProducer(t, rdb, promPool) + + // Saturate the sim — gate should close + // (query: 1 - clamp_max(vllm:num_requests_waiting / 5, 1) → 0 when waiting ≥ 5) + setSimWaitingRequests(t, 10) + defer setSimWaitingRequests(t, 0) + t.Log("Sim saturated (waiting-requests=10, gate should close)") + + // Give Prometheus time to scrape the new metric value + time.Sleep(10 * time.Second) + + // Enqueue a request — it should stay in the queue (gate closed) + reqID := fmt.Sprintf("prom-gate-%s", testRunID) + err := p.SubmitRequest(ctx, &asyncapi.RequestMessage{ + ID: reqID, + Created: time.Now().Unix(), + Deadline: time.Now().Add(5 * time.Minute).Unix(), + Payload: map[string]any{"model": testModel, "prompt": "Hello prom gate", "max_tokens": 5}, + Endpoint: "/v1/completions", + }) + if err != nil { + t.Fatalf("Failed to submit request: %v", err) + } + t.Logf("Enqueued request %s while gate is closed", reqID) + + // Verify no result arrives (gate closed) + time.Sleep(5 * time.Second) + queueDepth, _ := rdb.ZCard(ctx, promReqQueue).Result() + if queueDepth == 0 { + t.Fatal("Expected request in queue while gate is closed, but queue is empty") + } + t.Logf("Confirmed: request stuck in queue (depth=%d, gate closed)", queueDepth) + + // Clear saturation — gate should open + setSimWaitingRequests(t, 0) + t.Log("Sim idle (waiting-requests=0, gate should open)") + + // Wait for Prometheus to scrape the updated metric and dispatcher to react + deadline := time.After(60 * time.Second) + for { + select { + case <-deadline: + t.Fatal("Timed out waiting for dispatcher to drain queue after gate opened") + default: + } + depth, _ := rdb.ZCard(ctx, promReqQueue).Result() + if depth == 0 { + break + } + time.Sleep(1 * time.Second) + } + + // Poll result + pollCtx, pollCancel := context.WithTimeout(ctx, 15*time.Second) + defer pollCancel() + for { + result, err := p.GetResult(pollCtx) + if err != nil { + t.Fatalf("Failed to get result for %s: %v", reqID, err) + } + if result.ID == reqID { + t.Logf("Request %s completed after Prometheus gate opened", reqID) + return + } + t.Logf("Skipped stale result %s", result.ID) + } +} diff --git a/test/e2e/go.mod b/test/e2e/go.mod index b012aa024..fe31a8e23 100644 --- a/test/e2e/go.mod +++ b/test/e2e/go.mod @@ -1,15 +1,23 @@ module github.com/llm-d/llm-d-batch-gateway/test/e2e -go 1.25 +go 1.25.0 + +require ( + github.com/llm-d-incubation/llm-d-async/api v0.7.2 + github.com/llm-d-incubation/llm-d-async/producer v0.7.2 +) require ( github.com/openai/openai-go/v3 v3.24.0 + github.com/redis/go-redis/v9 v9.19.0 gopkg.in/yaml.v3 v3.0.1 ) require ( + github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/tidwall/gjson v1.18.0 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect + go.uber.org/atomic v1.11.0 // indirect ) diff --git a/test/e2e/go.sum b/test/e2e/go.sum index d33a2a4da..15d412fd8 100644 --- a/test/e2e/go.sum +++ b/test/e2e/go.sum @@ -1,5 +1,27 @@ +github.com/alicebob/miniredis/v2 v2.37.0 h1:RheObYW32G1aiJIj81XVt78ZHJpHonHLHW7OLIshq68= +github.com/alicebob/miniredis/v2 v2.37.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= +github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/llm-d-incubation/llm-d-async/api v0.7.2 h1:sf6iFDa5LpVoKDYiOOlsbnv+6Ykj5TA22yRqMdIbOfY= +github.com/llm-d-incubation/llm-d-async/api v0.7.2/go.mod h1:m2zJUwD/AZypJv8RPws3uvCnfQoNW7iv+hH9wPk/Xw4= +github.com/llm-d-incubation/llm-d-async/producer v0.7.2 h1:hhnLqi+MXq8kBjsQhnUCWYtTCG6uFTZLCH296CZMlpY= +github.com/llm-d-incubation/llm-d-async/producer v0.7.2/go.mod h1:IrClO4XwMtgLRnlkAqO1LDsuqmwtQjELDabT2f+5d40= github.com/openai/openai-go/v3 v3.24.0 h1:08x6GnYiB+AAejTo6yzPY8RkZMJQ8NpreiOyM5QfyYU= github.com/openai/openai-go/v3 v3.24.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/redis/go-redis/v9 v9.19.0 h1:XPVaaPSnG6RhYf7p+rmSa9zZfeVAnWsH5h3lxthOm/k= +github.com/redis/go-redis/v9 v9.19.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= @@ -10,6 +32,14 @@ github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= +github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= +go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= +go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= +golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=