Support RegisterCustomOpsLibrary via the Python API #4764

hariharans29 · 2020-08-12T04:50:14Z

Description: RegisterCustomOpsLibrary is a mechanism via which custom kernels in shared libraries can be used to support a model containing ops not yet in the ONNX opset. This change supports it via the Python API.

There are some challenges associated with this:

Freeing OrtCustomOpDomains associated with the custom kernels (

onnxruntime/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc

Line 120 in fdc5c30

OrtCustomOpDomain* domain = nullptr;

). Currently freeing is done by the user of C/C++ API via OrtReleaseCustomOpDomain(). In Python, just like in C/C++, we rely on the implementer of the shared library to release the OrtCustomOpDomains created within it (i.e.) we don't take ownership and manage the resource for the user.
Unloading the shared library
Currently the user has to manually unload the library via the provided handle while using the C/C++ API. In Python, we unload the library automatically (no action needed by the user) when the last session using the session options connected with the shared library goes out of scope.

TODO:

Update relevant documentation in the custom op md files
See how this plays with the shallow copy and the deep copying mechanism offered by the copy module (https://docs.python.org/3/library/copy.html). If we think it is unsafe to allow deep copying SessionOptions instances in Python (because they hold associated OrtCustomOpDomain pointers), explicit disallow it by overriding the method for now until we have support for it. Python aliases are okay and they are not copies. For example, so1 = Sessionoptions() so2 = so1. This only creates aliases to the same object and preserves the object until all the referenced aliases go out of scope and this is supported.
EDIT: It seems like currently SessionOptions isn't picklable. We need to investigate which instance attribute is causing the pickling issue and solve that. (Not sure if we even want to support pickling given that it could lead to security concerns). But given that it is not pickable, nobody will be able to shallow copy (copy.copy()) or deep copy (copy.deepcopy()) SessionOptions using the copy module and hence there is no action required.

CC: @wenbingl

Motivation and Context
Fix #3958

wenbingl · 2020-08-12T22:39:19Z

onnxruntime/test/python/onnxruntime_test_python.py

+        so2 = so1
+
+        # Model loading successfully indicates that the custom op node could be resolved successfully
+        sess1 = onnxrt.InferenceSession(custom_op_model, so1)


can we add custom_op_lib_paths as an optional arguments? So the user doesn't need to jump to another sessionOption to do that. And this looks more pythonic.

Do you want to support an arg in InferenceSession ? That is difficult to do.

Also, the user doesn't need to jump to another session options. There is only one session options.

so = ort.SessionOptions() so.register_custom_ops_library(shared_lib_path) sess = ort.InferenceSession(model_path, so)

I think this is very simple and it aligns with all the other language bindings. Also, session options can be shared across multiple sessions. If we tie it to Session, we have to keep referencing the library in all the Sessions.

I added another session options just to demonstrate that it is "re-useable"

It may not be difficult as you think.
In InferenceSession.init python function, just add a few lines as following
if custom_op_lib_paths:
so.register_custom_ops_library(custom_op_lib_paths)
...

The overall C++ and pybinding interface keeps same, the only difference is to make python wrapper more friendly to avoid everyone write couple of same code lines.

if custom_op_lib_paths: so.register_custom_ops_library(custom_op_lib_paths)

I find this hacky and not providing too much value (saves 2 lines of code while potentially causing maintenance hell) for several reasons:

What is so here then ? Construct a SessionOptions instance in InferenceSession on the fly just to hold the custom op libs ?

What if the user provides his/her own SessionOptions ? Then do we add the custom_op to that SessionOptions instance (and mutate something the user provided) or do we hold another "dummy" SessionOptions instance just to hold the custom_ops ?

I think this is unnecessary complications just to save so = Ort.SessionOptions so.register_ops_library(path). This is more clear in that it ties the shared library to the SessionOptions instance and thereby allowing a user to share it across sessions if he/she chooses to do so.

Also, by exposing the extra optional parameter in InferenceSession init, it does not keep it consistent with c/c++/c# - none of these language bindings will expose a Session constructor that takes a path to the shared library - they need to pass it in via SessionOptions as well. My suggestion is to stick with the same.

Is there any concern with instantiating a SessionOptions instance ?

Unless you think these two lines are problematic, I don't see what's the difference between putting it in ort code wrapper and the user code, right? And if you are the user, will it be helpful if there is a simpler API for you?

And No one expects C++ and Python interface are exactly same, only the parity features.

Is there any ETA on merging this PR? My concern is not a blocking issue here.

Thanks Wenbing.I think it is problematic for the concerns/ reasons listed above. Also, I don’t want to set the precedence of exposing an optional parameter in the InferenceSession constructor for every Session option because that defeats the purpose of SessionOptions.

The PR will be merged once I solve all build issues- working on it now.

hariharans29 · 2020-08-13T22:30:34Z

onnxruntime/core/session/custom_ops.cc

+    // If domain is empty, it is assumed to be part of the ONNX domain
+    if (domain->domain_[0]) {
+      // Add it to the DomainToVersion ONNX map if it doesn't already exist
+      // For example, two sessions using the same session_options should not add the same custom op domain to the version map twice


Fails a debug assert in the ONNX code if we try to add an item that already exists in the map. hence, added this guard.

hariharans29 · 2020-08-13T22:31:08Z

onnxruntime/python/onnxruntime_pybind_state.cc

+      // A non-nullptr indicates some error
+      // Free status and throw
+      platform_env.UnloadDynamicLibrary(library_handle_);
+      ::free(status);


Free the returned OrtStatus if not nullptr (error scenario)

hariharans29 · 2020-08-13T22:31:35Z

onnxruntime/python/onnxruntime_pybind_state.cc

+
+// Unload the library when the destructor is triggered
+CustomOpLibrary::~CustomOpLibrary() {
+  platform_env.UnloadDynamicLibrary(library_handle_);


unload the library when the corresponding struct that holds it goes out of scope

hariharans29 · 2020-08-13T22:32:56Z

onnxruntime/python/onnxruntime_pybind_state_common.h

+  ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CustomOpLibraries);
+
+ private:
+  std::vector<std::unique_ptr<CustomOpLibrary>> custom_op_libraries_;


CustomOpLibraries will "own" the Customlibrary. The session options that needs to load the shared library needn't hold onto it after it has gotten what it needs from the library.

onnxruntime/core/session/custom_ops.cc

onnxruntime/python/onnxruntime_pybind_state_common.h

onnxruntime/python/onnxruntime_pybind_state.cc

orttraining/orttraining/python/orttraining_pybind_state.cc

onnxruntime/python/onnxruntime_pybind_state.cc

onnxruntime/python/onnxruntime_pybind_state_common.h

orttraining/orttraining/python/orttraining_pybind_state.cc

docs/AddingCustomOp.md

saurabh-shandilya · 2020-09-07T23:24:41Z

@hariharans29 Thanks a lot for working on this. This is going to be very useful. When will this be available in the release branch or which release branch?

hariharans29 · 2020-09-08T02:57:16Z

Hi Saurabh- It should be in the upcoming release (due to release soon this month).

hariharans29 added 2 commits August 9, 2020 22:12

Initial commit

306a114

More changes

8bcc970

hariharans29 requested a review from a team as a code owner August 12, 2020 04:50

hariharans29 added 8 commits August 11, 2020 21:55

Resolve conflicts

b033785

Nit

ebb6987

Nit

a72aa6d

Fix warning

5a735ed

More changes

f65e504

Debug linux failure

cf70867

debug

08ebd49

debug 2

a1c53c9

wenbingl reviewed Aug 12, 2020

View reviewed changes

hariharans29 mentioned this pull request Aug 13, 2020

Re-enable the C/C++ API test for custom op libraries #4776

Closed

hariharans29 added 2 commits August 13, 2020 13:31

Debug

62064b8

More changes

6514806

hariharans29 commented Aug 13, 2020

View reviewed changes

hariharans29 added 3 commits August 13, 2020 19:03

More changes

aecbc45

Changes

14b0b36

Fix link

03569a0

hariharans29 changed the title ~~WIP: Support RegisterCustomOpsLibrary via the Python API~~ Support RegisterCustomOpsLibrary via the Python API Aug 14, 2020

hariharans29 added 4 commits August 18, 2020 19:17

Fix training build

bb4bf92

Merge master and resolve conflicts

29bb93f

Fix build

0f08762

Fix build

fa64d83

hariharans29 requested a review from jywu-msft August 19, 2020 04:31

hariharans29 added the api:Python label Aug 20, 2020

pranavsharma reviewed Aug 20, 2020

View reviewed changes

hariharans29 added 5 commits August 24, 2020 12:53

Merge remote-tracking branch 'origin/master' into CustomOpPythonApi

d46b900

PR feednack

537337c

Fix training build

dcc1963

Fix training build

3c554ad

Debug training build

adfa4e3

hariharans29 commented Aug 25, 2020

View reviewed changes

orttraining/orttraining/python/orttraining_pybind_state.cc Show resolved Hide resolved

pranavsharma reviewed Aug 25, 2020

View reviewed changes

onnxruntime/python/onnxruntime_pybind_state.cc Show resolved Hide resolved

onnxruntime/python/onnxruntime_pybind_state.cc Show resolved Hide resolved

hariharans29 added 12 commits August 25, 2020 13:14

a

7f562be

Fix training build

bc1e4d3

Revert some changes

ad8ff15

More changes

61450d4

Fix mac build

51d8f43

Fix build

e0977b4

More changes

bb8e0b1

Log

ac4f920

Nit

feb8501

Nits

82ef63d

refinement

1d0d7ed

Nit

57b4490

pranavsharma reviewed Aug 28, 2020

View reviewed changes

PR feedback

a841e6d

pranavsharma approved these changes Aug 28, 2020

View reviewed changes

hariharans29 merged commit 7045910 into master Aug 28, 2020

hariharans29 deleted the CustomOpPythonApi branch August 28, 2020 20:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support RegisterCustomOpsLibrary via the Python API #4764

Support RegisterCustomOpsLibrary via the Python API #4764

hariharans29 commented Aug 12, 2020 •

edited

Loading

wenbingl Aug 12, 2020

hariharans29 Aug 12, 2020 •

edited

Loading

hariharans29 Aug 12, 2020

wenbingl Aug 13, 2020

hariharans29 Aug 13, 2020 •

edited

Loading

wenbingl Aug 26, 2020

hariharans29 Aug 26, 2020

hariharans29 Aug 13, 2020

hariharans29 Aug 13, 2020

hariharans29 Aug 13, 2020

hariharans29 Aug 13, 2020

saurabh-shandilya commented Sep 7, 2020

hariharans29 commented Sep 8, 2020

Support RegisterCustomOpsLibrary via the Python API #4764

Support RegisterCustomOpsLibrary via the Python API #4764

Conversation

hariharans29 commented Aug 12, 2020 • edited Loading

Choose a reason for hiding this comment

hariharans29 Aug 12, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hariharans29 Aug 13, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

saurabh-shandilya commented Sep 7, 2020

hariharans29 commented Sep 8, 2020

hariharans29 commented Aug 12, 2020 •

edited

Loading

hariharans29 Aug 12, 2020 •

edited

Loading

hariharans29 Aug 13, 2020 •

edited

Loading