determine num gpus in python

Dockerfile

@@ -52,5 +52,5 @@ RUN mkdir -p /usr/local/bin
 COPY --chmod=755 start.sh /usr/local/bin/start.sh
 
 # Start the handler
-ENTRYPOINT ["/bin/bash", "/usr/local/bin/start.sh"] 
+#ENTRYPOINT ["/bin/bash", "/usr/local/bin/start.sh"] 
 CMD ["python3", "/src/handler.py"]

src/handler.py

@@ -3,6 +3,20 @@ import runpod
 from utils import JobInput
 from engine import vLLMEngine, OpenAIvLLMEngine
 
+# Detect number of visible GPUs
+gpu_count = torch.cuda.device_count()
+
+# Fallback to 1 if none detected
+if gpu_count < 1:
+    gpu_count = 1
+
+# Set the environment variable
+os.environ["TENSOR_PARALLEL_SIZE"] = str(gpu_count)
+
+print(f"Detected {gpu_count} GPU(s)")
+print(f"Set TENSOR_PARALLEL_SIZE={os.environ['TENSOR_PARALLEL_SIZE']}")
+
+
 vllm_engine = vLLMEngine()
 OpenAIvLLMEngine = OpenAIvLLMEngine(vllm_engine)