75 lines
2.7 KiB
Go
75 lines
2.7 KiB
Go
|
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
|
||
|
|
}
|