feat(lmeval): Add operator online and execution feature flag

trustyai-explainability · Dec 10, 2024 · 94afdbe · 94afdbe
1 parent 6f88955
commit 94afdbe
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 30 deletions.
diff --git a/config/base/params.env b/config/base/params.env
@@ -9,4 +9,5 @@ lmes-image-pull-policy=Always
 lmes-max-batch-size=24
 lmes-default-batch-size=8
 lmes-detect-device=true
-
+lmes-allow-online=true
+lmes-allow-code-execution=true
diff --git a/config/overlays/odh-kueue/params.env b/config/overlays/odh-kueue/params.env
@@ -9,3 +9,5 @@ lmes-image-pull-policy=Always
 lmes-max-batch-size=24
 lmes-default-batch-size=8
 lmes-detect-device=true
+lmes-allow-online=true
+lmes-allow-code-execution=true
diff --git a/config/overlays/odh/params.env b/config/overlays/odh/params.env
@@ -9,3 +9,5 @@ lmes-image-pull-policy=Always
 lmes-max-batch-size=24
 lmes-default-batch-size=8
 lmes-detect-device=true
+lmes-allow-online=true
+lmes-allow-code-execution=true
diff --git a/config/overlays/rhoai/params.env b/config/overlays/rhoai/params.env
@@ -9,3 +9,5 @@ lmes-image-pull-policy=Always
 lmes-max-batch-size=24
 lmes-default-batch-size=8
 lmes-detect-device=true
+lmes-allow-online=false
+lmes-allow-code-execution=false
diff --git a/controllers/lmes/config.go b/controllers/lmes/config.go
@@ -37,6 +37,8 @@ var Options *serviceOptions = &serviceOptions{
 	MaxBatchSize:        DefaultMaxBatchSize,
 	DetectDevice:        DefaultDetectDevice,
 	DefaultBatchSize:    DefaultBatchSize,
+	AllowOnline:         false,
+	AllowCodeExecution:  false,
 }
 
 type serviceOptions struct {
@@ -47,6 +49,8 @@ type serviceOptions struct {
 	MaxBatchSize        int
 	DefaultBatchSize    string
 	DetectDevice        bool
+	AllowOnline         bool
+	AllowCodeExecution  bool
 }
 
 func constructOptionsFromConfigMap(log *logr.Logger, configmap *corev1.ConfigMap) error {

diff --git a/controllers/lmes/constants.go b/controllers/lmes/constants.go
@@ -34,6 +34,8 @@ const (
 	MaxBatchSizeKey            = "lmes-max-batch-size"
 	DefaultBatchSizeKey        = "lmes-default-batch-size"
 	DetectDeviceKey            = "lmes-detect-device"
+	AllowOnline                = "lmes-allow-online"
+	AllowCodeExecution         = "lmes-allow-code-execution"
 	DefaultPodImage            = "quay.io/trustyai/ta-lmes-job:latest"
 	DefaultDriverImage         = "quay.io/trustyai/ta-lmes-driver:latest"
 	DefaultPodCheckingInterval = time.Second * 10

diff --git a/controllers/lmes/lmevaljob_controller.go b/controllers/lmes/lmevaljob_controller.go
@@ -68,6 +68,8 @@ var (
 		"DefaultBatchSize":    DefaultBatchSizeKey,
 		"MaxBatchSize":        MaxBatchSizeKey,
 		"DetectDevice":        DetectDeviceKey,
+		"AllowOnline":         AllowOnline,
+		"AllowCodeExecution":  AllowCodeExecution,
 	}
 
 	labelFilterPrefixes       = []string{}
@@ -713,41 +715,55 @@ func CreatePod(svcOpts *serviceOptions, job *lmesv1alpha1.LMEvalJob, log logr.Lo
 		volumes = append(volumes, outputPVC)
 	}
 
-	// Disable remote code execution by default
-	if job.Spec.AllowCodeExecution == nil || *job.Spec.AllowCodeExecution == false {
-		remoteCodeEnvVars := []corev1.EnvVar{
-			{
-				Name:  "TRUST_REMOTE_CODE",
-				Value: "0",
-			},
-			{
-				Name:  "HF_DATASETS_TRUST_REMOTE_CODE",
-				Value: "0",
-			},
+	remoteCodeEnvVars := []corev1.EnvVar{
+		{
+			Name:  "TRUST_REMOTE_CODE",
+			Value: "0",
+		},
+		{
+			Name:  "HF_DATASETS_TRUST_REMOTE_CODE",
+			Value: "0",
+		},
+	}
+
+	if job.Spec.AllowCodeExecution != nil && *job.Spec.AllowCodeExecution {
+		// Disable remote code execution by default
+
+		if !svcOpts.AllowCodeExecution {
+			log.Error(fmt.Errorf("code execution not allowed by the operator"), "change this setting and redeploy the operator")
+			envVars = append(envVars, remoteCodeEnvVars...)
 		}
+	} else {
 		envVars = append(envVars, remoteCodeEnvVars...)
 	}
 
+	offlineHuggingFaceEnvVars := []corev1.EnvVar{
+		{
+			Name:  "HF_DATASETS_OFFLINE",
+			Value: "1",
+		},
+		{
+			Name:  "HF_HUB_OFFLINE",
+			Value: "1",
+		},
+		{
+			Name:  "TRANSFORMERS_OFFLINE",
+			Value: "1",
+		},
+		{
+			Name:  "HF_EVALUATE_OFFLINE",
+			Value: "1",
+		},
+	}
+
 	// Enforce offline mode by default
-	if job.Spec.AllowOnline == nil || *job.Spec.AllowOnline == false {
-		offlineHuggingFaceEnvVars := []corev1.EnvVar{
-			{
-				Name:  "HF_DATASETS_OFFLINE",
-				Value: "1",
-			},
-			{
-				Name:  "HF_HUB_OFFLINE",
-				Value: "1",
-			},
-			{
-				Name:  "TRANSFORMERS_OFFLINE",
-				Value: "1",
-			},
-			{
-				Name:  "HF_EVALUATE_OFFLINE",
-				Value: "1",
-			},
+	if job.Spec.AllowOnline != nil && *job.Spec.AllowOnline {
+
+		if !svcOpts.AllowOnline {
+			log.Error(fmt.Errorf("online mode not allowed by the operator"), "change this setting and redeploy the operator")
+			envVars = append(envVars, offlineHuggingFaceEnvVars...)
 		}
+	} else {
 		envVars = append(envVars, offlineHuggingFaceEnvVars...)
 	}