fix broken cuda and rocm images (#263)

Co-authored-by: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com>

docker/cpu/Dockerfile

@@ -34,7 +34,7 @@ ARG TORCH_VERSION=stable
 
 RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
     pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu ; \
-elif [ "${TORCH_VERSION}" = "nighly" ]; then \
+elif [ "${TORCH_VERSION}" = "nightly" ]; then \
     pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu ; \
 else \
     pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu ; \

docker/cuda-ort/Dockerfile

@@ -33,8 +33,8 @@ ARG TORCH_VERSION=stable
 
 RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
     pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
-elif [ "${TORCH_VERSION}" = "nighly" ]; then \
-    pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+elif [ "${TORCH_VERSION}" = "nightly" ]; then \
+    pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/${TORCH_CUDA} ; \
 else \
     pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
 fi

docker/cuda/Dockerfile

@@ -28,14 +28,18 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 # Install PyTorch
 ARG TORCH_CUDA=cu124
-ARG TORCH_VERSION=stable
+ARG TORCH_RELEASE_TYPE=stable
+ARG TORCH_VERSION=""
 
-RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
+RUN if [ -n "${TORCH_VERSION}" ]; then \
+    # Install specific torch version if TORCH_VERSION is provided
+    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "stable" ]; then \
     pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
-elif [ "${TORCH_VERSION}" = "nighly" ]; then \
-    pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "nightly" ]; then \
+    pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/${TORCH_CUDA} ; \
 else \
-    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+    echo "Error: Invalid TORCH_RELEASE_TYPE. Must be 'stable', 'nightly', or specify a TORCH_VERSION." && exit 1 ; \
 fi
 
 # Install quantization libraries from source

docker/rocm/Dockerfile

@@ -31,16 +31,19 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-reco
 
 # Install PyTorch
 ARG TORCH_ROCM=rocm5.7
-ARG TORCH_VERSION=stable
+ARG TORCH_RELEASE_TYPE=stable
+ARG TORCH_VERSION=""
 
-RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
+RUN if [ -n "${TORCH_VERSION}" ]; then \
+    # Install specific torch version if TORCH_VERSION is provided
+    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_ROCM} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "stable" ]; then \
     pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_ROCM} ; \
-elif [ "${TORCH_VERSION}" = "nightly" ]; then \
+elif [ "${TORCH_RELEASE_TYPE}" = "nightly" ]; then \
     pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/${TORCH_ROCM} ; \
 else \
-    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_ROCM} ; \
+    echo "Error: Invalid TORCH_RELEASE_TYPE. Must be 'stable', 'nightly', or specify a TORCH_VERSION." && exit 1 ; \
 fi
-
 # Install quantization libraries from source
 ENV ROCM_VERSION=5.7
 ENV PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx1030;gfx1100"

scripts/install_quantization_libs.py

@@ -7,23 +7,6 @@
 EXTERNAL_REPOS_DIR = "external_repos"
 
 
-def process_setup_file(setup_file_path):
-    with open(setup_file_path, "r") as file:
-        setup_content = file.read()
-
-    # Use a regular expression to remove any line containing "torch=="
-    setup_content = re.sub(r'"torch==[^\"]+",', "", setup_content)
-
-    # Set IS_CPU_ONLY to False
-    setup_content = setup_content.replace(
-        "IS_CPU_ONLY = not torch.backends.mps.is_available() and not torch.cuda.is_available()", "IS_CPU_ONLY = False"
-    )
-
-    # Write the modified content back to setup.py
-    with open(setup_file_path, "w") as file:
-        file.write(setup_content)
-
-
 def clone_or_pull_repo(repo_url, repo_location_path):
     """Clone the repo if it doesn't exist; otherwise, pull the latest changes."""
     if os.path.exists(repo_location_path):
@@ -35,29 +18,42 @@ def clone_or_pull_repo(repo_url, repo_location_path):
         subprocess.run(f"git clone {repo_url} {repo_location_path}", shell=True, check=True)
 
 
+def process_setup_file_for_autoawq(setup_file_path):
+    print("Processing setup.py for AutoAWQ")
+
+    with open(setup_file_path, "r") as file:
+        setup_content = file.read()
+
+    # Use regex to match any line that starts with IS_CPU_ONLY = and modify it to IS_CPU_ONLY = False
+    setup_content = re.sub(r"(IS_CPU_ONLY\s*=\s*.*)", r"\1\nIS_CPU_ONLY = False", setup_content)
+
+    # Write the modified content back to setup.py
+    with open(setup_file_path, "w") as file:
+        file.write(setup_content)
+
+
 def install_autoawq_from_source():
     """Install the AutoAWQ and AutoAWQ_kernels packages from GitHub."""
     print("Installing AutoAWQ and AutoAWQ_kernels packages.")
-
     autoawq_repo_name = "AutoAWQ"
     autoawq_kernels_repo_name = "AutoAWQ_kernels"
 
     autoawq_repo_path = os.path.join(EXTERNAL_REPOS_DIR, autoawq_repo_name)
     kernels_repo_path = os.path.join(EXTERNAL_REPOS_DIR, autoawq_kernels_repo_name)
 
+    print("Installing AutoAWQ_kernels package.")
     clone_or_pull_repo(f"https://github.com/casper-hansen/{autoawq_kernels_repo_name}", kernels_repo_path)
-    kernels_setup_file_path = os.path.join(kernels_repo_path, "setup.py")
-    process_setup_file(kernels_setup_file_path)
     subprocess.run(
         f"cd {kernels_repo_path} && {sys.executable} -m pip install .",
         shell=True,
         check=True,
         env=os.environ,
     )
 
+    print("Installing AutoAWQ package.")
     clone_or_pull_repo(f"https://github.com/casper-hansen/{autoawq_repo_name}", autoawq_repo_path)
     autoawq_setup_file_path = os.path.join(autoawq_repo_path, "setup.py")
-    process_setup_file(autoawq_setup_file_path)
+    process_setup_file_for_autoawq(autoawq_setup_file_path)
     subprocess.run(
         f"cd {autoawq_repo_path} && {sys.executable} -m pip install .",
         shell=True,
@@ -75,8 +71,6 @@ def install_autogptq_from_source():
 
     clone_or_pull_repo("https://github.com/PanQiWei/AutoGPTQ.git", autogptq_repo_path)
     subprocess.run("pip install numpy gekko pandas", shell=True, check=True, env=os.environ)
-    autogptq_setup_file_path = os.path.join(autogptq_repo_path, "setup.py")
-    process_setup_file(autogptq_setup_file_path)
     subprocess.run(
         f"cd {autogptq_repo_path} && {sys.executable} -m pip install .",
         shell=True,