Skip to content

Commit fa0f511

Browse files
anmyachevanmyachev
authored
Revert "[WIN] Fix HTTP Error 404: Not Found" (#3348)
Signed-off-by: Anatoly Myachev <[email protected]>
1 parent 97fc7d9 commit fa0f511

File tree

1 file changed

+87
-81
lines changed

1 file changed

+87
-81
lines changed

‎python/setup.py‎

Lines changed: 87 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -318,9 +318,11 @@ def download_and_copy(name, src_func, dst_path, variable, version, url_func):
318318
system = platform.system()
319319
arch = platform.machine()
320320
# NOTE: This might be wrong for jetson if both grace chips and jetson chips return aarch64
321-
arch = {"arm64": "sbsa", "aarch64": "sbsa"}.get(arch, arch)
322-
supported = {"Linux": "linux", "Darwin": "linux", "Windows": "win"}
321+
arch = {"arm64": "sbsa", "aarch64": "sbsa", "AMD64": "x86_64"}.get(arch, arch)
322+
supported = {"Linux": "linux", "Darwin": "linux", "Windows": "windows"}
323323
url = url_func(supported[system], arch, version)
324+
if system == "Windows":
325+
url = url.replace("tar.xz", "zip")
324326
src_path = src_func(supported[system], arch, version)
325327
tmp_path = os.path.join(triton_cache_path, "nvidia", name) # path to cache the download
326328
dst_path = os.path.join(base_dir, os.pardir, "third_party", "nvidia", "backend", dst_path) # final binary path
@@ -333,8 +335,14 @@ def download_and_copy(name, src_func, dst_path, variable, version, url_func):
333335
download = download or curr_version.group(1) != version
334336
if download:
335337
print(f'downloading and extracting {url} ...')
336-
file = tarfile.open(fileobj=open_url(url), mode="r|*")
337-
file.extractall(path=tmp_path)
338+
if url.endswith(".zip"):
339+
with open_url(url) as response:
340+
file_bytes = BytesIO(response.read())
341+
with zipfile.ZipFile(file_bytes, "r") as file:
342+
file.extractall(path=tmp_path)
343+
else:
344+
file = tarfile.open(fileobj=open_url(url), mode="r|*")
345+
file.extractall(path=tmp_path)
338346
os.makedirs(os.path.split(dst_path)[0], exist_ok=True)
339347
print(f'copy {src_path} to {dst_path} ...')
340348
if os.path.isdir(src_path):
@@ -511,83 +519,81 @@ def build_extension(self, ext):
511519
NVIDIA_TOOLCHAIN_VERSION = json.load(nvidia_version_file)
512520

513521
exe_extension = sysconfig.get_config_var("EXE")
514-
# FIXME: urllib.error.HTTPError: HTTP Error 404: Not Found
515-
if platform.system() != "Windows":
516-
download_and_copy(
517-
name="nvcc",
518-
src_func=lambda system, arch, version: f"cuda_nvcc-{system}-{arch}-{version}-archive/bin/ptxas{exe_extension}",
519-
dst_path="bin/ptxas",
520-
variable="TRITON_PTXAS_PATH",
521-
version=NVIDIA_TOOLCHAIN_VERSION["ptxas"],
522-
url_func=lambda system, arch, version:
523-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/{system}-{arch}/cuda_nvcc-{system}-{arch}-{version}-archive.tar.xz",
524-
)
525-
# We download a separate ptxas for blackwell, since there are some bugs when using it for hopper
526-
download_and_copy(
527-
name="nvcc",
528-
src_func=lambda system, arch, version: f"cuda_nvcc-{system}-{arch}-{version}-archive/bin/ptxas{exe_extension}",
529-
dst_path="bin/ptxas-blackwell",
530-
variable="TRITON_PTXAS_PATH",
531-
version=NVIDIA_TOOLCHAIN_VERSION["ptxas-blackwell"],
532-
url_func=lambda system, arch, version:
533-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/{system}-{arch}/cuda_nvcc-{system}-{arch}-{version}-archive.tar.xz",
534-
)
535-
download_and_copy(
536-
name="cuobjdump",
537-
src_func=lambda system, arch, version:
538-
f"cuda_cuobjdump-{system}-{arch}-{version}-archive/bin/cuobjdump{exe_extension}",
539-
dst_path="bin/cuobjdump",
540-
variable="TRITON_CUOBJDUMP_PATH",
541-
version=NVIDIA_TOOLCHAIN_VERSION["cuobjdump"],
542-
url_func=lambda system, arch, version:
543-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cuobjdump/{system}-{arch}/cuda_cuobjdump-{system}-{arch}-{version}-archive.tar.xz",
544-
)
545-
download_and_copy(
546-
name="nvdisasm",
547-
src_func=lambda system, arch, version:
548-
f"cuda_nvdisasm-{system}-{arch}-{version}-archive/bin/nvdisasm{exe_extension}",
549-
dst_path="bin/nvdisasm",
550-
variable="TRITON_NVDISASM_PATH",
551-
version=NVIDIA_TOOLCHAIN_VERSION["nvdisasm"],
552-
url_func=lambda system, arch, version:
553-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvdisasm/{system}-{arch}/cuda_nvdisasm-{system}-{arch}-{version}-archive.tar.xz",
554-
)
555-
download_and_copy(
556-
name="nvcc",
557-
src_func=lambda system, arch, version: f"cuda_nvcc-{system}-{arch}-{version}-archive/include",
558-
dst_path="include",
559-
variable="TRITON_CUDACRT_PATH",
560-
version=NVIDIA_TOOLCHAIN_VERSION["cudacrt"],
561-
url_func=lambda system, arch, version:
562-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/{system}-{arch}/cuda_nvcc-{system}-{arch}-{version}-archive.tar.xz",
563-
)
564-
download_and_copy(
565-
name="cudart",
566-
src_func=lambda system, arch, version: f"cuda_cudart-{system}-{arch}-{version}-archive/include",
567-
dst_path="include",
568-
variable="TRITON_CUDART_PATH",
569-
version=NVIDIA_TOOLCHAIN_VERSION["cudart"],
570-
url_func=lambda system, arch, version:
571-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/{system}-{arch}/cuda_cudart-{system}-{arch}-{version}-archive.tar.xz",
572-
)
573-
download_and_copy(
574-
name="cupti",
575-
src_func=lambda system, arch, version: f"cuda_cupti-{system}-{arch}-{version}-archive/include",
576-
dst_path="include",
577-
variable="TRITON_CUPTI_INCLUDE_PATH",
578-
version=NVIDIA_TOOLCHAIN_VERSION["cupti"],
579-
url_func=lambda system, arch, version:
580-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/{system}-{arch}/cuda_cupti-{system}-{arch}-{version}-archive.tar.xz",
581-
)
582-
download_and_copy(
583-
name="cupti",
584-
src_func=lambda system, arch, version: f"cuda_cupti-{system}-{arch}-{version}-archive/lib",
585-
dst_path="lib/cupti",
586-
variable="TRITON_CUPTI_LIB_PATH",
587-
version=NVIDIA_TOOLCHAIN_VERSION["cupti"],
588-
url_func=lambda system, arch, version:
589-
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/{system}-{arch}/cuda_cupti-{system}-{arch}-{version}-archive.tar.xz",
590-
)
522+
download_and_copy(
523+
name="nvcc",
524+
src_func=lambda system, arch, version: f"cuda_nvcc-{system}-{arch}-{version}-archive/bin/ptxas{exe_extension}",
525+
dst_path="bin/ptxas",
526+
variable="TRITON_PTXAS_PATH",
527+
version=NVIDIA_TOOLCHAIN_VERSION["ptxas"],
528+
url_func=lambda system, arch, version:
529+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/{system}-{arch}/cuda_nvcc-{system}-{arch}-{version}-archive.tar.xz",
530+
)
531+
# We download a separate ptxas for blackwell, since there are some bugs when using it for hopper
532+
download_and_copy(
533+
name="nvcc",
534+
src_func=lambda system, arch, version: f"cuda_nvcc-{system}-{arch}-{version}-archive/bin/ptxas{exe_extension}",
535+
dst_path="bin/ptxas-blackwell",
536+
variable="TRITON_PTXAS_PATH",
537+
version=NVIDIA_TOOLCHAIN_VERSION["ptxas-blackwell"],
538+
url_func=lambda system, arch, version:
539+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/{system}-{arch}/cuda_nvcc-{system}-{arch}-{version}-archive.tar.xz",
540+
)
541+
download_and_copy(
542+
name="cuobjdump",
543+
src_func=lambda system, arch, version:
544+
f"cuda_cuobjdump-{system}-{arch}-{version}-archive/bin/cuobjdump{exe_extension}",
545+
dst_path="bin/cuobjdump",
546+
variable="TRITON_CUOBJDUMP_PATH",
547+
version=NVIDIA_TOOLCHAIN_VERSION["cuobjdump"],
548+
url_func=lambda system, arch, version:
549+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cuobjdump/{system}-{arch}/cuda_cuobjdump-{system}-{arch}-{version}-archive.tar.xz",
550+
)
551+
download_and_copy(
552+
name="nvdisasm",
553+
src_func=lambda system, arch, version:
554+
f"cuda_nvdisasm-{system}-{arch}-{version}-archive/bin/nvdisasm{exe_extension}",
555+
dst_path="bin/nvdisasm",
556+
variable="TRITON_NVDISASM_PATH",
557+
version=NVIDIA_TOOLCHAIN_VERSION["nvdisasm"],
558+
url_func=lambda system, arch, version:
559+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvdisasm/{system}-{arch}/cuda_nvdisasm-{system}-{arch}-{version}-archive.tar.xz",
560+
)
561+
download_and_copy(
562+
name="nvcc",
563+
src_func=lambda system, arch, version: f"cuda_nvcc-{system}-{arch}-{version}-archive/include",
564+
dst_path="include",
565+
variable="TRITON_CUDACRT_PATH",
566+
version=NVIDIA_TOOLCHAIN_VERSION["cudacrt"],
567+
url_func=lambda system, arch, version:
568+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/{system}-{arch}/cuda_nvcc-{system}-{arch}-{version}-archive.tar.xz",
569+
)
570+
download_and_copy(
571+
name="cudart",
572+
src_func=lambda system, arch, version: f"cuda_cudart-{system}-{arch}-{version}-archive/include",
573+
dst_path="include",
574+
variable="TRITON_CUDART_PATH",
575+
version=NVIDIA_TOOLCHAIN_VERSION["cudart"],
576+
url_func=lambda system, arch, version:
577+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/{system}-{arch}/cuda_cudart-{system}-{arch}-{version}-archive.tar.xz",
578+
)
579+
download_and_copy(
580+
name="cupti",
581+
src_func=lambda system, arch, version: f"cuda_cupti-{system}-{arch}-{version}-archive/include",
582+
dst_path="include",
583+
variable="TRITON_CUPTI_INCLUDE_PATH",
584+
version=NVIDIA_TOOLCHAIN_VERSION["cupti"],
585+
url_func=lambda system, arch, version:
586+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/{system}-{arch}/cuda_cupti-{system}-{arch}-{version}-archive.tar.xz",
587+
)
588+
download_and_copy(
589+
name="cupti",
590+
src_func=lambda system, arch, version: f"cuda_cupti-{system}-{arch}-{version}-archive/lib",
591+
dst_path="lib/cupti",
592+
variable="TRITON_CUPTI_LIB_PATH",
593+
version=NVIDIA_TOOLCHAIN_VERSION["cupti"],
594+
url_func=lambda system, arch, version:
595+
f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/{system}-{arch}/cuda_cupti-{system}-{arch}-{version}-archive.tar.xz",
596+
)
591597
backends = [*BackendInstaller.copy(["intel", "nvidia", "amd"]), *BackendInstaller.copy_externals()]
592598

593599

0 commit comments

Comments
 (0)