ce7cb702218693d08ea6d83464d2457c9250027c
This commit is contained in:
commit
34c5f8963c
15 changed files with 145336 additions and 0 deletions
213
services/vertex-batch/cloud.go
Normal file
213
services/vertex-batch/cloud.go
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type cloud struct {
|
||||
client *http.Client
|
||||
vertex string
|
||||
storage string
|
||||
project string
|
||||
bucket string
|
||||
poll time.Duration
|
||||
}
|
||||
|
||||
type job struct {
|
||||
Name string `json:"name"`
|
||||
State string `json:"state"`
|
||||
Completion struct {
|
||||
Success json.Number `json:"successfulCount"`
|
||||
} `json:"completionStats"`
|
||||
}
|
||||
|
||||
func (j job) terminal() bool {
|
||||
switch j.State {
|
||||
case "JOB_STATE_SUCCEEDED", "JOB_STATE_FAILED", "JOB_STATE_CANCELLED", "JOB_STATE_EXPIRED", "JOB_STATE_PARTIALLY_SUCCEEDED":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *cloud) exchange(ctx context.Context, method, endpoint string, body []byte, out any) error {
|
||||
req, err := http.NewRequestWithContext(ctx, method, endpoint, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cloud request: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return fmt.Errorf("cloud %s returned HTTP %d", method, resp.StatusCode)
|
||||
}
|
||||
if out == nil {
|
||||
_, err = io.Copy(io.Discard, resp.Body)
|
||||
return err
|
||||
}
|
||||
return json.NewDecoder(resp.Body).Decode(out)
|
||||
}
|
||||
|
||||
func (c *cloud) objects(ctx context.Context, prefix string) ([]string, error) {
|
||||
var names []string
|
||||
token := ""
|
||||
for {
|
||||
query := url.Values{"prefix": {prefix}}
|
||||
if token != "" {
|
||||
query.Set("pageToken", token)
|
||||
}
|
||||
var page struct {
|
||||
Items []struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"items"`
|
||||
Next string `json:"nextPageToken"`
|
||||
}
|
||||
if err := c.exchange(ctx, http.MethodGet, c.storage+"/storage/v1/b/"+url.PathEscape(c.bucket)+"/o?"+query.Encode(), nil, &page); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, item := range page.Items {
|
||||
names = append(names, item.Name)
|
||||
}
|
||||
if page.Next == "" {
|
||||
return names, nil
|
||||
}
|
||||
token = page.Next
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cloud) objectURL(name string) string {
|
||||
return c.storage + "/storage/v1/b/" + url.PathEscape(c.bucket) + "/o/" + url.PathEscape(name)
|
||||
}
|
||||
|
||||
func (c *cloud) run(ctx context.Context, model string, native json.RawMessage) (result json.RawMessage, err error) {
|
||||
var id [16]byte
|
||||
if _, err := rand.Read(id[:]); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prefix := hex.EncodeToString(id[:]) + "/"
|
||||
var current job
|
||||
defer func() {
|
||||
cleanup, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
|
||||
defer cancel()
|
||||
if current.Name != "" && !current.terminal() {
|
||||
if e := c.exchange(cleanup, http.MethodPost, c.vertex+"/v1/"+current.Name+":cancel", []byte("{}"), nil); e != nil && err == nil {
|
||||
err = e
|
||||
}
|
||||
}
|
||||
names, e := c.objects(cleanup, prefix)
|
||||
if e != nil && err == nil {
|
||||
err = e
|
||||
}
|
||||
for _, name := range names {
|
||||
if e := c.exchange(cleanup, http.MethodDelete, c.objectURL(name), nil, nil); e != nil && err == nil {
|
||||
err = e
|
||||
}
|
||||
}
|
||||
}()
|
||||
input := prefix + "input.jsonl"
|
||||
row, err := json.Marshal(struct {
|
||||
Request json.RawMessage `json:"request"`
|
||||
}{native})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
upload := c.storage + "/upload/storage/v1/b/" + url.PathEscape(c.bucket) + "/o?" + url.Values{"name": {input}, "uploadType": {"media"}}.Encode()
|
||||
if err = c.exchange(ctx, http.MethodPost, upload, append(row, '\n'), nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"displayName": strings.TrimSuffix(prefix, "/"), "model": "publishers/google/models/" + model,
|
||||
"inputConfig": map[string]any{"instancesFormat": "jsonl", "gcsSource": map[string]any{"uris": []string{"gs://" + c.bucket + "/" + input}}},
|
||||
"outputConfig": map[string]any{"predictionsFormat": "jsonl", "gcsDestination": map[string]any{"outputUriPrefix": "gs://" + c.bucket + "/" + prefix + "output/"}},
|
||||
})
|
||||
// A cancelled caller still needs the submitted job's name to cancel it.
|
||||
submit, cancelSubmit := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
|
||||
err = c.exchange(submit, http.MethodPost, c.vertex+"/v1/projects/"+url.PathEscape(c.project)+"/locations/global/batchPredictionJobs", body, ¤t)
|
||||
cancelSubmit()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err = ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ticker := time.NewTicker(c.poll)
|
||||
defer ticker.Stop()
|
||||
for !current.terminal() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-ticker.C:
|
||||
if err = c.exchange(ctx, http.MethodGet, c.vertex+"/v1/"+current.Name, nil, ¤t); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
if current.State != "JOB_STATE_SUCCEEDED" || current.Completion.Success != "1" {
|
||||
return nil, fmt.Errorf("batch did not produce one successful row (%s)", current.State)
|
||||
}
|
||||
names, err := c.objects(ctx, prefix+"output/")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var rows []prediction
|
||||
for _, name := range names {
|
||||
if !strings.HasSuffix(name, ".jsonl") {
|
||||
continue
|
||||
}
|
||||
output, err := c.readOutput(ctx, name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows = append(rows, output...)
|
||||
}
|
||||
if len(rows) != 1 || (len(rows[0].Status) != 0 && string(rows[0].Status) != `""`) || len(rows[0].Response) == 0 {
|
||||
return nil, errors.New("batch output is not one successful row")
|
||||
}
|
||||
return rows[0].Response, nil
|
||||
}
|
||||
|
||||
type prediction struct {
|
||||
Status json.RawMessage `json:"status"`
|
||||
Response json.RawMessage `json:"response"`
|
||||
}
|
||||
|
||||
func (c *cloud) readOutput(ctx context.Context, name string) ([]prediction, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.objectURL(name)+"?alt=media", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("batch output returned HTTP %d", resp.StatusCode)
|
||||
}
|
||||
decoder := json.NewDecoder(resp.Body)
|
||||
var rows []prediction
|
||||
for {
|
||||
var row prediction
|
||||
if err := decoder.Decode(&row); errors.Is(err, io.EOF) {
|
||||
return rows, nil
|
||||
} else if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows = append(rows, row)
|
||||
}
|
||||
}
|
||||
116
services/vertex-batch/cloud_test.go
Normal file
116
services/vertex-batch/cloud_test.go
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCloud(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name, state, success string
|
||||
cancel bool
|
||||
wantError bool
|
||||
}{
|
||||
{"success", "JOB_STATE_SUCCEEDED", "1", false, false},
|
||||
{"failed row", "JOB_STATE_SUCCEEDED", "0", false, true},
|
||||
{"failed job", "JOB_STATE_FAILED", "0", false, true},
|
||||
{"cancel", "JOB_STATE_PENDING", "0", true, true},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
objects := map[string]bool{}
|
||||
input, prefix := "", ""
|
||||
cancelled := false
|
||||
submitted := make(chan struct{})
|
||||
response := json.RawMessage(`{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}],"usageMetadata":{"totalTokenCount":16}}`)
|
||||
native := json.RawMessage(`{"contents":[{"role":"user","parts":[{"text":"synthetic"},{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}]}],"generationConfig":{"temperature":0.2}}`)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch {
|
||||
case strings.HasPrefix(r.URL.Path, "/upload/"):
|
||||
input = r.URL.Query().Get("name")
|
||||
objects[input] = true
|
||||
var row struct {
|
||||
Request json.RawMessage `json:"request"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&row); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if string(row.Request) != string(native) {
|
||||
t.Errorf("native body changed: %s", row.Request)
|
||||
}
|
||||
_, _ = fmt.Fprint(w, `{}`)
|
||||
case strings.HasSuffix(r.URL.Path, "/batchPredictionJobs"):
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Output struct {
|
||||
GCS struct {
|
||||
Prefix string `json:"outputUriPrefix"`
|
||||
} `json:"gcsDestination"`
|
||||
} `json:"outputConfig"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if body.Model != "publishers/google/models/gemini-test" {
|
||||
t.Error(body.Model)
|
||||
}
|
||||
prefix = strings.TrimPrefix(body.Output.GCS.Prefix, "gs://synthetic/")
|
||||
_, _ = fmt.Fprint(w, `{"name":"projects/synthetic/locations/global/batchPredictionJobs/1","state":"JOB_STATE_PENDING"}`)
|
||||
close(submitted)
|
||||
case strings.HasSuffix(r.URL.Path, ":cancel"):
|
||||
cancelled = true
|
||||
_, _ = fmt.Fprint(w, `{}`)
|
||||
case strings.Contains(r.URL.Path, "/batchPredictionJobs/"):
|
||||
objects[prefix+"predictions.jsonl"] = true
|
||||
_, _ = fmt.Fprintf(w, `{"name":"projects/synthetic/locations/global/batchPredictionJobs/1","state":%q,"completionStats":{"successfulCount":%q}}`, test.state, test.success)
|
||||
case r.Method == "DELETE":
|
||||
name := strings.SplitN(r.URL.Path, "/o/", 2)[1]
|
||||
delete(objects, name)
|
||||
w.WriteHeader(204)
|
||||
case r.URL.Query().Get("alt") == "media":
|
||||
_, _ = fmt.Fprintf(w, `{"status":"","response":%s}`, response)
|
||||
case strings.HasSuffix(r.URL.Path, "/o"):
|
||||
var items []map[string]string
|
||||
for name := range objects {
|
||||
if strings.HasPrefix(name, r.URL.Query().Get("prefix")) {
|
||||
items = append(items, map[string]string{"name": name})
|
||||
}
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"items": items})
|
||||
default:
|
||||
t.Error(r.Method, r.URL.String())
|
||||
w.WriteHeader(404)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := &cloud{client: server.Client(), vertex: server.URL, storage: server.URL, project: "synthetic", bucket: "synthetic", poll: time.Millisecond}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
if test.cancel {
|
||||
go func() { <-submitted; cancel() }()
|
||||
}
|
||||
result, err := client.run(ctx, "gemini-test", native)
|
||||
if (err != nil) != test.wantError {
|
||||
t.Fatalf("result %s, error %v", result, err)
|
||||
}
|
||||
if !test.wantError && string(result) != string(response) {
|
||||
t.Errorf("result changed: %s", result)
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(objects) != 0 || cancelled != test.cancel {
|
||||
t.Errorf("cleanup objects=%v cancelled=%v", objects, cancelled)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
136
services/vertex-batch/front.go
Normal file
136
services/vertex-batch/front.go
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type frontend struct {
|
||||
client *http.Client
|
||||
gateway string
|
||||
prefix string
|
||||
routePrefix string
|
||||
heartbeat time.Duration
|
||||
}
|
||||
|
||||
func (f frontend) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
data, err := io.ReadAll(r.Body)
|
||||
var body map[string]json.RawMessage
|
||||
if err != nil || json.Unmarshal(data, &body) != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var model string
|
||||
var streaming bool
|
||||
if json.Unmarshal(body["model"], &model) != nil || !strings.HasPrefix(model, f.prefix) || strings.TrimPrefix(model, f.prefix) == "" {
|
||||
http.Error(w, "invalid batch model", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if raw, ok := body["stream"]; ok {
|
||||
if json.Unmarshal(raw, &streaming) != nil {
|
||||
http.Error(w, "invalid stream flag", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
}
|
||||
body["model"], err = json.Marshal(f.routePrefix + strings.TrimPrefix(model, f.prefix))
|
||||
if err != nil {
|
||||
writeError(w)
|
||||
return
|
||||
}
|
||||
data, err = json.Marshal(body)
|
||||
if err != nil {
|
||||
writeError(w)
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithCancel(r.Context())
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, f.gateway, bytes.NewReader(data))
|
||||
if err != nil {
|
||||
writeError(w)
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", r.Header.Get("Authorization"))
|
||||
if streaming {
|
||||
f.stream(w, req)
|
||||
return
|
||||
}
|
||||
res, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
writeError(w)
|
||||
return
|
||||
}
|
||||
defer func() { _ = res.Body.Close() }()
|
||||
w.Header().Set("Content-Type", res.Header.Get("Content-Type"))
|
||||
w.WriteHeader(res.StatusCode)
|
||||
_, _ = io.Copy(w, res.Body)
|
||||
}
|
||||
|
||||
type fragment struct {
|
||||
data []byte
|
||||
end bool
|
||||
}
|
||||
|
||||
func (f frontend) receive(req *http.Request, incoming chan<- fragment) {
|
||||
send := func(part fragment) bool {
|
||||
select {
|
||||
case incoming <- part:
|
||||
return true
|
||||
case <-req.Context().Done():
|
||||
return false
|
||||
}
|
||||
}
|
||||
res, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
send(fragment{end: true})
|
||||
return
|
||||
}
|
||||
defer func() { _ = res.Body.Close() }()
|
||||
if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices {
|
||||
send(fragment{end: true})
|
||||
return
|
||||
}
|
||||
buffer := make([]byte, 4096)
|
||||
for {
|
||||
n, err := res.Body.Read(buffer)
|
||||
if !send(fragment{data: bytes.Clone(buffer[:n]), end: err != nil}) || err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (f frontend) stream(w http.ResponseWriter, req *http.Request) {
|
||||
incoming := make(chan fragment)
|
||||
go f.receive(req, incoming)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
flush := http.NewResponseController(w)
|
||||
if flush.Flush() != nil {
|
||||
return
|
||||
}
|
||||
ticker := time.NewTicker(f.heartbeat)
|
||||
defer ticker.Stop()
|
||||
var parser events
|
||||
for {
|
||||
var output []byte
|
||||
select {
|
||||
case <-req.Context().Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
output = []byte(": batch pending\n\n")
|
||||
case part := <-incoming:
|
||||
output = parser.feed(part.data, part.end)
|
||||
}
|
||||
if _, err := w.Write(output); err != nil {
|
||||
return
|
||||
}
|
||||
if flush.Flush() != nil || parser.done {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
116
services/vertex-batch/front_test.go
Normal file
116
services/vertex-batch/front_test.go
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func testFrontend(url string) frontend {
|
||||
return frontend{client: http.DefaultClient, gateway: url, prefix: "batch/", routePrefix: "native/", heartbeat: 10 * time.Millisecond}
|
||||
}
|
||||
|
||||
func TestFrontendPayload(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]json.RawMessage
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if string(body["model"]) != `"native/partner-model@001"` || string(body["tools"]) != `[{"opaque":true}]` || r.Header.Get("Authorization") != "Bearer synthetic" {
|
||||
t.Errorf("unexpected request: %s", body)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
_, _ = io.WriteString(w, `{"error":"synthetic failure"}`)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
req := httptest.NewRequest(http.MethodPost, "/front", strings.NewReader(`{"model":"batch/partner-model@001","tools":[{"opaque":true}]}`))
|
||||
req.Header.Set("Authorization", "Bearer synthetic")
|
||||
res := httptest.NewRecorder()
|
||||
testFrontend(upstream.URL).ServeHTTP(res, req)
|
||||
if res.Code != http.StatusBadGateway || res.Body.String() != `{"error":"synthetic failure"}` {
|
||||
t.Fatalf("unexpected response: %d %s", res.Code, res.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFrontendCancel(t *testing.T) {
|
||||
for _, headers := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "pending_headers", true: "pending_body"}[headers], func(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
cancelled := make(chan struct{})
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.Copy(io.Discard, r.Body)
|
||||
if headers {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
}
|
||||
close(entered)
|
||||
<-r.Context().Done()
|
||||
close(cancelled)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
defer upstream.CloseClientConnections()
|
||||
server := httptest.NewServer(testFrontend(upstream.URL))
|
||||
defer server.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, strings.NewReader(`{"model":"batch/gemini-3.8-flash","stream":true}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
res, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = res.Body.Close() }()
|
||||
select {
|
||||
case <-entered:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("upstream was not entered")
|
||||
}
|
||||
data := make([]byte, 1)
|
||||
if _, err := res.Body.Read(data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := res.Body.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case <-cancelled:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("disconnect did not cancel upstream")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFrontendStream(t *testing.T) {
|
||||
for _, complete := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "incomplete", true: "complete"}[complete], func(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
if complete {
|
||||
_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n")
|
||||
}
|
||||
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer upstream.Close()
|
||||
req := httptest.NewRequest(http.MethodPost, "/front", strings.NewReader(`{"model":"batch/gemini-3.8-flash","stream":true}`))
|
||||
res := httptest.NewRecorder()
|
||||
testFrontend(upstream.URL).ServeHTTP(res, req)
|
||||
output := res.Body.String()
|
||||
if complete {
|
||||
if !strings.Contains(output, "answer") || !strings.Contains(output, "[DONE]") {
|
||||
t.Fatal(output)
|
||||
}
|
||||
} else if !strings.Contains(output, "batch_error") || strings.Contains(output, "[DONE]") {
|
||||
t.Fatal(output)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
10
services/vertex-batch/go.mod
Normal file
10
services/vertex-batch/go.mod
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
module github.com/quorate/ms2k-internal/services/vertex-batch
|
||||
|
||||
go 1.26.4
|
||||
|
||||
require golang.org/x/oauth2 v0.37.0
|
||||
|
||||
require (
|
||||
cloud.google.com/go/compute/metadata v0.9.0 // indirect
|
||||
golang.org/x/sys v0.47.0 // indirect
|
||||
)
|
||||
8
services/vertex-batch/go.sum
Normal file
8
services/vertex-batch/go.sum
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs=
|
||||
cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
golang.org/x/oauth2 v0.37.0 h1:JUlcxA8oAtauLfiH8FX2/FkAWHAdi0QtGCGc+hofE98=
|
||||
golang.org/x/oauth2 v0.37.0/go.mod h1:IxwZNxUULJmpBFf9K/9NTMSIfZZuvuTy1gGxhigP/58=
|
||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
103
services/vertex-batch/http.go
Normal file
103
services/vertex-batch/http.go
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type batchRunner interface {
|
||||
run(context.Context, string, json.RawMessage) (json.RawMessage, error)
|
||||
}
|
||||
|
||||
type handler struct {
|
||||
runner batchRunner
|
||||
heartbeat time.Duration
|
||||
}
|
||||
|
||||
func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodGet && r.URL.Path == "/healthz" {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodPost || r.URL.Path != "/v1/chat/completions" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
model := r.Header.Get("x-batch-model")
|
||||
streaming := r.Header.Get("x-batch-stream")
|
||||
if model == "" || (streaming != "true" && streaming != "false") {
|
||||
http.Error(w, "invalid batch model or stream headers", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
native, err := io.ReadAll(r.Body)
|
||||
if err != nil || !json.Valid(native) {
|
||||
http.Error(w, "invalid Gemini request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithCancel(r.Context())
|
||||
type outcome struct {
|
||||
body json.RawMessage
|
||||
err error
|
||||
}
|
||||
done := make(chan outcome, 1)
|
||||
var workers sync.WaitGroup
|
||||
workers.Go(func() {
|
||||
body, err := h.runner.run(ctx, model, native)
|
||||
done <- outcome{body: body, err: err}
|
||||
})
|
||||
defer workers.Wait()
|
||||
defer cancel()
|
||||
if streaming != "true" {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case result := <-done:
|
||||
if result.err != nil {
|
||||
writeError(w)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write(result.body)
|
||||
return
|
||||
}
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
flush := http.NewResponseController(w)
|
||||
if err := flush.Flush(); err != nil {
|
||||
return
|
||||
}
|
||||
ticker := time.NewTicker(h.heartbeat)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if _, err := io.WriteString(w, ": batch pending\n\n"); err != nil {
|
||||
return
|
||||
}
|
||||
if err := flush.Flush(); err != nil {
|
||||
return
|
||||
}
|
||||
case result := <-done:
|
||||
if result.err != nil {
|
||||
slog.Error("batch failed", "kind", "cloud")
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, "data: "+string(result.body)+"\n\n")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func writeError(w http.ResponseWriter) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
_, _ = io.WriteString(w, `{"error":{"code":502,"status":"INTERNAL","message":"Batch did not produce a result"}}`)
|
||||
}
|
||||
84
services/vertex-batch/http_test.go
Normal file
84
services/vertex-batch/http_test.go
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type runnerFunc func(context.Context, string, json.RawMessage) (json.RawMessage, error)
|
||||
|
||||
func (f runnerFunc) run(ctx context.Context, model string, body json.RawMessage) (json.RawMessage, error) {
|
||||
return f(ctx, model, body)
|
||||
}
|
||||
|
||||
func TestHTTPStream(t *testing.T) {
|
||||
for _, failed := range []bool{false, true} {
|
||||
server := httptest.NewServer(handler{heartbeat: time.Millisecond, runner: runnerFunc(func(_ context.Context, model string, body json.RawMessage) (json.RawMessage, error) {
|
||||
if model != "partner-model@001" || string(body) != `{"contents":[]}` {
|
||||
t.Errorf("changed request %s %s", model, body)
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
if failed {
|
||||
return nil, errors.New("synthetic failure")
|
||||
}
|
||||
return json.RawMessage(`{"candidates":[{"finishReason":"STOP"}]}`), nil
|
||||
})})
|
||||
req, _ := http.NewRequest("POST", server.URL+"/v1/chat/completions", strings.NewReader(`{"contents":[]}`))
|
||||
req.Header.Set("x-batch-model", "partner-model@001")
|
||||
req.Header.Set("x-batch-stream", "true")
|
||||
response, err := server.Client().Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, err := io.ReadAll(response.Body)
|
||||
_ = response.Body.Close()
|
||||
server.Close()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(data), ": batch pending") || strings.Contains(string(data), "data:") == failed {
|
||||
t.Fatalf("failed=%v body=%s", failed, data)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCleanupWait(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
cancelled := make(chan struct{})
|
||||
cleanup := make(chan struct{})
|
||||
finished := make(chan struct{})
|
||||
h := handler{runner: runnerFunc(func(ctx context.Context, _ string, _ json.RawMessage) (json.RawMessage, error) {
|
||||
<-ctx.Done()
|
||||
close(cancelled)
|
||||
<-cleanup
|
||||
return nil, ctx.Err()
|
||||
})}
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"contents":[]}`)).WithContext(ctx)
|
||||
req.Header.Set("x-batch-model", "partner-model@001")
|
||||
req.Header.Set("x-batch-stream", "false")
|
||||
go func() {
|
||||
h.ServeHTTP(httptest.NewRecorder(), req)
|
||||
close(finished)
|
||||
}()
|
||||
cancel()
|
||||
<-cancelled
|
||||
select {
|
||||
case <-finished:
|
||||
t.Error("handler returned before cloud cleanup")
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
close(cleanup)
|
||||
select {
|
||||
case <-finished:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("handler did not finish after cleanup")
|
||||
}
|
||||
}
|
||||
77
services/vertex-batch/k8s/deployment.yaml
Normal file
77
services/vertex-batch/k8s/deployment.yaml
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: vertex-batch
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
replicas: 2
|
||||
selector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: vertex-batch
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app.kubernetes.io/name: vertex-batch
|
||||
app.kubernetes.io/instance: vertex-batch
|
||||
spec:
|
||||
terminationGracePeriodSeconds: 75
|
||||
automountServiceAccountToken: false
|
||||
imagePullSecrets:
|
||||
- name: quay-cluster-pull
|
||||
securityContext:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
runAsGroup: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
containers:
|
||||
- name: batch
|
||||
image: registry.quorate.internal/quorate/vertex-batch@sha256:a1416b98cc57ae27dde41df27ad696c933d6bad6906f241bc698fc48cf4c7a35
|
||||
args:
|
||||
- -gateway-url=http://agentgateway-proxy.agentgateway-system.svc.cluster.local/v1/chat/completions
|
||||
- -project=ms2k-personal
|
||||
- -bucket=ms2k-personal-vertex-batch
|
||||
- -credentials=/credentials/credentials.json
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: 8080
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
port: http
|
||||
lifecycle:
|
||||
preStop:
|
||||
sleep:
|
||||
seconds: 5
|
||||
resources:
|
||||
requests:
|
||||
cpu: 10m
|
||||
memory: 32Mi
|
||||
limits:
|
||||
memory: 128Mi
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop: [ALL]
|
||||
volumeMounts:
|
||||
- name: credentials
|
||||
mountPath: /credentials
|
||||
readOnly: true
|
||||
volumes:
|
||||
- name: credentials
|
||||
secret:
|
||||
secretName: agentgateway-vertex-ai
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: vertex-batch
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
selector:
|
||||
app.kubernetes.io/name: vertex-batch
|
||||
ports:
|
||||
- name: http
|
||||
port: 8080
|
||||
targetPort: http
|
||||
6
services/vertex-batch/k8s/kustomization.yaml
Normal file
6
services/vertex-batch/k8s/kustomization.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
apiVersion: kustomize.config.k8s.io/v1beta1
|
||||
kind: Kustomization
|
||||
resources:
|
||||
- deployment.yaml
|
||||
- routes.yaml
|
||||
- network-policy.yaml
|
||||
91
services/vertex-batch/k8s/network-policy.yaml
Normal file
91
services/vertex-batch/k8s/network-policy.yaml
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
apiVersion: cilium.io/v2
|
||||
kind: CiliumNetworkPolicy
|
||||
metadata:
|
||||
name: vertex-batch
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
endpointSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: vertex-batch
|
||||
ingress:
|
||||
- fromEntities: [host]
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
- fromEndpoints:
|
||||
- matchLabels:
|
||||
gateway.networking.k8s.io/gateway-name: agentgateway-proxy
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
egress:
|
||||
- toEndpoints:
|
||||
- matchLabels:
|
||||
gateway.networking.k8s.io/gateway-name: agentgateway-proxy
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "80"
|
||||
protocol: TCP
|
||||
- toEndpoints:
|
||||
- matchLabels:
|
||||
k8s:io.kubernetes.pod.namespace: kube-system
|
||||
k8s-app: kube-dns
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "53"
|
||||
protocol: UDP
|
||||
- port: "53"
|
||||
protocol: TCP
|
||||
- toEntities: [world]
|
||||
toPorts:
|
||||
- serverNames:
|
||||
- oauth2.googleapis.com
|
||||
- aiplatform.googleapis.com
|
||||
- storage.googleapis.com
|
||||
ports:
|
||||
- port: "443"
|
||||
protocol: TCP
|
||||
---
|
||||
apiVersion: cilium.io/v2
|
||||
kind: CiliumNetworkPolicy
|
||||
metadata:
|
||||
name: batch-gateway
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
endpointSelector:
|
||||
matchLabels:
|
||||
gateway.networking.k8s.io/gateway-name: agentgateway-proxy
|
||||
ingress:
|
||||
- fromEndpoints:
|
||||
- matchLabels:
|
||||
app.kubernetes.io/name: vertex-batch
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "80"
|
||||
protocol: TCP
|
||||
egress:
|
||||
- toEndpoints:
|
||||
- matchLabels:
|
||||
app.kubernetes.io/name: vertex-batch
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
---
|
||||
apiVersion: cilium.io/v2
|
||||
kind: CiliumClusterwideNetworkPolicy
|
||||
metadata:
|
||||
name: host-vertex-batch
|
||||
spec:
|
||||
nodeSelector: {}
|
||||
egress:
|
||||
- toEndpoints:
|
||||
- matchLabels:
|
||||
k8s:io.kubernetes.pod.namespace: agentgateway-system
|
||||
app.kubernetes.io/name: vertex-batch
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
88
services/vertex-batch/k8s/routes.yaml
Normal file
88
services/vertex-batch/k8s/routes.yaml
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
apiVersion: agentgateway.dev/v1alpha1
|
||||
kind: AgentgatewayBackend
|
||||
metadata:
|
||||
name: vertex-batch
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
ai:
|
||||
provider:
|
||||
vertexai:
|
||||
projectId: ms2k-personal
|
||||
region: global
|
||||
host: vertex-batch.agentgateway-system.svc.cluster.local
|
||||
port: 8080
|
||||
---
|
||||
apiVersion: gateway.networking.k8s.io/v1
|
||||
kind: HTTPRoute
|
||||
metadata:
|
||||
name: batch-front
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
parentRefs:
|
||||
- name: agentgateway-proxy
|
||||
rules:
|
||||
- matches:
|
||||
- method: POST
|
||||
path:
|
||||
type: Exact
|
||||
value: /v1/chat/completions
|
||||
headers:
|
||||
- name: x-model
|
||||
type: RegularExpression
|
||||
value: ^vertex_ai_batch/.+
|
||||
filters:
|
||||
- type: URLRewrite
|
||||
urlRewrite:
|
||||
path:
|
||||
type: ReplaceFullPath
|
||||
replaceFullPath: /front
|
||||
backendRefs:
|
||||
- name: vertex-batch
|
||||
port: 8080
|
||||
---
|
||||
apiVersion: gateway.networking.k8s.io/v1
|
||||
kind: HTTPRoute
|
||||
metadata:
|
||||
name: batch-native
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
parentRefs:
|
||||
- name: agentgateway-proxy
|
||||
rules:
|
||||
- matches:
|
||||
- method: POST
|
||||
path:
|
||||
type: Exact
|
||||
value: /v1/chat/completions
|
||||
headers:
|
||||
- name: x-model
|
||||
type: RegularExpression
|
||||
value: ^vertex_batch/.+
|
||||
backendRefs:
|
||||
- group: agentgateway.dev
|
||||
kind: AgentgatewayBackend
|
||||
name: vertex-batch
|
||||
---
|
||||
apiVersion: agentgateway.dev/v1alpha1
|
||||
kind: AgentgatewayPolicy
|
||||
metadata:
|
||||
name: batch-native
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
targetRefs:
|
||||
- group: gateway.networking.k8s.io
|
||||
kind: HTTPRoute
|
||||
name: batch-native
|
||||
traffic:
|
||||
transformation:
|
||||
request:
|
||||
set:
|
||||
- name: x-batch-stream
|
||||
value: 'string(has(json(request.body).stream) ? json(request.body).stream : false)'
|
||||
- name: x-batch-model
|
||||
value: 'json(request.body).model.stripPrefix("vertex_batch/")'
|
||||
backend:
|
||||
ai:
|
||||
transformations:
|
||||
- field: model
|
||||
expression: llmRequest.model.stripPrefix("vertex_batch/")
|
||||
75
services/vertex-batch/main.go
Normal file
75
services/vertex-batch/main.go
Normal file
|
|
@ -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
|
||||
}
|
||||
66
services/vertex-batch/stream.go
Normal file
66
services/vertex-batch/stream.go
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
var streamError = []byte("data: {\"error\":{\"message\":\"Batch stream ended without a completed response\",\"type\":\"batch_error\"}}\n\n")
|
||||
|
||||
type events struct {
|
||||
pending []byte
|
||||
complete bool
|
||||
done bool
|
||||
}
|
||||
|
||||
func (e *events) feed(body []byte, end bool) []byte {
|
||||
e.pending = append(e.pending, body...)
|
||||
var output []byte
|
||||
for !e.done {
|
||||
cut, size := bytes.Index(e.pending, []byte("\n\n")), 2
|
||||
if cr := bytes.Index(e.pending, []byte("\r\n\r\n")); cr >= 0 && (cut < 0 || cr < cut) {
|
||||
cut, size = cr, 4
|
||||
}
|
||||
if cut < 0 {
|
||||
break
|
||||
}
|
||||
event := e.pending[:cut+size]
|
||||
var data []byte
|
||||
for line := range bytes.SplitSeq(e.pending[:cut], []byte("\n")) {
|
||||
if bytes.HasPrefix(line, []byte("data:")) {
|
||||
if len(data) > 0 {
|
||||
data = append(data, '\n')
|
||||
}
|
||||
data = append(data, bytes.TrimSpace(line[5:])...)
|
||||
}
|
||||
}
|
||||
if bytes.Equal(data, []byte("[DONE]")) {
|
||||
e.done = true
|
||||
if !e.complete {
|
||||
output = append(output, streamError...)
|
||||
} else {
|
||||
output = append(output, event...)
|
||||
}
|
||||
} else {
|
||||
var chunk struct {
|
||||
Choices []struct {
|
||||
Finish *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if json.Unmarshal(data, &chunk) == nil {
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Finish != nil && *choice.Finish != "" {
|
||||
e.complete = true
|
||||
}
|
||||
}
|
||||
}
|
||||
output = append(output, event...)
|
||||
}
|
||||
e.pending = e.pending[cut+size:]
|
||||
}
|
||||
if end && !e.done {
|
||||
output = append(output, streamError...)
|
||||
e.done = true
|
||||
}
|
||||
return output
|
||||
}
|
||||
28
services/vertex-batch/stream_test.go
Normal file
28
services/vertex-batch/stream_test.go
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestEvents(t *testing.T) {
|
||||
finish := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n")
|
||||
usage := []byte("data: {\"choices\":[],\"usage\":{\"total_tokens\":16}}\n\n")
|
||||
done := []byte("data: [DONE]\n\n")
|
||||
input := append(append(append([]byte(nil), finish...), usage...), done...)
|
||||
for split := 0; split <= len(input); split++ {
|
||||
var parser events
|
||||
output := parser.feed(input[:split], false)
|
||||
output = append(output, parser.feed(input[split:], true)...)
|
||||
if !bytes.Equal(output, input) || !parser.done {
|
||||
t.Fatalf("split %d: %q", split, output)
|
||||
}
|
||||
}
|
||||
for _, input := range [][]byte{done, []byte("data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\ndata: [DONE]\n\n"), nil} {
|
||||
var parser events
|
||||
output := parser.feed(input, true)
|
||||
if !bytes.Contains(output, streamError) || bytes.Contains(output, done) {
|
||||
t.Fatalf("failed stream: %q", output)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue