Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use weights_only for load #3073

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Use weights_only for load #3073

wants to merge 4 commits into from

Conversation

kit1980
Copy link
Member

@kit1980 kit1980 commented Apr 3, 2024

Description

torch.load without weights_only parameter is unsafe. Explicitly set weights_only to False only if you trust the data you load and full pickle functionality is needed, otherwise set weights_only=True.

If weights_only=True doesn't work for some cases, then explicit weights_only=False should be used.

Found with https://github.com/pytorch-labs/torchfix/

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Feature/Issue validation/testing

This should be a no-op.

@msaroufim msaroufim requested a review from agunapal April 3, 2024 23:02
@kit1980
Copy link
Member Author

kit1980 commented Apr 3, 2024

The actual changes are in the first commit.
The rest are to fix unrelated pre-existing issues in the files to pass lint.

@kit1980 kit1980 marked this pull request as ready for review April 3, 2024 23:32
@agunapal
Copy link
Collaborator

agunapal commented Apr 4, 2024

CI should pass..not sure why it was failing.. Hopefully works this time

@agunapal
Copy link
Collaborator

agunapal commented Apr 4, 2024

CI should pass..not sure why it was failing.. Hopefully works this time

@msaroufim This needs some debugging. Its failing everytime

@msaroufim
Copy link
Member

msaroufim commented Apr 4, 2024

@agunapal that's fine I think this is important to look at regardless

@mreso
Copy link
Collaborator

mreso commented Apr 24, 2024

Seems like weights_only is not working in this test:

## Registering my_text_classifier_scripted_v3 model
2024-04-04T17:35:44,573 [DEBUG] epollEventLoopGroup-3-8 org.pytorch.serve.wlm.ModelVersionedRefs - Adding new version 1.0 for model my_text_classifier_scripted_v3
2024-04-04T17:35:44,573 [DEBUG] epollEventLoopGroup-3-8 org.pytorch.serve.wlm.ModelVersionedRefs - Setting default version to 1.0 for model my_text_classifier_scripted_v3
2024-04-04T17:35:44,573 [INFO ] epollEventLoopGroup-3-8 org.pytorch.serve.wlm.ModelManager - Model my_text_classifier_scripted_v3 loaded.
2024-04-04T17:35:44,573 [DEBUG] epollEventLoopGroup-3-8 org.pytorch.serve.wlm.ModelManager - updateModel: my_text_classifier_scripted_v3, count: 1
2024-04-04T17:35:44,574 [DEBUG] W-9009-my_text_classifier_scripted_v3_1.0 org.pytorch.serve.wlm.WorkerLifeCycle - Worker cmdline: [/opt/hostedtoolcache/Python/3.9.19/x64/bin/python, /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/model_service_worker.py, --sock-type, unix, --sock-name, /tmp/.ts.sock.9009, --metrics-config, /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/configs/metrics.yaml]
2024-04-04T17:35:45,804 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - s_name_part0=/tmp/.ts.sock, s_name_part1=9009, pid=9859
2024-04-04T17:35:45,804 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Listening on port: /tmp/.ts.sock.9009
2024-04-04T17:35:45,821 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Successfully loaded /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/configs/metrics.yaml.
2024-04-04T17:35:45,822 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - [PID]9859
2024-04-04T17:35:45,822 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Torch worker started.
2024-04-04T17:35:45,822 [DEBUG] W-9009-my_text_classifier_scripted_v3_1.0 org.pytorch.serve.wlm.WorkerThread - W-9009-my_text_classifier_scripted_v3_1.0 State change null -> WORKER_STARTED
2024-04-04T17:35:45,822 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Python runtime: 3.9.19
2024-04-04T17:35:45,822 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0 org.pytorch.serve.wlm.WorkerThread - Connecting to: /tmp/.ts.sock.9009
2024-04-04T17:35:45,822 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Connection accepted: /tmp/.ts.sock.9009.
2024-04-04T17:35:45,822 [DEBUG] W-9009-my_text_classifier_scripted_v3_1.0 org.pytorch.serve.wlm.WorkerThread - Flushing req.cmd LOAD repeats 1 to backend at: 1712252145822
2024-04-04T17:35:45,823 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0 org.pytorch.serve.wlm.WorkerThread - Looping backend response at: 1712252145823
2024-04-04T17:35:45,823 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - model_name: my_text_classifier_scripted_v3, batchSize: 1
2024-04-04T17:35:46,568 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - ONNX enabled
2024-04-04T17:35:46,568 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Torch TensorRT not enabled
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Backend worker process died.
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - Traceback (most recent call last):
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -   File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/model_service_worker.py", line 263, in <module>
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -     worker.run_server()
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -   File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/model_service_worker.py", line 231, in run_server
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -     self.handle_connection(cl_socket)
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -   File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/model_service_worker.py", line 194, in handle_connection
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -     service, result, code = self.load_model(msg)
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -   File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/model_service_worker.py", line 131, in load_model
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -     service = model_loader.load(
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -   File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/model_loader.py", line 143, in load
2024-04-04T17:35:46,619 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -     initialize_fn(service.context)
2024-04-04T17:35:46,620 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -   File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/ts/torch_handler/text_handler.py", line 59, in initialize
2024-04-04T17:35:46,620 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -     self.source_vocab = torch.load(
2024-04-04T17:35:46,620 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -   File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/torch/serialization.py", line 1025, in load
2024-04-04T17:35:46,620 [INFO ] epollEventLoopGroup-5-2 org.pytorch.serve.wlm.WorkerThread - 9009 Worker disconnected. WORKER_STARTED
2024-04-04T17:35:46,620 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG -     raise pickle.UnpicklingError(UNSAFE_MESSAGE + str(e)) from None
2024-04-04T17:35:46,620 [INFO ] W-9009-my_text_classifier_scripted_v3_1.0-stdout MODEL_LOG - _pickle.UnpicklingError: Weights only load failed. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution.Do it only if you get the file from a trusted source. WeightsUnpickler error: Unsupported class torchtext.vocab.Vocab

Will take a look later

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants