From 7bead9ef0f38a60341e5978ab957a44b2e9f85c3 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Wed, 2 Oct 2024 17:57:23 +0200
Subject: [PATCH] Attemp #1

---
 .github/workflows/build.yaml  |  4 +++-
 integration-tests/conftest.py | 24 +++++++++++++++---------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index d301ea94..99ec299e 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -197,5 +197,7 @@ jobs:
           # export EXTRA_PYTEST="${{ needs.build-and-push.outputs.extra_pytest }}"
           # export HF_TOKEN=${{ secrets.HF_TOKEN }}
           echo $DOCKER_IMAGE
-          python -c 'import docker; client = docker.from_env(); client.containers.run("registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-318bbc7-rocm", devices=["/dev/kfd", "/dev/dri"])'
+          # python -c 'import docker; client = docker.from_env(); client.containers.run("registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-318bbc7-rocm", devices=["/dev/kfd", "/dev/dri"], env={"MODEL_ID": "gpt2"}, detach=True)'
+          which pytest
+          pytest -s -vv integration-tests -k flash_llama_grammar_tools
           # pytest -s -vv integration-tests ${PYTEST_FLAGS} ${EXTRA_PYTEST}
diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
index 202d0f36..6e56be0c 100644
--- a/integration-tests/conftest.py
+++ b/integration-tests/conftest.py
@@ -507,9 +507,9 @@ def launcher(event_loop):
 
         if DOCKER_DEVICES:
             devices = DOCKER_DEVICES.split(",")
-            # visible = os.getenv("ROCR_VISIBLE_DEVICES")
-            # if visible:
-            #     env["ROCR_VISIBLE_DEVICES"] = visible
+            visible = os.getenv("ROCR_VISIBLE_DEVICES")
+            if visible:
+                env["ROCR_VISIBLE_DEVICES"] = visible
             device_requests = []
         else:
             devices = []
@@ -517,18 +517,24 @@ def launcher(event_loop):
                 docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
             ]
 
-        print("Starting docker")
-
-        args = ["sleep", "infinity"]
-
+        raise Exception(
+            f"""
+            Docoker image: {DOCKER_IMAGE}
+            args: {args}
+            container name: {container_name}
+            env: {env}
+            device_requests: {device_requests}
+            devices: {devices}
+        """
+        )
         container = client.containers.run(
             DOCKER_IMAGE,
             command=args,
-            # name=container_name,
+            name=container_name,
             environment=env,
             auto_remove=False,
             detach=True,
-            # device_requests=device_requests,
+            device_requests=device_requests,
             devices=devices,
             volumes=volumes,
             ports={"80/tcp": port},