Add Python 3.13 support

.github/workflows/build_image.yml

@@ -28,7 +28,7 @@ jobs:
     needs: get-version
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
       - uses: actions/checkout@v4
         with:

.github/workflows/pythonbuild.yml

@@ -27,9 +27,9 @@ jobs:
       - name: Set Python versions for run
         run: |
           if [[ ${{ github.event_name }} == "schedule" ]]; then
-            echo "python_versions=[\"3.10\",\"3.11\",\"3.12\"]" >> $GITHUB_ENV
+            echo "python_versions=[\"3.10\",\"3.11\",\"3.12\",\"3.13\"]" >> $GITHUB_ENV
           else
-            echo "python_versions=[\"3.12\"]" >> $GITHUB_ENV
+            echo "python_versions=[\"3.13\"]" >> $GITHUB_ENV
           fi
 
   build:
@@ -110,9 +110,6 @@ jobs:
           uv pip uninstall --system pandas pyarrow
           uv pip freeze
       - name: Run extras unit tests with coverage
-        # Skip this step if running on python 3.12 due to https://github.com/tensorflow/tensorflow/issues/62003
-        # and https://github.com/pytorch/pytorch/issues/110436
-        if: ${{ matrix.python-version != '3.12' }}
         run: |
           make unit_test_extras_codecov
       - name: Codecov
@@ -139,6 +136,9 @@ jobs:
             pandas: "pandas>=2.0.0"
           - pandas: "pandas<2.0.0"
             python-version: "3.12"
+          # pandas 1.x has no wheels for python 3.13
+          - pandas: "pandas<2.0.0"
+            python-version: "3.13"
 
     steps:
       - uses: actions/checkout@v4
@@ -389,6 +389,27 @@ jobs:
           # Segmentation fault on python 3.12: https://github.com/flyteorg/flyte/issues/5020
           - python-version: 3.12
             plugin-names: "flytekit-kf-pytorch"
+          # Python 3.13: the exclusions below mirror the python 3.12 ones. They have not been verified on 3.13
+          # (several of the underlying issues are closed by now); revisit both lists together.
+          - python-version: 3.13
+            plugin-names: "flytekit-airflow"
+          - python-version: 3.13
+            plugin-names: "flytekit-kf-pytorch"
+          - python-version: 3.13
+            plugin-names: "flytekit-modin"
+          - python-version: 3.13
+            plugin-names: "flytekit-onnx-pytorch"
+          - python-version: 3.13
+            plugin-names: "flytekit-ray"
+          - python-version: 3.13
+            plugin-names: "flytekit-vaex"
+          # The tensorflow releases that provide python 3.13 wheels require protobuf>=5, which conflicts with the
+          # "protobuf<5" pin used below for the plugin tests.
+          - python-version: 3.13
+            plugin-names: "flytekit-mlflow"
+          # flytekitplugins-duckdb declares python_requires<3.13 (duckdb<=1.0.0 pin)
+          - python-version: 3.13
+            plugin-names: "flytekit-duckdb"
     steps:
       - uses: actions/checkout@v4
       - name: "Clear action cache"

.github/workflows/pythonpublish.yml

@@ -82,6 +82,7 @@ jobs:
           - "3.10"
           - "3.11"
           - "3.12"
+          - "3.13"
     steps:
       - uses: actions/checkout@v4
         with:

dev-requirements.in

@@ -21,14 +21,16 @@ keyrings.alt
 setuptools_scm
 pytest-icdiff
 
-# Tensorflow is not available for python 3.12 yet: https://github.com/tensorflow/tensorflow/issues/62003
 tensorflow<=2.15.1; python_version<'3.12'
+# tensorflow>=2.16 ships python 3.12 wheels and >=2.20 ships python 3.13 wheels (linux x86_64/aarch64, macos arm64,
+# windows). With the protobuf pin below this resolves to tensorflow<2.20 on 3.12 and to the latest release on 3.13
+# (tensorflow>=2.20 requires protobuf>=5).
+tensorflow; python_version>='3.12'
 # Newer versions of torch bring in nvidia dependencies that are not present in windows, so
 # we put this constraint while we do not have per-environment requirements files
 torch<=1.12.1; python_version<'3.11'
-# pytorch 2 supports python 3.11
-# pytorch 2 does not support 3.12 yet: https://github.com/pytorch/pytorch/issues/110436
-torch; python_version<'3.12'
+# pytorch>=2.2 ships python 3.12 wheels and >=2.6 ships python 3.13 wheels
+torch; python_version>='3.11'
 pydantic
 
 # TODO: Currently, the python-magic library causes build errors on Windows due to its dependency on DLLs for libmagic.
@@ -40,8 +42,11 @@ python-magic; (platform_system=='Darwin' or platform_system=='Linux')
 # Google released a new major version of the protobuf library and once that started being used in the ecosystem at large,
 # including `googleapis-common-protos` we started seeing errors in CI, so let's constrain that for now.
 # The issue to support protobuf 5 is being tracked in https://github.com/flyteorg/flyte/issues/5448.
