add env var

.github/workflows/build.yaml

@@ -11,6 +11,10 @@ on:
     branches:
       - 'main'
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
 jobs:
   build-and-push-image:
     runs-on: ubuntu-latest

.github/workflows/tests.yaml

@@ -9,6 +9,10 @@ on:
       - "router/**"
       - "launcher/**"
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
 jobs:
   run_tests:
     runs-on: ubuntu-20.04

server/text_generation_server/models/__init__.py

@@ -1,5 +1,7 @@
+import os
 import torch
 
+from loguru import logger
 from transformers import AutoConfig
 from typing import Optional
 
@@ -14,9 +16,10 @@ from text_generation_server.models.t5 import T5Sharded
 
 try:
     from text_generation_server.models.flash_neox import FlashNeoX, FlashNeoXSharded
-
-    FLASH_NEOX = torch.cuda.is_available()
+    FLASH_NEOX = torch.cuda.is_available() and int(os.environ.get("FLASH_NEOX", 0)) == 1
 except ImportError:
+    if int(os.environ.get("FLASH_NEOX", 0)) == 1:
+        logger.exception("Could not import FlashNeoX")
     FLASH_NEOX = False
 
 __all__ = [