Unverified Commit 325aa3de authored by Ning Xie's avatar Ning Xie Committed by GitHub
Browse files

[Misc] local import code clean (#23420)


Signed-off-by: default avatarAndy Xie <andy.xning@gmail.com>
parent a073be6d
......@@ -292,7 +292,6 @@ class Worker(WorkerBase):
allocator = CuMemAllocator.get_instance()
context = allocator.use_memory_pool(tag="kv_cache")
else:
from contextlib import nullcontext
context = nullcontext()
with context:
self.model_runner.initialize_kv_cache(kv_cache_config)
......
......@@ -3,6 +3,7 @@
"""A GPU worker class."""
import gc
import os
from contextlib import nullcontext
from typing import Dict, List, Optional, Set, Tuple, Type, Union
import torch
......@@ -206,7 +207,6 @@ class Worker(LocalOrDistributedWorkerBase):
"used for one instance per process.")
context = allocator.use_memory_pool(tag="weights")
else:
from contextlib import nullcontext
context = nullcontext()
with context:
self.model_runner.load_model()
......@@ -330,7 +330,6 @@ class Worker(LocalOrDistributedWorkerBase):
allocator = CuMemAllocator.get_instance()
context = allocator.use_memory_pool(tag="kv_cache")
else:
from contextlib import nullcontext
context = nullcontext()
with context:
self._init_cache_engine()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment