Skip to content

[SPARK-53453][PYTHON][ML] Unblock 'torch<2.6.0' #52197

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.

Sign up for GitHub

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

Closed
zhengruifeng wants to merge 6 commits into from
Closed

[SPARK-53453][PYTHON][ML] Unblock 'torch<2.6.0' #52197

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e41e011
test
zhengruifeng Sep 2, 2025
640f899
test
zhengruifeng Sep 2, 2025
1b58279
test
zhengruifeng Sep 2, 2025
b2bfbb4
test
zhengruifeng Sep 2, 2025
ac26fc0
restore PR builder
zhengruifeng Sep 2, 2025
77222c5
more docker files
zhengruifeng Sep 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion dev/spark-test-image/python-311/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,6 @@ ARG CONNECT_PIP_PKGS="grpcio==1.67.0 grpcio-status==1.67.0 protobuf==5.29.1 goog
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.11
RUN python3.11 -m pip install --ignore-installed blinker>=1.6.2 # mlflow needs this
RUN python3.11 -m pip install $BASIC_PIP_PKGS unittest-xml-reporting $CONNECT_PIP_PKGS && \
python3.11 -m pip install 'torch<2.6.0' torchvision --index-url https://download.pytorch.org/whl/cpu && \
python3.11 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
Copy link
Copy Markdown
Contributor

@LuciferYang LuciferYang Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do other Dockerfiles need to be handled in separate pr?

Copy link
Copy Markdown
Contributor Author

@zhengruifeng zhengruifeng Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good catch

python3.11 -m pip install deepspeed torcheval && \
python3.11 -m pip cache purge
2 changes: 1 addition & 1 deletion python/pyspark/ml/connect/classification.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -381,7 +381,7 @@ def _save_core_model(self, path: str) -> None:
def _load_core_model(self, path: str) -> None:
import torch

lor_torch_model = torch.load(path)
lor_torch_model = torch.load(path, weights_only=False)
Copy link
Copy Markdown
Contributor

@bjornjorgensen bjornjorgensen Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need to have the weights_only=False . From there page "Loading un-trusted checkpoint with weights_only=False MUST never be done." https://github.com/pytorch/pytorch/security

Copy link
Copy Markdown
Contributor Author

@zhengruifeng zhengruifeng Sep 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The default value of weights_only is False in <2.6.0. And the default value was changed to True since 2.6.0.
This PR keeps the behavior.

Ideally, we should use weights_only=True, it will needs some investigation. Thanks for pointing it out.

self.torch_model = lor_torch_model[0]

def _get_extra_metadata(self) -> Dict[str, Any]:
Expand Down
3 changes: 2 additions & 1 deletion python/pyspark/ml/tests/connect/test_legacy_mode_classification.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,8 @@ def test_save_load(self):

# test saved torch model can be loaded by pytorch solely
lor_torch_model = torch.load(
os.path.join(local_model_path, "LogisticRegressionModel.torch")
os.path.join(local_model_path, "LogisticRegressionModel.torch"),
Copy link
Copy Markdown
Contributor

@bjornjorgensen bjornjorgensen Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dito but this is for a tests so I think its ok.

weights_only=False,
)

with torch.inference_mode():
Expand Down
⚡