Merge pull request #256 from doringeman/metrics-ua

metrics: always use default User-Agent and add action to model tracking

Author: doringeman

pkg/inference/scheduling/scheduler.go

@@ -231,7 +231,7 @@ func (s *Scheduler) handleOpenAIInference(w http.ResponseWriter, r *http.Request
 			return
 		}
 		// Non-blocking call to track the model usage.
-		s.tracker.TrackModel(model, r.UserAgent())
+		s.tracker.TrackModel(model, r.UserAgent(), "inference/"+backendMode.String())
 	}
 
 	modelID := s.modelManager.ResolveModelID(request.Model)
@@ -415,7 +415,7 @@ func (s *Scheduler) Configure(w http.ResponseWriter, r *http.Request) {
 
 	if model, err := s.modelManager.GetModel(configureRequest.Model); err == nil {
 		// Configure is called by compose for each model.
-		s.tracker.TrackModel(model, r.UserAgent())
+		s.tracker.TrackModel(model, r.UserAgent(), "configure/"+mode.String())
 	}
 	modelID := s.modelManager.ResolveModelID(configureRequest.Model)
 	if err := s.loader.setRunnerConfig(r.Context(), backend.Name(), modelID, mode, runnerConfig); err != nil {

pkg/metrics/metrics.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/docker/model-runner/pkg/distribution/types"
@@ -39,8 +40,11 @@ func NewTracker(httpClient *http.Client, log logging.Logger, userAgent string, d
 		client.Transport = http.DefaultTransport
 	}
 
+	userAgent = strings.TrimSpace(userAgent)
 	if userAgent == "" {
 		userAgent = "docker-model-runner"
+	} else {
+		userAgent = userAgent + " docker-model-runner"
 	}
 
 	if os.Getenv("DEBUG") == "1" {
@@ -59,24 +63,28 @@ func NewTracker(httpClient *http.Client, log logging.Logger, userAgent string, d
 	}
 }
 
-func (t *Tracker) TrackModel(model types.Model, userAgent string) {
+func (t *Tracker) TrackModel(model types.Model, userAgent, action string) {
 	if t.doNotTrack {
 		return
 	}
 
-	go t.trackModel(model, userAgent)
+	go t.trackModel(model, userAgent, action)
 }
 
-func (t *Tracker) trackModel(model types.Model, userAgent string) {
+func (t *Tracker) trackModel(model types.Model, userAgent, action string) {
 	tags := model.Tags()
 	t.log.Debugln("Tracking model:", tags)
 	if len(tags) == 0 {
 		return
 	}
-	ua := t.userAgent
+	parts := []string{t.userAgent}
 	if userAgent != "" {
-		ua += " " + userAgent
+		parts = append(parts, userAgent)
 	}
+	if action != "" {
+		parts = append(parts, action)
+	}
+	ua := strings.Join(parts, " ")
 	for _, tag := range tags {
 		ref, err := name.ParseReference(tag)
 		if err != nil {