Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
2bfdbf2a
Unverified
Commit
2bfdbf2a
authored
Dec 17, 2024
by
Tyler Michael Smith
Committed by
GitHub
Dec 16, 2024
Browse files
[V1][Core] Use weakref.finalize instead of atexit (#11242)
Signed-off-by:
Tyler Michael Smith
<
tyler@neuralmagic.com
>
parent
e88db68c
Changes
2
Show whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
5 additions
and
18 deletions
+5
-18
vllm/v1/engine/core_client.py
vllm/v1/engine/core_client.py
+2
-11
vllm/v1/executor/multiproc_executor.py
vllm/v1/executor/multiproc_executor.py
+3
-7
No files found.
vllm/v1/engine/core_client.py
View file @
2bfdbf2a
import
atexit
import
os
import
os
import
weakref
from
typing
import
List
,
Optional
from
typing
import
List
,
Optional
import
msgspec
import
msgspec
...
@@ -165,15 +165,9 @@ class MPClient(EngineCoreClient):
...
@@ -165,15 +165,9 @@ class MPClient(EngineCoreClient):
ready_path
=
ready_path
,
# type: ignore[misc]
ready_path
=
ready_path
,
# type: ignore[misc]
**
kwargs
,
**
kwargs
,
)
)
atexit
.
register
(
self
.
shutdown
)
self
.
_finalizer
=
weakref
.
finalize
(
self
,
self
.
shutdown
)
def
shutdown
(
self
):
def
shutdown
(
self
):
# During final garbage collection in process shutdown, atexit may be
# None.
if
atexit
:
# in case shutdown gets called via __del__ first
atexit
.
unregister
(
self
.
shutdown
)
# Shut down the zmq context.
# Shut down the zmq context.
self
.
ctx
.
destroy
(
linger
=
0
)
self
.
ctx
.
destroy
(
linger
=
0
)
...
@@ -197,9 +191,6 @@ class MPClient(EngineCoreClient):
...
@@ -197,9 +191,6 @@ class MPClient(EngineCoreClient):
os
.
remove
(
socket_file
)
os
.
remove
(
socket_file
)
self
.
proc_handle
=
None
self
.
proc_handle
=
None
def
__del__
(
self
):
self
.
shutdown
()
class
SyncMPClient
(
MPClient
):
class
SyncMPClient
(
MPClient
):
"""Synchronous client for multi-proc EngineCore."""
"""Synchronous client for multi-proc EngineCore."""
...
...
vllm/v1/executor/multiproc_executor.py
View file @
2bfdbf2a
import
atexit
import
os
import
os
import
pickle
import
pickle
import
signal
import
signal
import
sys
import
sys
import
time
import
time
import
weakref
from
dataclasses
import
dataclass
from
dataclasses
import
dataclass
from
enum
import
Enum
,
auto
from
enum
import
Enum
,
auto
from
multiprocessing.process
import
BaseProcess
from
multiprocessing.process
import
BaseProcess
...
@@ -37,7 +37,7 @@ class MultiprocExecutor(Executor):
...
@@ -37,7 +37,7 @@ class MultiprocExecutor(Executor):
def
__init__
(
self
,
vllm_config
:
VllmConfig
)
->
None
:
def
__init__
(
self
,
vllm_config
:
VllmConfig
)
->
None
:
# Call self.shutdown at exit to clean up
# Call self.shutdown at exit to clean up
# and ensure workers will be terminated.
# and ensure workers will be terminated.
atexit
.
register
(
self
.
shutdown
)
self
.
_finalizer
=
weakref
.
finalize
(
self
,
self
.
shutdown
)
self
.
vllm_config
=
vllm_config
self
.
vllm_config
=
vllm_config
self
.
parallel_config
=
vllm_config
.
parallel_config
self
.
parallel_config
=
vllm_config
.
parallel_config
...
@@ -195,14 +195,10 @@ class MultiprocExecutor(Executor):
...
@@ -195,14 +195,10 @@ class MultiprocExecutor(Executor):
os
.
remove
(
socket_path
)
os
.
remove
(
socket_path
)
def
shutdown
(
self
):
def
shutdown
(
self
):
if
atexit
:
# in case shutdown was called explicitly, we don't need to call it
# again
atexit
.
unregister
(
self
.
shutdown
)
"""Properly shut down the executor and its workers"""
"""Properly shut down the executor and its workers"""
if
getattr
(
self
,
'shutting_down'
,
False
):
if
getattr
(
self
,
'shutting_down'
,
False
):
self
.
shutting_down
=
True
self
.
shutting_down
=
True
for
w
in
self
.
workers
:
#TODO: not sure if needed
for
w
in
self
.
workers
:
w
.
worker_response_mq
=
None
w
.
worker_response_mq
=
None
self
.
_ensure_worker_termination
()
self
.
_ensure_worker_termination
()
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment