Skip to content

Commit

Permalink
Build pre-pruned CUDA/CuDNN tarballs
Browse files Browse the repository at this point in the history
  • Loading branch information
@nevillelyh
nevillelyh committed Jan 13, 2025
1 parent f7f03d8 commit 8b662df
Show file tree
Hide file tree
Showing 4 changed files with 157 additions and 20 deletions.
21 changes: 21 additions & 0 deletions script/build-cuda
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
#!/usr/bin/env bash

# Build pre-pruned CUDA & CuDNN tarballs

set -euo pipefail

MONOBASE_PYTHON='3.13'

cd "$(git rev-parse --show-toplevel)"

mkdir -p build/monobase build/cache

docker run --rm \
--user "$(id -u):$(id -g)" \
--volume "$PWD/src/monobase:/opt/r8/monobase" \
--volume "$PWD/build/monobase:/srv/r8/monobase" \
--volume "$PWD/build/cache:/var/cache/monobase" \
monobase:latest \
/opt/r8/monobase/run.sh \
monobase.cuda \
"$@"
2 changes: 2 additions & 0 deletions src/monobase/__main__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
python -m monobase.build
python -m monobase.cuda
python -m monobase.diff
python -m monobase.monogen
Expand Down
149 changes: 129 additions & 20 deletions src/monobase/cuda.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,15 @@
from dataclasses import dataclass

from monobase.urls import cuda_urls, cudnn_urls
from monobase.util import Version, mark_done, require_done_or_rm
from monobase.util import (
Version,
mark_done,
require_done_or_rm,
setup_logging,
)

# FIXME: set up delivery bucket
R8_PACKAGE_PREFIX = 'http://localhost:8000'

log = logging.getLogger(__name__)

Expand Down Expand Up @@ -65,17 +73,8 @@ def build_cudnns() -> dict[str, CuDNN]:
CUDNNS: dict[str, CuDNN] = build_cudnns()


def install_cuda(args: argparse.Namespace, version: str) -> str:
def build_cuda_tarball(args: argparse.Namespace, version: str) -> str:
cdir = os.path.join(args.prefix, 'cuda', f'cuda-{version}')
if require_done_or_rm(cdir):
log.info(f'CUDA {version} in {cdir} is complete')
return cdir

if args.skip_cuda:
os.makedirs(cdir, exist_ok=True)
mark_done(cdir, kind='cuda', version=version, skipped=True)
log.info(f'CUDA {version} skipped in {cdir}')
return cdir

cuda = CUDAS[version]
file = os.path.join(args.cache, cuda.filename)
Expand Down Expand Up @@ -115,24 +114,45 @@ def install_cuda(args: argparse.Namespace, version: str) -> str:
cmd = ['find', cdir, '-name', 'lib*.a', '-delete']
subprocess.run(cmd, check=True)

mark_done(cdir, kind='cuda', version=version, url=cuda.url)
log.info(f'CUDA {version} installed in {cdir}')
log.info(f'CUDA {version} build completed in {cdir}')
return cdir


def install_cudnn(args: argparse.Namespace, version: str, cuda_major: str) -> str:
key = f'{version}-cuda{cuda_major}'
cdir = os.path.join(args.prefix, 'cuda', f'cudnn-{key}')
def install_cuda(args: argparse.Namespace, version: str) -> str:
cdir = os.path.join(args.prefix, 'cuda', f'cuda-{version}')
if require_done_or_rm(cdir):
log.info(f'CuDNN {key} in {cdir} is complete')
log.info(f'CUDA {version} in {cdir} is complete')
return cdir

if args.skip_cuda:
os.makedirs(cdir, exist_ok=True)
mark_done(cdir, kind='cudnn', version=version, skipped=True)
log.info(f'CuDNN {key} skipped in {cdir}')
mark_done(cdir, kind='cuda', version=version, skipped=True)
log.info(f'CUDA {version} skipped in {cdir}')
return cdir

filename = f'monobase-cuda-{version}.tar.gz'
path = os.path.join(args.cache, filename)
url = f'file://{path}'
if not os.path.exists(path):
log.info(f'Downloading CUDA {version}...')
url = f'{R8_PACKAGE_PREFIX}/cuda/{filename}'
cmd = [f'{args.prefix}/bin/pget', '--pid-file', '/tmp/pget.pid', url, path]
subprocess.run(cmd, check=True)

log.info(f'Installing CUDA {version}...')
os.makedirs(cdir, exist_ok=True)
cmd = ['tar', '-xzf', path, '-C', cdir]
subprocess.run(cmd, check=True)

mark_done(cdir, kind='cuda', version=version, url=url)
log.info(f'CUDA {version} installed in {cdir}')
return cdir


def build_cudnn_tarball(args: argparse.Namespace, version: str, cuda_major: str) -> str:
key = f'{version}-cuda{cuda_major}'
cdir = os.path.join(args.prefix, 'cuda', f'cudnn-{key}')

cudnn = CUDNNS[key]
file = os.path.join(args.cache, cudnn.filename)
if not os.path.exists(file):
Expand All @@ -151,6 +171,95 @@ def install_cudnn(args: argparse.Namespace, version: str, cuda_major: str) -> st
cmd = ['tar', '-xf', file, '--strip-components=1', '--exclude=lib*.a', '-C', cdir]
subprocess.run(cmd, check=True)

