ce7cb702218693d08ea6d83464d2457c9250027c
This commit is contained in:
commit
34c5f8963c
1 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)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue