Rebuild ModelsManager and Scheduler routers with new CORS origins

Signed-off-by: Dorin Geman <[email protected]>

Author: doringeman

pkg/inference/models/manager.go

@@ -10,6 +10,7 @@ import (
 	"path"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/docker/model-distribution/distribution"
 	"github.com/docker/model-distribution/registry"
@@ -26,6 +27,8 @@ const (
 	maximumConcurrentModelPulls = 2
 )
 
+var lock sync.Mutex
+
 // Manager manages inference model pulls and storage.
 type Manager struct {
 	// log is the associated logger.
@@ -100,6 +103,20 @@ func NewManager(log logging.Logger, c ClientConfig, allowedOrigins []string) *Ma
 	return m
 }
 
+func (m *Manager) RebuildRoutes(allowedOrigins []string) {
+	lock.Lock()
+	defer lock.Unlock()
+	// Clear existing routes and re-register them.
+	m.router = http.NewServeMux()
+	// Register routes.
+	m.router.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
+		http.Error(w, "not found", http.StatusNotFound)
+	})
+	for route, handler := range m.routeHandlers(allowedOrigins) {
+		m.router.HandleFunc(route, handler)
+	}
+}
+
 func (m *Manager) routeHandlers(allowedOrigins []string) map[string]http.HandlerFunc {
 	handlers := map[string]http.HandlerFunc{
 		"POST " + inference.ModelsPrefix + "/create":                          m.handleCreateModel,
@@ -494,6 +511,8 @@ func (m *Manager) GetDiskUsage() (int64, error, int) {
 
 // ServeHTTP implement net/http.Handler.ServeHTTP.
 func (m *Manager) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	lock.Lock()
+	defer lock.Unlock()
 	m.router.ServeHTTP(w, r)
 }
 

pkg/inference/scheduling/scheduler.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/docker/model-distribution/distribution"
@@ -19,6 +20,8 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+var lock sync.Mutex
+
 // Scheduler is used to coordinate inference scheduling across multiple backends
 // and models.
 type Scheduler struct {
@@ -75,6 +78,20 @@ func NewScheduler(
 	return s
 }
 
+func (s *Scheduler) RebuildRoutes(allowedOrigins []string) {
+	lock.Lock()
+	defer lock.Unlock()
+	// Clear existing routes and re-register them.
+	s.router = http.NewServeMux()
+	// Register routes.
+	s.router.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
+		http.Error(w, "not found", http.StatusNotFound)
+	})
+	for route, handler := range s.routeHandlers(allowedOrigins) {
+		s.router.HandleFunc(route, handler)
+	}
+}
+
 func (s *Scheduler) routeHandlers(allowedOrigins []string) map[string]http.HandlerFunc {
 	openAIRoutes := []string{
 		"POST " + inference.InferencePrefix + "/{backend}/v1/chat/completions",
@@ -332,5 +349,7 @@ func (s *Scheduler) Unload(w http.ResponseWriter, r *http.Request) {
 
 // ServeHTTP implements net/http.Handler.ServeHTTP.
 func (s *Scheduler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	lock.Lock()
+	defer lock.Unlock()
 	s.router.ServeHTTP(w, r)
 }