initial Python 3.14 support (#2679)

Author: t-vi

.github/workflows/ci-testing.yml

@@ -37,6 +37,7 @@ jobs:
           - { os: "ubuntu-24.04", python-version: "3.11", requires: "latest" }
           - { os: "ubuntu-24.04", python-version: "3.12", requires: "latest" }
           - { os: "ubuntu-24.04", python-version: "3.13", requires: "latest" }
+          - { os: "ubuntu-24.04", python-version: "3.14", requires: "latest" }
         exclude:
           - { os: "windows-latest", suite: "ops" }
           - { os: "windows-latest", suite: "grads" }

pyproject.toml

@@ -14,7 +14,7 @@ authors = [
     { name = "Lightning AI", email = "[email protected]" },
 ]
 license = "Apache-2.0"
-requires-python = ">=3.10, <3.14"
+requires-python = ">=3.10, <3.15"
 keywords = ["deep learning", "AI", "compiler"]
 classifiers=[
     "Environment :: Console",

thunder/core/interpreter.py

@@ -1143,7 +1143,7 @@ def nexti(self, inst: dis.Instruction):
         if (3, 9) <= sys.version_info < (3, 11):
             if inst.starts_line is not None:
                 self.positions = Positions(inst.starts_line, inst.starts_line, 0, 999)
-        elif (3, 11) <= sys.version_info < (3, 14):
+        elif (3, 11) <= sys.version_info < (3, 15):
             if inst.positions is not None:
                 self.positions = inst.positions
         else:
@@ -3263,7 +3263,7 @@ def _async_gen_wrap_handler(inst: dis.Instruction, /, stack: InterpreterStack, *
 
 
 # https://docs.python.org/3.10/library/dis.html#opcode-BEFORE_ASYNC_WITH
-@register_opcode_handler("BEFORE_ASYNC_WITH")
+@register_opcode_handler("BEFORE_ASYNC_WITH", max_ver=(3, 13))
 def _before_async_with_handler(
     inst: dis.Instruction, /, stack: InterpreterStack, **kwargs
 ) -> None | INTERPRETER_SIGNALS:
@@ -3294,7 +3294,7 @@ def _before_async_with_handler(
 
 
 # https://docs.python.org/3.11/library/dis.html#opcode-BEFORE_WITH
-@register_opcode_handler("BEFORE_WITH", min_ver=(3, 11))
+@register_opcode_handler("BEFORE_WITH", min_ver=(3, 11), max_ver=(3, 13))
 def _before_with_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs) -> None | INTERPRETER_SIGNALS:
     mgr = stack.pop()
 
@@ -3349,6 +3349,7 @@ class BINARY_OP(enum.Enum):
     ISUB = 23
     ITRUEDIV = 24
     IXOR = 25
+    SUBSCR = 26
 
 
 def _binary_op(stack: InterpreterStack, op: BINARY_OP, a, b):
@@ -3384,6 +3385,9 @@ def _binary_op(stack: InterpreterStack, op: BINARY_OP, a, b):
     assert type(op) is BINARY_OP
     idx: int = op.value
 
+    if idx == BINARY_OP.SUBSCR.value:
+        return _binary_subscr(stack, a, b)
+
     res = Py_NULL()
     binop_name, *_ = ops[idx]
     _, left_method, right_method = ops[idx % BINARY_OP.IADD.value]
@@ -3631,11 +3635,15 @@ def impl(container, start, end):
 
 
 # https://docs.python.org/3.10/library/dis.html#opcode-BINARY_SUBSCR
-@register_opcode_handler("BINARY_SUBSCR")
+@register_opcode_handler("BINARY_SUBSCR", max_ver=(3, 13))
 def _binary_subscr_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs) -> None | INTERPRETER_SIGNALS:
     tos = stack.pop_wrapped()
     tos1 = stack.pop_wrapped()
 
+    return _binary_subscr(stack, tos1, tos)
+
+
+def _binary_subscr(stack, tos1, tos):
     def class_getitem_impl(cls, index):
         return cls.__class_getitem__(index)
 
@@ -3654,7 +3662,7 @@ def getitem_impl(obj, index):
 
 
 # https://docs.python.org/3.10/library/dis.html#opcode-BUILD_CONST_KEY_MAP
-@register_opcode_handler("BUILD_CONST_KEY_MAP")
+@register_opcode_handler("BUILD_CONST_KEY_MAP", max_ver=(3, 13))
 def _build_const_key_map_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs) -> None:
     assert type(inst.arg) is int
     count: int = inst.arg
@@ -3832,8 +3840,15 @@ def _call_function_handler(inst: dis.Instruction, /, stack: InterpreterStack, **
 def _call_function_ex_handler(
     inst: dis.Instruction, /, stack: InterpreterStack, **kwargs
 ) -> None | INTERPRETER_SIGNALS:
-    assert type(inst.arg) is int
-    kwargs = stack.pop_wrapped() if inst.arg & 0x01 else {}
+    if sys.version_info < (3, 14):
+        inst_arg = inst.arg
+    else:
+        inst_arg = 1
+    assert type(inst_arg) is int
+    kwargs = stack.pop_wrapped() if inst_arg & 0x01 else {}
+    if sys.version_info >= (3, 14) and wrapped_isinstance(kwargs, Py_NULL):
+        kwargs = {}
+
     assert wrapped_isinstance(kwargs, Mapping)
     args = stack.pop_wrapped()
     assert wrapped_isinstance(args, Iterable)
@@ -4314,7 +4329,9 @@ def _end_async_for_handler_3_11(
     **kwargs,
 ) -> None | INTERPRETER_SIGNALS:
     runtimectx: InterpreterRuntimeCtx = get_interpreterruntimectx()
-    assert inst.arg is None
+    if sys.version_info < (3, 14):
+        # 3.14+ has an (unused) int arg pointing to the END_SEND
+        assert inst.arg is None
 
     val = stack.pop()
     assert isinstance(val, BaseException)
@@ -4855,7 +4872,7 @@ def _list_to_tuple_handler(inst: dis.Instruction, /, stack: InterpreterStack, **
 
 
 # https://docs.python.org/3.13/library/dis.html#opcode-LOAD_ASSERTION_ERROR
-@register_opcode_handler("LOAD_ASSERTION_ERROR")
+@register_opcode_handler("LOAD_ASSERTION_ERROR", max_ver=(3, 13))
 def _load_assertion_error_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs) -> None:
     stack.append(wrap_const(AssertionError))
 
@@ -4889,7 +4906,35 @@ def _load_attr_handler_3_12(
         return load_method_helper(obj, name, stack)
 
 
-# https://docs.python.org/3.12/library/dis.html#opcode-LOAD_ATTR
+# https://docs.python.org/3.14/library/dis.html#opcode-LOAD_SPECIAL
+@register_opcode_handler("LOAD_SPECIAL", min_ver=(3, 14))
+def _load_special_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs) -> None | INTERPRETER_SIGNALS:
+    arg = inst.arg
+    mgr = stack.pop_wrapped()
+
+    type_name = type(unwrap(mgr)).__name__
+    special_methods = [
+        ("__enter__", f"'{type_name}' object does not support the context manager protocol (missed __enter__ method)"),
+        ("__exit__", f"'{type_name}' object does not support the context manager protocol (missed __exit__ method)"),
+        (
+            "__aenter__",
+            f"'{type_name}' object does not support the asynchronous context manager protocol (missed __aenter__ method)",
+        ),
+        (
+            "__aexit__",
+            f"'{type_name}' object does not support the asynchronous context manager protocol (missed __aexit__ method)",
+        ),
+    ]
+    name, error_msg = special_methods[arg]
+    res = load_method_helper(mgr, wrap_const(name), stack)
+    if res is INTERPRETER_SIGNALS.EXCEPTION_RAISED:
+        # clear previous error?
+        return do_raise(error_msg)
+
+    return res
+
+
+# https://docs.python.org/3.12/library/dis.html#opcode-LOAD_SUPER_ATTR
 @register_opcode_handler("LOAD_SUPER_ATTR", min_ver=(3, 12))
 def _load_super_attr_handler(
     inst: dis.Instruction, /, stack: InterpreterStack, co: CodeType, **kwargs
@@ -4944,6 +4989,23 @@ def _load_closure_handler(
     stack.append(val)
 
 
+# https://docs.python.org/3.14/library/dis.html#opcode-LOAD_COMMON_CONSTANT
+@register_opcode_handler("LOAD_COMMON_CONSTANT", min_ver=(3, 14))
+def _load_common_constant_handler(inst: dis.Instruction, /, stack: InterpreterStack, co: CodeType, **kwargs) -> None:
+    assert type(inst.arg) is int
+
+    common_consts = [
+        AssertionError,
+        NotImplementedError,
+        tuple,
+        all,
+        any,
+    ]
+    constant = wrap_const(common_consts[inst.arg])
+    constant = const_callback(constant)
+    stack.append(constant)
+
+
 # https://docs.python.org/3.10/library/dis.html#opcode-LOAD_CONST
 @register_opcode_handler("LOAD_CONST")
 def _load_const_handler(inst: dis.Instruction, /, stack: InterpreterStack, co: CodeType, **kwargs) -> None:
@@ -4955,6 +5017,16 @@ def _load_const_handler(inst: dis.Instruction, /, stack: InterpreterStack, co: C
     stack.append(constant)
 
 
+# https://docs.python.org/3.14/library/dis.html#opcode-LOAD_SMALL_INT
+@register_opcode_handler("LOAD_SMALL_INT", min_ver=(3, 14))
+def _load_small_int_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs) -> None:
+    assert type(inst.arg) is int
+
+    constant = wrap_const(inst.arg)
+    constant = const_callback(constant)
+    stack.append(constant)
+
+
 # https://docs.python.org/3.10/library/dis.html#opcode-LOAD_DEREF
 @register_opcode_handler("LOAD_DEREF")
 def _load_deref_handler(
@@ -4993,9 +5065,10 @@ def _load_deref_handler(
 # https://docs.python.org/3.10/library/dis.html#opcode-LOAD_FAST
 # https://docs.python.org/3.12/library/dis.html#opcode-LOAD_FAST_CHECK
 # LOAD_FAST for Python <3.12 is LOAD_FAST_CHECK
+@register_opcode_handler("LOAD_FAST_BORROW", min_ver=(3, 14))
 @register_opcode_handler("LOAD_FAST_CHECK", min_ver=(3, 12))
 @register_opcode_handler("LOAD_FAST")
-def _load_fast_check_handler(
+def _load_fast_check_load_fast_borrow_handler(
     inst: dis.Instruction, /, stack: InterpreterStack, co: CodeType, frame: InterpreterFrame, **kwargs
 ) -> None | INTERPRETER_SIGNALS:
     assert isinstance(inst.arg, int)
@@ -5019,8 +5092,9 @@ def _load_fast_check_handler(
 
 
 # https://docs.python.org/3.13/library/dis.html#opcode-LOAD_FAST_LOAD_FAST
+@register_opcode_handler("LOAD_FAST_BORROW_LOAD_FAST_BORROW", min_ver=(3, 14))
 @register_opcode_handler("LOAD_FAST_LOAD_FAST", min_ver=(3, 13))
-def _load_fast_load_fast_handler(
+def _load_fast_load_fast_load_fast_borrow_load_fast_borrow_handler(
     inst: dis.Instruction, /, stack: InterpreterStack, co: CodeType, frame: InterpreterFrame, **kwargs
 ) -> None | INTERPRETER_SIGNALS:
     assert isinstance(inst.arg, int)
@@ -5126,7 +5200,7 @@ def _load_global_handler(
 
 
 # https://docs.python.org/3.11/library/dis.html#opcode-LOAD_METHOD
-@register_opcode_handler("LOAD_METHOD")
+@register_opcode_handler("LOAD_METHOD", max_ver=(3, 13))
 def _load_method_handler(
     inst: dis.Instruction, /, stack: InterpreterStack, co: CodeType, **kwargs
 ) -> None | INTERPRETER_SIGNALS:
@@ -5360,6 +5434,13 @@ def _set_function_attribute_handler(
         stack.append(fn)
         return
 
+    if flag == 0x10:
+        annotate = stack.pop()
+        assert annotate is None or callable(annotate)
+        fn.__annotate__ = annotate
+        stack.append(fn)
+        return
+
     assert False, f"Flag value 0x{flag:x} unexpected in SET_FUNCTION_ATTRIBUTE"
 
 
@@ -5503,6 +5584,12 @@ def _match_sequence_handler(inst: dis.Instruction, /, stack: InterpreterStack, *
     stack.push(supported_sequence)
 
 
+# https://docs.python.org/3.14/library/dis.html#opcode-NOT_TAKEN
+@register_opcode_handler("NOT_TAKEN", min_ver=(3, 14))
+def _not_taken_handler(inst: dis.Instruction, /, **kwargs) -> None:
+    pass
+
+
 # https://docs.python.org/3.10/library/dis.html#opcode-NOP
 @register_opcode_handler("NOP")
 def _nop_handler(inst: dis.Instruction, /, **kwargs) -> None:
@@ -5777,6 +5864,12 @@ def _pop_top_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs
     stack.pop_wrapped()
 
 
+# https://docs.python.org/3.14/library/dis.html#opcode-POP_ITER
+@register_opcode_handler("POP_ITER", min_ver=(3, 14))
+def _pop_iter_handler(inst: dis.Instruction, /, stack: InterpreterStack, **kwargs) -> None:
+    stack.pop_wrapped()
+
+
 # Returns either
 def do_raise(exc: Any = Py_NULL(), cause: Any = Py_NULL()) -> Literal[INTERPRETER_SIGNALS.EXCEPTION_RAISED]:
     # Get the type and exception being raised
@@ -5944,7 +6037,7 @@ def _reraise_handler_3_11(
 
 
 # https://docs.python.org/3.12/library/dis.html#opcode-RETURN_CONST
-@register_opcode_handler("RETURN_CONST", min_ver=(3, 12))
+@register_opcode_handler("RETURN_CONST", min_ver=(3, 12), max_ver=(3, 13))
 def _return_const_handler(
     inst: dis.Instruction, /, co: CodeType, stack: InterpreterStack, **kwargs
 ) -> int | None | INTERPRETER_SIGNALS:
@@ -6576,7 +6669,7 @@ def _with_except_start_handler_3_10(
 
 
 # https://docs.python.org/3.11/library/dis.html#opcode-WITH_EXCEPT_START
-@register_opcode_handler("WITH_EXCEPT_START", min_ver=(3, 11))
+@register_opcode_handler("WITH_EXCEPT_START", min_ver=(3, 11), max_ver=(3, 13))
 def _with_except_start_handler_3_11(
     inst: dis.Instruction, *, inst_ptr: int, stack: InterpreterStack, try_stack: list[PyTryBlock], **kwargs
 ) -> None | INTERPRETER_SIGNALS:
@@ -6590,6 +6683,27 @@ def _with_except_start_handler_3_11(
     return check_and_append(stack, _interpret_call_with_unwrapping(exit_func, exc, val, tb))
 
 
+# https://docs.python.org/3.14/library/dis.html#opcode-WITH_EXCEPT_START
+@register_opcode_handler("WITH_EXCEPT_START", min_ver=(3, 14))
+def _with_except_start_handler_3_11(
+    inst: dis.Instruction, *, inst_ptr: int, stack: InterpreterStack, try_stack: list[PyTryBlock], **kwargs
+) -> None | INTERPRETER_SIGNALS:
+    # in 3.11 the exception representation changed to only val
+    val = stack[-1]
+    exc = type(val)
+    tb = val.__traceback__
+
+    assert isinstance(stack[-3], int)
+    exit_meth_or_func = stack[-5]
+    exit_meth_self_or_null = stack[-4]
+    if isinstance(exit_meth_self_or_null, Py_NULL):
+        return check_and_append(stack, _interpret_call_with_unwrapping(exit_meth_or_func, exc, val, tb))
+    else:
+        return check_and_append(
+            stack, _interpret_call_with_unwrapping(exit_meth_or_func, exit_meth_self_or_null, exc, val, tb)
+        )
+
+
 # https://docs.python.org/3.10/library/dis.html#opcode-YIELD_FROM
 @register_opcode_handler("YIELD_FROM", max_ver=(3, 10))
 def _yield_from_handler(
@@ -7157,7 +7271,7 @@ def _setup_frame_and_run_python_function(
         for i, (name, value) in enumerate(zip(code.co_freevars, closure)):
             local = freevar_callback(name, value, fn=wrapped_fn, idx=i)
             localsplus.append(local)
-    elif (3, 11) <= sys.version_info < (3, 14):
+    elif (3, 11) <= sys.version_info < (3, 15):
         assert len(code.co_varnames) == code.co_nlocals
         for n in code.co_varnames:
             local = locals_dict.get(n, Py_NULL())

thunder/tests/test_interpreter.py

@@ -340,8 +340,10 @@ def kwargsplat(**kwargs):
         assert is_jitting_with_raise() == jitting
         return foo(**kwargs)
 
-    assert any(i.opname == "CALL_FUNCTION_EX" and not i.arg & 1 for i in dis.get_instructions(argsplat))
-    assert any(i.opname == "CALL_FUNCTION_EX" and i.arg & 1 for i in dis.get_instructions(kwargsplat))
+    if sys.version_info < (3, 14):
+        # Python 3.14 has no arg in call function ex.
+        assert any(i.opname == "CALL_FUNCTION_EX" and not i.arg & 1 for i in dis.get_instructions(argsplat))
+        assert any(i.opname == "CALL_FUNCTION_EX" and i.arg & 1 for i in dis.get_instructions(kwargsplat))
 
     kwargs = {"a": 1, "b": 2}
 
@@ -357,6 +359,10 @@ def kwargsplat(**kwargs):
     assert_close(res2, jres2)
 
 
+@pytest.mark.skipif(
+    sys.version_info >= (3, 14),
+    reason="Python 3.14+ do not implement BUILD_CONST_KEY_MAP",
+)
 def test_build_const_key_map(jit):
     def fn1(a, b):
         return {"a": a, "b": b}
@@ -1074,7 +1080,7 @@ def foo(x):
 
     jfoo = jit(foo)
 
-    with pytest.raises(Exception, match=r"reduce\(\) takes no keyword arguments"):
+    with pytest.raises(Exception, match=r"reduce\(\) takes .* arguments"):  # Varies for 3.10-3.13, 3.14
         foo((1, 2, 3))
 
     with pytest.raises(Exception, match=r"got some positional-only arguments passed as keyword arguments"):