feat(batch): planner-batch-intergation test (vllm-project#2186)

* planner-batch-intergation
* planner--job-id-and-batch-ip-mapping
* planner-logging-refactor
* revert-mock-app
* address-review-comments

Signed-off-by: nwangfw <113055713+nwangfw@users.noreply.github.com>
Signed-off-by: Ning Wang <n.wang.chn@hotmail.com>

Author: nwangfw

apps/console/api/handler/job.go

@@ -16,33 +16,28 @@ limitations under the License.
 
 // JobHandler implements the Console BFF JobService:
 //
-//   - Calls the metadata service /v1/batches API via the official OpenAI Go
-//     SDK (openai-go v3). Talking to the metadata service through the SDK
-//     keeps it honest about being OpenAI-compatible — schema drift on the
-//     upstream side surfaces immediately as a deserialization or 4xx error.
+//   - Routes all job lifecycle calls (Enqueue / Get / List / Cancel)
+//     through the Planner. The Planner owns JobID -> MDS batch.ID
+//     translation; the handler never holds an MDS batch.ID.
 //   - Persists Console-owned fields (id, display name, created_by, future:
 //     organization, tags ...) in the local store.
-//   - Aggregates both sources into the wire-level *pb.Job returned to the UI.
+//   - Aggregates Planner Jobs into the wire-level *pb.Job returned to the UI.
 //
-// The AIBrix-only extension `aibrix.model_template` is passed via the SDK's
-// `option.WithJSONSet`, which is the OpenAI-recommended `extra_body` channel.
+// The AIBrix-only extension `aibrix.model_template` is forwarded to MDS by the
+// planner's BatchClient via the OpenAI SDK's `extra_body` channel.
 //
-// When the metadata service is unreachable the handler propagates the error
-// (codes.Unavailable). The frontend renders its mock fallback in that case.
+// When the planner / metadata service is unreachable the handler propagates
+// the error (codes.Unavailable). The frontend renders its mock fallback in
+// that case.
 package handler
 
 import (
-	"bytes"
 	"context"
-	"encoding/json"
 	"errors"
-	"fmt"
-	"io"
 	"net/http"
-	"strings"
 
+	"github.com/google/uuid"
 	"github.com/openai/openai-go/v3"
-	"github.com/openai/openai-go/v3/option"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"k8s.io/klog/v2"
@@ -53,75 +48,6 @@ import (
 	"github.com/vllm-project/aibrix/apps/console/api/store"
 )
 
-// loggingTransport dumps every request and response between the BFF and MDS
-// at klog -v=2 (verbose). Errors and non-2xx responses always log at info.
-// Enable with `--v=2` (or env KLOG_V=2) when debugging the OpenAI SDK payloads.
-type loggingTransport struct {
-	base http.RoundTripper
-}
-
-func (t *loggingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	verbose := klog.V(2).Enabled()
-
-	var reqBody []byte
-	if verbose && req.Body != nil {
-		var err error
-		reqBody, err = io.ReadAll(req.Body)
-		if err != nil {
-			return nil, fmt.Errorf("read request body for logging: %w", err)
-		}
-		req.Body = io.NopCloser(bytes.NewReader(reqBody))
-	}
-	if verbose {
-		klog.V(2).Infof("[BFF→MDS] %s %s\n%s", req.Method, req.URL.String(), prettyBody(reqBody))
-	}
-
-	resp, err := t.base.RoundTrip(req)
-	if err != nil {
-		klog.Warningf("[BFF→MDS] %s %s ERROR %v", req.Method, req.URL.String(), err)
-		return resp, err
-	}
-
-	if resp.StatusCode >= 400 || verbose {
-		respBody, rerr := io.ReadAll(resp.Body)
-		_ = resp.Body.Close()
-		if rerr != nil {
-			klog.Warningf("[BFF→MDS] %s %s -> %d (read body failed: %v)",
-				req.Method, req.URL.String(), resp.StatusCode, rerr)
-			return resp, nil
-		}
-		resp.Body = io.NopCloser(bytes.NewReader(respBody))
-		if resp.StatusCode >= 400 {
-			klog.Warningf("[BFF→MDS] %s %s -> %d\n%s", req.Method, req.URL.String(), resp.StatusCode, prettyBody(respBody))
-		} else {
-			klog.V(2).Infof("[BFF→MDS] %s %s -> %d\n%s", req.Method, req.URL.String(), resp.StatusCode, prettyBody(respBody))
-		}
-	}
-	return resp, nil
-}
-
-const maxLoggedBodyBytes = 8192
-
-// prettyBody indents JSON for readability and truncates oversized bodies.
-// Non-JSON bodies are returned as-is (also truncated).
-func prettyBody(b []byte) string {
-	if len(b) == 0 {
-		return "(empty)"
-	}
-	var pretty bytes.Buffer
-	if err := json.Indent(&pretty, b, "", "  "); err == nil {
-		out := pretty.Bytes()
-		if len(out) > maxLoggedBodyBytes {
-			return string(out[:maxLoggedBodyBytes]) + "\n...(truncated)"
-		}
-		return string(out)
-	}
-	if len(b) > maxLoggedBodyBytes {
-		return string(b[:maxLoggedBodyBytes]) + "...(truncated)"
-	}
-	return string(b)
-}
-
 // Console-owned fields we stash on the OpenAI batch.metadata map. Namespaced
 // to keep them out of user-supplied metadata's key space. The bare
 // "display_name" key is kept for backwards compatibility with batches
@@ -141,25 +67,15 @@ type JobHandler struct {
 
 	store                          store.Store
 	planner                        plannerapi.Planner
-	openai                         openai.Client
 	defaultModelDeploymentTemplate string
 	devMode                        bool
 }
 
 // NewJobHandler creates a JobHandler.
-func NewJobHandler(s store.Store, planner plannerapi.Planner, metadataServiceURL, defaultModelDeploymentTemplate string, devMode bool) *JobHandler {
-
-	baseURL := strings.TrimRight(metadataServiceURL, "/") + "/v1"
-	httpClient := &http.Client{Transport: &loggingTransport{base: http.DefaultTransport}}
-	client := openai.NewClient(
-		option.WithBaseURL(baseURL),
-		option.WithAPIKey("aibrix-console"),
-		option.WithHTTPClient(httpClient),
-	)
+func NewJobHandler(s store.Store, planner plannerapi.Planner, defaultModelDeploymentTemplate string, devMode bool) *JobHandler {
 	return &JobHandler{
 		store:                          s,
 		planner:                        planner,
-		openai:                         client,
 		defaultModelDeploymentTemplate: defaultModelDeploymentTemplate,
 		devMode:                        devMode,
 	}
@@ -193,8 +109,8 @@ func (h *JobHandler) ListJobs(ctx context.Context, req *pb.ListJobsRequest) (*pb
 	}
 
 	jobs := make([]*pb.Job, 0, len(resp.Data))
-	for _, batch := range resp.Data {
-		jobs = append(jobs, mergeJob(batch, nil))
+	for _, job := range resp.Data {
+		jobs = append(jobs, mergeJob(job, nil))
 	}
 	// SDK CursorPage exposes Data and HasMore. first_id / last_id ride along
 	// in the upstream JSON but are not surfaced as named fields; the UI
@@ -212,22 +128,22 @@ func (h *JobHandler) GetJob(ctx context.Context, req *pb.GetJobRequest) (*pb.Job
 	if req.Id == "" {
 		return nil, status.Error(codes.InvalidArgument, "id is required")
 	}
-	batch, err := h.planner.GetJob(ctx, req.Id)
+	job, err := h.planner.GetJob(ctx, req.Id)
 	if err != nil {
 		// Dev fallback: return the demo job if MDS is unreachable.
 		if h.devMode {
 			if dev, ok := h.store.(interface {
 				GetDemoJob(id string) (*pb.Job, bool)
 			}); ok {
-				if job, found := dev.GetDemoJob(req.Id); found {
+				if demoJob, found := dev.GetDemoJob(req.Id); found {
 					klog.Warningf("MDS unreachable, falling back to demo job %s: %v", req.Id, err)
-					return job, nil
+					return demoJob, nil
 				}
 			}
 		}
 		return nil, mapPlannerError(err, "get batch")
 	}
-	return mergeJob(batch, nil), nil
+	return mergeJob(job, nil), nil
 }
 
 // CreateJob calls POST /v1/batches with console-owned fields (display name,
@@ -253,6 +169,10 @@ func (h *JobHandler) CreateJob(ctx context.Context, req *pb.CreateJobRequest) (*
 		completionWindow = string(openai.BatchNewParamsCompletionWindow24h)
 	}
 
+	// Console-generated JobID. The queued planner will own a durable
+	// JobID -> BatchID map; until then the planner keeps it in-memory.
+	jobID := "job_" + uuid.NewString()
+
 	// Pack console-owned fields into batch.Metadata under the aibrix.console.*
 	// namespace. This keeps a single source of truth (MDS) for the e2e demo
 	// loop. The store-overlay path is parked (not deleted) for the future
@@ -289,32 +209,34 @@ func (h *JobHandler) CreateJob(ctx context.Context, req *pb.CreateJobRequest) (*
 	}
 
 	enqueueReq := &plannerapi.EnqueueRequest{
+		JobID:         jobID,
 		ModelTemplate: modelTemplate,
-		BatchPayload: plannerapi.BatchPayload{
+		BatchParams: openai.BatchNewParams{
 			InputFileID:      req.InputDataset,
-			Endpoint:         req.Endpoint,
-			CompletionWindow: completionWindow,
+			Endpoint:         openai.BatchNewParamsEndpoint(req.Endpoint),
+			CompletionWindow: openai.BatchNewParamsCompletionWindow(completionWindow),
 			Metadata:         metadata,
 		},
 	}
 
-	result, err := h.planner.Enqueue(ctx, enqueueReq)
+	job, err := h.planner.Enqueue(ctx, enqueueReq)
 	if err != nil {
 		return nil, mapPlannerError(err, "create batch")
 	}
-	return mergeJob(result.Batch, nil), nil
+	return mergeJob(job, nil), nil
 }
 
-// CancelJob proxies to POST /v1/batches/{id}/cancel and merges with store.
+// CancelJob routes through Planner.Cancel; the planner resolves JobID
+// to MDS batch.ID and forwards to /v1/batches/{id}/cancel.
 func (h *JobHandler) CancelJob(ctx context.Context, req *pb.CancelJobRequest) (*pb.Job, error) {
 	if req.Id == "" {
 		return nil, status.Error(codes.InvalidArgument, "id is required")
 	}
-	batch, err := h.openai.Batches.Cancel(ctx, req.Id)
+	job, err := h.planner.Cancel(ctx, req.Id)
 	if err != nil {
-		return nil, mapSDKError(err, "cancel batch")
+		return nil, mapPlannerError(err, "cancel batch")
 	}
-	return mergeJob(batch, nil), nil
+	return mergeJob(job, nil), nil
 }
 
 // currentUserEmail returns the authenticated user's email if available, else
@@ -372,64 +294,67 @@ func mapPlannerError(err error, op string) error {
 	return mapSDKError(err, op)
 }
 
-// mergeJob aggregates the OpenAI Batch state with optional Console overlay.
-// Console-owned fields (display name, created_by, template binding) are read
-// out of batch.metadata under the aibrix.console.* namespace; the overlay
-// argument is plumbing for the future store-backed reconcile path and is
-// expected to be nil today.
-func mergeJob(b *openai.Batch, overlay *pb.Job) *pb.Job {
+// mergeJob aggregates the planner's Job with optional Console overlay.
+// pb.Job.Id is set to the planner's JobID — the MDS batch.ID never reaches
+// this layer. Console-owned fields (display name, created_by, template
+// binding) are read out of batch.metadata under the aibrix.console.*
+// namespace; the overlay argument is plumbing for the future store-backed
+// reconcile path and is expected to be nil today.
+func mergeJob(v *plannerapi.Job, overlay *pb.Job) *pb.Job {
 	job := &pb.Job{}
-	if b != nil {
-		job.Id = b.ID
-		job.Object = string(b.Object)
-		job.Endpoint = b.Endpoint
-		job.Model = b.Model
-		job.InputDataset = b.InputFileID
-		job.CompletionWindow = b.CompletionWindow
-		job.Status = string(b.Status)
-		job.OutputDataset = b.OutputFileID
-		job.ErrorDataset = b.ErrorFileID
-		job.CreatedAt = b.CreatedAt
-		job.InProgressAt = b.InProgressAt
-		job.ExpiresAt = b.ExpiresAt
-		job.FinalizingAt = b.FinalizingAt
-		job.CompletedAt = b.CompletedAt
-		job.FailedAt = b.FailedAt
-		job.ExpiredAt = b.ExpiredAt
-		job.CancellingAt = b.CancellingAt
-		job.CancelledAt = b.CancelledAt
-		if len(b.Metadata) > 0 {
-			job.Metadata = map[string]string(b.Metadata)
-			// Console-owned fields. Prefer namespaced keys; fall back to the
-			// legacy bare "display_name" so batches created by older builds
-			// still surface their name.
-			if v := b.Metadata[metadataConsoleDisplayName]; v != "" {
-				job.Name = v
-			} else if v := b.Metadata[metadataDisplayName]; v != "" {
-				job.Name = v
-			}
-			if v := b.Metadata[metadataConsoleCreatedBy]; v != "" {
-				job.CreatedBy = v
-			}
-			if v := b.Metadata[metadataConsoleTemplateName]; v != "" {
-				job.ModelTemplateName = v
-			}
-			if v := b.Metadata[metadataConsoleTemplateVersion]; v != "" {
-				job.ModelTemplateVersion = v
+	if v != nil {
+		job.Id = v.JobID
+		if b := v.Batch; b != nil {
+			job.Object = string(b.Object)
+			job.Endpoint = b.Endpoint
+			job.Model = b.Model
+			job.InputDataset = b.InputFileID
+			job.CompletionWindow = b.CompletionWindow
+			job.Status = string(b.Status)
+			job.OutputDataset = b.OutputFileID
+			job.ErrorDataset = b.ErrorFileID
+			job.CreatedAt = b.CreatedAt
+			job.InProgressAt = b.InProgressAt
+			job.ExpiresAt = b.ExpiresAt
+			job.FinalizingAt = b.FinalizingAt
+			job.CompletedAt = b.CompletedAt
+			job.FailedAt = b.FailedAt
+			job.ExpiredAt = b.ExpiredAt
+			job.CancellingAt = b.CancellingAt
+			job.CancelledAt = b.CancelledAt
+			if len(b.Metadata) > 0 {
+				job.Metadata = map[string]string(b.Metadata)
+				// Console-owned fields. Prefer namespaced keys; fall back to the
+				// legacy bare "display_name" so batches created by older builds
+				// still surface their name.
+				if v := b.Metadata[metadataConsoleDisplayName]; v != "" {
+					job.Name = v
+				} else if v := b.Metadata[metadataDisplayName]; v != "" {
+					job.Name = v
+				}
+				if v := b.Metadata[metadataConsoleCreatedBy]; v != "" {
+					job.CreatedBy = v
+				}
+				if v := b.Metadata[metadataConsoleTemplateName]; v != "" {
+					job.ModelTemplateName = v
+				}
+				if v := b.Metadata[metadataConsoleTemplateVersion]; v != "" {
+					job.ModelTemplateVersion = v
+				}
 			}
-		}
-		if b.JSON.RequestCounts.Valid() {
-			job.RequestCounts = &pb.JobRequestCounts{
-				Total:     int32(b.RequestCounts.Total),
-				Completed: int32(b.RequestCounts.Completed),
-				Failed:    int32(b.RequestCounts.Failed),
+			if b.JSON.RequestCounts.Valid() {
+				job.RequestCounts = &pb.JobRequestCounts{
+					Total:     int32(b.RequestCounts.Total),
+					Completed: int32(b.RequestCounts.Completed),
+					Failed:    int32(b.RequestCounts.Failed),
+				}
 			}
-		}
-		if b.JSON.Usage.Valid() {
-			job.Usage = &pb.JobUsage{
-				InputTokens:  b.Usage.InputTokens,
-				OutputTokens: b.Usage.OutputTokens,
-				TotalTokens:  b.Usage.TotalTokens,
+			if b.JSON.Usage.Valid() {
+				job.Usage = &pb.JobUsage{
+					InputTokens:  b.Usage.InputTokens,
+					OutputTokens: b.Usage.OutputTokens,
+					TotalTokens:  b.Usage.TotalTokens,
+				}
 			}
 		}
 	}

apps/console/api/planner/api/errors.go

@@ -31,4 +31,10 @@ var (
 	// RM package) is wrapped with this sentinel so the planner stays
 	// decoupled from the concrete RM error vocabulary.
 	ErrInsufficientResources = errors.New("planner: insufficient resources")
+
+	// ErrJobNotFound indicates the JobID passed to GetJob/Cancel is not
+	// known to the planner. With Passthrough's in-memory map this fires
+	// for jobs created before the process started; the queued planner's
+	// durable index will narrow this to "truly unknown JobID" cases.
+	ErrJobNotFound = errors.New("planner: job not found")
 )