Consolidate GDS and MOFED modifiers

Signed-off-by: Evan Lezar <elezar@nvidia.com>

internal/modifier/gated.go

@@ -27,27 +27,43 @@ import (
 )
 
 const (
-	nvidiaGDSEnvvar = "NVIDIA_GDS"
+	nvidiaGDSEnvvar   = "NVIDIA_GDS"
+	nvidiaMOFEDEnvvar = "NVIDIA_MOFED"
 )
 
-// NewGDSModifier creates the modifiers for GDS devices.
+// NewFeatureGatedModifier creates the modifiers for optional features.
-// If the spec does not contain the NVIDIA_GDS=enabled environment variable no changes are made.
+// These include:
-func NewGDSModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
+//
+//	NVIDIA_GDS=enabled
+//	NVIDIA_MOFED=enabled
+//
+// If not devices are selected, no changes are made.
+func NewFeatureGatedModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
 	if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 {
 		logger.Infof("No modification required; no devices requested")
 		return nil, nil
 	}
 
-	if image.Getenv(nvidiaGDSEnvvar) != "enabled" {
+	var discoverers []discover.Discover
-		return nil, nil
-	}
 
 	driverRoot := cfg.NVIDIAContainerCLIConfig.Root
 	devRoot := cfg.NVIDIAContainerCLIConfig.Root
-	d, err := discover.NewGDSDiscoverer(logger, driverRoot, devRoot)
+
-	if err != nil {
+	if image.Getenv(nvidiaGDSEnvvar) == "enabled" {
-		return nil, fmt.Errorf("failed to construct discoverer for GDS devices: %v", err)
+		d, err := discover.NewGDSDiscoverer(logger, driverRoot, devRoot)
+		if err != nil {
+			return nil, fmt.Errorf("failed to construct discoverer for GDS devices: %w", err)
+		}
+		discoverers = append(discoverers, d)
 	}
 
-	return NewModifierFromDiscoverer(logger, d)
+	if image.Getenv(nvidiaMOFEDEnvvar) == "enabled" {
+		d, err := discover.NewMOFEDDiscoverer(logger, devRoot)
+		if err != nil {
+			return nil, fmt.Errorf("failed to construct discoverer for MOFED devices: %w", err)
+		}
+		discoverers = append(discoverers, d)
+	}
+
+	return NewModifierFromDiscoverer(logger, discover.Merge(discoverers...))
 }

internal/modifier/mofed.go

@@ -1,51 +0,0 @@
-/**
-# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
-#
-# 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 modifier
-
-import (
-	"fmt"
-
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/oci"
-)
-
-const (
-	nvidiaMOFEDEnvvar = "NVIDIA_MOFED"
-)
-
-// NewMOFEDModifier creates the modifiers for MOFED devices.
-// If the spec does not contain the NVIDIA_MOFED=enabled environment variable no changes are made.
-func NewMOFEDModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
-	if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 {
-		logger.Infof("No modification required; no devices requested")
-		return nil, nil
-	}
-
-	if image.Getenv(nvidiaMOFEDEnvvar) != "enabled" {
-		return nil, nil
-	}
-
-	d, err := discover.NewMOFEDDiscoverer(logger, cfg.NVIDIAContainerCLIConfig.Root)
-	if err != nil {
-		return nil, fmt.Errorf("failed to construct discoverer for MOFED devices: %v", err)
-	}
-
-	return NewModifierFromDiscoverer(logger, d)
-}

internal/runtime/runtime_factory.go

@@ -87,12 +87,7 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
 		return nil, err
 	}
 
-	gdsModifier, err := modifier.NewGDSModifier(logger, cfg, image)
+	featureModifier, err := modifier.NewFeatureGatedModifier(logger, cfg, image)
-	if err != nil {
-		return nil, err
-	}
-
-	mofedModifier, err := modifier.NewMOFEDModifier(logger, cfg, image)
 	if err != nil {
 		return nil, err
 	}
@@ -100,8 +95,7 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
 	modifiers := modifier.Merge(
 		modeModifier,
 		graphicsModifier,
-		gdsModifier,
+		featureModifier,
-		mofedModifier,
 	)
 	return modifiers, nil
 }