Fix env (docker) compat due to __file__ usage (#288)

This commit is contained in:
Qubitium
2024-03-13 13:02:48 +08:00
committed by GitHub
parent d9b3b01883
commit 92e2d74fd0

View File

@@ -4,6 +4,7 @@ import inspect
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
import importlib.resources
import numpy as np
import torch
@@ -31,10 +32,12 @@ global_server_args_dict: dict = None
@lru_cache()
def import_model_classes():
model_arch_name_to_cls = {}
for module_path in (Path(sglang.__file__).parent / "srt" / "models").glob("*.py"):
module = importlib.import_module(f"sglang.srt.models.{module_path.stem}")
if hasattr(module, "EntryClass"):
model_arch_name_to_cls[module.EntryClass.__name__] = module.EntryClass
for f in importlib.resources.files("sglang.srt.models").iterdir():
if f.name.endswith(".py"):
module_name = Path(f.name).with_suffix('')
module = importlib.import_module(f"sglang.srt.models.{module_name}")
if hasattr(module, "EntryClass"):
model_arch_name_to_cls[module.EntryClass.__name__] = module.EntryClass
return model_arch_name_to_cls