handler changes

Dockerfile

@@ -1,5 +1,5 @@
 # Base image
-FROM runpod/base:0.4.2-cuda11.8.0
+FROM runpod/base:0.4.2-cuda12.1.0
 
 ARG HUGGING_FACE_HUB_TOKEN
 
@@ -13,8 +13,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 # Add src files (Worker Template)
 ADD src .
 
-
-
 # Prepare argument for the model and tokenizer
 ARG MODEL_NAME=""
 ENV MODEL_NAME=$MODEL_NAME
@@ -28,6 +26,8 @@ ARG STREAMING=
 ENV STREAMING=$STREAMING
 ARG QUANTIZATION=
 ENV QUANTIZATION=$QUANTIZATION
+ARG MAX_CONCURRENCY=
+ENV MAX_CONCURRENCY=$MAX_CONCURRENCY
 
 ENV HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets"
 ENV HUGGINGFACE_HUB_CACHE="/runpod-volume/huggingface-cache/hub"
@@ -46,7 +46,7 @@ ENV MODEL_NAME=$MODEL_NAME \
 ENTRYPOINT ["/entrypoint.sh"]
 
 # Run the Python script to download the model
-RUN python3.11 -u /download_model.py --model_name $MODEL_NAME --model_revision $MODEL_REVISION --model_base_path $MODEL_BASE_PATH --hugging_face_hub_token $HUGGING_FACE_HUB_TOKEN
+RUN python3.11 -u /download_model.py --model_name $MODEL_NAME --model_revision $MODEL_REVISION --model_base_path $MODEL_BASE_PATH
 
 # Start the handler
-CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER QUANTIZATION=$QUANTIZATION python3.11 /handler.py
+CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER QUANTIZATION=$QUANTIZATION  python3.11 /handler.py

README.md

@@ -18,11 +18,14 @@ More information on receiving streaming responses from Serverless Endpoints can
 #### Optional:
 - `HUGGING_FACE_HUB_TOKEN`: Your Hugging Face token to access private or gated models. You can get your token [here](https://huggingface.co/settings/token).
 - `TOKENIZER`: The specified tokenizer to use. If you want to use the default tokenizer for the model, do not provide this docker argument at all.
-- `QUANTIZATION`: `awq` to use AWQ Quantization. Base model must be in AWQ format.
+- `QUANTIZATION`: `awq` to use AWQ Quantization (Base model must be in AWQ format). `squeezellm` for SqueezeLLM quantization - preliminary support.
 
 ### Compatible Models
 - LLaMA & LLaMA-2 
 - Mistral 
+- Yi
+- ChatGLM
+- Phi
 - MPT 
 - OPT 
 - Qwen 

src/download_model.py

@@ -6,6 +6,7 @@ from huggingface_hub import snapshot_download
 MODEL_NAME = os.environ.get('MODEL_NAME')
 MODEL_REVISION = os.environ.get('MODEL_REVISION', "main")
 MODEL_BASE_PATH = os.environ.get('MODEL_BASE_PATH', '/runpod-volume/')
+HUGGING_FACE_HUB_TOKEN = os.environ.get('HUGGING_FACE_HUB_TOKEN')
 
 
 if __name__ == "__main__":
@@ -13,7 +14,6 @@ if __name__ == "__main__":
     parser.add_argument('--model_name', type=str, default=MODEL_NAME)
     parser.add_argument('--model_revision', type=str, default=MODEL_REVISION)
     parser.add_argument('--model_base_path', type=str, default=MODEL_BASE_PATH)
-    parser.add_argument('--hugging_face_hub_token', type=str)
 
     args = parser.parse_args()
 
@@ -21,5 +21,4 @@ if __name__ == "__main__":
         args.model_name,
         revision=args.model_revision,
         local_dir=f"{args.model_base_path}{args.model_name.split('/')[1]}",
-        token=args.hugging_face_hub_token
     )

src/handler.py

