Send metrics on runner startup

Author: nevillelyh

internal/server/runner.go

@@ -120,7 +120,6 @@ func NewRunner(ipcUrl, uploadUrl string) *Runner {
 }
 
 func NewProcedureRunner(ipcUrl, uploadUrl, name, srcDir string) *Runner {
-	// Use srcDir as name
 	r := newRunner(name, ipcUrl, uploadUrl)
 	r.cmd.Dir = srcDir
 	return r
@@ -333,6 +332,12 @@ func (r *Runner) config() {
 		predictorName = c
 		// Default to 1 if not set in cog.yaml, regardless whether async predict or not
 		r.maxConcurrency = max(1, y.Concurrency.Max)
+
+		// Send metrics for normal single instance runner
+		// Do not send for multi-tenant procedure runners to reduce noise
+		if r.name == DefaultRunner {
+			go util.SendRunnerMetric(*y)
+		}
 	}
 	conf := PredictConfig{
 		ModuleName:     moduleName,

internal/util/metrics.go

@@ -0,0 +1,44 @@
+package util
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"os"
+)
+
+type MetricsPayload struct {
+	Source string         `json:"source,omitempty"`
+	Type   string         `json:"type,omitempty"`
+	Data   map[string]any `json:"data,omitempty"`
+}
+
+const MetricsEndpointEnv = "COG_METRICS_ENDPOINT"
+
+func SendRunnerMetric(yaml CogYaml) {
+	log := logger.Sugar()
+	endpoint := os.Getenv(MetricsEndpointEnv)
+	if endpoint == "" {
+		return
+	}
+	data := map[string]any{
+		"gpu":         yaml.Build.GPU,
+		"fast":        yaml.Build.Fast,
+		"cog_runtime": yaml.Build.CogRuntime,
+		"version":     Version(),
+	}
+	payload := MetricsPayload{
+		Source: "cog-runtime",
+		Type:   "runner",
+		Data:   data,
+	}
+	body, err := json.Marshal(payload)
+	if err != nil {
+		log.Errorw("failed to marshal payload", "error", err)
+		return
+	}
+	resp, err := http.DefaultClient.Post(endpoint, "application/json", bytes.NewBuffer(body))
+	if err != nil || resp.StatusCode != http.StatusOK {
+		log.Errorw("failed to send runner metrics", "error", err)
+	}
+}

internal/util/util.go

@@ -10,17 +10,28 @@ import (
 	"strings"
 	"time"
 
+	"github.com/replicate/go/logging"
+
 	"github.com/replicate/go/must"
 	"github.com/replicate/go/uuid"
 
 	"gopkg.in/yaml.v3"
 )
 
+var logger = logging.New("cog-util")
+
+type Build struct {
+	GPU        bool `yaml:"gpu"`
+	Fast       bool `yaml:"fast"`
+	CogRuntime bool `yaml:"cog_runtime"`
+}
+
 type Concurrency struct {
 	Max int `yaml:"max"`
 }
 
 type CogYaml struct {
+	Build       Build       `yaml:"build"`
 	Concurrency Concurrency `yaml:"concurrency"`
 	Predict     string      `yaml:"predict"`
 }