From 6f8687f0445cff1ae741377819ab4c35605338a6 Mon Sep 17 00:00:00 2001 From: Qubitium <417764+Qubitium@users.noreply.github.com> Date: Wed, 13 Mar 2024 02:25:42 +0000 Subject: [PATCH] Use importlib.resources instead of __file__ to resolve some env (docker) compat --- python/sglang/srt/managers/router/model_runner.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/managers/router/model_runner.py b/python/sglang/srt/managers/router/model_runner.py index 1f07286ed866..452979f8d7ca 100644 --- a/python/sglang/srt/managers/router/model_runner.py +++ b/python/sglang/srt/managers/router/model_runner.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from functools import lru_cache from pathlib import Path +import importlib.resources import numpy as np import torch @@ -29,10 +30,12 @@ @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