Python Session.run_async Causes Program Exit

[ivanthewebber]

Describe the issue

I am trying to utilize a callback and future with Session.run_async in order to asynchronously run inference without blocking the asyncio event loop. However, when I try to use Session.run_async the program exits silently.

Here is the relevant snippet of my code:

        fut = concurrent.futures.Future()
        def callback(outputs: np.ndarray, _, err: str):
            if err:
                fut.set_exception(RuntimeError(err))
            else:
                fut.set_result(outputs[0])

        # TODO FIXME calling thread exits after a second after this method even if debugger breakpoint
        self.session.run_async(
            self.config.get("output_names", None),
            vecs,
            callback,
            None,
            run_opts
        )

        return await fut

The documentation seems to be lacking, but hopefully this is something easy to work around.

To reproduce

Will add tomorrow.

Urgency

No response

Platform

Windows

OS Version

11

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.20.0

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[ivanthewebber]

There must be a dozen open issues for run_async, but there is not one working example I can find. What a shame. Perhaps Triton Server will work

[ivanthewebber]

I found the unit test for run_async:

onnxruntime/onnxruntime/test/python/onnxruntime_test_python.py

Line 704 in 1f70fc2

def test_run_async(self):

I found that the program only silently exits when RunOptions is provided; specifically my code was not setting any of the fields.

Here is a new unit test that could be added reproducing this problem, but keep in mind if the program silently exits the test may appear to have passed:

import concurrent.futures

import numpy as np

import onnxruntime as onnxrt

def test_run_async(self):
    fut = concurrent.futures.Future()
    output_expected = np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32)

    def callback(res: np.ndarray, fut: concurrent.futures.Future, err: str) -> None:
        if err:
            fut.set_exception(RuntimeError(err))
        else:
            fut.set_result(res)

    so = onnxrt.SessionOptions()
    so.intra_op_num_threads = 2

    sess = onnxrt.InferenceSession(get_name("mul_1.onnx"), so, providers=available_providers)

    x = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32)
    sess.run_async(
        ["Y"], 
        {"X": x}, 
        callback, 
        fut, 
        onnxrt.RunOptions() # this causes the program to exit
    )

    res = fut.result(10)  # timeout in 10 sec
    # e.g. can be used with asyncio.wrap_future(fut)
    self.assertTrue(fut.done())
    self.assertEqual(len(res), 1)
    np.testing.assert_allclose(output_expected, res, rtol=1e-05, atol=1e-08)

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

[ivanthewebber]

This bug has never been triaged and is easy to reproduce; can someone please comment if they have the same issue?