1  {-# LANGUAGE LambdaCase #-}
 2  {-# LANGUAGE OverloadedStrings #-}
 3  
 4  {- |
 5  nvidia-wheel-extract - Extract NVIDIA libraries from PyPI wheels
 6  
 7  This script downloads and extracts NVIDIA libraries from pypi.nvidia.com
 8  wheels, which have no redistribution restrictions.
 9  
10  Usage:
11    nvidia-wheel-extract <package> <output-dir>
12    nvidia-wheel-extract all <output-dir>
13  
14  Packages:
15    nccl       - NCCL 2.28.9
16    cudnn      - cuDNN 9.17.0.29
17    tensorrt   - TensorRT 10.14.1.48
18    cutensor   - cuTensor 2.4.1
19    cusparselt - cuSPARSELt 0.8.1
20    all        - Extract all packages
21  
22  Examples:
23    nvidia-wheel-extract nccl ./nccl-sdk
24    nvidia-wheel-extract all ./nvidia-libs
25  -}
26  module Main where
27  
28  import Aleph.Script hiding (FilePath)
29  import qualified Aleph.Script.Nvidia.Versions as V
30  import qualified Aleph.Script.Nvidia.Wheel as Wheel
31  import qualified Data.Text as T
32  import System.Environment (getArgs)
33  
34  main :: IO ()
35  main = do
36      args <- getArgs
37      case args of
38          [pkg, outDir] -> script $ verbosely $ extractPackage pkg (fromText $ pack outDir)
39          _ -> printHelp
40  
41  printHelp :: IO ()
42  printHelp = do
43      putStrLn "nvidia-wheel-extract - Extract NVIDIA libraries from PyPI wheels"
44      putStrLn ""
45      putStrLn "Usage:"
46      putStrLn "  nvidia-wheel-extract <package> <output-dir>"
47      putStrLn ""
48      putStrLn "Packages:"
49      putStrLn $ "  nccl       - NCCL " ++ T.unpack V.ncclVersion
50      putStrLn $ "  cudnn      - cuDNN " ++ T.unpack V.cudnnVersion
51      putStrLn $ "  tensorrt   - TensorRT " ++ T.unpack V.tensorrtVersion
52      putStrLn $ "  cutensor   - cuTensor " ++ T.unpack V.cutensorVersion
53      putStrLn $ "  cusparselt - cuSPARSELt " ++ T.unpack V.cusparseltVersion
54      putStrLn "  all        - Extract all packages"
55      putStrLn ""
56      putStrLn "Examples:"
57      putStrLn "  nvidia-wheel-extract nccl ./nccl-sdk"
58      putStrLn "  nvidia-wheel-extract all ./nvidia-libs"
59  
60  extractPackage :: String -> FilePath -> Sh ()
61  extractPackage pkg outDir = case pkg of
62      "nccl" -> extractOne Wheel.nccl outDir
63      "cudnn" -> extractOne Wheel.cudnn outDir
64      "tensorrt" -> extractOne Wheel.tensorrt outDir
65      "cutensor" -> extractOne Wheel.cutensor outDir
66      "cusparselt" -> extractOne Wheel.cusparselt outDir
67      "all" -> Wheel.extractAllWheels V.currentPlatform outDir
68      _ -> do
69          echoErr $ "Unknown package: " <> pack pkg
70          echoErr "Use one of: nccl, cudnn, tensorrt, cutensor, cusparselt, all"
71          exit 1
72  
73  extractOne :: Wheel.WheelDef -> FilePath -> Sh ()
74  extractOne def outDir =
75      case Wheel.forPlatform V.currentPlatform def of
76          Just spec -> do
77              mkdirP outDir
78              Wheel.extractToNixLayout spec outDir
79              echoErr $ ":: Done! Extracted to " <> toTextIgnore outDir
80          Nothing -> do
81              echoErr $ ":: Error: " <> Wheel.defName def <> " not available for this platform"
82              exit 1