Python 3.12 support (#10548)

Oneflow-Inc · Aug 3, 2024 · f0d13a6 · f0d13a6
1 parent 7355fc1
commit f0d13a6
Show file tree

Hide file tree

Showing 30 changed files with 51 additions and 9 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -170,6 +170,7 @@ jobs:
           nightly-date: ${{ needs.generate-build-matrix.outputs.formatted_date }}
           use-nvidia-wheels: ${{ matrix.entry !='cu112' }}
           python-versions: |
+            3.12
             3.11
             3.10
             3.9
@@ -195,6 +196,7 @@ jobs:
           nightly-date: ${{ needs.generate-build-matrix.outputs.formatted_date }}
           use-nvidia-wheels: ${{ matrix.entry !='cu112' }}
           python-versions: |
+            3.12
             3.11
             3.10
             3.9
@@ -219,6 +221,7 @@ jobs:
           nightly: ${{ inputs.is_priv || github.event_name == 'schedule' || github.ref == 'refs/heads/release/add_nightly_date_index'}}
           nightly-date: ${{ needs.generate-build-matrix.outputs.formatted_date }}
           python-versions: |
+            3.12
             3.11
             3.10
             3.9

diff --git a/ci/fixed-dev-requirements.txt b/ci/fixed-dev-requirements.txt
@@ -1,2 +1,3 @@
-numpy==1.22.1 ; python_version >= "3.10"
+numpy==1.26.4 ; python_version >= "3.12"
+numpy==1.22.1 ; python_version >= "3.10" and python_version < "3.12"
 numpy==1.21.6 ; python_version >= "3.7" and python_version < "3.10"
diff --git a/oneflow/api/python/autograd/autograd_function_state.h b/oneflow/api/python/autograd/autograd_function_state.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_AUTOGRAD_AUTOGRAD_FUNCTION_STATE_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <pybind11/pybind11.h>
 
 #include "oneflow/core/framework/op_expr_grad_function.h"

diff --git a/oneflow/api/python/caster/size.h b/oneflow/api/python/caster/size.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_CASTER_SIZE_H_
 #include <type_traits>
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <pybind11/pybind11.h>
 
 #include "oneflow/api/python/framework/size.h"

diff --git a/oneflow/api/python/exception/exception.h b/oneflow/api/python/exception/exception.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_COMMON_EXCEPTION_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <pybind11/pybind11.h>
 
 #include "oneflow/core/common/exception.h"

diff --git a/oneflow/api/python/framework/memory_format.h b/oneflow/api/python/framework/memory_format.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FRAMEWORK_MEMORY_FORMAT_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <pybind11/pybind11.h>
 
 #include "oneflow/core/common/memory_format.pb.h"

diff --git a/oneflow/api/python/framework/size.h b/oneflow/api/python/framework/size.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FRAMEWORK_SIZE_H_
 #include <type_traits>
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <pybind11/pybind11.h>
 #include "oneflow/core/common/shape.h"
 

diff --git a/oneflow/api/python/framework/tensor.cpp b/oneflow/api/python/framework/tensor.cpp
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <pybind11/pybind11.h>
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include "oneflow/api/python/exception/exception.h"
 #include "oneflow/api/python/framework/size.h"
 #include "oneflow/api/python/framework/tensortype.h"

diff --git a/oneflow/api/python/framework/tensor.h b/oneflow/api/python/framework/tensor.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FRAMEWORK_TENSOR_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/core/framework/tensor.h"
 

diff --git a/oneflow/api/python/framework/tensor_functions.cpp b/oneflow/api/python/framework/tensor_functions.cpp
@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include "oneflow/api/python/exception/exception.h"
 #include "oneflow/api/python/framework/size.h"
 #include "oneflow/api/python/framework/tensor_functions_util.h"

diff --git a/oneflow/api/python/framework/tensor_functions_util.h b/oneflow/api/python/framework/tensor_functions_util.h
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <string>
 #include "oneflow/api/python/exception/exception.h"
 #include "oneflow/api/python/functional/common.h"
@@ -71,4 +72,4 @@ PyObject* PyParseArgs(PyObject* args, PyObject* kwargs, const char* func_name,
 }
 
 }  // namespace one
-}  // namespace oneflow
+}  // namespace oneflow
diff --git a/oneflow/api/python/framework/tensortype.cpp b/oneflow/api/python/framework/tensortype.cpp
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <pybind11/pybind11.h>
 #include "oneflow/api/python/framework/tensor.h"
 #include "oneflow/api/python/framework/tensortype.h"

diff --git a/oneflow/api/python/framework/tensortype.h b/oneflow/api/python/framework/tensortype.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FRAMEWORK_TENSORTYPE_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include "oneflow/core/framework/dtype.h"
 #include "oneflow/core/framework/device.h"
 

diff --git a/oneflow/api/python/framework/typeinfo.h b/oneflow/api/python/framework/typeinfo.h
@@ -18,6 +18,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FRAMEWORK_TYPEINFO_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include "oneflow/core/common/throw.h"
 #include "oneflow/core/framework/dtype.h"
 

diff --git a/oneflow/api/python/functional/indexing.h b/oneflow/api/python/functional/indexing.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FUNCTIONAL_INDEXING_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/api/python/functional/common.h"
 #include "oneflow/core/common/maybe.h"

diff --git a/oneflow/api/python/functional/python_arg.h b/oneflow/api/python/functional/python_arg.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <pybind11/pybind11.h>
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/core/common/throw.h"
 #include "oneflow/api/python/functional/value_types.h"

diff --git a/oneflow/api/python/functional/python_arg_parser.h b/oneflow/api/python/functional/python_arg_parser.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FUNCTIONAL_PYTHON_ARG_PARSER_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/api/python/functional/function_def.h"
 #include "oneflow/api/python/functional/python_arg.h"

