diff --git a/src/rocker/cli.py b/src/rocker/cli.py index 70d2cfbc..314692d9 100644 --- a/src/rocker/cli.py +++ b/src/rocker/cli.py @@ -37,6 +37,7 @@ def main(): parser.add_argument('--nocleanup', action='store_true', help='do not remove the docker container when stopped') parser.add_argument('--persist-image', action='store_true', help='do not remove the docker image when stopped', default=False) #TODO(tfoote) Add a name to it if persisting parser.add_argument('--pull', action='store_true') + parser.add_argument('--dive', action='store_true', help='Run Dive for introspection after building the image') # 310 parser.add_argument('--version', action='version', version='%(prog)s ' + get_rocker_version()) diff --git a/src/rocker/nvidia_extension.py b/src/rocker/nvidia_extension.py index 3644c220..ac764309 100644 --- a/src/rocker/nvidia_extension.py +++ b/src/rocker/nvidia_extension.py @@ -188,6 +188,21 @@ def __init__(self): self.supported_distros = ['Ubuntu', 'Debian GNU/Linux'] self.supported_versions = ['20.04', '22.04', '24.04', '11', '12'] # Debian 11 and 12 +# ADDED: Function to check if CUDA is installed + def is_cuda_installed(self): + """Check if CUDA is already installed on the system.""" + try: + result = subprocess.run( + ["nvidia-smi"], # Use `nvidia-smi` to check for CUDA + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True + ) + return result.returncode == 0 # Return True if CUDA is installed + except (FileNotFoundError, subprocess.CalledProcessError): + # If `nvidia-smi` is not found or fails, CUDA is not installed + return False + def get_environment_subs(self, cliargs={}): if not self._env_subs: self._env_subs = {} @@ -222,7 +237,14 @@ def get_preamble(self, cliargs): # preamble = pkgutil.get_data('rocker', 'templates/%s_preamble.Dockerfile.em' % self.name).decode('utf-8') # return empy_expand(preamble, self.get_environment_subs(cliargs)) +# CHANGED: Modified `get_snippet()` to skip installation if CUDA is detected def get_snippet(self, cliargs): + """Generate Dockerfile snippet for installing CUDA.""" + if self.is_cuda_installed(): + # ADDED: Skip installation if CUDA is already installed + print("CUDA is already installed. Skipping installation.") + return "" # No Dockerfile snippet if CUDA is installed + print("CUDA not found. Proceeding with installation.") snippet = pkgutil.get_data('rocker', 'templates/%s_snippet.Dockerfile.em' % self.name).decode('utf-8') return empy_expand(snippet, self.get_environment_subs(cliargs))