quorate/services/vertex-batch/cloud.go

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, &current)
	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, &current); 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)
	}
}