4
4
import subprocess
5
5
import uuid
6
6
import gc
7
+ from functools import cache
7
8
from platform import uname
8
9
from typing import List , Tuple , Union
9
10
from packaging .version import parse , Version
@@ -120,6 +121,7 @@ def is_hip() -> bool:
120
121
return torch .version .hip is not None
121
122
122
123
124
+ @cache
123
125
def is_neuron () -> bool :
124
126
try :
125
127
import transformers_neuronx
@@ -128,6 +130,7 @@ def is_neuron() -> bool:
128
130
return transformers_neuronx is not None
129
131
130
132
133
+ @cache
131
134
def get_max_shared_memory_bytes (gpu : int = 0 ) -> int :
132
135
"""Returns the maximum shared memory per thread block in bytes."""
133
136
# NOTE: This import statement should be executed lazily since
@@ -151,6 +154,7 @@ def random_uuid() -> str:
151
154
return str (uuid .uuid4 ().hex )
152
155
153
156
157
+ @cache
154
158
def in_wsl () -> bool :
155
159
# Reference: https://github.com/microsoft/WSL/issues/4071
156
160
return "microsoft" in " " .join (uname ()).lower ()
@@ -225,6 +229,7 @@ def set_cuda_visible_devices(device_ids: List[int]) -> None:
225
229
os .environ ["CUDA_VISIBLE_DEVICES" ] = "," .join (map (str , device_ids ))
226
230
227
231
232
+ @cache
228
233
def get_nvcc_cuda_version () -> Optional [Version ]:
229
234
cuda_home = os .environ .get ('CUDA_HOME' )
230
235
if not cuda_home :
0 commit comments