Spaces:
Sleeping
Sleeping
| """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) | |