@@ -18,12 +18,13 @@ STREAMING = os.environ.get('STREAMING', False) == 'True'
 TOKENIZER = os.environ.get('TOKENIZER', None)
 USE_FULL_METRICS = os.environ.get('USE_FULL_METRICS', True)
 DTYPE = "auto"
-USE_HF_CHAT_TEMPLATE = os.environ.get('USE_HF_CHAT_TEMPLATE', False) == 'True'
+MAX_CONCURRENCY = os.environ.get('MAX_CONCURRENCY', 200)
+TOTAL_RUNNING_JOBS = 0
 
 # Set up quantization-related parameters
 QUANTIZATION = os.environ.get('QUANTIZATION', None)
 
-if type(QUANTIZATION) is str and QUANTIZATION.lower() != "awq":
+if type(QUANTIZATION) is str and QUANTIZATION.lower() in ["awq", "squeezellm"]:
     QUANTIZATION = None
     print("Invalid quantization parameter. Using default value of None.")
 else:
@@ -32,7 +33,7 @@ else:
 if not MODEL_NAME:
     print("Error: The model has not been provided.")
 
-if len(TOKENIZER) == 0:
+if not TOKENIZER or len(TOKENIZER) == 0:
     print("Error: The tokenizer has not been provided. Defaulting to MODEL_NAME.")
 
 # Tensor parallelism
@@ -62,10 +63,9 @@ llm.engine._log_system_stats = lambda x, y: vllm_log_system_stats(
     llm.engine, x, y)
 
 
-def concurrency_controller() -> bool:
-    # Calculate pending sequences
-    total_queued_sequences = len(llm.engine.scheduler.waiting)
-    return total_queued_sequences > 0
+def concurrency_controller() -> int:
+    global TOTAL_RUNNING_JOBS
+    return MAX_CONCURRENCY - TOTAL_RUNNING_JOBS
 
 
 def prepare_metrics() -> dict:
@@ -144,7 +144,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
     This is the handler function that will be called by the serverless worker.
     '''
     print("Job received by handler: {}".format(job))
-
+    global TOTAL_RUNNING_JOBS
     # Retrieve the job input.
     job_input = job['input']
 
@@ -163,6 +163,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
 
     # Send request to VLLM
     request_id = random_uuid()
+    TOTAL_RUNNING_JOBS += 1
     results_generator = llm.generate(prompt, sampling_params, request_id)
 
     # Keep track of the stream's information to perform the appropriate chunking.
@@ -289,13 +290,15 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
         # Yield the output
         yield ret
 
+        TOTAL_RUNNING_JOBS -= 1
+
 
 async def handler(job: dict) -> dict[str, list]:
     '''
     This is the handler function that will be called by the serverless worker.
     '''
     print("Job received by handler: {}".format(job))
-
+    global TOTAL_RUNNING_JOBS
     # Retrieve the job input.
     job_input = job['input']
     # Create the prompt using the template.
@@ -310,6 +313,7 @@ async def handler(job: dict) -> dict[str, list]:
 
     # Send request to VLLM
     request_id = random_uuid()
+    TOTAL_RUNNING_JOBS += 1
     results_generator = llm.generate(prompt, sampling_params, request_id)
 
     # Get the final generated output
@@ -347,6 +351,9 @@ async def handler(job: dict) -> dict[str, list]:
         "input_tokens": runpod_metrics['input_tokens'],
         "output_tokens": runpod_metrics['output_tokens']
     }
+
+    TOTAL_RUNNING_JOBS -= 1
+
     return ret
 
 

src/templates.py

@@ -1,17 +0,0 @@
-class Template():
-    def __init__(self, template_method):
-        self.template_method = template_method
-
-    def __call__(self, prompt):
-        return self.template_method(prompt)
-
-
-LLAMA2_TEMPLATE = Template(
-    lambda prompt: """SYSTEM: You are a helpful assistant.
-USER: {}
-ASSISTANT: """.format(prompt)
-)
-
-DEFAULT_TEMPLATE = Template(
-    lambda prompt: prompt
-)