@@ -125,24 +125,26 @@ curl localhost:8080/v1/chat/completions \
### Streaming with JavaScript
First, we need to install the `@huggingface/inference` library.
-`npm install @huggingface/inference`
-If you're using the free Inference API, you can use `HfInference`. If you're using inference endpoints, you can use `HfInferenceEndpoint`.
+```bash
+npm install @huggingface/inference
+```
+
+Whether you use Inference Providers (our serverless API), or Inference Endpoints, you can call `InferenceClient`.
-We can create a `HfInferenceEndpoint` providing our endpoint URL and credential.
```js
-import { HfInferenceEndpoint } from '@huggingface/inference'
+import { InferenceClient } from '@huggingface/inference';
-const hf = new HfInferenceEndpoint('https://your_endpoint.endpoints.huggingface.cloud/', 'hf_YOUR_TOKEN')
+const client = new InferenceClient('hf_YOUR_TOKEN', { endpointUrl: 'https://your_endpoint.endpoints.huggingface.cloud/' });
// prompt
-const prompt = 'What can you do in Nuremberg, Germany? Give me 3 Tips'
+const prompt = 'What can you do in Nuremberg, Germany? Give me 3 Tips';
-const stream = hf.textGenerationStream({ inputs: prompt })
+const stream = client.textGenerationStream({ inputs: prompt });
for await (const r of stream) {
// yield the generated token
- process.stdout.write(r.token.text)
+ process.stdout.write(r.token.text);
}
```
diff --git a/docs/source/installation_amd.md b/docs/source/installation_amd.md
index 3c9c0eecf05..df4abb3b43c 100644
--- a/docs/source/installation_amd.md
+++ b/docs/source/installation_amd.md
@@ -11,7 +11,7 @@ volume=$PWD/data # share a volume with the Docker container to avoid downloading
docker run --rm -it --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
--device=/dev/kfd --device=/dev/dri --group-add video \
--ipc=host --shm-size 256g --net host -v $volume:/data \
- ghcr.io/huggingface/text-generation-inference:3.0.0-rocm \
+ ghcr.io/huggingface/text-generation-inference:3.3.5-rocm \
--model-id $model
```
@@ -39,6 +39,6 @@ The custom kernel supports bf16 and fp16 data types, block size of 16, head size
## Unsupported features
-The following features are currently not supported in the ROCm version of TGI, and the supported may be extended in the future:
+The following features are currently not supported in the ROCm version of TGI, and the support may be extended in the future:
* Loading [AWQ](https://huggingface.co/docs/transformers/quantization#awq) checkpoints.
* Kernel for sliding window attention (Mistral)
diff --git a/docs/source/installation_gaudi.md b/docs/source/installation_gaudi.md
index 1ddf2b47225..51aa667dcad 100644
--- a/docs/source/installation_gaudi.md
+++ b/docs/source/installation_gaudi.md
@@ -1,3 +1,3 @@
# Using TGI with Intel Gaudi
-Check out this [repository](https://github.com/huggingface/tgi-gaudi) to serve models with TGI on Gaudi and Gaudi2 with [Optimum Habana](https://huggingface.co/docs/optimum/habana/index).
+You can use TGI on Intel Gaudi using the [TGI gaudi backend](https://huggingface.co/docs/text-generation-inference/backends/gaudi).
diff --git a/docs/source/installation_inferentia.md b/docs/source/installation_inferentia.md
index 0394e6ded37..bfd0f657754 100644
--- a/docs/source/installation_inferentia.md
+++ b/docs/source/installation_inferentia.md
@@ -1,3 +1,3 @@
# Using TGI with Inferentia
-Check out this [guide](https://github.com/huggingface/optimum-neuron/tree/main/text-generation-inference) on how to serve models with TGI on Inferentia2.
+You can use TGI on AWS Trainium and Inferentia platforms using the [TGI neuron backend](https://huggingface.co/docs/text-generation-inference/backends/neuron).
diff --git a/docs/source/installation_intel.md b/docs/source/installation_intel.md
index a5578e769c0..60b0bcc056f 100644
--- a/docs/source/installation_intel.md
+++ b/docs/source/installation_intel.md
@@ -12,7 +12,7 @@ volume=$PWD/data # share a volume with the Docker container to avoid downloading
docker run --rm --privileged --cap-add=sys_nice \
--device=/dev/dri \
--ipc=host --shm-size 1g --net host -v $volume:/data \
- ghcr.io/huggingface/text-generation-inference:3.0.0-intel-xpu \
+ ghcr.io/huggingface/text-generation-inference:3.3.5-intel-xpu \
--model-id $model --cuda-graphs 0
```
@@ -29,7 +29,7 @@ volume=$PWD/data # share a volume with the Docker container to avoid downloading
docker run --rm --privileged --cap-add=sys_nice \
--device=/dev/dri \
--ipc=host --shm-size 1g --net host -v $volume:/data \
- ghcr.io/huggingface/text-generation-inference:3.0.0-intel-cpu \
+ ghcr.io/huggingface/text-generation-inference:3.3.5-intel-cpu \
--model-id $model --cuda-graphs 0
```
diff --git a/docs/source/installation_nvidia.md b/docs/source/installation_nvidia.md
index d954481eda5..37cb841c8b3 100644
--- a/docs/source/installation_nvidia.md
+++ b/docs/source/installation_nvidia.md
@@ -11,7 +11,7 @@ model=teknium/OpenHermes-2.5-Mistral-7B
volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run
docker run --gpus all --shm-size 64g -p 8080:80 -v $volume:/data \
- ghcr.io/huggingface/text-generation-inference:3.0.0 \
+ ghcr.io/huggingface/text-generation-inference:3.3.5 \
--model-id $model
```
diff --git a/docs/source/installation_tpu.md b/docs/source/installation_tpu.md
new file mode 100644
index 00000000000..559e83aa7dd
--- /dev/null
+++ b/docs/source/installation_tpu.md
@@ -0,0 +1,3 @@
+# Using TGI with Google TPUs
+
+Check out this [guide](https://huggingface.co/docs/optimum-tpu) on how to serve models with TGI on TPUs.
diff --git a/docs/source/multi_backend_support.md b/docs/source/multi_backend_support.md
new file mode 100644
index 00000000000..997503a4f19
--- /dev/null
+++ b/docs/source/multi_backend_support.md
@@ -0,0 +1,16 @@
+# Multi-backend support
+
+TGI (Text Generation Inference) offers flexibility by supporting multiple backends for serving large language models (LLMs).
+With multi-backend support, you can choose the backend that best suits your needs,
+whether you prioritize performance, ease of use, or compatibility with specific hardware. API interaction with
+TGI remains consistent across backends, allowing you to switch between them seamlessly.
+
+**Supported backends:**
+* **TGI CUDA backend**: This high-performance backend is optimized for NVIDIA GPUs and serves as the default option
+ within TGI. Developed in-house, it boasts numerous optimizations and is used in production by various projects, including those by Hugging Face.
+* **[TGI TRTLLM backend](./backends/trtllm)**: This backend leverages NVIDIA's TensorRT library to accelerate LLM inference.
+ It utilizes specialized optimizations and custom kernels for enhanced performance.
+ However, it requires a model-specific compilation step for each GPU architecture.
+* **[TGI Llamacpp backend](./backends/llamacpp)**: This backend facilitates the deployment of large language models
+ (LLMs) by integrating [llama.cpp][llama.cpp], an advanced inference engine optimized for both CPU and GPU computation.
+* **[TGI Neuron backend](./backends/neuron)**: This backend leverages the [AWS Neuron SDK](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) to allow the deployment of large language models (LLMs) on [AWS Trainium and Inferentia chips](https://aws.amazon.com/ai/machine-learning/trainium/).
diff --git a/docs/source/quicktour.md b/docs/source/quicktour.md
index d1f3efa4f40..bd8495c52e2 100644
--- a/docs/source/quicktour.md
+++ b/docs/source/quicktour.md
@@ -11,7 +11,7 @@ model=teknium/OpenHermes-2.5-Mistral-7B
volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run
docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data \
- ghcr.io/huggingface/text-generation-inference:3.0.0 \
+ ghcr.io/huggingface/text-generation-inference:3.3.5 \
--model-id $model
```
@@ -96,7 +96,7 @@ curl 127.0.0.1:8080/generate \
To see all possible deploy flags and options, you can use the `--help` flag. It's possible to configure the number of shards, quantization, generation parameters, and more.
```bash
-docker run ghcr.io/huggingface/text-generation-inference:3.0.0 --help
+docker run ghcr.io/huggingface/text-generation-inference:3.3.5 --help
```
diff --git a/docs/source/reference/api_reference.md b/docs/source/reference/api_reference.md
index 42a777039f3..7d21eca70a8 100644
--- a/docs/source/reference/api_reference.md
+++ b/docs/source/reference/api_reference.md
@@ -163,7 +163,7 @@ hub = {
# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
- image_uri=get_huggingface_llm_image_uri("huggingface",version="3.0.0"),
+ image_uri=get_huggingface_llm_image_uri("huggingface",version="3.3.5"),
env=hub,
role=role,
)
diff --git a/docs/source/reference/launcher.md b/docs/source/reference/launcher.md
index 159b22e7352..5b7321b73a3 100644
--- a/docs/source/reference/launcher.md
+++ b/docs/source/reference/launcher.md
@@ -58,8 +58,6 @@ Options:
Quantization method to use for the model. It is not necessary to specify this option for pre-quantized models, since the quantization method is read from the model configuration.
Marlin kernels will be used automatically for GPTQ/AWQ models.
-
- [env: QUANTIZE=]
Possible values:
- awq: 4 bit quantization. Requires a specific AWQ quantized model:
. Should replace GPTQ models wherever possible because of the better latency
@@ -72,6 +70,8 @@ Options:
- bitsandbytes-nf4: Bitsandbytes 4bit. Can be applied on any model, will cut the memory requirement by 4x, but it is known that the model will be much slower to run than the native f16
- bitsandbytes-fp4: Bitsandbytes 4bit. nf4 should be preferred in most cases but maybe this one has better perplexity performance for you model
- fp8: [FP8](https://developer.nvidia.com/blog/nvidia-arm-and-intel-publish-fp8-specification-for-standardization-as-an-interchange-format-for-ai/) (e4m3) works on H100 and above This dtype has native ops should be the fastest if available. This is currently not the fastest because of local unpacking + padding to satisfy matrix multiplication limitations
+
+ [env: QUANTIZE=]
```
## SPECULATE
@@ -198,7 +198,7 @@ Options:
For `max_batch_total_tokens=1000`, you could fit `10` queries of `total_tokens=100` or a single query of `1000` tokens.
- Overall this number should be the largest possible amount that fits the remaining memory (after the model is loaded). Since the actual memory overhead depends on other parameters like if you're using quantization, flash attention or the model implementation, text-generation-inference cannot infer this number automatically.
+ Overall this number should be the largest possible amount that fits the remaining memory (after the model is loaded). Since the actual memory overhead depends on other parameters like if you're using quantization, flash attention or the model implementation, text-generation-inference infers this number automatically if not provided ensuring that the value is as large as possible.
[env: MAX_BATCH_TOTAL_TOKENS=]
@@ -251,6 +251,15 @@ Options:
[env: PORT=]
[default: 3000]
+```
+## PROMETHEUS_PORT
+```shell
+ -p, --prometheus-port
+ The Prometheus port to listen on
+
+ [env: PROMETHEUS_PORT=]
+ [default: 9000]
+
```
## SHARD_UDS_PATH
```shell
@@ -447,14 +456,14 @@ Options:
```shell
--usage-stats
Control if anonymous usage stats are collected. Options are "on", "off" and "no-stack" Defaul is on
-
- [env: USAGE_STATS=]
- [default: on]
Possible values:
- on: Default option, usage statistics are collected anonymously
- off: Disables all collection of usage statistics
- no-stack: Doesn't send the error stack trace or error type, but allows sending a crash event
+
+ [env: USAGE_STATS=]
+ [default: on]
```
## PAYLOAD_LIMIT
@@ -477,6 +486,15 @@ Options:
[env: ENABLE_PREFILL_LOGPROBS=]
+```
+## GRACEFUL_TERMINATION_TIMEOUT
+```shell
+ -g, --graceful-termination-timeout
+ Change timeout of graceful termination of the TGI server
+
+ [env: GRACEFUL_TERMINATION_TIMEOUT=]
+ [default: 90]
+
```
## HELP
```shell
diff --git a/docs/source/supported_models.md b/docs/source/supported_models.md
index 0f39ff28e42..5d855a5b905 100644
--- a/docs/source/supported_models.md
+++ b/docs/source/supported_models.md
@@ -4,14 +4,19 @@
Text Generation Inference enables serving optimized models. The following sections list which models (VLMs & LLMs) are supported.
- [Deepseek V2](https://huggingface.co/deepseek-ai/DeepSeek-V2)
+- [Deepseek V3](https://huggingface.co/deepseek-ai/DeepSeek-V3)
- [Idefics 2](https://huggingface.co/HuggingFaceM4/idefics2-8b) (Multimodal)
+- [Idefics 3](https://huggingface.co/HuggingFaceM4/Idefics3-8B-Llama3) (Multimodal)
- [Llava Next (1.6)](https://huggingface.co/llava-hf/llava-v1.6-vicuna-13b-hf) (Multimodal)
- [Llama](https://huggingface.co/collections/meta-llama/llama-31-669fc079a0c406a149a5738f)
+- [Llama4](https://huggingface.co/collections/meta-llama/llama-31-669fc079a0c406a149a5738f)
- [Phi 3](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct)
- [Granite](https://huggingface.co/ibm-granite/granite-3.0-8b-instruct)
- [Gemma](https://huggingface.co/google/gemma-7b)
- [PaliGemma](https://huggingface.co/google/paligemma-3b-pt-224)
- [Gemma2](https://huggingface.co/collections/google/gemma-2-release-667d6600fd5220e7b967f315)
+- [Gemma3](https://huggingface.co/collections/google/gemma-3-release-67c6c6f89c4f76621268bb6d)
+- [Gemma3 Text](https://huggingface.co/collections/google/gemma-3-release-67c6c6f89c4f76621268bb6d)
- [Cohere](https://huggingface.co/CohereForAI/c4ai-command-r-plus)
- [Dbrx](https://huggingface.co/databricks/dbrx-instruct)
- [Mamba](https://huggingface.co/state-spaces/mamba-2.8b-slimpj)
@@ -25,6 +30,7 @@ Text Generation Inference enables serving optimized models. The following sectio
- [StarCoder 2](https://huggingface.co/bigcode/starcoder2-15b-instruct-v0.1)
- [Qwen 2](https://huggingface.co/collections/Qwen/qwen2-6659360b33528ced941e557f)
- [Qwen 2 VL](https://huggingface.co/collections/Qwen/qwen2-vl-66cee7455501d7126940800d)
+- [Qwen 2.5 VL](https://huggingface.co/collections/Qwen/qwen25-66e81a666513e518adb90d9e)
- [Opt](https://huggingface.co/facebook/opt-6.7b)
- [T5](https://huggingface.co/google/flan-t5-xxl)
- [Galactica](https://huggingface.co/facebook/galactica-120b)
diff --git a/docs/source/usage_statistics.md b/docs/source/usage_statistics.md
index d3878b53ccb..dbf36c77921 100644
--- a/docs/source/usage_statistics.md
+++ b/docs/source/usage_statistics.md
@@ -3,7 +3,7 @@
Text Generation Inference collects anonymous usage statistics to help us improve the service. The collected data is used to improve TGI and to understand what causes failures. The data is collected transparently and any sensitive information is omitted.
-Data is sent twice, once on server startup and once when server stops. Also, usage statistics are only enabled when TGI is running in docker to avoid collecting data then TGI runs directly on the host machine.
+Usage statistics are collected only when TGI is running in a Docker container. This prevents data collection when TGI is run directly on the host machine. The collected data includes startup and shutdown events, as well as a heartbeat signal sent every 15 minutes.
## What data is collected
diff --git a/flake.lock b/flake.lock
index ec87d569231..e57990c89a1 100644
--- a/flake.lock
+++ b/flake.lock
@@ -102,17 +102,17 @@
"flake-parts": "flake-parts_3",
"nix-test-runner": "nix-test-runner_3",
"nixpkgs": [
- "tgi-nix",
+ "hf-nix",
"nixpkgs"
],
"pre-commit-hooks": "pre-commit-hooks_3"
},
"locked": {
- "lastModified": 1732039290,
- "narHash": "sha256-LQKY7bShf2H9kJouxa9ZspfdrulnZF9o4kLTqGqCDYM=",
+ "lastModified": 1739473963,
+ "narHash": "sha256-ItAhpjNUzEWd/cgZVyW/jvoGbCec4TK29e1Mnmn1oJE=",
"owner": "nix-community",
"repo": "crate2nix",
- "rev": "9ff208ce7f5a482272b1bcefbe363c772d7ff914",
+ "rev": "be31feae9a82c225c0fd1bdf978565dc452a483a",
"type": "github"
},
"original": {
@@ -305,11 +305,11 @@
},
"flake-compat_4": {
"locked": {
- "lastModified": 1696426674,
- "narHash": "sha256-kvjfFW7WAETZlt09AgDn1MrtKzP7t90Vf7vypd3OL1U=",
+ "lastModified": 1733328505,
+ "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=",
"owner": "edolstra",
"repo": "flake-compat",
- "rev": "0f9255e01c2351cc7d116c072cb317785dd33b33",
+ "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec",
"type": "github"
},
"original": {
@@ -579,6 +579,26 @@
"type": "github"
}
},
+ "hf-nix": {
+ "inputs": {
+ "flake-compat": "flake-compat_4",
+ "flake-utils": "flake-utils_7",
+ "nixpkgs": "nixpkgs_6"
+ },
+ "locked": {
+ "lastModified": 1747919133,
+ "narHash": "sha256-VvF1naQOvv7yulQ5/cDiaxkNxlh1Y84QMZnderv1szk=",
+ "owner": "huggingface",
+ "repo": "hf-nix",
+ "rev": "9c71e026d6c7c8588ef85a5f7c77f57d598e038c",
+ "type": "github"
+ },
+ "original": {
+ "owner": "huggingface",
+ "repo": "hf-nix",
+ "type": "github"
+ }
+ },
"nix-filter": {
"locked": {
"lastModified": 1731533336,
@@ -718,16 +738,16 @@
},
"nixpkgs_6": {
"locked": {
- "lastModified": 1732034459,
- "narHash": "sha256-Zais/zMRuJdlALidkUgEuasXOd37ZZLqkPkF9bIYSrY=",
+ "lastModified": 1747820358,
+ "narHash": "sha256-fTqsZsUX6M3yeEvgyQvXcbGmT2CaRVyVwsi8eK29Oj4=",
"owner": "danieldk",
"repo": "nixpkgs",
- "rev": "40280e7bf9743cdf563494db4ece2a43aa674fa8",
+ "rev": "d3c1681180717528068082103bf323147de6ab0b",
"type": "github"
},
"original": {
"owner": "danieldk",
- "ref": "outlines-v0.1.4-tgi",
+ "ref": "cudatoolkit-12.9-kernel-builder",
"repo": "nixpkgs",
"type": "github"
}
@@ -836,28 +856,28 @@
"inputs": {
"crate2nix": "crate2nix",
"flake-utils": "flake-utils_6",
+ "hf-nix": "hf-nix",
"nix-filter": "nix-filter",
"nixpkgs": [
- "tgi-nix",
+ "hf-nix",
"nixpkgs"
],
- "rust-overlay": "rust-overlay",
- "tgi-nix": "tgi-nix"
+ "rust-overlay": "rust-overlay"
}
},
"rust-overlay": {
"inputs": {
"nixpkgs": [
- "tgi-nix",
+ "hf-nix",
"nixpkgs"
]
},
"locked": {
- "lastModified": 1732242723,
- "narHash": "sha256-NWI8csIK0ujFlFuEXKnoc+7hWoCiEtINK9r48LUUMeU=",
+ "lastModified": 1743993291,
+ "narHash": "sha256-u8GHvduU1gCtoFXvTS/wGjH1ouv5S/GRGq6MAT+sG/k=",
"owner": "oxalica",
"repo": "rust-overlay",
- "rev": "a229311fcb45b88a95fdfa5cecd8349c809a272a",
+ "rev": "0cb3c8979c65dc6a5812dfe67499a8c7b8b4325b",
"type": "github"
},
"original": {
@@ -970,26 +990,6 @@
"repo": "default",
"type": "github"
}
- },
- "tgi-nix": {
- "inputs": {
- "flake-compat": "flake-compat_4",
- "flake-utils": "flake-utils_7",
- "nixpkgs": "nixpkgs_6"
- },
- "locked": {
- "lastModified": 1732218602,
- "narHash": "sha256-BElslL34KjOJCFMPkNtilOz6S/7iY7Vd72FNbRRWKDY=",
- "owner": "huggingface",
- "repo": "text-generation-inference-nix",
- "rev": "f79638ac4e420e661321261744e745a3a747e182",
- "type": "github"
- },
- "original": {
- "owner": "huggingface",
- "repo": "text-generation-inference-nix",
- "type": "github"
- }
}
},
"root": "root",
diff --git a/flake.nix b/flake.nix
index 83cedfa620f..b5b13cad8bf 100644
--- a/flake.nix
+++ b/flake.nix
@@ -2,15 +2,15 @@
inputs = {
crate2nix = {
url = "github:nix-community/crate2nix";
- inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
+ inputs.nixpkgs.follows = "hf-nix/nixpkgs";
};
nix-filter.url = "github:numtide/nix-filter";
- tgi-nix.url = "github:huggingface/text-generation-inference-nix";
- nixpkgs.follows = "tgi-nix/nixpkgs";
+ hf-nix.url = "github:huggingface/hf-nix";
+ nixpkgs.follows = "hf-nix/nixpkgs";
flake-utils.url = "github:numtide/flake-utils";
rust-overlay = {
url = "github:oxalica/rust-overlay";
- inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
+ inputs.nixpkgs.follows = "hf-nix/nixpkgs";
};
};
outputs =
@@ -21,7 +21,7 @@
nixpkgs,
flake-utils,
rust-overlay,
- tgi-nix,
+ hf-nix,
}:
flake-utils.lib.eachDefaultSystem (
system:
@@ -33,10 +33,10 @@
};
pkgs = import nixpkgs {
inherit system;
- inherit (tgi-nix.lib) config;
+ inherit (hf-nix.lib) config;
overlays = [
rust-overlay.overlays.default
- tgi-nix.overlays.default
+ hf-nix.overlays.default
(import nix/overlay.nix)
];
};
@@ -44,9 +44,24 @@
benchmark = cargoNix.workspaceMembers.text-generation-benchmark.build.override {
inherit crateOverrides;
};
- launcher = cargoNix.workspaceMembers.text-generation-launcher.build.override {
- inherit crateOverrides;
- };
+ launcher =
+ let
+ launcherUnwrapped = cargoNix.workspaceMembers.text-generation-launcher.build.override {
+ inherit crateOverrides;
+ };
+ packagePath =
+ with pkgs.python3.pkgs;
+ makePythonPath [
+ torch
+ ];
+ in
+ pkgs.writeShellApplication {
+ name = "text-generation-launcher";
+ text = ''
+ PYTHONPATH="${packagePath}" ${launcherUnwrapped}/bin/text-generation-launcher "$@"
+ '';
+ };
+
router =
let
routerUnwrapped = cargoNix.workspaceMembers.text-generation-router-v3.build.override {
@@ -161,11 +176,15 @@
'';
};
- dockerImage = pkgs.callPackage nix/docker.nix {
+ # Use plain nixpkgs without overlays for dockerTools. dockerTools
+ # uses a Python package for computing the layers from the transitive
+ # closure. However, this needs a lot of rebuilds due to our overlay.
+
+ dockerImage = nixpkgs.legacyPackages.${system}.callPackage nix/docker.nix {
text-generation-inference = default;
};
- dockerImageStreamed = pkgs.callPackage nix/docker.nix {
+ dockerImageStreamed = nixpkgs.legacyPackages.${system}.callPackage nix/docker.nix {
text-generation-inference = default;
stream = true;
};
diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
index c9c477665a3..9cc334168aa 100644
--- a/integration-tests/conftest.py
+++ b/integration-tests/conftest.py
@@ -1,4 +1,18 @@
+pytest_plugins = [
+ "fixtures.neuron.service",
+ "fixtures.neuron.export_models",
+ "fixtures.gaudi.service",
+]
# ruff: noqa: E402
+from _pytest.fixtures import SubRequest
+from huggingface_hub.inference._generated.types.chat_completion import (
+ ChatCompletionStreamOutput,
+ ChatCompletionOutput,
+)
+from openai.types.chat.chat_completion_chunk import (
+ ChatCompletionChunk as OAIChatCompletionChunk,
+)
+from openai.types.completion import Completion as OAICompletion
import requests
@@ -10,13 +24,13 @@ def request(self, *args, **kwargs):
requests.sessions.Session = SessionTimeoutFix
+import warnings
import asyncio
import contextlib
import json
import math
import os
import random
-import shutil
import subprocess
import sys
import tempfile
@@ -30,7 +44,6 @@ def request(self, *args, **kwargs):
from aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError
from docker.errors import NotFound
from syrupy.extensions.json import JSONSnapshotExtension
-
from text_generation import AsyncClient
from text_generation.types import (
BestOfSequence,
@@ -56,45 +69,124 @@ def pytest_addoption(parser):
parser.addoption(
"--release", action="http://github.com/store_true", default=False, help="run release tests"
)
+ parser.addoption(
+ "--neuron", action="http://github.com/store_true", default=False, help="run neuron tests"
+ )
+ parser.addoption(
+ "--gaudi", action="http://github.com/store_true", default=False, help="run gaudi tests"
+ )
+ parser.addoption(
+ "--gaudi-all-models",
+ action="http://github.com/store_true",
+ default=False,
+ help="Run tests for all models instead of just the default subset",
+ )
def pytest_configure(config):
config.addinivalue_line("markers", "release: mark test as a release-only test")
+ config.addinivalue_line("markers", "neuron: mark test as a neuron test")
def pytest_collection_modifyitems(config, items):
- if config.getoption("--release"):
- # --release given in cli: do not skip release tests
- return
- skip_release = pytest.mark.skip(reason="need --release option to run")
+ selectors = []
+ if not config.getoption("--release"):
+ # --release not given in cli: skip release tests
+ def skip_release(item):
+ if "release" in item.keywords:
+ item.add_marker(pytest.mark.skip(reason="need --release option to run"))
+
+ selectors.append(skip_release)
+
+ if config.getoption("--gaudi"):
+
+ def skip_not_gaudi(item):
+ if "gaudi" not in item.keywords:
+ item.add_marker(pytest.mark.skip(reason="requires --gaudi to run"))
+
+ selectors.append(skip_not_gaudi)
+ else:
+
+ def skip_gaudi(item):
+ if "gaudi" in item.keywords:
+ item.add_marker(pytest.mark.skip(reason="requires --gaudi to run"))
+
+ selectors.append(skip_gaudi)
+
+ if config.getoption("--neuron"):
+
+ def skip_not_neuron(item):
+ if "neuron" not in item.keywords:
+ item.add_marker(
+ pytest.mark.skip(reason="incompatible with --neuron option")
+ )
+
+ selectors.append(skip_not_neuron)
+ else:
+
+ def skip_neuron(item):
+ if "neuron" in item.keywords:
+ item.add_marker(pytest.mark.skip(reason="requires --neuron to run"))
+
+ selectors.append(skip_neuron)
+
for item in items:
- if "release" in item.keywords:
- item.add_marker(skip_release)
+ for selector in selectors:
+ selector(item)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def container_log(request: SubRequest):
+ error_log = request.getfixturevalue("error_log")
+ assert error_log is not None
+ yield
+ if request.session.testsfailed:
+ error_log.seek(0)
+ print(error_log.read(), file=sys.stderr)
+ else:
+ error_log.truncate(0)
+ error_log.seek(0)
class ResponseComparator(JSONSnapshotExtension):
rtol = 0.2
ignore_logprob = False
- def serialize(
+ def _serialize(
self,
data,
- *,
- include=None,
- exclude=None,
- matcher=None,
):
if (
isinstance(data, Response)
or isinstance(data, ChatComplete)
or isinstance(data, ChatCompletionChunk)
or isinstance(data, ChatCompletionComplete)
+ or isinstance(data, Completion)
+ or isinstance(data, OAIChatCompletionChunk)
+ or isinstance(data, OAICompletion)
):
data = data.model_dump()
+ elif isinstance(data, ChatCompletionStreamOutput) or isinstance(
+ data, ChatCompletionOutput
+ ):
+ data = dict(data)
+ elif isinstance(data, List):
+ data = [self._serialize(d) for d in data]
+ elif isinstance(data, dict):
+ return data
+ else:
+ raise RuntimeError(f"Unexpected data {type(data)} : {data}")
+ return data
- if isinstance(data, List):
- data = [d.model_dump() for d in data]
-
+ def serialize(
+ self,
+ data,
+ *,
+ include=None,
+ exclude=None,
+ matcher=None,
+ ):
+ data = self._serialize(data)
data = self._filter(
data=data,
depth=0,
@@ -103,7 +195,8 @@ def serialize(
include=include,
matcher=matcher,
)
- return json.dumps(data, indent=2, ensure_ascii=False, sort_keys=False) + "\n"
+ data = json.dumps(data, indent=2, ensure_ascii=False, sort_keys=False) + "\n"
+ return data
def matches(
self,
@@ -119,7 +212,7 @@ def _convert_data(data):
if isinstance(data, Dict):
if "choices" in data:
data["choices"] = list(
- sorted(data["choices"], key=lambda x: x["index"])
+ sorted(data["choices"], key=lambda x: int(x["index"]))
)
choices = data["choices"]
if isinstance(choices, List) and len(choices) >= 1:
@@ -132,7 +225,7 @@ def _convert_data(data):
return Response(**data)
if isinstance(data, List):
return [_convert_data(d) for d in data]
- raise NotImplementedError
+ raise NotImplementedError(f"Data: {data}")
def eq_token(token: Token, other: Token) -> bool:
return (
@@ -230,7 +323,25 @@ def eq_chat_complete(response: ChatComplete, other: ChatComplete) -> bool:
def eq_chat_complete_chunk(
response: ChatCompletionChunk, other: ChatCompletionChunk
) -> bool:
- return response.choices[0].delta.content == other.choices[0].delta.content
+ if response.choices:
+ if response.choices[0].delta.content is not None:
+ return (
+ response.choices[0].delta.content
+ == other.choices[0].delta.content
+ )
+ elif response.choices[0].delta.tool_calls is not None:
+ return (
+ response.choices[0].delta.tool_calls
+ == other.choices[0].delta.tool_calls
+ )
+ else:
+ raise RuntimeError(
+ f"Invalid empty chat chunk {response} vs {other}"
+ )
+ elif response.usage is not None:
+ return response.usage == other.usage
+ else:
+ raise RuntimeError(f"Invalid empty chat {response} vs {other}")
def eq_response(response: Response, other: Response) -> bool:
return response.generated_text == other.generated_text and eq_details(
@@ -245,6 +356,9 @@ def eq_response(response: Response, other: Response) -> bool:
if not isinstance(snapshot_data, List):
snapshot_data = [snapshot_data]
+ if len(serialized_data) == 0:
+ return len(snapshot_data) == len(serialized_data)
+
if isinstance(serialized_data[0], Completion):
return len(snapshot_data) == len(serialized_data) and all(
[eq_completion(r, o) for r, o in zip(serialized_data, snapshot_data)]
@@ -278,8 +392,10 @@ class IgnoreLogProbResponseComparator(ResponseComparator):
class LauncherHandle:
- def __init__(self, port: int):
- self.client = AsyncClient(f"http://localhost:{port}", timeout=30)
+ def __init__(self, port: int, error_log):
+ with warnings.catch_warnings(action="http://github.com/ignore"):
+ self.client = AsyncClient(f"http://localhost:{port}", timeout=30)
+ self.error_log = error_log
def _inner_health(self):
raise NotImplementedError
@@ -288,6 +404,8 @@ async def health(self, timeout: int = 60):
assert timeout > 0
for _ in range(timeout):
if not self._inner_health():
+ self.error_log.seek(0)
+ print(self.error_log.read(), file=sys.stderr)
raise RuntimeError("Launcher crashed")
try:
@@ -295,12 +413,14 @@ async def health(self, timeout: int = 60):
return
except (ClientConnectorError, ClientOSError, ServerDisconnectedError):
time.sleep(1)
+ self.error_log.seek(0)
+ print(self.error_log.read(), file=sys.stderr)
raise RuntimeError("Health check failed")
class ContainerLauncherHandle(LauncherHandle):
- def __init__(self, docker_client, container_name, port: int):
- super(ContainerLauncherHandle, self).__init__(port)
+ def __init__(self, docker_client, container_name, port: int, error_log):
+ super().__init__(port, error_log)
self.docker_client = docker_client
self.container_name = container_name
@@ -310,8 +430,8 @@ def _inner_health(self) -> bool:
class ProcessLauncherHandle(LauncherHandle):
- def __init__(self, process, port: int):
- super(ProcessLauncherHandle, self).__init__(port)
+ def __init__(self, process, port: int, error_log):
+ super().__init__(port, error_log)
self.process = process
def _inner_health(self) -> bool:
@@ -333,15 +453,14 @@ def ignore_logprob_response_snapshot(snapshot):
return snapshot.use_extension(IgnoreLogProbResponseComparator)
-@pytest.fixture(scope="module")
-def event_loop():
- loop = asyncio.get_event_loop()
- yield loop
- loop.close()
+@pytest.fixture(scope="session")
+def error_log():
+ with tempfile.TemporaryFile("w+") as tmp:
+ yield tmp
-@pytest.fixture(scope="module")
-def launcher(event_loop):
+@pytest.fixture(scope="session")
+async def launcher(error_log):
@contextlib.contextmanager
def local_launcher(
model_id: str,
@@ -354,6 +473,7 @@ def local_launcher(
kv_cache_dtype: Optional[str] = None,
revision: Optional[str] = None,
max_input_length: Optional[int] = None,
+ max_input_tokens: Optional[int] = None,
max_batch_prefill_tokens: Optional[int] = None,
max_total_tokens: Optional[int] = None,
lora_adapters: Optional[List[str]] = None,
@@ -402,6 +522,9 @@ def local_launcher(
if max_input_length:
args.append("--max-input-length")
args.append(str(max_input_length))
+ if max_input_tokens:
+ args.append("--max-input-tokens")
+ args.append(str(max_input_tokens))
if max_batch_prefill_tokens:
args.append("--max-batch-prefill-tokens")
args.append(str(max_batch_prefill_tokens))
@@ -425,22 +548,19 @@ def local_launcher(
if attention is not None:
env["ATTENTION"] = attention
- with tempfile.TemporaryFile("w+") as tmp:
+ # with tempfile.TemporaryFile("w+") as tmp:
# We'll output stdout/stderr to a temporary file. Using a pipe
# cause the process to block until stdout is read.
- with subprocess.Popen(
- args,
- stdout=tmp,
- stderr=subprocess.STDOUT,
- env=env,
- ) as process:
- yield ProcessLauncherHandle(process, port)
+ with subprocess.Popen(
+ args,
+ stdout=error_log,
+ stderr=subprocess.STDOUT,
+ env=env,
+ ) as process:
+ yield ProcessLauncherHandle(process, port, error_log=error_log)
- process.terminate()
- process.wait(60)
-
- tmp.seek(0)
- shutil.copyfileobj(tmp, sys.stderr)
+ process.terminate()
+ process.wait(60)
if not use_flash_attention:
del env["USE_FLASH_ATTENTION"]
@@ -558,6 +678,7 @@ def docker_launcher(
docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
]
+ client.api.timeout = 1000
container = client.containers.run(
DOCKER_IMAGE,
command=args,
@@ -569,12 +690,25 @@ def docker_launcher(
devices=devices,
volumes=volumes,
ports={"80/tcp": port},
- healthcheck={"timeout": int(60 * 1e9), "retries": 2}, # 60s
+ healthcheck={"timeout": int(180 * 1e9), "retries": 2}, # 60s
shm_size="1G",
)
+ def pipe():
+ for log in container.logs(stream=True):
+ log = log.decode("utf-8")
+ error_log.write(log)
+
+ # Start looping to pipe the logs
+ import threading
+
+ t = threading.Thread(target=pipe, args=())
+ t.start()
+
try:
- yield ContainerLauncherHandle(client, container.name, port)
+ yield ContainerLauncherHandle(
+ client, container.name, port, error_log=error_log
+ )
if not use_flash_attention:
del env["USE_FLASH_ATTENTION"]
@@ -585,9 +719,6 @@ def docker_launcher(
except NotFound:
pass
- container_output = container.logs().decode("utf-8")
- print(container_output, file=sys.stderr)
-
finally:
try:
container.remove()
diff --git a/integration-tests/fixtures/gaudi/service.py b/integration-tests/fixtures/gaudi/service.py
new file mode 100644
index 00000000000..f4f43691cd9
--- /dev/null
+++ b/integration-tests/fixtures/gaudi/service.py
@@ -0,0 +1,311 @@
+import asyncio
+import contextlib
+import os
+import shlex
+import subprocess
+import sys
+import threading
+import time
+from tempfile import TemporaryDirectory
+from typing import List
+import socket
+
+import docker
+import pytest
+from aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError
+from docker.errors import NotFound
+import logging
+from huggingface_hub import AsyncInferenceClient, TextGenerationOutput
+import huggingface_hub
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
+ stream=sys.stdout,
+)
+logger = logging.getLogger(__file__)
+
+# Use the latest image from the local docker build
+DOCKER_IMAGE = os.getenv("DOCKER_IMAGE", "tgi-gaudi")
+DOCKER_VOLUME = os.getenv("DOCKER_VOLUME", None)
+HF_TOKEN = huggingface_hub.get_token()
+
+assert (
+ HF_TOKEN is not None
+), "HF_TOKEN is not set, please set it as some models are gated and thus the test will fail without it"
+
+if DOCKER_VOLUME is None:
+ logger.warning(
+ "DOCKER_VOLUME is not set, this will lead to the tests redownloading the models on each run, consider setting it to speed up testing"
+ )
+
+LOG_LEVEL = os.getenv("LOG_LEVEL", "info")
+
+BASE_ENV = {
+ "HF_HUB_ENABLE_HF_TRANSFER": "1",
+ "LOG_LEVEL": LOG_LEVEL,
+ "HF_TOKEN": os.getenv("HF_TOKEN", None),
+}
+
+
+HABANA_RUN_ARGS = {
+ "runtime": "habana",
+ "ipc_mode": "host",
+ "cap_add": ["sys_nice"],
+}
+
+
+def stream_container_logs(container, test_name):
+ """Stream container logs in a separate thread."""
+ try:
+ for log in container.logs(stream=True, follow=True):
+ print(
+ f"[TGI Server Logs - {test_name}] {log.decode('utf-8')}",
+ end="",
+ file=sys.stderr,
+ flush=True,
+ )
+ except Exception as e:
+ logger.error(f"Error streaming container logs: {str(e)}")
+
+
+class TestClient(AsyncInferenceClient):
+ def __init__(self, service_name: str, base_url: str):
+ super().__init__(model=base_url)
+ self.service_name = service_name
+
+
+class LauncherHandle:
+ def __init__(self, service_name: str, port: int):
+ self.client = TestClient(service_name, f"http://localhost:{port}")
+
+ def _inner_health(self):
+ raise NotImplementedError
+
+ async def health(self, timeout: int = 60):
+ assert timeout > 0
+ start_time = time.time()
+ logger.info(f"Starting health check with timeout of {timeout}s")
+
+ for attempt in range(timeout):
+ if not self._inner_health():
+ logger.error("Launcher crashed during health check")
+ raise RuntimeError("Launcher crashed")
+
+ try:
+ await self.client.text_generation("test", max_new_tokens=1)
+ elapsed = time.time() - start_time
+ logger.info(f"Health check passed after {elapsed:.1f}s")
+ return
+ except (ClientConnectorError, ClientOSError, ServerDisconnectedError) as e:
+ if attempt == timeout - 1:
+ logger.error(f"Health check failed after {timeout}s: {str(e)}")
+ raise RuntimeError(f"Health check failed: {str(e)}")
+ if attempt % 10 == 0 and attempt != 0: # Only log every 10th attempt
+ logger.debug(
+ f"Connection attempt {attempt}/{timeout} failed: {str(e)}"
+ )
+ time.sleep(1)
+ except Exception as e:
+ logger.error(f"Unexpected error during health check: {str(e)}")
+ # Get full traceback for debugging
+ import traceback
+
+ logger.error(f"Full traceback:\n{traceback.format_exc()}")
+ raise
+
+
+class ContainerLauncherHandle(LauncherHandle):
+ def __init__(self, docker_client, container_name, port: int):
+ service_name = container_name # Use container name as service name
+ super(ContainerLauncherHandle, self).__init__(service_name, port)
+ self.docker_client = docker_client
+ self.container_name = container_name
+
+ def _inner_health(self) -> bool:
+ try:
+ container = self.docker_client.containers.get(self.container_name)
+ status = container.status
+ if status not in ["running", "created"]:
+ logger.warning(f"Container status is {status}")
+ # Get container logs for debugging
+ logs = container.logs().decode("utf-8")
+ logger.debug(f"Container logs:\n{logs}")
+ return status in ["running", "created"]
+ except Exception as e:
+ logger.error(f"Error checking container health: {str(e)}")
+ return False
+
+
+class ProcessLauncherHandle(LauncherHandle):
+ def __init__(self, process, port: int):
+ service_name = "process" # Use generic name for process launcher
+ super(ProcessLauncherHandle, self).__init__(service_name, port)
+ self.process = process
+
+ def _inner_health(self) -> bool:
+ return self.process.poll() is None
+
+
+@pytest.fixture(scope="module")
+def data_volume():
+ tmpdir = TemporaryDirectory()
+ yield tmpdir.name
+ try:
+ # Cleanup the temporary directory using sudo as it contains root files created by the container
+ subprocess.run(shlex.split(f"sudo rm -rf {tmpdir.name}"), check=True)
+ except subprocess.CalledProcessError as e:
+ logger.error(f"Error cleaning up temporary directory: {str(e)}")
+
+
+@pytest.fixture(scope="module")
+def gaudi_launcher():
+ @contextlib.contextmanager
+ def docker_launcher(
+ model_id: str,
+ test_name: str,
+ tgi_args: List[str] = None,
+ env_config: dict = None,
+ ):
+ logger.info(
+ f"Starting docker launcher for model {model_id} and test {test_name}"
+ )
+
+ # Get a random available port
+ def get_free_port():
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+ s.bind(("", 0))
+ s.listen(1)
+ port = s.getsockname()[1]
+ return port
+
+ port = get_free_port()
+ logger.debug(f"Using port {port}")
+
+ client = docker.from_env()
+
+ container_name = f"tgi-gaudi-test-{test_name.replace('/', '-')}"
+
+ try:
+ container = client.containers.get(container_name)
+ logger.info(
+ f"Stopping existing container {container_name} for test {test_name}"
+ )
+ container.stop()
+ container.wait()
+ container.remove()
+ logger.info(f"Removed existing container {container_name}")
+ except NotFound:
+ pass
+ except Exception as e:
+ logger.error(f"Error handling existing container: {str(e)}")
+
+ if tgi_args is None:
+ tgi_args = []
+ else:
+ tgi_args = tgi_args.copy()
+
+ env = BASE_ENV.copy()
+
+ # Add model_id to env
+ env["MODEL_ID"] = model_id
+
+ # Add env config that is defined in the fixture parameter
+ if env_config is not None:
+ env.update(env_config.copy())
+
+ volumes = []
+ if DOCKER_VOLUME:
+ volumes = [f"{DOCKER_VOLUME}:/data"]
+ logger.debug(f"Using volume {volumes}")
+
+ try:
+ logger.debug(f"Using command {tgi_args}")
+ logger.info(f"Creating container with name {container_name}")
+
+ logger.debug(f"Using environment {env}")
+ logger.debug(f"Using volumes {volumes}")
+ logger.debug(f"HABANA_RUN_ARGS {HABANA_RUN_ARGS}")
+
+ # Log equivalent docker run command for debugging, this is not actually executed
+ container = client.containers.run(
+ DOCKER_IMAGE,
+ command=tgi_args,
+ name=container_name,
+ environment=env,
+ detach=True,
+ volumes=volumes,
+ ports={"80/tcp": port},
+ **HABANA_RUN_ARGS,
+ )
+
+ logger.info(f"Container {container_name} started successfully")
+
+ # Start log streaming in a background thread
+ log_thread = threading.Thread(
+ target=stream_container_logs,
+ args=(container, test_name),
+ daemon=True, # This ensures the thread will be killed when the main program exits
+ )
+ log_thread.start()
+
+ # Add a small delay to allow container to initialize
+ time.sleep(2)
+
+ # Check container status after creation
+ status = container.status
+ logger.debug(f"Initial container status: {status}")
+ if status not in ["running", "created"]:
+ logs = container.logs().decode("utf-8")
+ logger.error(f"Container failed to start properly. Logs:\n{logs}")
+
+ yield ContainerLauncherHandle(client, container.name, port)
+
+ except Exception as e:
+ logger.error(f"Error starting container: {str(e)}")
+ # Get full traceback for debugging
+ import traceback
+
+ logger.error(f"Full traceback:\n{traceback.format_exc()}")
+ raise
+ finally:
+ try:
+ container = client.containers.get(container_name)
+ logger.info(f"Stopping container {container_name}")
+ container.stop()
+ container.wait()
+
+ container_output = container.logs().decode("utf-8")
+ print(container_output, file=sys.stderr)
+
+ container.remove()
+ logger.info(f"Container {container_name} removed successfully")
+ except NotFound:
+ pass
+ except Exception as e:
+ logger.warning(f"Error cleaning up container: {str(e)}")
+
+ return docker_launcher
+
+
+@pytest.fixture(scope="module")
+def gaudi_generate_load():
+ async def generate_load_inner(
+ client: AsyncInferenceClient, prompt: str, max_new_tokens: int, n: int
+ ) -> List[TextGenerationOutput]:
+ try:
+ futures = [
+ client.text_generation(
+ prompt,
+ max_new_tokens=max_new_tokens,
+ details=True,
+ decoder_input_details=True,
+ )
+ for _ in range(n)
+ ]
+ return await asyncio.gather(*futures)
+ except Exception as e:
+ logger.error(f"Error generating load: {str(e)}")
+ raise
+
+ return generate_load_inner
diff --git a/integration-tests/fixtures/neuron/export_models.py b/integration-tests/fixtures/neuron/export_models.py
new file mode 100644
index 00000000000..beee2ba7127
--- /dev/null
+++ b/integration-tests/fixtures/neuron/export_models.py
@@ -0,0 +1,274 @@
+import copy
+import logging
+import sys
+from tempfile import TemporaryDirectory
+
+import huggingface_hub
+import pytest
+import docker
+import hashlib
+import os
+import tempfile
+
+from docker.errors import NotFound
+
+
+TEST_ORGANIZATION = "optimum-internal-testing"
+TEST_CACHE_REPO_ID = f"{TEST_ORGANIZATION}/neuron-testing-cache"
+HF_TOKEN = huggingface_hub.get_token()
+
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="[%(asctime)s] %(levelname)s [%(filename)s.%(funcName)s:%(lineno)d] %(message)s",
+ stream=sys.stdout,
+)
+logger = logging.getLogger(__file__)
+
+
+# All model configurations below will be added to the neuron_model_config fixture
+MODEL_CONFIGURATIONS = {
+ "llama": {
+ "model_id": "unsloth/Llama-3.2-1B-Instruct",
+ "export_kwargs": {
+ "batch_size": 4,
+ "sequence_length": 2048,
+ "num_cores": 2,
+ "auto_cast_type": "fp16",
+ },
+ },
+ "qwen2": {
+ "model_id": "Qwen/Qwen2.5-0.5B",
+ "export_kwargs": {
+ "batch_size": 4,
+ "sequence_length": 4096,
+ "num_cores": 2,
+ "auto_cast_type": "fp16",
+ },
+ },
+ "qwen3": {
+ "model_id": "Qwen/Qwen3-1.7B",
+ "export_kwargs": {
+ "batch_size": 4,
+ "sequence_length": 4096,
+ "num_cores": 2,
+ "auto_cast_type": "bf16",
+ },
+ },
+ "granite": {
+ "model_id": "ibm-granite/granite-3.1-2b-instruct",
+ "export_kwargs": {
+ "batch_size": 4,
+ "sequence_length": 4096,
+ "num_cores": 2,
+ "auto_cast_type": "bf16",
+ },
+ },
+ "phi3": {
+ "model_id": "microsoft/Phi-3-mini-4k-instruct",
+ "export_kwargs": {
+ "batch_size": 4,
+ "sequence_length": 4096,
+ "num_cores": 2,
+ "auto_cast_type": "bf16",
+ },
+ },
+}
+
+
+def get_neuron_backend_hash():
+ import subprocess
+
+ res = subprocess.run(
+ ["git", "rev-parse", "--show-toplevel"], capture_output=True, text=True
+ )
+ root_dir = res.stdout.split("\n")[0]
+
+ def get_sha(path):
+ res = subprocess.run(
+ ["git", "ls-tree", "HEAD", f"{root_dir}/{path}"],
+ capture_output=True,
+ text=True,
+ )
+ # Output of the command is in the form '040000 tree|blob \t\n'
+ sha = res.stdout.split("\t")[0].split(" ")[-1]
+ return sha.encode()
+
+ # We hash both the neuron backends directory and Dockerfile and create a smaller hash out of that
+ m = hashlib.sha256()
+ m.update(get_sha("backends/neuron"))
+ m.update(get_sha("Dockerfile.neuron"))
+ return m.hexdigest()[:10]
+
+
+def get_neuron_model_name(config_name: str):
+ return f"neuron-tgi-testing-{config_name}-{get_neuron_backend_hash()}"
+
+
+def get_tgi_docker_image():
+ docker_image = os.getenv("DOCKER_IMAGE", None)
+ if docker_image is None:
+ client = docker.from_env()
+ images = client.images.list(filters={"reference": "text-generation-inference"})
+ if not images:
+ raise ValueError(
+ "No text-generation-inference image found on this host to run tests."
+ )
+ docker_image = images[0].tags[0]
+ return docker_image
+
+
+def maybe_export_model(config_name, model_config):
+ """Export a neuron model for the specified test configuration.
+
+ If the neuron model has not already been compiled and pushed to the hub, it is
+ exported by a custom image built on the fly from the base TGI image.
+ This makes sure the exported model and image are aligned and avoids introducing
+ neuron specific imports in the test suite.
+
+ Args:
+ config_name (`str`):
+ Used to identify test configurations
+ model_config (`str`):
+ The model configuration for export (includes the original model id)
+ """
+ neuron_model_name = get_neuron_model_name(config_name)
+ neuron_model_id = f"{TEST_ORGANIZATION}/{neuron_model_name}"
+ hub = huggingface_hub.HfApi()
+ if hub.repo_exists(neuron_model_id):
+ logger.info(
+ f"Skipping model export for config {config_name} as {neuron_model_id} already exists"
+ )
+ return neuron_model_id
+
+ client = docker.from_env()
+
+ env = {"LOG_LEVEL": "info", "CUSTOM_CACHE_REPO": TEST_CACHE_REPO_ID}
+ if HF_TOKEN is not None:
+ env["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
+ env["HF_TOKEN"] = HF_TOKEN
+
+ # Create a sub-image to export the model to workaround docker dind issues preventing
+ # to share a volume from the container running tests
+ model_id = model_config["model_id"]
+ export_kwargs = model_config["export_kwargs"]
+ base_image = get_tgi_docker_image()
+ export_image = f"neuron-tgi-tests-{config_name}-export-img"
+ logger.info(f"Building temporary image {export_image} from {base_image}")
+ with tempfile.TemporaryDirectory() as context_dir:
+ # Create entrypoint
+ model_path = "/data/neuron_model"
+ export_command = (
+ f"optimum-cli export neuron -m {model_id} --task text-generation"
+ )
+ for kwarg, value in export_kwargs.items():
+ export_command += f" --{kwarg} {str(value)}"
+ export_command += f" {model_path}"
+ entrypoint_content = f"""#!/bin/sh
+ {export_command}
+ huggingface-cli repo create --organization {TEST_ORGANIZATION} {neuron_model_name}
+ huggingface-cli upload {TEST_ORGANIZATION}/{neuron_model_name} {model_path} --exclude *.bin *.safetensors
+ optimum-cli neuron cache synchronize --repo_id {TEST_CACHE_REPO_ID}
+ """
+ with open(os.path.join(context_dir, "entrypoint.sh"), "wb") as f:
+ f.write(entrypoint_content.encode("utf-8"))
+ f.flush()
+ # Create Dockerfile
+ docker_content = f"""
+ FROM {base_image}
+ COPY entrypoint.sh /export-entrypoint.sh
+ RUN chmod +x /export-entrypoint.sh
+ ENTRYPOINT ["/export-entrypoint.sh"]
+ """
+ with open(os.path.join(context_dir, "Dockerfile"), "wb") as f:
+ f.write(docker_content.encode("utf-8"))
+ f.flush()
+ image, logs = client.images.build(
+ path=context_dir, dockerfile=f.name, tag=export_image
+ )
+ logger.info("Successfully built image %s", image.id)
+ logger.debug("Build logs %s", logs)
+
+ try:
+ client.containers.run(
+ export_image,
+ environment=env,
+ auto_remove=True,
+ detach=False,
+ devices=["/dev/neuron0"],
+ shm_size="1G",
+ )
+ logger.info(f"Successfully exported model for config {config_name}")
+ except Exception as e:
+ logger.exception(f"An exception occurred while running container: {e}.")
+ pass
+ finally:
+ # Cleanup the export image
+ logger.info("Cleaning image %s", image.id)
+ try:
+ image.remove(force=True)
+ except NotFound:
+ pass
+ except Exception as e:
+ logger.error("Error while removing image %s, skipping", image.id)
+ logger.exception(e)
+ return neuron_model_id
+
+
+def maybe_export_models():
+ for config_name, model_config in MODEL_CONFIGURATIONS.items():
+ maybe_export_model(config_name, model_config)
+
+
+@pytest.fixture(scope="session", params=MODEL_CONFIGURATIONS.keys())
+def neuron_model_config(request):
+ """Expose a pre-trained neuron model
+
+ The fixture first makes sure the following model artifacts are present on the hub:
+ - exported neuron model under optimum-internal-testing/neuron-testing-- 
+