Add model exists check (#353)

vllm-project · Nov 7, 2024 · e3e5b4e · e3e5b4e
1 parent 7cfb63b
commit e3e5b4e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go
@@ -400,6 +400,15 @@ func (c *Cache) GetModelsForPod(podName string) (map[string]struct{}, error) {
 	return models, nil
 }
 
+func (c *Cache) CheckModelExists(modelName string) bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	_, ok := c.ModelToPodMapping[modelName]
+
+	return ok
+}
+
 func (c *Cache) GetPodMetric(podName, metricName string) (float64, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()

diff --git a/pkg/plugins/gateway/gateway.go b/pkg/plugins/gateway/gateway.go
@@ -213,11 +213,11 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestID string, req *e
 			"error processing request body"), targetPodIP
 	}
 
-	if model, ok = jsonMap["model"].(string); !ok || model == "" {
+	if model, ok = jsonMap["model"].(string); !ok || model == "" || !s.cache.CheckModelExists(model) {
 		return generateErrorResponse(envoyTypePb.StatusCode_InternalServerError,
 			[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
-				Key: "x-no-model", RawValue: []byte("true")}}},
-			"no model in request body"), targetPodIP
+				Key: "x-no-model", RawValue: []byte(model)}}},
+			"no model in request body or model does not exist"), targetPodIP
 	}
 
 	headers := []*configPb.HeaderValueOption{}