Skip to content

Commit acae017

Browse files
authored
Add _unsafe_reset_threadpool to pybindings
Differential Revision: D71023514 Pull Request resolved: #9196
1 parent 7ac7270 commit acae017

File tree

4 files changed

+20
-0
lines changed

4 files changed

+20
-0
lines changed

extension/pybindings/portable_lib.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
_load_for_executorch_from_buffer, # noqa: F401
4646
_load_for_executorch_from_bundled_program, # noqa: F401
4747
_reset_profile_results, # noqa: F401
48+
_unsafe_reset_threadpool, # noqa: F401
4849
BundledModule, # noqa: F401
4950
ExecuTorchModule, # noqa: F401
5051
MethodMeta, # noqa: F401

extension/pybindings/pybindings.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#include <executorch/extension/data_loader/buffer_data_loader.h>
2424
#include <executorch/extension/data_loader/mmap_data_loader.h>
2525
#include <executorch/extension/memory_allocator/malloc_memory_allocator.h>
26+
#include <executorch/extension/threadpool/threadpool.h>
2627
#include <executorch/runtime/backend/interface.h>
2728
#include <executorch/runtime/core/data_loader.h>
2829
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -1064,6 +1065,14 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
10641065
"_reset_profile_results",
10651066
[]() { EXECUTORCH_RESET_PROFILE_RESULTS(); },
10661067
call_guard);
1068+
m.def(
1069+
"_unsafe_reset_threadpool",
1070+
[](int num_threads) {
1071+
executorch::extension::threadpool::get_threadpool()
1072+
->_unsafe_reset_threadpool(num_threads);
1073+
},
1074+
py::arg("num_threads"),
1075+
call_guard);
10671076

10681077
py::class_<PyModule>(m, "ExecuTorchModule")
10691078
.def("load_bundled_input", &PyModule::load_bundled_input, call_guard)

extension/pybindings/pybindings.pyi

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,3 +264,12 @@ def _reset_profile_results() -> None:
264264
This API is experimental and subject to change without notice.
265265
"""
266266
...
267+
268+
@experimental("This API is experimental and subject to change without notice.")
269+
def _unsafe_reset_threadpool(num_threads: int) -> None:
270+
"""
271+
.. warning::
272+
273+
This API is experimental and subject to change without notice.
274+
"""
275+
...

shim_et/xplat/executorch/extension/pybindings/pybindings.bzl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ def executorch_pybindings(python_module_name, srcs = [], cppdeps = [], visibilit
5454
],
5555
deps = [
5656
"//executorch/runtime/core:core",
57+
"//executorch/extension/threadpool:threadpool",
5758
] + cppdeps,
5859
external_deps = [
5960
"pybind11",

0 commit comments

Comments
 (0)