Use weights_only for load

[kit1980]

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.

[kit1980]

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

[agunapal]

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

[agunapal]

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

@msaroufim This needs some debugging. Its failing everytime

[msaroufim]

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

[mreso]

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