ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/main.go

@@ -0,0 +1,75 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"flag"
+	"fmt"
+	"log/slog"
+	"net"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google"
+)
+
+func main() {
+	if err := run(); err != nil {
+		slog.Error("vertex-batch stopped", "error", err)
+		os.Exit(1)
+	}
+}
+
+func run() error {
+	httpAddr := flag.String("http", ":8080", "HTTP listen address")
+	gateway := flag.String("gateway-url", "", "Gateway native conversion URL")
+	prefix := flag.String("prefix", "vertex_ai_batch/", "Client model prefix")
+	routePrefix := flag.String("route-prefix", "vertex_batch/", "Gateway native route prefix")
+	project := flag.String("project", "", "Vertex project")
+	bucket := flag.String("bucket", "", "GCS batch bucket")
+	credentials := flag.String("credentials", "", "Google service account JSON file")
+	vertex := flag.String("vertex-url", "https://aiplatform.googleapis.com", "Vertex API URL")
+	storage := flag.String("storage-url", "https://storage.googleapis.com", "GCS API URL")
+	poll := flag.Duration("poll", 10*time.Second, "Batch poll interval")
+	heartbeat := flag.Duration("heartbeat", 15*time.Second, "SSE keepalive interval")
+	flag.Parse()
+	if *gateway == "" || *prefix == "" || *routePrefix == "" || *project == "" || *bucket == "" || *credentials == "" || *poll <= 0 || *heartbeat <= 0 {
+		return errors.New("gateway-url, prefixes, project, bucket, credentials and positive intervals are required")
+	}
+	data, err := os.ReadFile(*credentials)
+	if err != nil {
+		return fmt.Errorf("read credentials file: %w", err)
+	}
+	config, err := google.JWTConfigFromJSON(data, "https://www.googleapis.com/auth/cloud-platform")
+	if err != nil {
+		return errors.New("invalid service account configuration")
+	}
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer cancel()
+	client := oauth2.NewClient(context.Background(), config.TokenSource(context.Background()))
+	backend := &cloud{client: client, vertex: *vertex, storage: *storage, project: *project, bucket: *bucket, poll: *poll}
+	mux := http.NewServeMux()
+	mux.Handle("POST /front", frontend{client: http.DefaultClient, gateway: *gateway, prefix: *prefix, routePrefix: *routePrefix, heartbeat: *heartbeat})
+	mux.Handle("/", handler{runner: backend, heartbeat: *heartbeat})
+	httpServer := &http.Server{Addr: *httpAddr, Handler: mux, BaseContext: func(net.Listener) context.Context { return ctx }}
+	stopped := make(chan error, 1)
+	go func() { stopped <- httpServer.ListenAndServe() }()
+	select {
+	case err = <-stopped:
+	case <-ctx.Done():
+	}
+	cancel()
+	shutdown, stop := context.WithTimeout(context.Background(), 65*time.Second)
+	defer stop()
+	if e := httpServer.Shutdown(shutdown); e != nil {
+		return e
+	}
+	if errors.Is(err, http.ErrServerClosed) {
+		return nil
+	}
+	return err
+}