run on mounted disk

predibase · Nov 8, 2024 · 24576f7 · 24576f7
1 parent 78cd63d
commit 24576f7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/.github/workflows/integration_tests2.yaml b/.github/workflows/integration_tests2.yaml
@@ -33,7 +33,12 @@ jobs:
       run: |
         echo "HUGGING_FACE_HUB_TOKEN=${{ secrets.HUGGING_FACE_HUB_TOKEN }}" >> $GITHUB_ENV
 
-    - name: Run Mistral tests
+    - name: Run LLM tests
       run: |
         cd integration-tests-2
         pytest test_base_llms.py -v --capture=tee-sys
+
+    # - name: Run Embedding tests
+    #   run: |
+    #     cd integration-tests-2
+    #     pytest test_embeddings.py -v --capture=tee-sys
diff --git a/integration-tests-2/utils/docker_runner.py b/integration-tests-2/utils/docker_runner.py
@@ -1,5 +1,4 @@
 import docker
-import tempfile
 import time
 import os
 import requests
@@ -20,8 +19,7 @@ def start_container(self, model_config: Dict[str, Any]) -> None:
             for key, value in model_config["docker_args"].items():
                 cmd.extend([f"--{key.replace('_', '-')}", str(value)])
 
-        data_temp_dir = tempfile.mkdtemp()
-        print("Using temporary directory for data: ", data_temp_dir)
+        data_dir = os.getenv("DATA_DIR", "/integration-test-data")
         # Start container
         self.container = self.client.containers.run(
             "ghcr.io/predibase/lorax:main",
@@ -31,7 +29,7 @@ def start_container(self, model_config: Dict[str, Any]) -> None:
             },
             device_requests=[docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])],
             ports={"80": 8080},
-            volumes={data_temp_dir: {"bind": "/data", "mode": "rw"}},
+            volumes={data_dir: {"bind": "/data", "mode": "rw"}},
             shm_size="1g",
             detach=True,
             name=f"lorax-test-{model_config['name']}-{time.time()}",