add conftest.py

Signed-off-by: Ylang Tsou <ylangt@google.com>

Author: Ylang Tsou

docker/Dockerfile.pypi

@@ -15,27 +15,19 @@ RUN apt-get update && apt-get install -y \
 WORKDIR /workspace/tpu_inference
 COPY . .
 RUN export TPU_INFERENCE_VERSION=$(pip index versions tpu-inference --pre 2>/dev/null | grep -oE "[0-9]+\.[0-9]+\.[0-9]+\.dev[0-9]+" | head -n 1) && \
-    echo -n "$TPU_INFERENCE_VERSION" > /tmp/tpu_inference_version && \
-    echo "Installing tpu-inference version: $TPU_INFERENCE_VERSION"
+    echo -n "${TPU_INFERENCE_VERSION}" > /tmp/tpu_inference_version
 
 # Clone vLLM
-WORKDIR /workspace/vllm_build
+WORKDIR /workspace/vllm
 ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
-RUN git clone $VLLM_REPO /workspace/vllm_build
+RUN git clone $VLLM_REPO /workspace/vllm
 RUN pip install build
 RUN export TPU_INFERENCE_VERSION=$(cat /tmp/tpu_inference_version) && \
-    echo "Using tpu-inference version: ${TPU_INFERENCE_VERSION}" && \
     sed -i "s/^tpu-inference==.*/tpu-inference==${TPU_INFERENCE_VERSION}/" requirements/tpu.txt && \
-    echo "Starting build.sh" && \
     bash tools/vllm-tpu/build.sh ${TPU_INFERENCE_VERSION}
 
 # Install vllm-tpu from wheel
-RUN echo "pip installing vllm-tpu" && \
-    pip install --no-cache-dir dist/*.whl
-# RUN export TPU_INFERENCE_VERSION=$(cat /tmp/tpu_inference_version) && \
-  #   export PIP_VERSION=$(echo "$TPU_INFERENCE_VERSION" | sed 's/^v//') && \
-    # echo "Installing vllm-tpu version: $PIP_VERSION" && \
-    # python3 -m pip install -v --no-cache-dir "vllm-tpu==$PIP_VERSION"
+RUN pip install --no-cache-dir dist/*.whl
 
 # Install test dependencies
 RUN python3 -m pip install -e tests/vllm_test_utils
@@ -46,8 +38,7 @@ RUN python3 -m pip install --no-cache-dir \
     pytest-cov \
     tblib
 
-# Preventing shadowing
-# WORKDIR /workspace
-# RUN mv /workspace/vllm /workspace/vllm_backup
+# Set environment variable to use site-packages
+ENV FORCE_USE_SITE_PACKAGES=1
 
 CMD ["/bin/bash"]

tests/conftest.py

@@ -0,0 +1,36 @@
+import os
+import site
+import sys
+
+
+def pytest_configure(config):
+    if os.environ.get("FORCE_USE_SITE_PACKAGES", "0") != "1":
+        print("[DEBUG] Running in Repo Mode. Using local source code.",
+              file=sys.stderr)
+        return
+
+    print("\n[DEBUG] conftest.py is fixing sys.path...", file=sys.stderr)
+
+    local_vllm_path = os.path.abspath("/workspace/vllm")
+
+    removed = False
+    for path in list(sys.path):
+        if local_vllm_path in os.path.abspath(path):
+            sys.path.remove(path)
+            removed = True
+            print(f"[DEBUG] Removed local source from path: {path}",
+                  file=sys.stderr)
+
+    site_packages = site.getsitepackages()
+    for sp in site_packages:
+        if sp not in sys.path:
+            sys.path.insert(0, sp)
+        else:
+            sys.path.remove(sp)
+            sys.path.insert(0, sp)
+
+    if removed:
+        print("Path fix applied successfully.", file=sys.stderr)
+    else:
+        print("Local vllm path not found in sys.path (maybe already clean?)",
+              file=sys.stderr)