diff --git a/integration-tests-2/utils/docker_runner.py b/integration-tests-2/utils/docker_runner.py index 0c43bc1c..a9ed3958 100644 --- a/integration-tests-2/utils/docker_runner.py +++ b/integration-tests-2/utils/docker_runner.py @@ -1,5 +1,5 @@ import docker -import json +import tempfile import time import os import requests @@ -20,6 +20,8 @@ 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) # Start container self.container = self.client.containers.run( "ghcr.io/predibase/lorax:main", @@ -29,7 +31,7 @@ def start_container(self, model_config: Dict[str, Any]) -> None: }, device_requests=[docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])], ports={"80": 8080}, - volumes={f"{os.getcwd()}/data": {"bind": "/data", "mode": "rw"}}, + volumes={data_temp_dir: {"bind": "/data", "mode": "rw"}}, shm_size="1g", detach=True, name=f"lorax-test-{model_config['name']}-{time.time()}", @@ -55,19 +57,3 @@ def stop_container(self) -> None: self.container.stop() self.container.remove() self.container = None - - -if __name__ == "__main__": - with open("configs.json", "r") as f: - model_configs = json.load(f) - runner = DockerModelRunner() - runner.start_container(model_configs[0]) - runner.wait_for_healthy() - response = requests.post( - "http://localhost:8080/generate", - json={"inputs": model_configs[0]["test_prompt"], "parameters": {"max_new_tokens": 10}}, - ) - response.raise_for_status() - print(response.json()) - runner.stop_container() - print("Done")