From 3a062c106253c5add1b3ad64fb37ff78ef508dea Mon Sep 17 00:00:00 2001
From: Swati Gupta <swatig@nvidia.com>
Date: Mon, 28 Oct 2024 20:45:52 +0000
Subject: [PATCH] install systemd

---
 ubuntu22.04/Dockerfile                | 2 ++
 ubuntu22.04/nvidia-driver             | 8 ++++++--
 ubuntu22.04/precompiled/Dockerfile    | 2 ++
 ubuntu22.04/precompiled/nvidia-driver | 8 ++++++--
 4 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/ubuntu22.04/Dockerfile b/ubuntu22.04/Dockerfile
index 2c690fb3..193492ea 100644
--- a/ubuntu22.04/Dockerfile
+++ b/ubuntu22.04/Dockerfile
@@ -93,6 +93,8 @@ RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | h
     curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     nvidia-imex --version && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends systemd && \
     rm -rf /tmp/nvidia-imex_${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb; fi
 
 WORKDIR  /drivers
diff --git a/ubuntu22.04/nvidia-driver b/ubuntu22.04/nvidia-driver
index d48c445e..c1fced6b 100755
--- a/ubuntu22.04/nvidia-driver
+++ b/ubuntu22.04/nvidia-driver
@@ -189,7 +189,7 @@ _create_driver_package() (
 
 # Check if imex node config is present
 _assert_imex_node_config() {
-    [ -f /etc/nvidia-imex/node_config.cfg ]] || return 1
+    [ -f /etc/nvidia-imex/nodes_config.cfg ] || return 1
 }
 
 _assert_nvswitch_system() {
@@ -333,7 +333,9 @@ _load_driver() {
 
     if _assert_imex_node_config; then
 	echo "Starting NVIDIA imex daemon..."
+	systemctl enable nvidia-imex.service
 	systemctl start nvidia-imex.service
+	systemctl status nvidia-imex.service
     fi
 
     return 0
@@ -394,9 +396,11 @@ _unload_driver() {
     fi
 
     # stop nvidia-imex daemon
-    if [ systemctl is-active --quiet nvidia-imex.service ]; then
+    if systemctl is-active --quiet nvidia-imex.service; then
 	echo "Stopping NVIDIA imex daemon..."
 	systemctl stop nvidia-imex.service
+	systemctl status nvidia-imex.service
+	systemctl disable nvidia-imex.service
     fi
 
     echo "Unloading NVIDIA driver kernel modules..."
diff --git a/ubuntu22.04/precompiled/Dockerfile b/ubuntu22.04/precompiled/Dockerfile
index 64bc2a83..0846e049 100644
--- a/ubuntu22.04/precompiled/Dockerfile
+++ b/ubuntu22.04/precompiled/Dockerfile
@@ -66,6 +66,8 @@ RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | h
     curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     nvidia-imex --version && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends systemd && \
     rm -rf /tmp/nvidia-imex_${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb; fi
 
 # update pkg cache and download pkgs for driver module installation during runtime.
diff --git a/ubuntu22.04/precompiled/nvidia-driver b/ubuntu22.04/precompiled/nvidia-driver
index 5253b550..503269b7 100755
--- a/ubuntu22.04/precompiled/nvidia-driver
+++ b/ubuntu22.04/precompiled/nvidia-driver
@@ -23,7 +23,7 @@ _update_package_cache() {
 }
 
 _assert_imex_node_config() {
-    [ -f /etc/nvidia-imex/node_config.cfg ]] || return 1
+    [ -f /etc/nvidia-imex/nodes_config.cfg ] || return 1
 }
 
 _assert_nvswitch_system() {
@@ -135,7 +135,9 @@ _load_driver() {
 
     if _assert_imex_node_config; then
 	echo "Starting NVIDIA imex daemon..."
+	systemctl enable nvidia-imex.service
 	systemctl start nvidia-imex.service
+	systemctl status nvidia-imex.service
     fi
 
     return 0
@@ -180,9 +182,11 @@ _unload_driver() {
         fi
     fi
 
-    if [ systemctl is-active --quiet nvidia-imex.service ]; then
+    if systemctl is-active --quiet nvidia-imex.service; then
 	echo "Stopping NVIDIA imex daemon..."
 	systemctl stop nvidia-imex.service
+	systemctl status nvidia-imex.service
+	systemctl disable nvidia-imex.service
     fi
 
     echo "Unloading NVIDIA driver kernel modules..."