diff --git a/cmd/nvidia-container-runtime/modifier/experimental.go b/cmd/nvidia-container-runtime/modifier/experimental.go index 9693d05b..399f9534 100644 --- a/cmd/nvidia-container-runtime/modifier/experimental.go +++ b/cmd/nvidia-container-runtime/modifier/experimental.go @@ -18,7 +18,6 @@ package modifier import ( "fmt" - "path/filepath" "github.com/NVIDIA/nvidia-container-toolkit/internal/config" "github.com/NVIDIA/nvidia-container-toolkit/internal/discover" @@ -69,9 +68,9 @@ func newExperimentalModifierFromDiscoverer(logger *logrus.Logger, d discover.Dis // Modify applies the required modifications to the incomming OCI spec. These modifications // are applied in-place. func (m experimental) Modify(spec *specs.Spec) error { - err := m.assertSpecIsCompatible(spec) + err := nvidiaContainerRuntimeHookRemover{m.logger}.Modify(spec) if err != nil { - return fmt.Errorf("OCI specification cannot be modified: %v", err) + return fmt.Errorf("failed to remove existing hooks: %v", err) } specEdits, err := edits.NewSpecEdits(m.logger, m.discoverer) @@ -81,38 +80,3 @@ func (m experimental) Modify(spec *specs.Spec) error { return specEdits.Modify(spec) } - -func (m experimental) assertSpecIsCompatible(spec *specs.Spec) error { - if spec == nil { - return nil - } - - if spec.Hooks == nil { - return nil - } - - if hookPath := findStableHook(spec.Hooks.Prestart); hookPath != "" { - return fmt.Errorf("spec already contains required 'prestart' hook: %v", hookPath) - } - - return nil -} - -// findStableHook checks the list of OCI hooks for the nvidia-container-runtime-hook -// or nvidia-container-toolkit hook. These are included, for example, by the non-experimental -// nvidia-container-runtime or docker when specifying the --gpus flag. -func findStableHook(hooks []specs.Hook) string { - lookFor := map[string]bool{ - nvidiaContainerRuntimeHookExecuable: true, - nvidiaContainerToolkitExecutable: true, - } - - for _, h := range hooks { - base := filepath.Base(h.Path) - if lookFor[base] { - return h.Path - } - } - - return "" -} diff --git a/cmd/nvidia-container-runtime/modifier/hook_remover.go b/cmd/nvidia-container-runtime/modifier/hook_remover.go new file mode 100644 index 00000000..ac3354d7 --- /dev/null +++ b/cmd/nvidia-container-runtime/modifier/hook_remover.go @@ -0,0 +1,79 @@ +/** +# 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" + "path/filepath" + + "github.com/NVIDIA/nvidia-container-toolkit/internal/oci" + "github.com/opencontainers/runtime-spec/specs-go" + "github.com/sirupsen/logrus" +) + +// nvidiaContainerRuntimeHookRemover is a spec modifer that detects and removes inserted nvidia-container-runtime hooks +type nvidiaContainerRuntimeHookRemover struct { + logger *logrus.Logger +} + +var _ oci.SpecModifier = (*nvidiaContainerRuntimeHookRemover)(nil) + +// Modify removes any NVIDIA Container Runtime hooks from the provided spec +func (m nvidiaContainerRuntimeHookRemover) Modify(spec *specs.Spec) error { + if spec == nil { + return nil + } + + if spec.Hooks == nil { + return nil + } + + var updateRequired bool + newPrestart := make([]specs.Hook, 0, len(spec.Hooks.Prestart)) + + for _, hook := range spec.Hooks.Prestart { + if isNVIDIAContainerRuntimeHook(&hook) { + m.logger.Infof("Removing hook %v", hook) + updateRequired = true + continue + } + newPrestart = append(newPrestart, hook) + } + + if updateRequired { + // TODO: Once we have updated the hook implementation to give an error if invoked incorrectly, we will update the spec hooks here instead of just logging. + // We can then also use a boolean to track whether this is required instead of storing the removed hooks + // spec.Hooks.Prestart = newPrestart + m.logger.Debugf("Updating 'prestart' hooks to %v", newPrestart) + return fmt.Errorf("spec already contains required 'prestart' hook") + } + + return nil +} + +// isNVIDIAContainerRuntimeHook checks if the provided hook is an nvidia-container-runtime-hook +// or nvidia-container-toolkit hook. These are included, for example, by the non-experimental +// nvidia-container-runtime or docker when specifying the --gpus flag. +func isNVIDIAContainerRuntimeHook(hook *specs.Hook) bool { + lookFor := map[string]bool{ + nvidiaContainerRuntimeHookExecuable: true, + nvidiaContainerToolkitExecutable: true, + } + base := filepath.Base(hook.Path) + + return lookFor[base] +}