[Fix] Register custom ops only if they exist (#13321)

This commit is contained in:
Lianmin Zheng
2025-11-15 17:48:17 -08:00
committed by GitHub
parent 13366843a4
commit db7299aa30
6 changed files with 47 additions and 34 deletions

View File

@@ -683,7 +683,6 @@ class Engine(EngineBase):
def _set_envs_and_config(server_args: ServerArgs):
# Set global environments
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
if "NCCL_CUMEM_ENABLE" not in os.environ or server_args.enable_symm_mem:
os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem))
if (
@@ -756,10 +755,13 @@ def _set_envs_and_config(server_args: ServerArgs):
def _init_tokenizer_manager(
server_args: ServerArgs, port_args: PortArgs
server_args: ServerArgs,
port_args: PortArgs,
TokenizerManagerClass: Optional[TokenizerManager] = None,
) -> TokenizerManager:
# Launch tokenizer process
tokenizer_manager = TokenizerManager(server_args, port_args)
TokenizerManagerClass = TokenizerManagerClass or TokenizerManager
tokenizer_manager = TokenizerManagerClass(server_args, port_args)
# Initialize templates
template_manager = TemplateManager()