More Inference Endpoints features and fixes

.github/workflows/tpu-tgi-release.yml

@@ -3,6 +3,9 @@ name: Release
 on:
   release:
     types: [published]
+  push:
+    branches:
+      - debug-tgi-ie-pt2
 
 jobs:
   docker:

optimum/tpu/xla_mp_comm.py

@@ -11,6 +11,8 @@ def __init__(self, manager: mp.Manager):
         self.root_command = manager.list()
         self.agent_ready = manager.Event()
         self.output_data = manager.list()
+        self.agent_error = manager.Event()
+        self.agent_error.clear()
 
     def send(self, command: int, *args) -> ListProxy:
         """Send a command and arguments to the agents and wait for the response.
@@ -30,6 +32,8 @@ def send(self, command: int, *args) -> ListProxy:
         self.root_bell.set()
         # wait again until agent is ready, meaning command has been processed
         self.agent_ready.wait()
+        if self.agent_error.is_set():
+            raise RuntimeError("Error on one of threads, stopping.")
         ret = self.output_data
         return ret
 
@@ -41,6 +45,7 @@ def __init__(self, root_mailbox: RootMailbox):
         self.root_command = root_mailbox.root_command
         self.agent_ready = root_mailbox.agent_ready
         self.output_data = root_mailbox.output_data
+        self.agent_error = root_mailbox.agent_error
 
     def receive(self) -> ListProxy:
         """Wait for a command from the root process and return it.

text-generation-inference/docker/entrypoint.sh

@@ -5,13 +5,26 @@
 ulimit -l 68719476736
 
 # Hugging Face Hub related
+if [[ -z "${BATCH_SIZE}" ]]; then
+  BATCH_SIZE=4
+fi
+export BATCH_SIZE="${BATCH_SIZE}"
+
+if [[ -z "${JSON_OUTPUT_DISABLE}" ]]; then
+  JSON_OUTPUT_DISABLE=--json-output
+else
+  JSON_OUTPUT_DISABLE=""
+fi
+export JSON_OUTPUT_DISABLE="${JSON_OUTPUT_DISABLE}"
+
 if [[ -z "${MODEL_ID}" ]]; then
   echo "MODEL_ID must be set"
   exit 1
 fi
 export MODEL_ID="${MODEL_ID}"
 
 text-generation-launcher --port 8080 \
-  --max-batch-size 4 \
+  --max-batch-size ${BATCH_SIZE} \
+  ${JSON_OUTPUT_DISABLE} \
   --model-id ${MODEL_ID}