From ed1161db36f6cd5c1b4a31445bce6691e53b7e90 Mon Sep 17 00:00:00 2001
From: Tariq Ibrahim <tibrahim@nvidia.com>
Date: Sun, 5 Jul 2026 11:22:30 -0700
Subject: [PATCH] [CDI Hooks] add ability to specify OCI hook type

Signed-off-by: Tariq Ibrahim <tibrahim@nvidia.com>
---
 internal/discover/hooks.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/internal/discover/hooks.go b/internal/discover/hooks.go
index fb16f8587..a31e097fc 100644
--- a/internal/discover/hooks.go
+++ b/internal/discover/hooks.go
@@ -26,6 +26,18 @@ import (
 // A HookName represents a supported CDI hooks.
 type HookName string
 
+// An OCIHookType is the OCI lifecycle stage at which a CDI hook runs.
+type OCIHookType string
+
+const (
+	// OCIHookTypeCreateRuntime runs in the runtime namespace before the container is created.
+	OCIHookTypeCreateRuntime = OCIHookType(cdi.CreateRuntimeHook)
+	// OCIHookTypeCreateContainer runs in the container mount namespace before pivot_root.
+	OCIHookTypeCreateContainer = OCIHookType(cdi.CreateContainerHook)
+	// OCIHookTypeStartContainer runs in the container namespace after the container is started.
+	OCIHookTypeStartContainer = OCIHookType(cdi.StartContainerHook)
+)
+
 const (
 	// AllHooks is a special hook name that allows all hooks to be matched.
 	AllHooks = HookName("all")
@@ -195,13 +207,22 @@ func (c cdiHookCreator) Create(name HookName, args ...string) *Hook {
 	}
 
 	return &Hook{
-		Lifecycle: cdi.CreateContainerHook,
+		Lifecycle: string(c.getOCIHookType(name)),
 		Path:      c.nvidiaCDIHookPath,
 		Args:      append(c.requiredArgs(name), c.transformArgs(name, args...)...),
 		Env:       []string{fmt.Sprintf("NVIDIA_CTK_DEBUG=%v", c.debugLogging)},
 	}
 }
 
+func (c cdiHookCreator) getOCIHookType(name HookName) OCIHookType {
+	switch name {
+	case CreateSymlinksHook, ChmodHook, DisableDeviceNodeModificationHook, EnableCudaCompatHook, UpdateLDCacheHook:
+		return OCIHookTypeCreateContainer
+	default:
+		return OCIHookTypeCreateContainer
+	}
+}
+
 func (c cdiHookCreator) isDisabled(name HookName, args ...string) bool {
 	disabled, ok := c.disabledHooks[name]
 	if ok {