Fix vllm batch docker image (#463)

* Fix vllm batch docker image

* try again with 0.2.5

Author: yunfeng-scale

model-engine/model_engine_server/inference/batch_inference/Dockerfile_vllm

@@ -6,10 +6,15 @@ RUN apt-get update && \
     rm -rf /var/lib/apt/lists/* && \
     apt-get clean
 
-RUN pip uninstall torch -y
 COPY model-engine/model_engine_server/inference/batch_inference/requirements.txt /workspace/requirements.txt
 RUN pip install -r requirements.txt
 
+RUN pip uninstall torch -y
+RUN pip install torch==2.1.1 --index-url https://download.pytorch.org/whl/cu121
+
+RUN pip uninstall xformers -y
+RUN pip install xformers==0.0.23 --index-url https://download.pytorch.org/whl/cu121
+
 RUN wget https://github.com/peak/s5cmd/releases/download/v2.2.1/s5cmd_2.2.1_Linux-64bit.tar.gz
 RUN tar -xvzf s5cmd_2.2.1_Linux-64bit.tar.gz
 

model-engine/model_engine_server/inference/batch_inference/requirements.txt

@@ -1,5 +1,3 @@
-ray==2.6.3
-#git+https://github.com/vllm-project/vllm.git@4b61c6b669e368c6850531815940d9a542b9f223#egg=vllm
 vllm==0.2.5
 pydantic==1.10.13
 boto3==1.34.15

model-engine/model_engine_server/inference/batch_inference/sample_config.json

@@ -8,4 +8,4 @@
         "labels": {"team": "my_team"}
     },
     "data_parallelism":2
-}
+}

model-engine/model_engine_server/inference/batch_inference/vllm_batch.py

@@ -17,6 +17,7 @@
 
 CONFIG_FILE = os.getenv("CONFIG_FILE")
 AWS_REGION = os.getenv("AWS_REGION", "us-west-2")
+MODEL_WEIGHTS_FOLDER = os.getenv("MODEL_WEIGHTS_FOLDER", "./model_weights")
 
 os.environ["AWS_PROFILE"] = os.getenv("S3_WRITE_AWS_PROFILE", "default")
 
@@ -118,15 +119,15 @@ async def batch_inference():
     request = CreateBatchCompletionsRequest.parse_file(CONFIG_FILE)
 
     if request.model_config.checkpoint_path is not None:
-        download_model(request.model_config.checkpoint_path, "./model_weights")
+        download_model(request.model_config.checkpoint_path, MODEL_WEIGHTS_FOLDER)
 
     content = request.content
     if content is None:
         with smart_open.open(request.input_data_path, "r") as f:
             content = CreateBatchCompletionsRequestContent.parse_raw(f.read())
 
     model = (
-        "./model_weights" if request.model_config.checkpoint_path else request.model_config.model
+        MODEL_WEIGHTS_FOLDER if request.model_config.checkpoint_path else request.model_config.model
     )
 
     results_generators = await generate_with_vllm(request, content, model, job_index)