Add HPU model runner & HPU input batch unit tests (#44)

Signed-off-by: Konrad Zawora <[email protected]>

Author: kzawora-intel

.github/workflows/pre-merge.yaml

@@ -12,7 +12,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  pre_merge_hpu_test:
+  pre_merge_hpu_test_build:
     runs-on: ucb-vllm-cicd-g2
     permissions:
       contents: read # Required to checkout code and read history
@@ -73,15 +73,42 @@ jobs:
 
           EOF
           echo "Docker image built successfully."
+  hpu_unit_tests:
+    runs-on: ucb-vllm-cicd-g2
+    needs: pre_merge_hpu_test_build
+    steps:
+      - name: Run pytest in tests/unit_tests
+        run: |
+          EXITCODE=1
+          remove_docker_containers() { docker rm -f hpu-plugin-v1-test-unit-tests || true; }
+          trap 'remove_docker_containers; exit $EXITCODE;' EXIT
+          remove_docker_containers
 
+          echo "Running HPU plugin v1 unit tests"
+          docker run --rm --runtime=habana --name=hpu-plugin-v1-test-unit-tests --network=host \
+            -e HABANA_VISIBLE_DEVICES=all \
+            -e HF_HOME=/workspace/hf_cache \
+            -v /mnt/hf_cache:/workspace/hf_cache \
+            hpu-plugin-v1-test-env-pre-merge \
+            /bin/bash -c "pytest -vvv /workspace/vllm-gaudi/tests/unit_tests"
+          
+          EXITCODE=$?
+          echo "Test script exited with code: $EXITCODE"
+
+  hpu_e2e_tests:
+    runs-on: ucb-vllm-cicd-g2
+    # This is a final job that runs after the build and unit tests 
+    # Unit tests are supposed to fail fast if anything goes wrong, removing the need for this job to run.
+    needs: [pre_merge_hpu_test_build, hpu_unit_tests]
+    steps:
       - name: Run test scripts
         run: |
           EXITCODE=1
-          remove_docker_containers() { docker rm -f hpu-plugin-v1-test-pre-merge || true; }
+          remove_docker_containers() { docker rm -f hpu-plugin-v1-e2e-tests || true; }
           trap 'remove_docker_containers; exit $EXITCODE;' EXIT
           remove_docker_containers
 
-          echo "Running HPU plugin v1 test"
+          echo "Running HPU plugin v1 e2e tests"
           docker run --rm --runtime=habana --name=hpu-plugin-v1-test-pre-merge --network=host \
             -e HABANA_VISIBLE_DEVICES=all \
             -e HF_HOME=/workspace/hf_cache \
@@ -91,3 +118,11 @@ jobs:
           
           EXITCODE=$?
           echo "Test script exited with code: $EXITCODE"
+
+  pre_merge_hpu_test:
+    runs-on: ucb-vllm-cicd-g2
+    needs: [hpu_unit_tests, hpu_e2e_tests]
+    # This job is required to pass for pre-merge CI. By itself it does nothing, and will only pass if all jobs specified in "needs" list pass.
+    steps:
+      - name: Succeeded if all previous jobs passed
+        run: echo "All previous jobs passed." 

tests/unit_tests/conftest.py

@@ -0,0 +1,20 @@
+from vllm.distributed import (cleanup_dist_env_and_memory,
+                              init_distributed_environment,
+                              initialize_model_parallel)
+import pytest
+import tempfile
+
+
+@pytest.fixture
+def dist_init():
+    temp_file = tempfile.mkstemp()[1]
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        distributed_init_method=f"file://{temp_file}",
+        local_rank=0,
+        backend="hccl",
+    )
+    initialize_model_parallel(1, 1)
+    yield
+    cleanup_dist_env_and_memory()