huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-05-01 23:12:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix integration-tests config for docker runt .

Browse Source
This commit is contained in:
Nicolas Patry 2024-06-06 19:44:22 +02:00
parent 81704d28e8
commit fa05db296a
1 changed files with 15 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
integration-tests/conftest.py
View File

@ -34,6 +34,7 @@ from text_generation.types import (
DOCKER_IMAGE = os.getenv("DOCKER_IMAGE", None) DOCKER_IMAGE = os.getenv("DOCKER_IMAGE", None)
HUGGING_FACE_HUB_TOKEN = os.getenv("HUGGING_FACE_HUB_TOKEN", None) HUGGING_FACE_HUB_TOKEN = os.getenv("HUGGING_FACE_HUB_TOKEN", None)
DOCKER_VOLUME = os.getenv("DOCKER_VOLUME", "/data") DOCKER_VOLUME = os.getenv("DOCKER_VOLUME", "/data")
DOCKER_DEVICES = os.getenv("DOCKER_DEVICES")
class ResponseComparator(JSONSnapshotExtension): class ResponseComparator(JSONSnapshotExtension):
@ -453,6 +454,18 @@ def launcher(event_loop):
if DOCKER_VOLUME: if DOCKER_VOLUME:
volumes = [f"{DOCKER_VOLUME}:/data"] volumes = [f"{DOCKER_VOLUME}:/data"]
if DOCKER_DEVICES:
devices = DOCKER_DEVICES.split(",")
visible = os.getenv("ROCR_VISIBLE_DEVICES")
if visible:
env["ROCR_VISIBLE_DEVICES"] = visible
device_requests = []
else:
devices = []
device_requests = [
docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
]
container = client.containers.run( container = client.containers.run(
DOCKER_IMAGE, DOCKER_IMAGE,
command=args, command=args,
@ -460,9 +473,8 @@ def launcher(event_loop):
environment=env, environment=env,
auto_remove=False, auto_remove=False,
detach=True, detach=True,
device_requests=[ device_requests=device_requests,
docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]]) devices=devices,
],
volumes=volumes, volumes=volumes,
ports={"80/tcp": port}, ports={"80/tcp": port},
shm_size="1G", shm_size="1G",