diff --git a/oneflow/api/python/functional/python_return_types.h b/oneflow/api/python/functional/python_return_types.h
@@ -21,6 +21,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_FUNCTIONAL_PYTHON_RETURN_TYPES_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <string>
 #include <sstream>
 #include <structmember.h>

diff --git a/oneflow/api/python/functional/tensor_api.cpp b/oneflow/api/python/functional/tensor_api.cpp
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <memory>
 
 #include "oneflow/api/python/utils/tensor_utils.h"

diff --git a/oneflow/api/python/functional/value_types.h b/oneflow/api/python/functional/value_types.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include <complex>
 #include <memory>
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/core/common/data_type.pb.h"
 #include "oneflow/core/common/maybe.h"

diff --git a/oneflow/api/python/utils/tensor_utils.h b/oneflow/api/python/utils/tensor_utils.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_UTILS_TENSOR_UTILS_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <pybind11/functional.h>

diff --git a/oneflow/extension/python/py_compute.h b/oneflow/extension/python/py_compute.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef ONEFLOW_EXTENSION_PYTHON_PY_COMPUTE_H_
 #define ONEFLOW_EXTENSION_PYTHON_PY_COMPUTE_H_
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include "oneflow/core/framework/framework.h"
 
 namespace oneflow {

diff --git a/oneflow/extension/python/py_kernel_registry.h b/oneflow/extension/python/py_kernel_registry.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <string>
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include "oneflow/core/common/maybe.h"
 
 namespace oneflow {

diff --git a/oneflow/extension/stack/python/custom_eval_frame.c b/oneflow/extension/stack/python/custom_eval_frame.c
@@ -21,6 +21,7 @@ limitations under the License.
 
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
+#undef _PyGC_FINALIZED
 #include <frameobject.h>
 #include <pystate.h>
 // see https://bugs.python.org/issue35886

diff --git a/oneflow/extension/stack/python/custom_eval_frame.h b/oneflow/extension/stack/python/custom_eval_frame.h
@@ -21,6 +21,7 @@ extern "C" {
 #endif
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #if PY_VERSION_HEX >= 0x03090000
 typedef PyObject* (*PyFrameEvalFunc)(struct _ts*, struct _frame*, int);

diff --git a/python/oneflow/framework/check_point_v2.py b/python/oneflow/framework/check_point_v2.py
@@ -29,7 +29,6 @@
     IO,
     BinaryIO,
 )
-from typing_extensions import TypeAlias
 from pathlib import Path
 import pickle
 import json
@@ -58,10 +57,10 @@
 PROTOCOL_VERSION = 1
 ONEFLOW_MAGIC_KEY = "__oneflow__"
 
-MAP_LOCATION: TypeAlias = Optional[
+MAP_LOCATION = Optional[
     Union[Callable[[Tensor, str], Tensor], flow.device, str, flow.placement]
 ]
-FILE_LIKE: TypeAlias = Union[os.PathLike, BinaryIO, IO[bytes], Path]
+FILE_LIKE = Union[os.PathLike, BinaryIO, IO[bytes], Path]
 
 
 class _opener(object):

diff --git a/python/setup.py b/python/setup.py
@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from __future__ import absolute_import
 
 import argparse
@@ -96,11 +97,16 @@ def has_ext_modules(self):
 
 
 def get_oneflow_internal_so_path():
-    import imp
-
-    fp, pathname, description = imp.find_module("_oneflow_internal", ["oneflow"])
+    import importlib
+
+    suffixes = importlib.machinery.EXTENSION_SUFFIXES
+    loader = importlib.machinery.ExtensionFileLoader
+    lazy_loader = importlib.util.LazyLoader.factory(loader)
+    finder = importlib.machinery.FileFinder("oneflow", (lazy_loader, suffixes))
+    spec = finder.find_spec("_oneflow_internal")
+    pathname = spec.origin
     assert os.path.isfile(pathname)
-    return os.path.relpath(pathname, "oneflow")
+    return os.path.basename(pathname)
 
 
 package_data = {"oneflow": [get_oneflow_internal_so_path()] + include_files}

diff --git a/tools/functional/generate_dispatch_stateful_ops.py b/tools/functional/generate_dispatch_stateful_ops.py
@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import re
 import argparse
@@ -52,6 +53,7 @@
 #define ONEFLOW_API_PYTHON_FUNCTIONAL_GENERATED_DISPATCH_OP_API_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/core/common/optional.h"
 #include "oneflow/core/common/scalar.h"
@@ -109,6 +111,7 @@
     + """
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/api/python/of_api_registry.h"
 #include "oneflow/api/python/functional/common.h"

diff --git a/tools/functional/generate_functional_api.py b/tools/functional/generate_functional_api.py
@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import re
 import argparse
@@ -101,6 +102,7 @@
     + """
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 namespace oneflow {{
 namespace one {{
@@ -117,6 +119,7 @@
     + """
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/api/python/of_api_registry.h"
 #include "oneflow/api/python/exception/exception.h"

diff --git a/tools/functional/generate_tensor_api.py b/tools/functional/generate_tensor_api.py
@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import re
 import argparse
@@ -52,6 +53,7 @@
 #define ONEFLOW_API_PYTHON_FUNCTIONAL_GENERATED_TENSOR_API_H_
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/core/common/optional.h"
 #include "oneflow/core/common/scalar.h"
@@ -95,6 +97,7 @@
     + """
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 namespace oneflow {{
 namespace one {{
@@ -111,6 +114,7 @@
     + """
 
 #include <Python.h>
+#undef _PyGC_FINALIZED
 
 #include "oneflow/api/python/of_api_registry.h"
 #include "oneflow/api/python/exception/exception.h"