follow up changes in the upstream; and update test scripts

Signed-off-by: Juncheng Gu <jcgu@google.com>

Author: juncgu-google

.buildkite/features/KV_Cache_Offload.yml

@@ -9,12 +9,7 @@ steps:
     commands:
       - |
         .buildkite/scripts/run_in_docker.sh \
-          python3 -m pytest -s -v /workspace/tpu_inference/tests/distributed/offload/tpu_offload_connector_scheduler_test.py \
-          /workspace/tpu_inference/tests/distributed/offload/tpu_offload_connector_worker_test.py \
-          /workspace/tpu_inference/tests/distributed/offload/tpu_offload_cpu_backend_test.py \
-          /workspace/tpu_inference/tests/distributed/offload/tpu_offload_manager_test.py \
-          /workspace/tpu_inference/tests/distributed/offload/tpu_offload_utils_test.py \
-          /workspace/tpu_inference/tests/distributed/offload/tpu_offload_accuracy_test.py
+          python3 -m pytest -s -v /workspace/tpu_inference/tests/distributed/offload/
   - label: "Record correctness test result for KV Cache Offload"
     key: "record_KV_Cache_Offload_CorrectnessTest"
     depends_on: "KV_Cache_Offload_CorrectnessTest"

examples/offload/gke/pod_tpu_host_offload_unit_tests.yaml

@@ -17,12 +17,7 @@ spec:
     command:
     - /bin/bash
     - -c
-    - "pytest -sv tests/distributed/offload/tpu_offload_cpu_backend_test.py"
-    - "pytest -sv tests/distributed/offload/tpu_offload_connector_worker_test.py"
-    - "pytest -sv tests/distributed/offload/tpu_offload_connector_scheduler_test.py"
-    - "pytest -sv tests/distributed/offload/tpu_offload_utils_test.py"
-    - "pytest -sv tests/distributed/offload/tpu_offload_manager_test.py"
-    - "pytest -sv tests/distributed/offload/tpu_offload_accuracy_test.py"
+    - "pytest -sv tests/distributed/offload/"
     env:
     - name: HUGGING_FACE_HUB_TOKEN
       valueFrom:

tests/distributed/offload/tpu_offload_connector_worker_test.py

@@ -94,6 +94,7 @@ def setUp(self):
     def tearDown(self):
         super().tearDown()
         cc.reset_cache()
+        jax.clear_caches()
 
     def create_mesh(self, axis_shapes, axis_names):
         """Creates a JAX device mesh with the default device order."""

tpu_inference/worker/tpu_worker.py

@@ -289,7 +289,8 @@ def determine_available_memory(self) -> int:
 
                 kv_cache_specs = self.model_runner.get_kv_cache_spec()
                 num_layers = len(kv_cache_specs)
-                vllm_page_size_bytes = get_uniform_page_size(kv_cache_specs)
+                vllm_page_size_bytes = get_uniform_page_size(
+                    list(kv_cache_specs.values()))
                 stage_buffer_size_bytes = staging_buffer_pages * num_layers * vllm_page_size_bytes
 
                 total_hbm_avail = total_hbm_avail - stage_buffer_size_bytes