mirror/OpenLLM
1
0
Fork 0
mirror of https://github.com/bentoml/OpenLLM.git synced 2026-03-07 16:47:13 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix: device imports using strategies (#584)

Browse Source 
* fix: device imports using strategies

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

* chore: support trust_remote_code for vLLM runners

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

---------

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
This commit is contained in:
Aaron Pham
2023-11-08 05:10:50 -05:00
committed by GitHub
parent 6d81fbbccd
commit 85a7243ac3
6 changed files with 16 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
openllm-python/src/openllm/utils/__init__.py
View File

@@ -6,6 +6,7 @@ we won't ensure backward compatibility for these functions. So use with caution.
from __future__ import annotations
import typing as t
import functools
import openllm_core
if t.TYPE_CHECKING:
@@ -14,7 +15,16 @@ if t.TYPE_CHECKING:
def generate_labels(llm: openllm.LLM[t.Any, t.Any]) -> dict[str, t.Any]:
return {'backend': llm.__llm_backend__, 'framework': 'openllm', 'model_name': llm.config['model_name'], 'architecture': llm.config['architecture'], 'serialisation': llm._serialisation}
__all__ = ['generate_labels']
def available_devices() -> tuple[str, ...]:
"""Return available GPU under system. Currently only supports NVIDIA GPUs."""
from .._strategies import NvidiaGpuResource
return tuple(NvidiaGpuResource.from_system())
@functools.lru_cache(maxsize=1)
def device_count() -> int:
return len(available_devices())
__all__ = ['generate_labels', 'available_devices', 'device_count']
def __dir__() -> t.Sequence[str]:
return sorted(__all__)