From 3539ea37e25c14ce3bab1d80016bd9912f4ee539 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Wed, 5 Jun 2024 15:52:09 +0200 Subject: [PATCH] Making it work ? --- .github/workflows/build.yaml | 1 + integration-tests/conftest.py | 18 +++++++++++++++--- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index f60a2f59..059d267d 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -39,4 +39,5 @@ jobs: run: | export DOCKER_IMAGE=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:2.0.4-rocm export HUGGING_FACE_HUB_TOKEN=${{ secrets.HUGGING_FACE_HUB_TOKEN }} + export DEVICES=/dev/kfd,/dev/dri python -m pytest -s -vv integration-tests/models/test_flash_gpt2.py diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index 2ef85da6..a11bf781 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -34,6 +34,7 @@ from text_generation.types import ( DOCKER_IMAGE = os.getenv("DOCKER_IMAGE", None) HUGGING_FACE_HUB_TOKEN = os.getenv("HUGGING_FACE_HUB_TOKEN", None) DOCKER_VOLUME = os.getenv("DOCKER_VOLUME", "/data") +DOCKER_DEVICES = os.getenv("DOCKER_DEVICES") class ResponseComparator(JSONSnapshotExtension): @@ -453,16 +454,27 @@ def launcher(event_loop): if DOCKER_VOLUME: 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( DOCKER_IMAGE, command=args, name=container_name, environment=env, auto_remove=False, + devices=devices, + device_requests=device_requests, detach=True, - device_requests=[ - docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]]) - ], volumes=volumes, ports={"80/tcp": port}, shm_size="1G",