1  {-# LANGUAGE OverloadedStrings #-}
 2  
 3  {- |
 4  Run a command in a namespace with GPU device access.
 5  
 6  Usage: gpu-run COMMAND [ARGS...]
 7  
 8  Like fhs-run but with:
 9  - Full /dev access (including GPU devices)
10  - /sys mounted (for GPU info)
11  - /run/current-system for nvidia drivers
12  -}
13  module Main where
14  
15  import Aleph.Script hiding (FilePath)
16  import qualified Aleph.Script.Tools.Bwrap as Bwrap
17  import Data.Function ((&))
18  import qualified Data.List as L
19  import System.Environment (getArgs)
20  
21  main :: IO ()
22  main = do
23      args <- getArgs
24      case args of
25          [] -> script $ do
26              echoErr "Usage: gpu-run COMMAND [ARGS...]"
27              echoErr ""
28              echoErr "Run COMMAND in a namespace with GPU device access."
29              echoErr ""
30              echoErr "Examples:"
31              echoErr "  gpu-run nvidia-smi"
32              echoErr "  gpu-run python3 -c 'import torch; print(torch.cuda.is_available())'"
33              exit 1
34          cmd -> script $ do
35              homeDir <- getEnvDefault "HOME" "/root"
36              cwd <- pwd
37  
38              let sandbox =
39                      Bwrap.defaults
40                          -- Core system
41                          & Bwrap.roBind "/nix/store" "/nix/store"
42                          & Bwrap.devBind "/dev" "/dev"
43                          & Bwrap.proc "/proc"
44                          & Bwrap.roBind "/sys" "/sys"
45                          & Bwrap.tmpfs "/tmp"
46                          -- NixOS specifics for GPU
47                          & Bwrap.roBind "/run/current-system" "/run/current-system"
48                          -- Network/SSL
49                          & Bwrap.roBind "/etc/resolv.conf" "/etc/resolv.conf"
50                          & Bwrap.roBind "/etc/ssl" "/etc/ssl"
51                          -- User directories (read-write)
52                          & Bwrap.bind (unpack homeDir) (unpack homeDir)
53                          & Bwrap.bind cwd cwd
54                          & Bwrap.chdir cwd
55                          -- Environment
56                          & Bwrap.setenv "PATH" "/run/current-system/sw/bin:/nix/store/bin:/usr/local/bin:/usr/bin:/bin"
57                          & Bwrap.setenv "HOME" homeDir
58                          & Bwrap.dieWithParent
59  
60              Bwrap.exec sandbox (map pack cmd)