[CI/CD] add neuron docker and ci test scripts (#3571)

2024-04-18 15:26:01 -07:00
parent 87fa80c91f
commit cd2f63fb36
6 changed files with 103 additions and 4 deletions
--- a/setup.py
+++ b/setup.py
@ -204,7 +204,8 @@ def _is_neuron() -> bool:
        subprocess.run(["neuron-ls"], capture_output=True, check=True)
    except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
        torch_neuronx_installed = False
-    return torch_neuronx_installed
+    return torch_neuronx_installed or os.environ.get("VLLM_BUILD_WITH_NEURON",
+                                                     False)


 def _is_cpu() -> bool: