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): ...@@ -292,7 +292,6 @@ class Worker(WorkerBase):
allocator = CuMemAllocator.get_instance() allocator = CuMemAllocator.get_instance()
context = allocator.use_memory_pool(tag="kv_cache") context = allocator.use_memory_pool(tag="kv_cache")
else: else:
from contextlib import nullcontext
context = nullcontext() context = nullcontext()
with context: with context:
self.model_runner.initialize_kv_cache(kv_cache_config) self.model_runner.initialize_kv_cache(kv_cache_config)
......
...@@ -3,6 +3,7 @@ ...@@ -3,6 +3,7 @@
"""A GPU worker class.""" """A GPU worker class."""
import gc import gc
import os import os
from contextlib import nullcontext
from typing import Dict, List, Optional, Set, Tuple, Type, Union from typing import Dict, List, Optional, Set, Tuple, Type, Union
import torch import torch
...@@ -206,7 +207,6 @@ class Worker(LocalOrDistributedWorkerBase): ...@@ -206,7 +207,6 @@ class Worker(LocalOrDistributedWorkerBase):
"used for one instance per process.") "used for one instance per process.")
context = allocator.use_memory_pool(tag="weights") context = allocator.use_memory_pool(tag="weights")
else: else:
from contextlib import nullcontext
context = nullcontext() context = nullcontext()
with context: with context:
self.model_runner.load_model() self.model_runner.load_model()
...@@ -330,7 +330,6 @@ class Worker(LocalOrDistributedWorkerBase): ...@@ -330,7 +330,6 @@ class Worker(LocalOrDistributedWorkerBase):
allocator = CuMemAllocator.get_instance() allocator = CuMemAllocator.get_instance()
context = allocator.use_memory_pool(tag="kv_cache") context = allocator.use_memory_pool(tag="kv_cache")
else: else:
from contextlib import nullcontext
context = nullcontext() context = nullcontext()
with context: with context:
self._init_cache_engine() 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