unset env NVIDIA_VISIBLE_DEVICES when gpushare is enabled (#1273)

* unset env NVIDIA_VISIBLE_DEVICES when gpushare is enabled

Signed-off-by: Yi Chen <github@chenyicn.net>

* Group constants into one const block

Signed-off-by: Yi Chen <github@chenyicn.net>

---------

Signed-off-by: Yi Chen <github@chenyicn.net>

Author: ChenYi015

pkg/argsbuilder/serving.go

@@ -24,15 +24,16 @@ import (
 	"strings"
 	"time"
 
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/kubeflow/arena/pkg/apis/config"
 	"github.com/kubeflow/arena/pkg/apis/types"
+	"github.com/kubeflow/arena/pkg/common"
 	"github.com/kubeflow/arena/pkg/util"
-	log "github.com/sirupsen/logrus"
-	"github.com/spf13/cobra"
-	"k8s.io/apimachinery/pkg/api/errors"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 type ServingArgsBuilder struct {
@@ -685,12 +686,26 @@ func (s *ServingArgsBuilder) disabledNvidiaENVWithNoneGPURequest() error {
 	if s.args.Envs == nil {
 		s.args.Envs = map[string]string{}
 	}
+
+	// Handle cloud vendor special configurations first.
+	if s.hasAliyunGPUConfig() {
+		return nil
+	}
+
+	// Handle general GPU logic.
 	if s.args.GPUCount == 0 && s.args.GPUMemory == 0 && s.args.GPUCore == 0 {
-		s.args.Envs["NVIDIA_VISIBLE_DEVICES"] = "void"
+		s.args.Envs[common.ENV_NVIDIA_VISIBLE_DEVICES] = "void"
 	}
+
 	return nil
 }
 
+// hasAliyunGPUConfig check whether args has labels/devices related to Aliyun GPU config
+func (s *ServingArgsBuilder) hasAliyunGPUConfig() bool {
+	return (s.args.Labels != nil && s.args.Labels[common.LABEL_ALIYUN_COM_GPU_COUNT] != "") ||
+		(s.args.Devices != nil && s.args.Devices[common.DEVICE_ALIYUN_COM_GPU_MEM] != "")
+}
+
 func (s *ServingArgsBuilder) checkGPUCore() error {
 	if s.args.GPUCore%5 != 0 {
 		return fmt.Errorf("GPUCore should be the multiple of 5")

pkg/argsbuilder/serving_distributed.go

@@ -19,8 +19,10 @@ import (
 	"reflect"
 	"strings"
 
-	"github.com/kubeflow/arena/pkg/apis/types"
 	"github.com/spf13/cobra"
+
+	"github.com/kubeflow/arena/pkg/apis/types"
+	"github.com/kubeflow/arena/pkg/common"
 )
 
 type DistributedServingArgsBuilder struct {
@@ -139,7 +141,7 @@ func (s *DistributedServingArgsBuilder) setNvidiaENV() error {
 	// Since master and worker share the same envs, but they may have
 	// different gpu resource, we delete the NVIDIA_VISIBLE_DEVICES env
 	// and set it in helm chart manually
-	delete(s.args.Envs, "NVIDIA_VISIBLE_DEVICES")
+	delete(s.args.Envs, common.ENV_NVIDIA_VISIBLE_DEVICES)
 	return nil
 }
 

pkg/argsbuilder/submit.go

@@ -23,11 +23,13 @@ import (
 	"strconv"
 	"strings"
 
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+
 	"github.com/kubeflow/arena/pkg/apis/config"
 	"github.com/kubeflow/arena/pkg/apis/types"
+	"github.com/kubeflow/arena/pkg/common"
 	"github.com/kubeflow/arena/pkg/util"
-	log "github.com/sirupsen/logrus"
-	"github.com/spf13/cobra"
 )
 
 type SubmitArgsBuilder struct {
@@ -628,12 +630,26 @@ func (s *SubmitArgsBuilder) disabledNvidiaENVWithNoneGPURequest() error {
 	if s.args.Envs == nil {
 		s.args.Envs = map[string]string{}
 	}
+
+	// Handle cloud vendor special configurations first.
+	if s.hasAliyunGPUConfig() {
+		return nil
+	}
+
+	// Handle general GPU logic.
 	if s.args.GPUCount == 0 {
-		s.args.Envs["NVIDIA_VISIBLE_DEVICES"] = "void"
+		s.args.Envs[common.ENV_NVIDIA_VISIBLE_DEVICES] = "void"
 	}
+
 	return nil
 }
 
+// hasAliyunGPUConfig check whether args has labels/devices related to Aliyun GPU config.
+func (s *SubmitArgsBuilder) hasAliyunGPUConfig() bool {
+	return (s.args.Labels != nil && s.args.Labels[common.LABEL_ALIYUN_COM_GPU_COUNT] != "") ||
+		(s.args.Devices != nil && s.args.Devices[common.DEVICE_ALIYUN_COM_GPU_MEM] != "")
+}
+
 func (s *SubmitArgsBuilder) setJobInfoToEnv() error {
 	if s.args.Envs == nil {
 		s.args.Envs = map[string]string{}

pkg/common/constants.go

@@ -0,0 +1,25 @@
+//
+// Copyright 2025 The Kubeflow Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//       http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package common
+
+const (
+	LABEL_ALIYUN_COM_GPU_COUNT = "aliyun.com/gpu-count"
+
+	ENV_NVIDIA_VISIBLE_DEVICES = "NVIDIA_VISIBLE_DEVICES"
+
+	DEVICE_ALIYUN_COM_GPU_MEM = "aliyun.com/gpu-mem"
+)