Skip to content

GitLab

Unverified Commit e99ba957 authored by Aaron Hao's avatar Aaron Hao Committed by GitHub
Browse files

[BUG] Fixing Weight Sync unit test (#34841) 


Signed-off-by: default avatarahao-anyscale <ahao@anyscale.com>
parent 64ac1395
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 0 deletions
tests/entrypoints/weight_transfer/test_weight_transfer_llm.py
View file @ e99ba957
...@@ -124,6 +124,8 @@ def test_init_weight_transfer_engine_calls_engine(): ...@@ -124,6 +124,8 @@ def test_init_weight_transfer_engine_calls_engine():
if torch.cuda.device_count() < 1: if torch.cuda.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test") pytest.skip("Need at least 1 GPU for this test")
# Run in-process so mock.patch works (spawn won't inherit the mock)
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
# Enable insecure serialization to allow pickling functions for collective_rpc # Enable insecure serialization to allow pickling functions for collective_rpc
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
...@@ -171,6 +173,8 @@ def test_update_weights_calls_engine(): ...@@ -171,6 +173,8 @@ def test_update_weights_calls_engine():
if torch.cuda.device_count() < 1: if torch.cuda.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test") pytest.skip("Need at least 1 GPU for this test")
# Run in-process so mock.patch works (spawn won't inherit the mock)
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
# Enable insecure serialization to allow pickling functions for collective_rpc # Enable insecure serialization to allow pickling functions for collective_rpc
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
...@@ -228,6 +232,8 @@ def test_full_weight_transfer_flow(): ...@@ -228,6 +232,8 @@ def test_full_weight_transfer_flow():
if torch.cuda.device_count() < 1: if torch.cuda.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test") pytest.skip("Need at least 1 GPU for this test")
# Run in-process so mock.patch works (spawn won't inherit the mock)
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
# Enable insecure serialization to allow pickling functions for collective_rpc # Enable insecure serialization to allow pickling functions for collective_rpc
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
......
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