-protobuf<5
-types-protobuf<5
+# The root cause was fixed in googleapis-common-protos>=1.63.2. On python 3.13 the pin is lifted, because the
+# tensorflow releases that support 3.13 require protobuf>=5 (flytekit itself does not cap protobuf).
+protobuf<5; python_version<'3.13'
+types-protobuf<5; python_version<'3.13'
+types-protobuf; python_version>='3.13'
 
 types-croniter
 types-decorator

flytekit/configuration/default_images.py

@@ -13,6 +13,7 @@ class PythonVersion(enum.Enum):
     PYTHON_3_10 = (3, 10)
     PYTHON_3_11 = (3, 11)
     PYTHON_3_12 = (3, 12)
+    PYTHON_3_13 = (3, 13)
 
 
 class DefaultImages(object):
@@ -26,6 +27,7 @@ class DefaultImages(object):
         PythonVersion.PYTHON_3_10: "cr.flyte.org/flyteorg/flytekit:py3.10-",
         PythonVersion.PYTHON_3_11: "cr.flyte.org/flyteorg/flytekit:py3.11-",
         PythonVersion.PYTHON_3_12: "cr.flyte.org/flyteorg/flytekit:py3.12-",
+        PythonVersion.PYTHON_3_13: "cr.flyte.org/flyteorg/flytekit:py3.13-",
     }
 
     @classmethod

flytekit/extras/pytorch/checkpoint.py

@@ -9,6 +9,7 @@
 
 from flytekit.core.context_manager import FlyteContext
 from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError
+from flytekit.extras.pytorch.native import load_torch_object
 from flytekit.models.core import types as _core_types
 from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar
 from flytekit.models.types import LiteralType
@@ -118,8 +119,11 @@ def to_python_value(
         else:
             map_location = torch.device("cpu")
 
-        # load checkpoint from a file
-        return typing.cast(PyTorchCheckpoint, torch.load(local_path, map_location=map_location))
+        # load checkpoint from a file; the hyperparameters can be arbitrary objects, so the weights-only loader
+        # cannot be used
+        return typing.cast(
+            PyTorchCheckpoint, load_torch_object(local_path, map_location=map_location, weights_only=False)
+        )
 
     def guess_python_type(self, literal_type: LiteralType) -> Type[PyTorchCheckpoint]:
         if (

flytekit/extras/pytorch/native.py

@@ -1,5 +1,6 @@
+import inspect
 import pathlib
-from typing import Type, TypeVar
+from typing import Any, Optional, Type, TypeVar
 
 import torch
 
@@ -12,7 +13,25 @@
 T = TypeVar("T")
 
 
+def load_torch_object(path: str, map_location: Any, weights_only: Optional[bool] = None) -> Any:
+    """
+    Wrapper around ``torch.load``.
+
+    ``weights_only=None`` keeps torch's own default, which is ``True`` since torch 2.6: only tensors, containers and
+    primitives are unpickled. ``weights_only=False`` unpickles arbitrary objects and therefore executes pickle code
+    from the blob. flytekit needs that for ``nn.Module`` and checkpoint objects; the trust model is the same as for
+    ``FlytePickle``, the data was written by flytekit itself in an upstream task. The keyword does not exist before
+    torch 1.13, in which case it is omitted.
+    """
+    if weights_only is not None and "weights_only" in inspect.signature(torch.load).parameters:
+        return torch.load(path, map_location=map_location, weights_only=weights_only)
+    return torch.load(path, map_location=map_location)
+
+
 class PyTorchTypeTransformer(TypeTransformer[T]):
+    # Passed to ``load_torch_object``; ``None`` keeps torch's default.
+    WEIGHTS_ONLY: Optional[bool] = None
+
     def get_literal_type(self, t: Type[T]) -> LiteralType:
         return LiteralType(
             blob=_core_types.BlobType(
@@ -63,7 +82,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
             map_location = torch.device("cpu")
 
         # load pytorch tensor/module from a file
-        return torch.load(local_path, map_location=map_location)
+        return load_torch_object(local_path, map_location=map_location, weights_only=self.WEIGHTS_ONLY)
 
 
 class PyTorchTensorTransformer(PyTorchTypeTransformer[torch.Tensor]):
@@ -85,6 +104,8 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[torch.Tensor]:
 
 class PyTorchModuleTransformer(PyTorchTypeTransformer[torch.nn.Module]):
     PYTORCH_FORMAT = "PyTorchModule"
+    # A whole ``nn.Module`` object is stored, which torch's weights-only loader refuses to unpickle.
+    WEIGHTS_ONLY = False
 
     def __init__(self):
         super().__init__(name="PyTorch Module", t=torch.nn.Module)

plugins/flytekit-openai/flytekitplugins/openai/batch/task.py

@@ -64,6 +64,7 @@ class OpenAIFileDefaultImages(DefaultImages):
         PythonVersion.PYTHON_3_10: "cr.flyte.org/flyteorg/flytekit:py3.10-openai-batch-",
         PythonVersion.PYTHON_3_11: "cr.flyte.org/flyteorg/flytekit:py3.11-openai-batch-",
         PythonVersion.PYTHON_3_12: "cr.flyte.org/flyteorg/flytekit:py3.12-openai-batch-",
+        PythonVersion.PYTHON_3_13: "cr.flyte.org/flyteorg/flytekit:py3.13-openai-batch-",
     }
 
 

plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py

@@ -26,6 +26,7 @@ class SQLAlchemyDefaultImages(DefaultImages):
         PythonVersion.PYTHON_3_10: "cr.flyte.org/flyteorg/flytekit:py3.10-sqlalchemy-",
         PythonVersion.PYTHON_3_11: "cr.flyte.org/flyteorg/flytekit:py3.11-sqlalchemy-",
         PythonVersion.PYTHON_3_12: "cr.flyte.org/flyteorg/flytekit:py3.12-sqlalchemy-",
+        PythonVersion.PYTHON_3_13: "cr.flyte.org/flyteorg/flytekit:py3.13-sqlalchemy-",
     }
 
 

pyproject.toml

@@ -9,7 +9,7 @@ authors = [{ name = "Flyte Contributors", email = "admin@flyte.org" }]
 description = "Flyte SDK for Python"
 license = { text = "Apache-2.0" }
 readme = { file = "README.md", content-type = "text/markdown" }
-requires-python = ">=3.10,<3.13"
+requires-python = ">=3.10,<3.14"
 dependencies = [
     # Please maintain an alphabetical order in the following list
     "adlfs>=2023.3.0",
@@ -62,6 +62,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Topic :: Scientific/Engineering",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Software Development",

tests/flytekit/unit/configuration/test_image_config.py

@@ -16,6 +16,7 @@
     [
         (PythonVersion.PYTHON_3_9, "cr.flyte.org/flyteorg/flytekit:py3.9-latest"),
         (PythonVersion.PYTHON_3_10, "cr.flyte.org/flyteorg/flytekit:py3.10-latest"),
+        (PythonVersion.PYTHON_3_13, "cr.flyte.org/flyteorg/flytekit:py3.13-latest"),
     ],
 )
 def test_defaults(python_version_enum, expected_image_string):
@@ -26,6 +27,7 @@ def test_defaults(python_version_enum, expected_image_string):
     "python_version_enum, flytekit_version, expected_image_string",
     [
         (PythonVersion.PYTHON_3_9, "v0.32.0", "cr.flyte.org/flyteorg/flytekit:py3.9-0.32.0"),
+        (PythonVersion.PYTHON_3_13, "v1.17.0", "cr.flyte.org/flyteorg/flytekit:py3.13-1.17.0"),
     ],
 )
 def test_set_both(python_version_enum, flytekit_version, expected_image_string):

tests/flytekit/unit/extras/pytorch/test_load.py

@@ -0,0 +1,40 @@
+from unittest import mock
+
+import torch
+
+from flytekit.extras.pytorch import PyTorchModuleTransformer, PyTorchTensorTransformer
+from flytekit.extras.pytorch.native import load_torch_object
+
+
+def test_load_torch_object_keeps_torch_default():
+    with mock.patch("torch.load", autospec=True) as m:
+        load_torch_object("model.pt", map_location="cpu")
+    m.assert_called_once_with("model.pt", map_location="cpu")
+
+
+def test_load_torch_object_passes_weights_only():
+    with mock.patch("torch.load", autospec=True) as m:
+        load_torch_object("model.pt", map_location="cpu", weights_only=False)
+    m.assert_called_once_with("model.pt", map_location="cpu", weights_only=False)
+
+
+def test_load_torch_object_omits_weights_only_for_old_torch():
+    def legacy_load(f, map_location=None):
+        return "loaded"
+
+    with mock.patch("torch.load", new=legacy_load):
+        assert load_torch_object("model.pt", map_location="cpu", weights_only=False) == "loaded"
+
+
+def test_transformers_weights_only_setting():
+    # tensors load fine with torch's safe default, whole modules do not
+    assert PyTorchTensorTransformer.WEIGHTS_ONLY is None
+    assert PyTorchModuleTransformer.WEIGHTS_ONLY is False
+
+
+def test_module_round_trip_uses_full_unpickling(tmp_path):
+    module = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.BatchNorm1d(2))
+    path = str(tmp_path / "module.pt")
+    torch.save(module, path)
+    loaded = load_torch_object(path, map_location="cpu", weights_only=PyTorchModuleTransformer.WEIGHTS_ONLY)
+    assert isinstance(loaded, torch.nn.Sequential)