ModelProxy supports the "jobType" field of inference jobs (#118)

* support job Type and inference job parameters

* update

* update

* update

* update

* update

---------

Co-authored-by: Yuting Jiang <[email protected]>

Author: zhogu

src/model-proxy/config/model-proxy.yaml

@@ -5,4 +5,4 @@ service_type: "common"
 
 port: 9999
 retry: 5
-modelkey: "123"
+concurrency: 10

src/model-proxy/deploy/model-proxy.yaml.template

@@ -27,11 +27,12 @@ spec:
         args:
           - "--port={{ cluster_cfg['model-proxy']['port'] }}"
           - "--retry={{ cluster_cfg['model-proxy']['retry'] }}"
-          - "--modelkey={{ cluster_cfg['model-proxy']['modelkey'] }}"
           - "--logdir=/usr/local/ltp/model-proxy/logs"
         env:
         - name: REST_SERVER_URI
           value: {{ cluster_cfg["rest-server"]["uri"] }}
+        - name: CONCURRENCY
+          value: "{{ cluster_cfg['model-proxy']['concurrency'] }}"
         volumeMounts:
         {%- if cluster_cfg['model-proxy']['log_pvc'] %}
         - name: model-proxy-log-storage

src/model-proxy/src/main.go

@@ -15,14 +15,12 @@ var (
 	port       int
 	maxRetries int = 5 // default value
 	logFileDir string
-	modelKey   string
 )
 
 func init() {
 	flag.IntVar(&port, "port", 9999, "port for the proxy server")
 	flag.IntVar(&maxRetries, "retry", 5, "max retries for the request to the model server")
 	flag.StringVar(&logFileDir, "logdir", "./logs", "path to the log file directory")
-	flag.StringVar(&modelKey, "modelkey", "", "model key for requesting model serving jobs")
 }
 
 func main() {
@@ -33,7 +31,6 @@ func main() {
 			Host:       "0.0.0.0",
 			Port:       port,
 			MaxRetries: maxRetries,
-			ModelKey:   modelKey,
 		},
 		Log: &types.Log{
 			LogStorage: &types.LogStorage{

src/model-proxy/src/proxy/authenticator.go

@@ -7,6 +7,8 @@ import (
 	"log"
 	"net/http"
 	"strings"
+
+	"modelproxy/types"
 )
 
 // obfuscateToken returns a truncated identifier for safely logging tokens.
@@ -18,61 +20,57 @@ func obfuscateToken(token string) string {
 }
 
 type RestServerAuthenticator struct {
-	// rest-server token => model names => model urls
-	tokenToModels map[string]map[string][]string
-	modelKey      string
+	// rest-server token => model names => model service list
+	tokenToModels map[string]map[string][]*types.BaseSpec
 }
 
-func NewRestServerAuthenticator(tokenToModels map[string]map[string][]string, modelKey string) *RestServerAuthenticator {
-	if tokenToModels == nil {
-		tokenToModels = make(map[string]map[string][]string)
-	}
+func NewRestServerAuthenticator() *RestServerAuthenticator {
 	return &RestServerAuthenticator{
-		tokenToModels: tokenToModels,
-		modelKey:      modelKey,
+		tokenToModels: make(map[string]map[string][]*types.BaseSpec),
 	}
 }
 
-func (ra *RestServerAuthenticator) UpdateTokenModels(token string, model2Url map[string][]string) {
+// UpdateTokenModels updates the model mapping for a given token
+func (ra *RestServerAuthenticator) UpdateTokenModels(token string, model2Service map[string][]*types.BaseSpec) {
 	if ra.tokenToModels == nil {
-		ra.tokenToModels = make(map[string]map[string][]string)
+		ra.tokenToModels = make(map[string]map[string][]*types.BaseSpec)
 	}
-	ra.tokenToModels[token] = model2Url
+	ra.tokenToModels[token] = model2Service
 }
 
 // Check if the request is authenticated and return the available model urls
-func (ra *RestServerAuthenticator) AuthenticateReq(req *http.Request, reqBody map[string]interface{}) (bool, []string) {
+func (ra *RestServerAuthenticator) AuthenticateReq(req *http.Request, reqBody map[string]interface{}) (bool, []*types.BaseSpec) {
 	token := req.Header.Get("Authorization")
 	token = strings.Replace(token, "Bearer ", "", 1)
 	//  read request body
 	model, ok := reqBody["model"].(string)
 	if !ok {
 		log.Printf("[-] Error: 'model' field missing or not a string in request body")
-		return false, []string{}
+		return false, nil
 	}
 	availableModels, ok := ra.tokenToModels[token]
 	if !ok {
 		// request to RestServer to get the models
 		log.Printf("[-] Error: token %s not found in the authenticator\n", obfuscateToken(token))
-		availableModels, err := GetJobModelsMapping(req, ra.modelKey)
+		availableModels, err := GetJobModelsMapping(req)
 		if err != nil {
 			log.Printf("[-] Error: failed to get models for token %s: %v\n", obfuscateToken(token), err)
-			return false, []string{}
+			return false, nil
 		}
 		ra.tokenToModels[token] = availableModels
 	}
 	if len(availableModels) == 0 {
 		log.Printf("[-] Error: no models found")
-		return false, []string{}
+		return false, nil
 	}
 	if model == "" {
 		log.Printf("[-] Error: model is empty")
-		return false, []string{}
+		return false, nil
 	}
 	for m, v := range availableModels {
 		if m == model {
 			return true, v
 		}
 	}
-	return false, []string{}
+	return false, nil
 }

src/model-proxy/src/proxy/load_balancer.go

@@ -28,16 +28,16 @@ func NewUrlPoller(url string, bsl types.BaseSpecList) *UrlPoller {
 	}
 }
 
-func NewUrlPollerWithKey(url string, modelUrls []string, modelKey string) *UrlPoller {
-	if len(modelUrls) == 0 {
+func NewUrlPollerWithKey(url string, modelServices []*types.BaseSpec) *UrlPoller {
+	if len(modelServices) == 0 {
 		return nil
 	}
-	bsl := make(types.BaseSpecList, 0, len(modelUrls))
-	for _, v := range modelUrls {
+	bsl := make(types.BaseSpecList, 0, len(modelServices))
+	for _, v := range modelServices {
 		bsl = append(bsl, &types.BaseSpecStatistic{
 			BaseSpec: &types.BaseSpec{
-				URL: v,
-				Key: modelKey,
+				URL: v.URL,
+				Key: v.Key,
 			},
 		})
 	}