Some batch inference improvements (#460)

* Some batch inference improvements

* fix unit test

* coverage

* integration test

* fix

Author: yunfeng-scale

docs/guides/completions.md

@@ -120,6 +120,55 @@ async def main():
 asyncio.run(main())
 ```
 
+## Batch completions
+
+The Python client also supports batch completions. Batch completions supports distributing data to multiple workers to accelerate inference. It also tries to maximize throughput so the completions should finish quite a bit faster than hitting models through HTTP. Use [Completion.batch_complete](../../api/python_client/#llmengine.completion.Completion.batch_complete) to utilize batch completions.
+
+Some examples of batch completions:
+
+=== "Batch completions with prompts in the request"
+```python
+from llmengine import Completion
+from llmengine.data_types import CreateBatchCompletionsModelConfig, CreateBatchCompletionsRequestContent
+
+content = CreateBatchCompletionsRequestContent(
+    prompts=["What is deep learning", "What is a neural network"],
+    max_new_tokens=10,
+    temperature=0.0
+)
+
+response = Completion.batch_create(
+    output_data_path="s3://my-path",
+    model_config=CreateBatchCompletionsModelConfig(
+        model="llama-2-7b",
+        checkpoint_path="s3://checkpoint-path",
+        labels={"team":"my-team", "product":"my-product"}
+    ),
+    content=content
+)
+print(response.job_id)
+```
+
+=== "Batch completions with prompts in a file and with 2 parallel jobs"
+```python
+from llmengine import Completion
+from llmengine.data_types import CreateBatchCompletionsModelConfig, CreateBatchCompletionsRequestContent
+
+# Store CreateBatchCompletionsRequestContent data into input file "s3://my-input-path"
+
+response = Completion.batch_create(
+    input_data_path="s3://my-input-path",
+    output_data_path="s3://my-output-path",
+    model_config=CreateBatchCompletionsModelConfig(
+        model="llama-2-7b",
+        checkpoint_path="s3://checkpoint-path",
+        labels={"team":"my-team", "product":"my-product"}
+    ),
+    data_parallelism=2
+)
+print(response.job_id)
+```
+
 ## Which model should I use?
 
 See the [Model Zoo](../../model_zoo) for more information on best practices for which model to use for Completions.

model-engine/model_engine_server/inference/batch_inference/Dockerfile_vllm

@@ -1,7 +1,7 @@
 FROM nvcr.io/nvidia/pytorch:23.09-py3
 
 RUN apt-get update && \
-    apt-get install -y dumb-init && \
+    apt-get install -y dumb-init psmisc && \
     apt-get autoremove -y && \
     rm -rf /var/lib/apt/lists/* && \
     apt-get clean

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

@@ -33,6 +33,7 @@ def download_model(checkpoint_path, final_weights_folder):
     # Need to override these env vars so s5cmd uses AWS_PROFILE
     env["AWS_ROLE_ARN"] = ""
     env["AWS_WEB_IDENTITY_TOKEN_FILE"] = ""
+    # nosemgrep
     process = subprocess.Popen(
         s5cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
     )
@@ -193,6 +194,7 @@ async def generate_with_vllm(request, content, model, job_index):
         tensor_parallel_size=request.model_config.num_shards,
         seed=request.model_config.seed or 0,
         disable_log_requests=True,
+        gpu_memory_utilization=0.8,  # To avoid OOM errors when there's host machine GPU usage
     )
 
     llm = AsyncLLMEngine.from_engine_args(engine_args)
@@ -220,5 +222,33 @@ async def generate_with_vllm(request, content, model, job_index):
     return results_generators
 
 
+def get_gpu_free_memory():  # pragma: no cover
+    """Get GPU free memory using nvidia-smi."""
+    try:
+        output = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"]
+        ).decode("utf-8")
+        gpu_memory = [int(x) for x in output.strip().split("\n")]
+        return gpu_memory
+    except subprocess.CalledProcessError:
+        return None
+
+
+def check_unknown_startup_memory_usage():  # pragma: no cover
+    """Check for unknown memory usage at startup."""
+    gpu_free_memory = get_gpu_free_memory()
+    if gpu_free_memory is not None:
+        min_mem = min(gpu_free_memory)
+        max_mem = max(gpu_free_memory)
+        if max_mem - min_mem > 10:
+            print(
+                f"WARNING: Unbalanced GPU memory usage at start up. This may cause OOM. Memory usage per GPU in MB: {gpu_free_memory}."
+            )
+            # nosemgrep
+            output = subprocess.check_output(["fuser -v /dev/nvidia*"], shell=True).decode("utf-8")
+            print(f"Processes using GPU: {output}")
+
+
 if __name__ == "__main__":
+    check_unknown_startup_memory_usage()
     asyncio.run(batch_inference())

model-engine/model_engine_server/inference/vllm/vllm_server.py

@@ -114,6 +114,7 @@ def check_unknown_startup_memory_usage():
             print(
                 f"WARNING: Unbalanced GPU memory usage at start up. This may cause OOM. Memory usage per GPU in MB: {gpu_free_memory}."
             )
+            # nosemgrep
             output = subprocess.check_output(["fuser -v /dev/nvidia*"], shell=True).decode("utf-8")
             print(f"Processes using GPU: {output}")
 

model-engine/model_engine_server/infra/services/image_cache_service.py

@@ -3,7 +3,7 @@
 
 import pytz
 from model_engine_server.common.config import hmi_config
-from model_engine_server.common.env_vars import GIT_TAG
+from model_engine_server.common.env_vars import CIRCLECI, GIT_TAG
 from model_engine_server.core.config import infra_config
 from model_engine_server.core.loggers import logger_name, make_logger
 from model_engine_server.domain.entities import GpuType, ModelEndpointInfraState
@@ -69,17 +69,38 @@ def _cache_finetune_llm_images(
         )
 
         istio_image = DockerImage("gcr.io/istio-release/proxyv2", "1.15.0")
-        tgi_image = DockerImage(
-            f"{infra_config().docker_repo_prefix}/{hmi_config.tgi_repository}", "0.9.3-launch_s3"
+        tgi_image_110 = DockerImage(
+            f"{infra_config().docker_repo_prefix}/{hmi_config.tgi_repository}", "1.1.0"
         )
-        tgi_image_2 = DockerImage(
-            f"{infra_config().docker_repo_prefix}/{hmi_config.tgi_repository}", "0.9.4"
+        vllm_image_027 = DockerImage(
+            f"{infra_config().docker_repo_prefix}/{hmi_config.vllm_repository}", "0.2.7"
+        )
+        vllm_image_032 = DockerImage(
+            f"{infra_config().docker_repo_prefix}/{hmi_config.vllm_repository}", "0.3.2"
+        )
+        latest_tag = (
+            self.docker_repository.get_latest_image_tag(
+                f"{infra_config().docker_repo_prefix}/{hmi_config.batch_inference_vllm_repository}"
+            )
+            if not CIRCLECI
+            else "fake_docker_repository_latest_image_tag"
+        )
+        vllm_batch_image_latest = DockerImage(
+            f"{infra_config().docker_repo_prefix}/{hmi_config.batch_inference_vllm_repository}",
+            latest_tag,
         )
         forwarder_image = DockerImage(
             f"{infra_config().docker_repo_prefix}/launch/gateway", GIT_TAG
         )
 
-        for llm_image in [istio_image, tgi_image, tgi_image_2, forwarder_image]:
+        for llm_image in [
+            istio_image,
+            tgi_image_110,
+            vllm_image_027,
+            vllm_image_032,
+            vllm_batch_image_latest,
+            forwarder_image,
+        ]:
             if self.docker_repository.is_repo_name(
                 llm_image.repo
             ) and not self.docker_repository.image_exists(llm_image.tag, llm_image.repo):

model-engine/tests/unit/infra/services/test_image_cache_service.py

@@ -52,14 +52,29 @@ async def test_caching_finetune_llm_images(
     gateway: Any = fake_image_cache_service.image_cache_gateway
 
     istio_image = DockerImage("gcr.io/istio-release/proxyv2", "1.15.0")
-    tgi_image = DockerImage(
-        f"{infra_config().docker_repo_prefix}/{hmi_config.tgi_repository}", "0.9.3-launch_s3"
+    tgi_image_110 = DockerImage(
+        f"{infra_config().docker_repo_prefix}/{hmi_config.tgi_repository}", "1.1.0"
     )
-    tgi_image_2 = DockerImage(
-        f"{infra_config().docker_repo_prefix}/{hmi_config.tgi_repository}", "0.9.4"
+    vllm_image_027 = DockerImage(
+        f"{infra_config().docker_repo_prefix}/{hmi_config.vllm_repository}", "0.2.7"
+    )
+    vllm_image_032 = DockerImage(
+        f"{infra_config().docker_repo_prefix}/{hmi_config.vllm_repository}", "0.3.2"
+    )
+    latest_tag = "fake_docker_repository_latest_image_tag"
+    vllm_batch_image_latest = DockerImage(
+        f"{infra_config().docker_repo_prefix}/{hmi_config.batch_inference_vllm_repository}",
+        latest_tag,
     )
     forwarder_image = DockerImage(f"{infra_config().docker_repo_prefix}/launch/gateway", GIT_TAG)
 
     for key in ["a10", "a100"]:
-        for llm_image in [istio_image, tgi_image, tgi_image_2, forwarder_image]:
+        for llm_image in [
+            istio_image,
+            tgi_image_110,
+            vllm_image_027,
+            vllm_image_032,
+            vllm_batch_image_latest,
+            forwarder_image,
+        ]:
             assert f"{llm_image.repo}:{llm_image.tag}" in gateway.cached_images[key]