feat(docker): no support for --gpus flag in

Author: Parity Agent

cmd/amd-container-runtime-hook/hook.go

@@ -0,0 +1,146 @@
+/**
+# Copyright (c) Advanced Micro Devices, Inc. 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 main
+
+import (
+	"github.com/ROCm/container-toolkit/internal/amdgpu"
+	gpuTracker "github.com/ROCm/container-toolkit/internal/gpu-tracker"
+	"github.com/opencontainers/runtime-spec/specs-go"
+)
+
+// ociHookHandler wraps OCI functionality for hook-specific operations
+type ociHookHandler struct {
+	spec        *specs.Spec
+	bundlePath  string
+	containerId string
+
+	reserveGPUs func(string, string) ([]int, error)
+}
+
+func newOCIHandlerFromSpec(spec *specs.Spec, bundlePath, containerId string) (*ociHookHandler, error) {
+	gpuTracker, err := gpuTracker.New()
+	if err != nil {
+		return nil, err
+	}
+
+	return &ociHookHandler{
+		spec:        spec,
+		bundlePath:  bundlePath,
+		containerId: containerId,
+		reserveGPUs: gpuTracker.ReserveGPUs,
+	}, nil
+}
+
+func (h *ociHookHandler) InjectGPUs() error {
+	// Extract AMD_VISIBLE_DEVICES from spec environment
+	amdDevices, err := h.getRequestedDevices()
+	if err != nil {
+		return err
+	}
+
+	if len(amdDevices) == 0 {
+		return nil // No devices requested
+	}
+
+	// Get all GPUs
+	gpus, err := amdgpu.GetAMDGPUs()
+	if err != nil {
+		return err
+	}
+
+	// Add each requested GPU
+	for _, idx := range amdDevices {
+		if idx >= len(gpus) {
+			continue // Skip invalid indices
+		}
+		for _, drmDev := range gpus[idx].DrmDevices {
+			gpu, err := amdgpu.GetAMDGPU(drmDev)
+			if err != nil {
+				return err
+			}
+			if err := h.addGPUDevice(gpu); err != nil {
+				return err
+			}
+		}
+	}
+
+	// Add /dev/kfd
+	kfd, err := amdgpu.GetAMDGPU("/dev/kfd")
+	if err != nil {
+		return err
+	}
+	return h.addGPUDevice(kfd)
+}
+
+func (h *ociHookHandler) getRequestedDevices() ([]int, error) {
+	if h.spec.Process == nil {
+		return []int{}, nil
+	}
+
+	for _, env := range h.spec.Process.Env {
+		if len(env) > 20 && env[:20] == "AMD_VISIBLE_DEVICES=" {
+			gpuList := env[20:]
+			return h.reserveGPUs(gpuList, h.containerId)
+		}
+		if len(env) > 16 && env[:16] == "DOCKER_RESOURCE_" {
+			gpuList := env[16:]
+			idx := 16
+			for idx < len(env) && env[idx] != '=' {
+				idx++
+			}
+			if idx < len(env) {
+				gpuList = env[idx+1:]
+				return h.reserveGPUs(gpuList, h.containerId)
+			}
+		}
+	}
+	return []int{}, nil
+}
+
+func (h *ociHookHandler) addGPUDevice(gpu amdgpu.AMDGPU) error {
+	// Add to devices list
+	dev := specs.LinuxDevice{
+		Path:     gpu.Path,
+		Type:     gpu.DevType,
+		Major:    gpu.Major,
+		Minor:    gpu.Minor,
+		FileMode: &gpu.FileMode,
+		GID:      &gpu.Gid,
+		UID:      &gpu.Uid,
+	}
+
+	if h.spec.Linux == nil {
+		h.spec.Linux = &specs.Linux{}
+	}
+	h.spec.Linux.Devices = append(h.spec.Linux.Devices, dev)
+
+	// Add to cgroup resources
+	rdev := specs.LinuxDeviceCgroup{
+		Allow:  gpu.Allow,
+		Type:   gpu.DevType,
+		Major:  &gpu.Major,
+		Minor:  &gpu.Minor,
+		Access: gpu.Access,
+	}
+
+	if h.spec.Linux.Resources == nil {
+		h.spec.Linux.Resources = &specs.LinuxResources{}
+	}
+	h.spec.Linux.Resources.Devices = append(h.spec.Linux.Resources.Devices, rdev)
+
+	return nil
+}

cmd/amd-container-runtime-hook/main.go

@@ -0,0 +1,151 @@
+/**
+# Copyright (c) Advanced Micro Devices, Inc. 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 main
+
+import (
+	"encoding/json"
+	"flag"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	gpuTracker "github.com/ROCm/container-toolkit/internal/gpu-tracker"
+	"github.com/ROCm/container-toolkit/internal/logger"
+	"github.com/opencontainers/runtime-spec/specs-go"
+)
+
+var (
+	debugFlag   = flag.Bool("debug", false, "enable debug logging")
+	versionFlag = flag.Bool("version", false, "print version and exit")
+)
+
+// HookState represents the container state passed to the hook via stdin
+type HookState struct {
+	Pid    int    `json:"pid,omitempty"`
+	Bundle string `json:"bundle"`
+}
+
+func main() {
+	flag.Parse()
+
+	if *versionFlag {
+		fmt.Printf("amd-container-runtime-hook version 1.0.0\n")
+		return
+	}
+
+	logger.Init(*debugFlag)
+	logger.Log.Printf("AMD Container Runtime Hook started")
+
+	// OCI hooks receive container state on stdin
+	var state HookState
+	decoder := json.NewDecoder(os.Stdin)
+	if err := decoder.Decode(&state); err != nil {
+		logger.Log.Printf("Failed to decode container state: %v", err)
+		os.Exit(1)
+	}
+
+	logger.Log.Printf("Processing container in bundle: %s", state.Bundle)
+
+	// Load OCI spec
+	specPath := filepath.Join(state.Bundle, "config.json")
+	spec, err := loadSpec(specPath)
+	if err != nil {
+		logger.Log.Printf("Failed to load spec from %s: %v", specPath, err)
+		os.Exit(1)
+	}
+
+	// Check if GPU devices are requested
+	if !needsGPUs(spec) {
+		logger.Log.Printf("No AMD GPU devices requested, exiting")
+		return
+	}
+
+	// Create OCI handler with spec already loaded
+	containerId := filepath.Base(state.Bundle)
+	ociHandler, err := newOCIHandlerFromSpec(spec, state.Bundle, containerId)
+	if err != nil {
+		logger.Log.Printf("Failed to create OCI handler: %v", err)
+		releaseGPUs(containerId)
+		os.Exit(1)
+	}
+
+	// Add GPU devices (without cleanup hook - Docker manages hook lifecycle)
+	if err := ociHandler.InjectGPUs(); err != nil {
+		logger.Log.Printf("Failed to inject GPUs: %v", err)
+		releaseGPUs(containerId)
+		os.Exit(1)
+	}
+
+	// Write modified spec back
+	if err := writeSpec(specPath, spec); err != nil {
+		logger.Log.Printf("Failed to write spec to %s: %v", specPath, err)
+		releaseGPUs(containerId)
+		os.Exit(1)
+	}
+
+	logger.Log.Printf("Successfully injected AMD GPU devices")
+}
+
+func loadSpec(path string) (*specs.Spec, error) {
+	file, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	var spec specs.Spec
+	if err := json.NewDecoder(file).Decode(&spec); err != nil {
+		return nil, err
+	}
+	return &spec, nil
+}
+
+func writeSpec(path string, spec *specs.Spec) error {
+	file, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	encoder := json.NewEncoder(file)
+	encoder.SetIndent("", "  ")
+	return encoder.Encode(spec)
+}
+
+func needsGPUs(spec *specs.Spec) bool {
+	if spec.Process == nil {
+		return false
+	}
+	for _, env := range spec.Process.Env {
+		if len(env) >= 20 && env[:20] == "AMD_VISIBLE_DEVICES=" {
+			return true
+		}
+		if len(env) >= 16 && env[:16] == "DOCKER_RESOURCE_" {
+			return true
+		}
+	}
+	return false
+}
+
+func releaseGPUs(containerId string) {
+	tracker, err := gpuTracker.New()
+	if err != nil {
+		logger.Log.Printf("Failed to create GPU tracker for cleanup: %v", err)
+		return
+	}
+	tracker.ReleaseGPUs(containerId)
+}

cmd/amd-ctk/runtime/runtime.go

@@ -17,6 +17,9 @@
 package runtime
 
 import (
+	"fmt"
+	"os"
+
 	"github.com/ROCm/container-toolkit/cmd/amd-ctk/runtime/configure"
 	"github.com/urfave/cli/v2"
 )
@@ -31,7 +34,69 @@ func AddNewCommand() *cli.Command {
 
 	runtimeCmd.Subcommands = []*cli.Command{
 		configure.AddNewCommand(),
+		configureHookCommand(),
 	}
 
 	return &runtimeCmd
 }
+
+func configureHookCommand() *cli.Command {
+	return &cli.Command{
+		Name:  "configure-hooks",
+		Usage: "Configure OCI prestart hooks for --gpus flag support",
+		Flags: []cli.Flag{
+			&cli.StringFlag{
+				Name:  "hook-path",
+				Usage: "Path to amd-container-runtime-hook binary",
+				Value: "/usr/bin/amd-container-runtime-hook",
+			},
+			&cli.StringFlag{
+				Name:  "config-path",
+				Usage: "Path to Docker daemon.json",
+				Value: "/etc/docker/daemon.json",
+			},
+			&cli.BoolFlag{
+				Name:  "remove",
+				Usage: "Remove hook configuration",
+			},
+		},
+		Action: func(c *cli.Context) error {
+			hookPath := c.String("hook-path")
+			configPath := c.String("config-path")
+			remove := c.Bool("remove")
+
+			if remove {
+				fmt.Println("Removing hook configuration not yet implemented")
+				return nil
+			}
+
+			// Verify hook binary exists
+			if _, err := os.Stat(hookPath); os.IsNotExist(err) {
+				return fmt.Errorf("hook binary not found at %s", hookPath)
+			}
+
+			fmt.Printf("Hook configuration:\n")
+			fmt.Printf("  Hook binary: %s\n", hookPath)
+			fmt.Printf("  Config file: %s\n", configPath)
+			fmt.Printf("\nTo enable --gpus flag support, add this to %s:\n\n", configPath)
+			fmt.Printf(`{
+  "runtimes": {
+    "amd": {
+      "path": "amd-container-runtime",
+      "runtimeArgs": []
+    }
+  },
+  "hooks": {
+    "prestart": [
+      {
+        "path": "%s"
+      }
+    ]
+  }
+}
+`, hookPath)
+			fmt.Printf("\nThen restart Docker: sudo systemctl restart docker\n")
+			return nil
+		},
+	}
+}