mark_done(cdir, kind='cudnn', version=version, url=cudnn.url)
log.info(f'CuDNN {key} build completed in {cdir}')
return cdir


def install_cudnn(args: argparse.Namespace, version: str, cuda_major: str) -> str:
key = f'{version}-cuda{cuda_major}'
cdir = os.path.join(args.prefix, 'cuda', f'cudnn-{key}')
if require_done_or_rm(cdir):
log.info(f'CuDNN {key} in {cdir} is complete')
return cdir

if args.skip_cuda:
os.makedirs(cdir, exist_ok=True)
mark_done(cdir, kind='cudnn', version=version, skipped=True)
log.info(f'CuDNN {key} skipped in {cdir}')
return cdir

filename = f'monobase-cudnn-{key}.tar.gz'
path = os.path.join(args.cache, filename)
url = f'file://{path}'
if not os.path.exists(path):
log.info(f'Downloading CuDNN {key}...')
url = f'{R8_PACKAGE_PREFIX}/cudnn/{filename}'
cmd = [f'{args.prefix}/bin/pget', '--pid-file', '/tmp/pget.pid', url, path]
subprocess.run(cmd, check=True)

log.info(f'Installing CuDNN {key}...')
os.makedirs(cdir, exist_ok=True)
cmd = ['tar', '-xzf', path, '-C', cdir]
subprocess.run(cmd, check=True)

mark_done(cdir, kind='cudnn', version=version, url=url)
log.info(f'CuDNN {key} installed in {cdir}')
return cdir


parser = argparse.ArgumentParser(description='Build monobase environment')
parser.add_argument(
'--prefix',
metavar='PATH',
default='/srv/r8/monobase',
help='prefix for monobase',
)
parser.add_argument(
'--cache',
metavar='PATH',
default='/var/cache/monobase',
help='cache for monobase',
)

if __name__ == '__main__':
setup_logging()

args = parser.parse_args()

def tar(d, f):
# https://www.gnu.org/software//tar/manual/html_section/Reproducibility.html
tar_flags = [
'--sort=name',
'--format=posix',
'--pax-option=exthdr.name=%d/PaxHeaders/%f',
'--pax-option=delete=atime,delete=ctime',
'--clamp-mtime',
'--mtime=0',
'--numeric-owner',
'--owner=0',
'--group=0',
'--mode=go+u,go-w',
]
tar_env = os.environ.copy()
tar_env['LC_ALL'] = 'C'
tar_env['TZ'] = 'UTC'
tar_env['GZIPFLAGS'] = '--no-name --best'
cmd = ['tar', '-C', d] + tar_flags + ['-czf', f] + os.listdir(d)
subprocess.run(cmd, check=True, env=tar_env)
shutil.rmtree(d, ignore_errors=True)

for k in CUDAS.keys():
tf = os.path.join(args.cache, f'monobase-cuda-{k}.tar.gz')
if os.path.exists(tf):
continue
cdir = build_cuda_tarball(args, k)
log.info(f'Creating CUDA tarball {tf}...')
tar(cdir, tf)

for k, v in CUDNNS.items():
tf = os.path.join(args.cache, f'monobase-cudnn-{k}.tar.gz')
if os.path.exists(tf):
continue
cdir = build_cudnn_tarball(args, str(v.cudnn_version), str(v.cuda_major))
log.info(f'Creating CuDNN tarball {tf}...')
tar(cdir, tf)
5 changes: 5 additions & 0 deletions src/monobase/urls.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ def getenv_or(key: str, default: str) -> str:
)

cuda_urls = [
f'{cuda_prefix}/cuda_12.6.3_560.35.05_linux.run',
f'{cuda_prefix}/cuda_12.6.2_560.35.03_linux.run',
f'{cuda_prefix}/cuda_12.6.1_560.35.03_linux.run',
f'{cuda_prefix}/cuda_12.6.0_560.28.03_linux.run',
Expand Down Expand Up @@ -50,6 +51,10 @@ def getenv_or(key: str, default: str) -> str:
)

cudnn_urls = [
f'{cudnn_prefix}/cudnn-linux-x86_64-9.6.0.74_cuda12-archive.tar.xz',
f'{cudnn_prefix}/cudnn-linux-x86_64-9.6.0.74_cuda11-archive.tar.xz',
f'{cudnn_prefix}/cudnn-linux-x86_64-9.5.1.17_cuda12-archive.tar.xz',
f'{cudnn_prefix}/cudnn-linux-x86_64-9.5.1.17_cuda11-archive.tar.xz',
f'{cudnn_prefix}/cudnn-linux-x86_64-9.5.0.50_cuda12-archive.tar.xz',
f'{cudnn_prefix}/cudnn-linux-x86_64-9.5.0.50_cuda11-archive.tar.xz',
f'{cudnn_prefix}/cudnn-linux-x86_64-9.4.0.58_cuda12-archive.tar.xz',
Expand Down

0 comments on commit 8b662df

Please sign in to comment.