Spaces:
Sleeping
Sleeping
File size: 4,960 Bytes
c83d2d6 3a68bf0 c83d2d6 3a68bf0 c83d2d6 3a68bf0 c83d2d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """Build the pinned native SPAC extensions on the Space CPU at startup."""
from __future__ import annotations
import hashlib
import importlib.util
import os
from pathlib import Path
import shutil
import site
import subprocess
import sys
import tarfile
import requests
ROOT=Path(__file__).resolve().parent
CACHE=ROOT/'.cache/zerogpu-native'
CUDA=CACHE/'cuda-12.8'
REDIST='https://developer.download.nvidia.com/compute/cuda/redist/'
ARCHIVES=[
('cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-12.8.93-archive.tar.xz','9961b3484b6b71314063709a4f9529654f96782ad39e72bf1e00f070db8210d3'),
('cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.8.90-archive.tar.xz','8d566b5fe745c46842dc16945cf36686227536decd2302c372be86da37faca68'),
('cuda_cccl/linux-x86_64/cuda_cccl-linux-x86_64-12.8.90-archive.tar.xz','0740e9e01e4f15e17c5ab8d68bba4f8ec0eb6b84edccba4ac45112d2d2174e4b'),
]
def download(url,path,sha=None):
path.parent.mkdir(parents=True,exist_ok=True)
if not path.exists():
session=requests.Session();session.trust_env=False
with session.get(url,stream=True,timeout=(15,120)) as response:
response.raise_for_status()
temporary=path.with_suffix(path.suffix+'.partial')
with temporary.open('wb') as stream:
for chunk in response.iter_content(4*1024*1024):stream.write(chunk)
temporary.replace(path)
if sha and hashlib.sha256(path.read_bytes()).hexdigest()!=sha:
raise RuntimeError('CUDA toolkit archive checksum mismatch')
def extract(path,folder):
folder.mkdir(parents=True,exist_ok=True)
with tarfile.open(path) as archive:
for member in archive.getmembers():
if not (folder/member.name).resolve().is_relative_to(folder.resolve()):
raise RuntimeError('Unsafe archive member')
archive.extractall(folder)
def ensure_native():
for key in list(os.environ):
if 'proxy' in key.lower():os.environ.pop(key)
CUDA.mkdir(parents=True,exist_ok=True)
if not (CUDA/'bin/nvcc').exists():
for relative,sha in ARCHIVES:
path=CACHE/Path(relative).name
print('CUDA_TOOLCHAIN_COMPONENT',path.name,flush=True)
download(REDIST+relative,path,sha)
folder=CACHE/path.name.replace('.tar.xz','')
extract(path,folder)
source=next(folder.iterdir())
shutil.copytree(source,CUDA,dirs_exist_ok=True,symlinks=True)
os.environ['CUDA_HOME']=str(CUDA)
os.environ['PATH']=str(CUDA/'bin')+os.pathsep+os.environ.get('PATH','')
os.environ['TORCH_CUDA_ARCH_LIST']='12.0+PTX'
os.environ['FORCE_CUDA']='1'
os.environ['MAX_JOBS']='2'
includes=[str(CUDA/'include')];libraries=[str(CUDA/'lib'),str(CUDA/'lib64')]
for base in site.getsitepackages():
nvidia=Path(base)/'nvidia'
includes.extend(str(p) for p in nvidia.glob('*/include'))
libraries.extend(str(p) for p in nvidia.glob('*/lib'))
for key,values in [('CPATH',includes),('CPLUS_INCLUDE_PATH',includes),('LIBRARY_PATH',libraries),('LD_LIBRARY_PATH',libraries)]:
os.environ[key]=os.pathsep.join(values+[os.environ.get(key,'')])
wheels=CACHE/'wheels'
wheels.mkdir(parents=True,exist_ok=True)
if importlib.util.find_spec('torchsparse') is None:
revision='385f5ce8718fcae93540511b7f5832f4e71fd835'
archive=CACHE/'torchsparse.tar.gz'
download('https://codeload.github.com/mit-han-lab/torchsparse/tar.gz/'+revision,archive)
folder=CACHE/'torchsparse-source';extract(archive,folder)
source=next(folder.iterdir())
print('BUILDING_TORCHSPARSE_FOR_BLACKWELL',flush=True)
setup=source/'setup.py'
setup.write_text(setup.read_text().replace('["-g", "-O3"','["-g0", "-O3"'))
subprocess.run([sys.executable,'-m','pip','wheel','--no-deps','--no-build-isolation','-w',str(wheels),str(source)],check=True)
wheel=next(wheels.glob('torchsparse-*.whl'))
subprocess.run([sys.executable,'-m','pip','install','--no-deps',str(wheel)],check=True)
if importlib.util.find_spec('spac_sdf_cuda') is None:
print('BUILDING_SDF_KERNEL_FOR_BLACKWELL',flush=True)
setup=CACHE/'setup_sdf.py'
sources=[str(ROOT/'spac/_kernels/src/udf_kernel.cu'),str(ROOT/'spac/_kernels/src/udf_cuda.cpp')]
setup.write_text('from setuptools import setup\nfrom torch.utils.cpp_extension import BuildExtension,CUDAExtension\nsetup(name="spac-sdf-cuda",version="0.1.0",ext_modules=[CUDAExtension("spac_sdf_cuda",'+repr(sources)+',extra_compile_args={"cxx":["-g0"],"nvcc":["-O3"]})],cmdclass={"build_ext":BuildExtension})\n')
subprocess.run([sys.executable,str(setup),'bdist_wheel','--dist-dir',str(wheels)],cwd=ROOT,check=True)
wheel=next(wheels.glob('spac_sdf_cuda-*.whl'))
subprocess.run([sys.executable,'-m','pip','install','--no-deps',str(wheel)],check=True)
print('NATIVE_EXTENSIONS_READY',flush=True)
|