Merge remote-tracking branch 'upstream/main'

moulalis · moulalis · commit acbbd0be5b38 · 2026-07-16T00:30:37.000Z
diff --git a/.gitignore b/.gitignore
@@ -11,3 +11,4 @@ benchmarks/results/*
 !benchmarks/results/.gitkeep
 .dispatcher-port-forward.pid
 .dispatcher-sim-port-forward.pid
+*.DS_Store
diff --git a/internal/processor/worker/executor.go b/internal/processor/worker/executor.go
@@ -35,10 +35,16 @@ import (
 	"github.com/go-logr/logr"
 	"github.com/google/uuid"
 
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	semconv "go.opentelemetry.io/otel/semconv/v1.26.0"
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/llm-d/llm-d-batch-gateway/internal/processor/metrics"
 	"github.com/llm-d/llm-d-batch-gateway/internal/shared/openai"
 	batch_types "github.com/llm-d/llm-d-batch-gateway/internal/shared/types"
 	"github.com/llm-d/llm-d-batch-gateway/internal/util/logging"
+	uotel "github.com/llm-d/llm-d-batch-gateway/internal/util/otel"
 	httpclient "github.com/llm-d/llm-d-batch-gateway/pkg/clients/http"
 	"github.com/llm-d/llm-d-batch-gateway/pkg/clients/inference"
 )
@@ -180,6 +186,10 @@ func (ep *executionProgress) counts() *openai.BatchRequestCounts {
 // SLO expiry or SIGTERM. Its sole purpose is to let the drain phase distinguish user cancel from
 // SLO expiry.
 func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx context.Context, params *jobExecutionParams) (*openai.BatchRequestCounts, error) {
+	// requestAbortCtx is derived from sloCtx (not ctx), so it doesn't carry
+	// the execute-job span. Graft it so process-model spans nest correctly.
+	requestAbortCtx = trace.ContextWithSpan(requestAbortCtx, trace.SpanFromContext(ctx))
+
 	logger := logr.FromContextOrDiscard(ctx)
 	logger.V(logging.INFO).Info("Starting execution: executing job")
 
@@ -193,6 +203,11 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte
 		return nil, fmt.Errorf("failed to read model map: %w", err)
 	}
 
+	uotel.SetAttr(ctx,
+		attribute.Int(uotel.AttrModelCount, len(modelMap.SafeToModel)),
+		attribute.Int64(uotel.AttrInputLineCount, modelMap.LineCount),
+	)
+
 	// Early SLO check: if the deadline already fired before execution begins (e.g. SLO expired
 	// during ingestion), skip dispatch entirely. No output file is written since no requests
 	// were dispatched, but error.jsonl may already contain model_not_found entries from
@@ -279,10 +294,13 @@ 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) {
+			modelCtx, modelSpan := uotel.StartSpan(requestAbortCtx, "process-model",
+				trace.WithAttributes(semconv.GenAiRequestModel(modelID)),
+			)
 			var err error
 			if p.asyncInference != nil {
 				err = p.processModelAsync(
-					requestAbortCtx,
+					modelCtx,
 					ctx,
 					sloCtx,
 					userCancelCtx,
@@ -295,7 +313,7 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte
 				)
 			} else {
 				err = p.processModel(
-					requestAbortCtx,
+					modelCtx,
 					ctx,
 					sloCtx,
 					userCancelCtx,
@@ -307,6 +325,11 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte
 					tenantID,
 				)
 			}
+			if err != nil && !errors.Is(err, errExpired) && !errors.Is(err, errCancelled) && !errors.Is(err, errShutdown) {
+				modelSpan.RecordError(err)
+				modelSpan.SetStatus(codes.Error, "process-model failed")
+			}
+			modelSpan.End()
 			// 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
@@ -422,6 +445,7 @@ func (p *Processor) processModel(
 	if err != nil {
 		return fmt.Errorf("model setup failed: read plan for model %s: %w", modelID, err)
 	}
+	uotel.SetAttr(requestAbortCtx, attribute.Int(uotel.AttrRequestCount, len(entries)))
 
 	logger.V(logging.INFO).Info("Processing requests for a model", "numEntries", len(entries))
 
@@ -563,6 +587,7 @@ func (p *Processor) processModelAsync(
 	if err != nil {
 		return fmt.Errorf("model setup failed: read plan for model %s: %w", modelID, err)
 	}
+	uotel.SetAttr(requestAbortCtx, attribute.Int(uotel.AttrRequestCount, len(entries)))
 
 	logger.V(logging.INFO).Info("Processing requests for model (async)", "numEntries", len(entries))
 
diff --git a/internal/processor/worker/job_runner.go b/internal/processor/worker/job_runner.go
@@ -168,12 +168,17 @@ func (p *Processor) runJob(ctx context.Context, params *jobExecutionParams) {
 	go p.heartbeat(heartbeatCtx, params.jobItem.ID, requestAbortFn)
 
 	// ingestion: pre-process job (rejects unregistered-model requests early)
-	if err := p.preProcessJob(ctx, sloCtx, userCancelCtx, params.jobInfo); err != nil {
+	ingestCtx, ingestSpan := uotel.StartSpan(ctx, "ingest-and-plan")
+	err = p.preProcessJob(ingestCtx, sloCtx, userCancelCtx, params.jobInfo)
+	if err != nil && !errors.Is(err, errExpired) && !errors.Is(err, errCancelled) && !errors.Is(err, errShutdown) {
+		ingestSpan.RecordError(err)
+		ingestSpan.SetStatus(codes.Error, "pre-process failed")
+	}
+	ingestSpan.End()
+	if err != nil {
 		// errExpired, errCancelled, and errShutdown are expected terminal states, not system errors.
 		if !errors.Is(err, errExpired) && !errors.Is(err, errCancelled) && !errors.Is(err, errShutdown) {
 			logger.Error(err, "Pre-processing failed")
-			span.RecordError(err)
-			span.SetStatus(codes.Error, "pre-process failed")
 		}
 		p.handleJobError(ctx, params, err)
 		// No RecordJobProcessingDuration here: preprocessing is ingestion (parsing, plan
@@ -195,15 +200,16 @@ func (p *Processor) runJob(ctx context.Context, params *jobExecutionParams) {
 	transitionedToInProgress = true
 
 	// execution: execute inference requests
+	execCtx, execSpan := uotel.StartSpan(ctx, "execute-job")
 	var execErr error
-	requestCounts, execErr = p.executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx, params)
+	requestCounts, execErr = p.executeJob(execCtx, sloCtx, userCancelCtx, requestAbortCtx, params)
+	if execErr != nil && !errors.Is(execErr, errExpired) && !errors.Is(execErr, errCancelled) && !errors.Is(execErr, errShutdown) {
+		execSpan.RecordError(execErr)
+		execSpan.SetStatus(codes.Error, "execution failed")
+	}
+	execSpan.End()
 	params.requestCounts = requestCounts
 	if execErr != nil {
-		// errExpired, errCancelled, and errShutdown are expected terminal states, not system errors.
-		if !errors.Is(execErr, errExpired) && !errors.Is(execErr, errCancelled) && !errors.Is(execErr, errShutdown) {
-			span.RecordError(execErr)
-			span.SetStatus(codes.Error, "execution failed")
-		}
 		p.handleJobError(ctx, params, execErr)
 		// Record processing duration for any job that ran (partially or fully).
 		// executeJob always returns non-nil counts alongside its sentinel errors
diff --git a/internal/processor/worker/preprocessor.go b/internal/processor/worker/preprocessor.go
@@ -32,10 +32,13 @@ import (
 	"github.com/go-logr/logr"
 	"github.com/google/uuid"
 
+	"go.opentelemetry.io/otel/attribute"
+
 	"github.com/llm-d/llm-d-batch-gateway/internal/processor/metrics"
 	"github.com/llm-d/llm-d-batch-gateway/internal/shared/openai"
 	batch_types "github.com/llm-d/llm-d-batch-gateway/internal/shared/types"
 	"github.com/llm-d/llm-d-batch-gateway/internal/util/logging"
+	uotel "github.com/llm-d/llm-d-batch-gateway/internal/util/otel"
 	"github.com/llm-d/llm-d-batch-gateway/pkg/clients/inference"
 )
 
@@ -221,7 +224,16 @@ func (p *Processor) preProcessJob(ctx, sloCtx, userCancelCtx context.Context, jo
 		return fmt.Errorf("write model map: %w", err)
 	}
 
-	metrics.RecordPlanBuildDuration(time.Since(planBuildStart), metrics.GetSizeBucket(int(lineCount)))
+	sizeBucket := metrics.GetSizeBucket(int(lineCount))
+	metrics.RecordPlanBuildDuration(time.Since(planBuildStart), sizeBucket)
+
+	uotel.SetAttr(ctx,
+		attribute.Int64(uotel.AttrInputLineCount, lineCount),
+		attribute.Int(uotel.AttrModelCount, len(modelToSafe)),
+		attribute.Int64(uotel.AttrRejectedCount, rejectedCount),
+		attribute.String(uotel.AttrSizeBucket, sizeBucket),
+	)
+
 	modelCounts := make(map[string]int, len(modelToSafe))
 	for model, safe := range modelToSafe {
 		modelCounts[model] = len(acc.entries[safe])
diff --git a/internal/util/otel/otel.go b/internal/util/otel/otel.go
@@ -47,6 +47,11 @@ const (
 	AttrRequestTotal     = "batch.request.total"
 	AttrRequestCompleted = "batch.request.completed"
 	AttrRequestFailed    = "batch.request.failed"
+	AttrModelCount       = "batch.model.count"
+	AttrRequestCount     = "batch.request.count"
+	AttrInputLineCount   = "batch.input.line_count"
+	AttrRejectedCount    = "batch.input.rejected_count"
+	AttrSizeBucket       = "batch.size_bucket"
 )
 
 // baseLoggerKey stores the logger captured before the first trace enrichment.
diff --git a/test/e2e/batches_test.go b/test/e2e/batches_test.go
@@ -74,11 +74,11 @@ func doTestBatchCancel(t *testing.T) {
 	var lines []string
 	for i := 1; i <= 5; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"fast-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"sim-model","max_tokens":1,"messages":[{"role":"user","content":"Hi %d"}]}}`, i, i))
+			`{"custom_id":"fast-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":1,"messages":[{"role":"user","content":"Hi %d"}]}}`, i, testSimModel, i))
 	}
 	for i := 1; i <= 20; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"slow-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"sim-model","max_tokens":200,"messages":[{"role":"user","content":"Tell me a long story %d"}]}}`, i, i))
+			`{"custom_id":"slow-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"Tell me a long story %d"}]}}`, i, testSimModel, i))
 	}
 	slowJSONL := strings.Join(lines, "\n")
 	fileID := mustCreateFile(t, fmt.Sprintf("test-batch-cancel-%s.jsonl", testRunID), slowJSONL)
@@ -371,11 +371,6 @@ func doTestBatchMixedSuccessFailure(t *testing.T) {
 func doTestPassThroughHeaders(t *testing.T) {
 	t.Helper()
 
-	// Verify processor logs contain the pass-through header names
-	if !testKubectlAvailable {
-		t.Skip("kubectl not available, skipping processor log verification")
-	}
-
 	// Create batch with pass-through headers
 	fileID := mustCreateFile(t, fmt.Sprintf("test-pass-through-headers-%s.jsonl", testRunID), testJSONL)
 
@@ -405,6 +400,13 @@ func doTestPassThroughHeaders(t *testing.T) {
 		t.Errorf("expected empty error_file_id, got %q", finalBatch.ErrorFileID)
 	}
 
+	// Verify processor logs contain the pass-through header names.
+	// Skip when kubectl is unavailable — the batch completion above already
+	// validates the core pass-through functionality.
+	if !testKubectlAvailable {
+		t.Skip("kubectl not available, skipping processor log verification")
+	}
+
 	out, err := exec.Command("kubectl", "logs",
 		"-l", fmt.Sprintf("app.kubernetes.io/instance=%s,app.kubernetes.io/component=processor", testHelmRelease),
 		"-n", testNamespace,
@@ -519,7 +521,7 @@ func doTestBatchExpiration(t *testing.T) {
 	var blockerLines []string
 	for i := 1; i <= 50; i++ {
 		blockerLines = append(blockerLines, fmt.Sprintf(
-			`{"custom_id":"blocker-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"Block %d"}]}}`, i, testModel, i))
+			`{"custom_id":"blocker-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"Block %d"}]}}`, i, testSimModel, i))
 	}
 	blockerFileID := mustCreateFile(t, fmt.Sprintf("test-expiration-blocker-%s.jsonl", testRunID), strings.Join(blockerLines, "\n"))
 	blockerBatchID := mustCreateBatch(t, blockerFileID)
@@ -545,7 +547,7 @@ func doTestBatchExpiration(t *testing.T) {
 	var lines []string
 	for i := 1; i <= numRequests; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"expire-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"Expire %d"}]}}`, i, testModel, i))
+			`{"custom_id":"expire-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"Expire %d"}]}}`, i, testSimModel, i))
 	}
 	fileID := mustCreateFile(t, fmt.Sprintf("test-batch-expiration-%s.jsonl", testRunID), strings.Join(lines, "\n"))
 
@@ -604,7 +606,7 @@ func doTestCancelIdempotentRetry(t *testing.T) {
 	var lines []string
 	for i := 1; i <= 20; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"retry-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testModel, i))
+			`{"custom_id":"retry-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testSimModel, i))
 	}
 	fileID := mustCreateFile(t, fmt.Sprintf("test-cancel-retry-%s.jsonl", testRunID), strings.Join(lines, "\n"))
 	batchID := mustCreateBatch(t, fileID)
@@ -740,11 +742,11 @@ func doTestProgressPolling(t *testing.T) {
 	var lines []string
 	for i := 1; i <= 5; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"fast-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":1,"messages":[{"role":"user","content":"fast %d"}]}}`, i, testModel, i))
+			`{"custom_id":"fast-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":1,"messages":[{"role":"user","content":"fast %d"}]}}`, i, testSimModel, i))
 	}
 	for i := 1; i <= 15; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"slow-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testModel, i))
+			`{"custom_id":"slow-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testSimModel, i))
 	}
 	fileID := mustCreateFile(t, fmt.Sprintf("test-progress-%s.jsonl", testRunID), strings.Join(lines, "\n"))
 	batchID := mustCreateBatch(t, fileID)
diff --git a/test/e2e/e2e_test.go b/test/e2e/e2e_test.go
@@ -37,6 +37,7 @@ var (
 	testHelmRelease       = getEnvOrDefault("TEST_HELM_RELEASE", "batch-gateway")
 	testPostgresqlRelease = getEnvOrDefault("TEST_POSTGRESQL_RELEASE", "postgresql")
 	testRedisRelease      = getEnvOrDefault("TEST_REDIS_RELEASE", "redis")
+	testDBPod             = getEnvOrDefault("TEST_DB_POD", "")
 
 	// testDBClientType and testExchangeClientType are detected from Helm
 	// releases at startup; see detectDBClientType / detectExchangeClientType.
@@ -46,8 +47,14 @@ var (
 	testRunID = fmt.Sprintf("%d", time.Now().UnixNano())
 
 	// testModel is the model name used in batch input; configurable via TEST_MODEL env var.
-	testModel           = getEnvOrDefault("TEST_MODEL", "sim-model")
-	testModelB          = getEnvOrDefault("TEST_MODEL_B", "sim-model-b")
+	testModel  = getEnvOrDefault("TEST_MODEL", "sim-model")
+	testModelB = getEnvOrDefault("TEST_MODEL_B", "sim-model-b")
+	// testSimModel is the model name used in timing-dependent tests that require
+	// controlled latency (cancel, expiration, progress polling, orphan recovery,
+	// graceful shutdown). It should always point to a simulator with predictable
+	// inter-token latency. Defaults to testModel for dev-deploy where all models
+	// are simulated.
+	testSimModel        = getEnvOrDefault("TEST_SIM_MODEL", testModel)
 	testModel429        = getEnvOrDefault("TEST_MODEL_429", "sim-model-429")
 	testModelAlwaysFail = getEnvOrDefault("TEST_MODEL_ALWAYS_FAIL", "sim-model-always-fail")
 	testModelAIMD       = getEnvOrDefault("TEST_MODEL_AIMD", "sim-model-aimd")
diff --git a/test/e2e/gc_test.go b/test/e2e/gc_test.go
@@ -47,10 +47,15 @@ func expireInDB(t *testing.T, table, id string) {
 func expireInPostgresql(t *testing.T, table, id string) {
 	t.Helper()
 
+	podName := testDBPod
+	if podName == "" {
+		podName = fmt.Sprintf("%s-0", testPostgresqlRelease)
+	}
+
 	sql := fmt.Sprintf("UPDATE %s SET expiry = 1 WHERE id = '%s'", table, id)
 	cmd := fmt.Sprintf(`PGPASSWORD="$(cat "$POSTGRES_PASSWORD_FILE")" psql -U postgres -d postgres -c %q`, sql)
 	out, err := exec.Command("kubectl", "exec",
-		fmt.Sprintf("%s-0", testPostgresqlRelease),
+		podName,
 		"-n", testNamespace,
 		"--", "bash", "-c", cmd,
 	).CombinedOutput()
diff --git a/test/e2e/orphan_recovery_test.go b/test/e2e/orphan_recovery_test.go
@@ -68,7 +68,7 @@ func doTestHardCrashOrphanRecovery(t *testing.T) {
 	var lines []string
 	for i := 1; i <= 50; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"orphan-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testModel, i))
+			`{"custom_id":"orphan-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testSimModel, i))
 	}
 	fileID := mustCreateFile(t, fmt.Sprintf("test-orphan-recovery-%s.jsonl", testRunID), strings.Join(lines, "\n"))
 	batchID := mustCreateBatch(t, fileID)
@@ -115,11 +115,11 @@ func doTestCancellingOrphanRecovery(t *testing.T) {
 	var lines []string
 	for i := 1; i <= 5; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"cancel-orphan-fast-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":1,"messages":[{"role":"user","content":"Hi %d"}]}}`, i, testModel, i))
+			`{"custom_id":"cancel-orphan-fast-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":1,"messages":[{"role":"user","content":"Hi %d"}]}}`, i, testSimModel, i))
 	}
 	for i := 1; i <= 20; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"cancel-orphan-slow-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testModel, i))
+			`{"custom_id":"cancel-orphan-slow-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testSimModel, i))
 	}
 	fileID := mustCreateFile(t, fmt.Sprintf("test-cancelling-orphan-%s.jsonl", testRunID), strings.Join(lines, "\n"))
 	batchID := mustCreateBatch(t, fileID)
diff --git a/test/e2e/processor_graceful_shutdown_test.go b/test/e2e/processor_graceful_shutdown_test.go
@@ -62,7 +62,7 @@ func doTestPodDeleteMidJob(t *testing.T) {
 	var lines []string
 	for i := 1; i <= 10; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"pod-del-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testModel, i))
+			`{"custom_id":"pod-del-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testSimModel, i))
 	}
 	fileID := mustCreateFile(t, fmt.Sprintf("test-pod-delete-%s.jsonl", testRunID), strings.Join(lines, "\n"))
 	batchID := mustCreateBatch(t, fileID)
@@ -124,7 +124,7 @@ func doTestRollingRestartReEnqueue(t *testing.T) {
 	var lines []string
 	for i := 1; i <= 10; i++ {
 		lines = append(lines, fmt.Sprintf(
-			`{"custom_id":"restart-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testModel, i))
+			`{"custom_id":"restart-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testSimModel, i))
 	}
 	fileID := mustCreateFile(t, fmt.Sprintf("test-rolling-restart-%s.jsonl", testRunID), strings.Join(lines, "\n"))
 	batchID := mustCreateBatch(t, fileID)
