Procedure support

cmd/cog/main.go

@@ -19,14 +19,6 @@ import (
 	"github.com/replicate/cog-runtime/internal/util"
 )
 
-type ServerConfig struct {
-	Host                  string `ff:"long: host, default: 0.0.0.0, usage: HTTP server host"`
-	Port                  int    `ff:"long: port, default: 5000, usage: HTTP server port"`
-	WorkingDir            string `ff:"long: working-dir, nodefault, usage: working directory"`
-	AwaitExplicitShutdown bool   `ff:"long: await-explicit-shutdown, default: false, usage: await explicit shutdown"`
-	UploadUrl             string `ff:"long: upload-url, nodefault, usage: output file upload URL"`
-}
-
 var logger = logging.New("cog")
 
 func schemaCommand() *ff.Command {
@@ -44,7 +36,7 @@ func schemaCommand() *ff.Command {
 				log.Errorw("failed to read cog.yaml", "err", err)
 				return err
 			}
-			m, c, err := y.PredictModuleAndClass()
+			m, c, err := y.PredictModuleAndPredictor()
 			if err != nil {
 				log.Errorw("failed to parse predict", "err", err)
 				return err
@@ -58,25 +50,15 @@ func schemaCommand() *ff.Command {
 func serverCommand() *ff.Command {
 	log := logger.Sugar()
 
-	var cfg ServerConfig
+	cfg := &server.Config{}
 	flags := ff.NewFlagSet("server")
-	must.Do(flags.AddStruct(&cfg))
+	must.Do(flags.AddStruct(cfg))
 
 	return &ff.Command{
 		Name:  "server",
 		Usage: "server [FLAGS]",
 		Flags: flags,
 		Exec: func(ctx context.Context, args []string) error {
-			workingDir := cfg.WorkingDir
-			if workingDir == "" {
-				workingDir = must.Get(os.MkdirTemp("", "cog-server-"))
-			}
-			log.Infow("configuration",
-				"working-dir", workingDir,
-				"await-explicit-shutdown", cfg.AwaitExplicitShutdown,
-				"upload-url", cfg.UploadUrl,
-			)
-
 			ctx, cancel := context.WithCancel(ctx)
 			go func() {
 				ch := make(chan os.Signal, 1)
@@ -88,13 +70,18 @@ func serverCommand() *ff.Command {
 
 			addr := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)
 			log.Infow("starting Cog HTTP server", "addr", addr)
-			r := server.NewRunner(workingDir, cfg.AwaitExplicitShutdown, cfg.UploadUrl)
+			r, err := server.NewRunner(cfg)
+			if err != nil {
+				log.Fatalw("failed to build runner", "err", err)
+			}
 			must.Do(r.Start())
 			s := server.NewServer(addr, r)
 			go func() {
 				<-ctx.Done()
 				must.Do(r.Shutdown())
-				must.Do(s.Shutdown(ctx))
+				if err := s.Shutdown(ctx); err != nil && !errors.Is(err, context.Canceled) {
+					panic(err)
+				}
 			}()
 			if err := s.ListenAndServe(); errors.Is(err, http.ErrServerClosed) {
 				if r.ExitCode() == 0 {
@@ -125,7 +112,7 @@ func testCommand() *ff.Command {
 				log.Errorw("failed to read cog.yaml", "err", err)
 				return err
 			}
-			m, c, err := y.PredictModuleAndClass()
+			m, c, err := y.PredictModuleAndPredictor()
 			if err != nil {
 				log.Errorw("failed to parse predict", "err", err)
 				return err

go.mod

@@ -1,11 +1,13 @@
 module github.com/replicate/cog-runtime
 
 go 1.23.3
+
 require (
 	github.com/gabriel-vasile/mimetype v1.4.8
 	github.com/peterbourgon/ff/v4 v4.0.0-alpha.4
 	github.com/replicate/go v0.0.0-20241101110715-45e9ae8c2040
 	github.com/stretchr/testify v1.10.0
+	golang.org/x/sync v0.13.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 

go.sum

@@ -30,6 +30,8 @@ go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
 go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
 golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
 golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
+golang.org/x/sync v0.13.0 h1:AauUjRAJ9OSnvULf/ARrrVywoJDy0YS2AwQ98I37610=
+golang.org/x/sync v0.13.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

internal/server/config.go

@@ -0,0 +1,10 @@
+package server
+
+type Config struct {
+	Host                  string `ff:"long: host, default: 0.0.0.0, usage: HTTP server host"`
+	Port                  int    `ff:"long: port, default: 5000, usage: HTTP server port"`
+	WorkingDir            string `ff:"long: working-dir, nodefault, usage: working directory"`
+	AwaitExplicitShutdown bool   `ff:"long: await-explicit-shutdown, default: false, usage: await explicit shutdown"`
+	UploadUrl             string `ff:"long: upload-url, nodefault, usage: output file upload URL"`
+	UseProcedureMode      bool   `ff:"long: use-procedure-mode, default: false, usage: use procedure-specific runner"`
+}

internal/server/http.go

@@ -13,7 +13,7 @@ var (
 	ErrSetupFailed = errors.New("setup failed")
 )
 
-func NewServer(addr string, runner *Runner) *http.Server {
+func NewServer(addr string, runner Runner) *http.Server {
 	handler := Handler{runner: runner}
 	serveMux := http.NewServeMux()
 	serveMux.HandleFunc("GET /{$}", handler.Root)

internal/server/pending_prediction.go

@@ -0,0 +1,97 @@
+package server
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"slices"
+	"sync"
+	"time"
+)
+
+type PendingPrediction struct {
+	request     PredictionRequest
+	response    PredictionResponse
+	lastUpdated time.Time
+	inputPaths  []string
+	mu          sync.Mutex
+	c           chan *PredictionResponse
+}
+
+func (pr *PendingPrediction) appendLogLine(line string) {
+	pr.mu.Lock()
+	defer pr.mu.Unlock()
+	pr.response.Logs += fmt.Sprintln(line)
+}
+
+func (pr *PendingPrediction) sendWebhook(event WebhookEvent) error {
+	pr.mu.Lock()
+	defer pr.mu.Unlock()
+
+	if pr.request.Webhook == "" {
+		return nil
+	}
+
+	if len(pr.request.WebhookEventsFilter) > 0 && !slices.Contains(pr.request.WebhookEventsFilter, event) {
+		return nil
+	}
+
+	if event == WebhookLogs || event == WebhookOutput {
+		if time.Since(pr.lastUpdated) < 500*time.Millisecond {
+			return nil
+		}
+
+		pr.lastUpdated = time.Now()
+	}
+
+	log := logger.Sugar()
+	log.Infow("sending webhook", "url", pr.request.Webhook, "response", pr.response)
+
+	bodyBytes, err := json.Marshal(pr.response)
+	if err != nil {
+		log.Errorw("failed to marshal pending prediction response", "err", err)
+
+		return err
+	}
+
+	req, err := http.NewRequest(http.MethodPost, pr.request.Webhook, bytes.NewBuffer(bodyBytes))
+	if err != nil {
+		log.Errorw("failed to build pending prediction webhook request", "err", err)
+
+		return err
+	}
+
+	req.Header.Add("Content-Type", "application/json")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		log.Errorw("failed to send webhook", "error", err)
+
+		return err
+	} else if resp.StatusCode != 200 {
+		defer resp.Body.Close()
+
+		respBody, err := io.ReadAll(resp.Body)
+		if err != nil {
+			log.Errorw("failed to read pending prediction webhook response body", "err", err)
+
+			return err
+		}
+
+		log.Errorw("failed to send pending prediction webhook", "code", resp.StatusCode, "body", string(respBody))
+	}
+
+	return nil
+}
+
+func (pr *PendingPrediction) sendResponse() {
+	if pr.c == nil {
+		logger.Warn("attempted pending prediction response with nil channel")
+
+		return
+	}
+
+	pr.c <- &pr.response
+}