tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1#!/usr/bin/env nix-shell
2#! nix-shell -i bash -p nix jq
3
4# shellcheck shell=bash
5
6set -euo pipefail
7
8mkRedistUrlRelativePath() {
9 local -r cudaMajorMinorVersion=${1:?}
10 local -r tensorrtMajorMinorPatchBuildVersion=${2:?}
11 local -r redistSystem=${3:?}
12
13 local -r tensorrtMajorMinorPatchVersion="$(echo "$tensorrtMajorMinorPatchBuildVersion" | cut -d. -f1-3)"
14 local -r tensorrtMinorVersion="$(echo "$tensorrtMajorMinorPatchVersion" | cut -d. -f2)"
15
16 local archiveDir=""
17 local archiveExtension=""
18 local osName=""
19 local platformName=""
20 case "$redistSystem" in
21 linux-aarch64) archiveDir="tars" && archiveExtension="tar.gz" && osName="l4t" && platformName="aarch64-gnu" ;;
22 linux-sbsa)
23 archiveDir="tars" && archiveExtension="tar.gz" && platformName="aarch64-gnu"
24 # 10.0-10.3 use Ubuntu 22.40
25 # 10.4-10.6 use Ubuntu 24.04
26 # 10.7+ use Linux
27 case "$tensorrtMinorVersion" in
28 0 | 1 | 2 | 3) osName="Ubuntu-22.04" ;;
29 4 | 5 | 6) osName="Ubuntu-24.04" ;;
30 *) osName="Linux" ;;
31 esac
32 ;;
33 linux-x86_64) archiveDir="tars" && archiveExtension="tar.gz" && osName="Linux" && platformName="x86_64-gnu" ;;
34 windows-x86_64)
35 archiveExtension="zip" && platformName="win10"
36 # Windows info is different for 10.0.*
37 case "$tensorrtMinorVersion" in
38 0) archiveDir="zips" && osName="Windows10" ;;
39 *) archiveDir="zip" && osName="Windows" ;;
40 esac
41 ;;
42 *)
43 echo "mkRedistUrlRelativePath: Unsupported redistSystem: $redistSystem" >&2
44 exit 1
45 ;;
46 esac
47
48 local -r relativePath="tensorrt/$tensorrtMajorMinorPatchVersion/$archiveDir/TensorRT-${tensorrtMajorMinorPatchBuildVersion}.${osName}.${platformName}.cuda-${cudaMajorMinorVersion}.${archiveExtension}"
49 echo "$relativePath"
50}
51
52getNixStorePath() {
53 local -r relativePath=${1:?}
54 local -r jsonBlob="$(nix store prefetch-file --json "https://developer.nvidia.com/downloads/compute/machine-learning/$relativePath")"
55 if [[ -z $jsonBlob ]]; then
56 echo "getNixStorePath: Failed to fetch jsonBlob for relativePath: $relativePath" >&2
57 exit 1
58 fi
59 local -r storePath="$(echo "$jsonBlob" | jq -cr '.storePath')"
60 echo "$storePath"
61}
62
63printOutput() {
64 local -r cudaMajorMinorVersion=${1:?}
65 local -r redistSystem=${2:?}
66 local -r md5Hash=${3:?}
67 local -r relativePath=${4:?}
68 local -r sha256Hash=${5:?}
69 local -r size=${6:?}
70
71 local -r cudaVariant="cuda$(echo "$cudaMajorMinorVersion" | cut -d. -f1)"
72
73 # Echo everything to stdout using JQ to format the output as JSON
74 jq \
75 --raw-output \
76 --sort-keys \
77 --null-input \
78 --arg redistSystem "$redistSystem" \
79 --arg cudaVariant "$cudaVariant" \
80 --arg md5Hash "$md5Hash" \
81 --arg relativePath "$relativePath" \
82 --arg sha256Hash "$sha256Hash" \
83 --arg size "$size" \
84 '{
85 $redistSystem: {
86 $cudaVariant: {
87 md5: $md5Hash,
88 relative_path: $relativePath,
89 sha256: $sha256Hash,
90 size: $size
91 }
92 }
93 }'
94}
95
96main() {
97 local -r cudaMajorMinorVersion=${1:?}
98 if [[ ! $cudaMajorMinorVersion =~ ^[0-9]+\.[0-9]+$ ]]; then
99 echo "main: Invalid cudaMajorMinorVersion: $cudaMajorMinorVersion" >&2
100 exit 1
101 fi
102
103 local -r tensorrtMajorMinorPatchBuildVersion=${2:?}
104 if [[ ! $tensorrtMajorMinorPatchBuildVersion =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
105 echo "main: Invalid tensorrtMajorMinorPatchBuildVersion: $tensorrtMajorMinorPatchBuildVersion" >&2
106 exit 1
107 fi
108
109 local -r redistSystem=${3:?}
110 case "$redistSystem" in
111 linux-aarch64) ;;
112 linux-sbsa) ;;
113 linux-x86_64) ;;
114 windows-x86_64) ;;
115 *)
116 echo "main: Unsupported redistSystem: $redistSystem" >&2
117 exit 1
118 ;;
119 esac
120
121 local -r relativePath="$(mkRedistUrlRelativePath "$cudaMajorMinorVersion" "$tensorrtMajorMinorPatchBuildVersion" "$redistSystem")"
122 local -r storePath="$(getNixStorePath "$relativePath")"
123 echo "main: storePath: $storePath" >&2
124 local -r md5Hash="$(nix hash file --type md5 --base16 "$storePath")"
125 local -r sha256Hash="$(nix hash file --type sha256 --base16 "$storePath")"
126 local -r size="$(du -b "$storePath" | cut -f1)"
127
128 printOutput "$cudaMajorMinorVersion" "$redistSystem" "$md5Hash" "$relativePath" "$sha256Hash" "$size"
129}
130
131main "$@"