styling

.github/workflows/quality.yaml

@@ -2,6 +2,11 @@ name: Quality Checks
 
 on:
   push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
 
 concurrency:
   cancel-in-progress: true

.github/workflows/security.yml

@@ -2,6 +2,11 @@ name: Security Checks
 
 on:
   push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
 
 concurrency:
   cancel-in-progress: true

docker/cpu/Dockerfile

@@ -20,7 +20,6 @@ FROM ubuntu:${UBUNTU_VERSION}
 ENV DEBIAN_FRONTEND noninteractive
 ENV PATH="/home/user/.local/bin:${PATH}"
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    libgl1-mesa-dev libglib2.0-0 \
     sudo build-essential git bash-completion numactl \
     python3.10 python3-pip python3.10-dev google-perftools && \
     apt-get clean && rm -rf /var/lib/apt/lists/* && \
@@ -30,12 +29,15 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 ENV LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4:/usr/local/lib/libiomp5.so"
 
 # Install PyTorch
-ARG TORCH_VERSION=stable
+ARG TORCH_VERSION=""
+ARG TORCH_RELEASE_TYPE=stable
 
-RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
+RUN if [ -n "${TORCH_VERSION}" ]; then \
+    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "stable" ]; then \
     pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu ; \
-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/cpu ; \
 else \
-    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu ; \
+    echo "Error: Invalid TORCH_RELEASE_TYPE. Must be 'stable', 'nightly', or specify a TORCH_VERSION." && exit 1 ; \
 fi

docker/cuda-ort/Dockerfile

@@ -39,5 +39,7 @@ else \
     pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
 fi
 
+# Install torch-ort and onnxruntime-training
 ENV TORCH_CUDA_ARCH_LIST="5.0 6.0 7.0 7.5 8.0 8.6 9.0+PTX"
+
 RUN pip install --no-cache-dir torch-ort onnxruntime-training && python -m torch_ort.configure

docker/cuda/Dockerfile

@@ -27,12 +27,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     pip install --no-cache-dir --upgrade pip setuptools wheel requests
 
 # Install PyTorch
+ARG TORCH_VERSION=""
 ARG TORCH_CUDA=cu124
 ARG TORCH_RELEASE_TYPE=stable
-ARG TORCH_VERSION=""
 
 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} ; \

docker/rocm/Dockerfile

@@ -30,12 +30,11 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-reco
     cd /opt/rocm/share/amd_smi && pip install .
 
 # Install PyTorch
+ARG TORCH_VERSION=""
 ARG TORCH_ROCM=rocm5.7
 ARG TORCH_RELEASE_TYPE=stable
-ARG TORCH_VERSION=""
 
 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} ; \
@@ -44,6 +43,7 @@ elif [ "${TORCH_RELEASE_TYPE}" = "nightly" ]; then \
 else \
     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"

llm_perf/update_llm_perf_cpu_pytorch.py

@@ -28,7 +28,7 @@
 
 if os.getenv("MACHINE", None) is None and os.getenv("SUBSET", None) is None:
     PUSH_REPO_ID = f"optimum-benchmark/llm-perf-{BACKEND}-{HARDWARE}-debug"
-    CANONICAL_PRETRAINED_OPEN_LLM_LIST = ["gpt2"]
+    CANONICAL_PRETRAINED_OPEN_LLM_LIST = ["gpt2"]  # noqa: F811
     SUBSET = "unquantized"
 elif os.getenv("MACHINE", None) is not None and os.getenv("SUBSET", None) is not None:
     PUSH_REPO_ID = f"optimum-benchmark/llm-perf-{BACKEND}-{HARDWARE}-{SUBSET}-{MACHINE}"

llm_perf/update_llm_perf_cuda_pytorch.py

@@ -19,7 +19,7 @@
 
 if os.getenv("MACHINE", None) is None and os.getenv("SUBSET", None) is None:
     PUSH_REPO_ID = "optimum-benchmark/llm-perf-pytorch-cuda-debug"
-    CANONICAL_PRETRAINED_OPEN_LLM_LIST = ["gpt2"]
+    CANONICAL_PRETRAINED_OPEN_LLM_LIST = ["gpt2"]  # noqa: F811
     SUBSET = "unquantized"
 elif os.getenv("MACHINE", None) is not None and os.getenv("SUBSET", None) is not None:
     PUSH_REPO_ID = f"optimum-benchmark/llm-perf-pytorch-cuda-{SUBSET}-{MACHINE}"

optimum_benchmark/trackers/latency.py

@@ -121,8 +121,9 @@ def __init__(self, device: str, backend: str):
         self.device = device
         self.backend = backend
         self.is_asynchronous = self.backend == "pytorch" and self.device == "cuda"
-        self.is_distributed = (self.backend != "vllm" and
-                               is_torch_distributed_available() and torch.distributed.is_initialized())
+        self.is_distributed = (
+            self.backend != "vllm" and is_torch_distributed_available() and torch.distributed.is_initialized()
+        )
 
         if self.is_asynchronous:
             LOGGER.info("\t+ Tracking latency using Pytorch CUDA events")

tests/test_cli.py

@@ -7,7 +7,7 @@
 
 from optimum_benchmark.logging_utils import run_subprocess_and_log_stream_output
 
-LOGGER = getLogger("test")
+LOGGER = getLogger("test-cli")
 
 
 FORCE_SERIAL = os.environ.get("FORCE_SERIAL", "0") == "1"