Use map as param in customize function in scheduler

SeldonIO · Jan 28, 2025 · 81e2fa1 · 81e2fa1
1 parent 9a9bd1e
commit 81e2fa1
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/scheduler/pkg/agent/repository/mlserver/mlserver.go b/scheduler/pkg/agent/repository/mlserver/mlserver.go
@@ -166,7 +166,7 @@ func (m *MLServerRepositoryHandler) setExtraParameters(
 	pipelineRef *string,
 	envoyHost string,
 	envoyPort int,
-	customize func(*ModelSettings) error,
+	customize func(map[string]interface{}) error,
 ) error {
 	workers := 0
 	settingsPath := filepath.Join(modelRepoPath, mlserverConfigFilename)
@@ -204,7 +204,7 @@ func (m *MLServerRepositoryHandler) setExtraParameters(
 
 	// custom logic for specific caller
 	if customize != nil {
-		if err := customize(ms); err != nil {
+		if err := customize(ms.Parameters.Extra); err != nil {
 			return err
 		}
 	}
@@ -226,8 +226,8 @@ func (m *MLServerRepositoryHandler) SetExplainer(modelRepoPath string, explainer
 			explainerSpec.PipelineRef,
 			envoyHost,
 			envoyPort,
-			func(ms *ModelSettings) error {
-				ms.Parameters.Extra[explainerTypeKey] = &explainerSpec.Type
+			func(extra map[string]interface{}) error {
+				extra[explainerTypeKey] = &explainerSpec.Type
 				return nil
 			},
 		)