Add metrics.predict_time to response

replicate · Dec 13, 2024 · bd8d0f0 · bd8d0f0
1 parent 061bf30
commit bd8d0f0
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/internal/server/runner.go b/internal/server/runner.go
@@ -397,6 +397,10 @@ func (r *Runner) handleResponses() {
 			log.Infow("prediction processing", "id", pr.request.Id, "status", pr.response.Status)
 			pr.sendWebhook(WebhookOutput)
 		} else if pr.response.Status.IsCompleted() {
+			if pr.response.Status == PredictionSucceeded {
+				t := util.ParseTime(pr.response.CompletedAt).Sub(util.ParseTime(pr.response.StartedAt)).Seconds()
+				pr.response.Metrics = PredictionMetrics{PredictTime: t}
+			}
 			log.Infow("prediction completed", "id", pr.request.Id, "status", pr.response.Status)
 			pr.sendWebhook(WebhookCompleted)
 			pr.sendResponse()

diff --git a/internal/server/types.go b/internal/server/types.go
@@ -85,14 +85,19 @@ type PredictionRequest struct {
 	OutputFilePrefix    string         `json:"output_file_prefix,omitempty"`
 }
 
+type PredictionMetrics struct {
+	PredictTime float64 `json:"predict_time"`
+}
+
 type PredictionResponse struct {
-	Input       any              `json:"input"`
-	Output      any              `json:"output"`
-	Id          string           `json:"id"`
-	CreatedAt   string           `json:"created_at"`
-	StartedAt   string           `json:"started_at"`
-	CompletedAt string           `json:"completed_at"`
-	Logs        string           `json:"logs,omitempty"`
-	Error       string           `json:"error,omitempty"`
-	Status      PredictionStatus `json:"status,omitempty"`
+	Input       any               `json:"input"`
+	Output      any               `json:"output"`
+	Id          string            `json:"id"`
+	CreatedAt   string            `json:"created_at"`
+	StartedAt   string            `json:"started_at"`
+	CompletedAt string            `json:"completed_at"`
+	Logs        string            `json:"logs,omitempty"`
+	Error       string            `json:"error,omitempty"`
+	Status      PredictionStatus  `json:"status,omitempty"`
+	Metrics     PredictionMetrics `json:"metrics,omitempty"`
 }
diff --git a/internal/util/util.go b/internal/util/util.go
@@ -48,9 +48,15 @@ func PredictionId() string {
 	return encoding.EncodeToString(shuffle)
 }
 
+const TimeLayout = "2006-01-02T15:04:05.999999-07:00"
+
 func NowIso() string {
 	// Python: datetime.now(tz=timezone.utc).isoformat()
-	return time.Now().UTC().Format("2006-01-02T15:04:05.999999-07:00")
+	return time.Now().UTC().Format(TimeLayout)
+}
+
+func ParseTime(t string) time.Time {
+	return must.Get(time.Parse(TimeLayout, t))
 }
 
 func JoinLogs(logs []string) string {