perf: Python binding inference performance improvement #426

kthui · 2025-01-11T02:47:56Z

What does the PR do?

Refactor infer() and async_infer() APIs to handle memory allocation and callbacks internally in C++, and only expose the basic interface for the Python iterator to fetch responses.

Checklist

Commit Type:

Check the conventional commit type
box here and add the label to the github PR.

Related PRs:

triton-inference-server/server#7949

Where should the reviewer start?

Start with the tritonserver_pybind.cc for the interface change, then move on to _model.py on how the Python iterator interacts with the interface. Finally, move on to the _request.py and _response.py on how they support the Python iterator.

For testing, start with test_binding.py and test_api.py, and then _tensor.py on the DLPack limitation regarding bytes.

Test plan:

Existing L0_python_api is sufficient for catching any regression from this performance improvement. It is modified to test from the new interface.

CI Pipeline ID: 23514700

Caveats:

User are no longer able to specify custom:

request release callback
response allocator
response callback

Currently, only CPU memory output is supported at the binding level, so GPU memory output will involve an extra D2H copy at the backend and a H2D copy at the frontend. This will be resolved as a follow-up.

The test_stop failure will have to be triaged and fixed as a follow-up.

Background

N/A

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

N/A

* fix request tensor lifecycle * remove prints for benchmarking * remove one h2h response copy * add allocator delete * schedule future.set_result to be called in event loop

* Use tensor object in output * Enable infer() to use the improved binding * Pass the C pointer directly to Python * Move output device if differ from request setting * Copyright and pre-commit

* Fix py_future object lifecycle * Fix request released after complete final

nnshah1 · 2025-01-15T21:28:25Z

python/test/test_api.py

@@ -137,6 +137,7 @@ def test_memory_fallback_to_cpu(self, server_options):

        tritonserver.default_memory_allocators[tritonserver.MemoryType.GPU] = allocator


should this be removed - or changed in some way to indicate the allocator is internal ....

yes, updated the test to indicate the allocator is internal, and it will always use CPU memory regardless of the backend memory preference.

nnshah1 · 2025-01-15T21:30:43Z

python/test/test_api.py

@@ -164,6 +165,7 @@ def test_memory_allocator_exception(self, server_options):
            ):
                pass

+    @pytest.mark.skip(reason="Skipping test, infer no longer use allocator")


we should keep this but refactor - if user requests output memory type gpu and that is not supported by the internal allocator - we would still want to raise an exception during inference

yes, we should raise an exception if the output memory type specified on the request is not supported, but currently the bindings does not accept a requested output memory type, so I think we can skip this test for now and add proper testing after adding support for allocating GPU memory.

hmm, but if we fail because say cupy is not available? Can we still make sure the right error gets propagated?

yes, the "test_unsupported_memory_type" is repurposed for testing moving outputs to unsupported memory type

nnshah1 · 2025-01-15T21:32:06Z

python/test/test_api.py

@@ -418,6 +420,9 @@ def test_ready(self, server_options):
        server = tritonserver.Server(server_options).start()
        assert server.ready()

+    @pytest.mark.skip(
+        reason="Skipping test, some request/response object may not be released which may cause server stop to fail"


Can we use xfail instead of skip?

@rmccorm4 - what do you think - as we had spent time on fixing this - how much an issue is this?

yes, switched to xfail.

python/test/test_binding.py

python/tritonserver/_c/tritonserver_pybind.cc

* Indicate allocator is internal on test_memory_fallback_to_cpu() * Remove test_memory_allocator_exception() as infer no longer use custom allocators * Update reason for skipping test_unsupported_memory_type() * Mark test_stop() with xfail instead of skip

python/tritonserver/_c/tritonserver_pybind.cc

nnshah1 · 2025-01-16T00:47:47Z

python/tritonserver/_c/tritonserver_pybind.cc

+      TRITONSERVER_MemoryType* actual_memory_type,
+      int64_t* actual_memory_type_id)
+  {
+    *buffer = malloc(byte_size * sizeof(uint8_t));


are there tritonserver apis we should be using for memory allocation instead of directly calling malloc?

I don't think there is one, because there are multiple calls to malloc in the core code base for different purpose, i.e.
https://github.com/triton-inference-server/core/blob/r24.12/src/backend_memory_manager.cc#L98
https://github.com/triton-inference-server/core/blob/r24.12/src/backend_model_instance.cc#L73
https://github.com/triton-inference-server/core/blob/r24.12/src/cache_manager.cc#L95

There is one on the server repo, but it is our of reach from the core binding:
https://github.com/triton-inference-server/server/blob/r24.12/src/memory_alloc.cc#L96

@GuanLuo to double check - should we be using the backend_memory_manager ? Reason I ask is that these / and the server already do support the different types of memory gpu / cpu / cpu pinned - and we would want that kind of support here if we can get it - and we should reuse it ...

What is in server certainly won't work, it is tightly related to the server frontend for optimization which is why allocator is up for the user to define.

You probably can reuse the one for backend but that is kind of strange. Plus the one in backend may not be as well crafted as you expected.. Please evaluate whether it may sense to proceed in this way.

Let's make that a todo as we bring in support for more memory types - we should look to see if we can reuse a common memory allocator in the core - of the ones backend_memeory_manager seems the closest - but agree is odd to reuse here directly.

nnshah1 · 2025-01-16T00:49:16Z

python/tritonserver/_c/tritonserver_pybind.cc

+    *actual_memory_type = TRITONSERVER_MEMORY_CPU;
+    *actual_memory_type_id = 0;
+    // The response allocator needs to be kept alive until the allocated memory
+    // is released via the release function.


is this true? - question: why not an allocator singleton for cpu - and reuse it for all requests?

yes, the core will need the allocator opaque object for finding the set release function. The allocator opaque object is basically this class that stores the release function pointer.

yes, updated to use singleton allocator object for all instances of the request wrapper.

nnshah1 · 2025-01-16T00:50:20Z

python/tritonserver/_c/tritonserver_pybind.cc

+      int64_t memory_type_id)
+  {
+    free(buffer);
+    // Release ownership of the response allocator.


thinking a singleton could be simpler and avoid alloc / dealloc of the allocator object itself - if possible.

yes, updated to use a singleton response allocator.

nnshah1 · 2025-01-16T00:51:30Z

python/tritonserver/_c/tritonserver_pybind.cc

+    // Release ownership of the response allocator.
+    std::unique_ptr<std::shared_ptr<struct TRITONSERVER_ResponseAllocator>>
+        allocator_ptr(reinterpret_cast<
+                      std::shared_ptr<struct TRITONSERVER_ResponseAllocator>*>(


actually - what do we use the buffer_userp for?

it was used to store the shared pointer to the response allocator object and passed to the release function, so the reference count can be increased/decreased as malloc/free that enables the allocator to be destructed only after the last allocated response memory is freed.

but it is no longer necessary with the singleton allocator.

nnshah1 · 2025-01-16T01:06:27Z

python/tritonserver/_c/tritonserver_pybind.cc

+      if (response_future_.get() != nullptr) {
+        throw AlreadyExistsError("cannot call GetNextResponse concurrently");
+      }
+      response_future_.reset(new py::object(py_future));


not sure I follow the logic here - why do we create a new object based on the passed in object? why not a reference on the original object?

second: how does the swap below work?

I tried holding a raw pointer to the original object, but the pointer can no longer be dereferenced (segmentation fault) after the GetNextResponse() returned, so the only safe way of holding on to the future object after GetNextResponse() returns is to increment the reference count on the future object, which is achieved by "copying" it to a new object.

the reason for using a unique pointer is to have a mechanism for checking if there is a future object pending to be set with the next response without accessing any Python variable, to avoid holding the GIL when doing the simple check. it can be achieved with a simple bool flag, but I think it is safer to have one variable than two to avoid setting one variable but forgetting about the other one.

~~Experiment: Increment the reference count on the py_future object, and see if the object can be remained valid after the function return.~~

TODO: The new will always allocate on heap, can we keep the variable on the rts, to avoid new/delete for every response?

Can we have some kind of small memory leak/growth pytest test that checks memory before, does some inferences, and checks memory after?

The implementation is updated to remove the pointer wrapper around response_future_, and more comments clarifying the flow among GetNextResponse() and AddNextResponse(): Avoid extra malloc when tracking response_future_
cc @nnshah1 @rmccorm4

Can we have some kind of small memory leak/growth pytest test that checks memory before, does some inferences, and checks memory after?

Do you think something like this makes sense? I think there will always be some small growth, but we can set the acceptable growth on the test case and have the test fail if it grows beyond what we set.

More discussion here: #426 (comment)

Nice find! Looks worth a try to me - I see it's used by quite a few big projects - https://github.com/search?q=pytest-memray&type=code

pydantic

datadog

huggingface

...

I'm curious how much time overhead it adds to run.

but we can set the acceptable growth on the test case and have the test fail if it grows beyond what we set.

Let's make sure whatever threshold we set isn't tailored to the specific test case or number of iterations, and doesn't get hit if we increased the iterations on the test case or something

nnshah1 · 2025-01-16T01:08:35Z

python/tritonserver/_c/tritonserver_pybind.cc

+    {
+      py::gil_scoped_acquire gil;
+      std::unique_ptr<py::object> response_future_local(nullptr);
+      response_future_.swap(response_future_local);


not sure I follow this - we create a new py object from null, swap it with the stored one (releasing it?) - then set result on the nullptr based object?...

Or ... I guess swap works the opposite way :)

we swap the null ptr for the one stored in the future ....

the guarantee here that the GetNextResponse() will not be called and receive a new future object (vs this future object) is this future object is not yet set with a result, so as soon as this future object is set with a result, we need to be ready for the next future object from GetNextResponse() that may take the place of this future object. thus, the solution here is to move the current future object from the class to local variable, making the class variable ready for the next future object before the current future object is set with a result.

python/tritonserver/_c/tritonserver_pybind.cc

nnshah1 · 2025-01-16T01:14:17Z

python/tritonserver/_api/_model.py

-            inference_request.response_queue,
-            raise_on_error,
-        )
+        self._server.infer_async(request)


how do we pass in the requested memory type to the C++ side?

I think we can add a new binding function to the request wrapper that allows setting a requested output memory type to the request wrapper object, then the response allocation function can access the requested memory type from the request wrapper object and allocate the correct memory type.

let's make sure the stories to follow up on customizing allocation are in place and scheduled. Priority would be to specify the memory type (device id, etc.)

yes, DLIS-7824.

GuanLuo · 2025-01-29T11:41:12Z

python/tritonserver/_c/tritonserver_pybind.cc

+    // The request wrapper needs to be kept alive until the release callback is
+    // invoked, so the Triton request can be deleted after the core/backend is
+    // done with it.
+    std::unique_ptr<std::shared_ptr<PyInferenceRequest>> request_ptr(


IIUC, the unique pointer is used here to make sure the resource is being cleaned up if exception is thrown below.. But if that is the intention, I don't think the current code does the job as the ownership is released before calling into TRITONSERVER_InferenceRequestSetReleaseCallback which is not safe guarding the exception.

GuanLuo · 2025-01-29T11:41:57Z

python/tritonserver/_c/tritonserver_pybind.cc

+      struct TRITONSERVER_InferenceRequest* request, const uint32_t flags,
+      void* userp)
+  {
+    std::unique_ptr<std::shared_ptr<PyInferenceRequest>> request_ptr(


Yea, at this point just reinterpret cast and delete should be fine.

GuanLuo · 2025-01-29T11:56:53Z

python/tritonserver/_c/tritonserver_pybind.cc

+      ThrowIfError(TRITONSERVER_ResponseAllocatorSetBufferAttributesFunction(
+          allocator_.get(), ResponseAllocatorBufferAttributesFn));
+    });
+  }


In this case, I think you can move the allocator initialization and the related callback functions out to a separate class and make it a singleton, which then this SetResponseAllocator function simply grabs the Triton allocator pointer from the singleton

Yes, this is a good idea. I think we can proceed with the refactoring when we add GPU memory support into the allocators, which the class abstraction will be necessary to keep the function block short. DLIS-7824

GuanLuo · 2025-01-29T11:58:08Z

python/tritonserver/_c/tritonserver_pybind.cc

-    auto cr = reinterpret_cast<CallbackResource*>(userp);
-    cr->release_fn(cr->request, flags, cr->user_object);
-    delete cr;
+    ThrowIfError(TRITONSERVER_BufferAttributesSetMemoryType(


Should return error as that is what the caller expects

GuanLuo · 2025-01-29T12:05:16Z

python/tritonserver/_c/tritonserver_pybind.cc

+      TRITONSERVER_MemoryType* actual_memory_type,
+      int64_t* actual_memory_type_id)
+  {
+    *buffer = malloc(byte_size * sizeof(uint8_t));


What is in server certainly won't work, it is tightly related to the server frontend for optimization which is why allocator is up for the user to define.

You probably can reuse the one for backend but that is kind of strange. Plus the one in backend may not be as well crafted as you expected.. Please evaluate whether it may sense to proceed in this way.

kthui · 2025-02-04T18:15:05Z

What is in server certainly won't work, it is tightly related to the server frontend for optimization which is why allocator is up for the user to define.

You probably can reuse the one for backend but that is kind of strange. Plus the one in backend may not be as well crafted as you expected.. Please evaluate whether it may sense to proceed in this way.

I think we will proceed with a static class in the binding that is singleton and exclusive for the binding, which as support for both CPU and GPU allocation and is optimized for the binding use-case. DLIS-7824

rmccorm4 · 2025-02-07T07:12:22Z

python/test/test_api.py

+    # increase is not proportional to the increase in inference requests.
+    @pytest.mark.parametrize("infer_repeat", [200, 500, 1000, 8000])
+    def test_inference_leak(self, server_options, infer_repeat):
+        min_mem_inc, max_mem_inc = 20000, 120000


are these values in bytes? 20KB and 120KB?

yes, those are in bytes, so 20KB and 120KB

Can you add _bytes suffix to the var names?

Sure, added the _bytes suffix. Add _bytes suffix to mem size vars

rmccorm4 · 2025-02-07T07:13:51Z

python/test/test_api.py

+    # max_infer_repeat / min_infer_repeat >> max_mem_inc / min_mem_inc, so the memory
+    # increase is not proportional to the increase in inference requests.


If there is no per-request/response memory leaks, shouldn't sending ~200 requests and ~8000 requests with the same concurrency/batching see roughly the same amount of memory usage before/after?

Do you have more detailed results/findings from local experiments that led you to these numbers and methodology?

I found the memory usage increment will start higher and converge to around 40KB locally on my machine as the test progresses, regardless of the requests count, i.e.

requests = 200; memory increment = 71193

requests = 500; memory increment = 60405

requests = 1000; memory increment = 40133

requests = 8000; memory increment = 38519

and in the opposite order:

requests = 8000; memory increment = 83592

requests = 1000; memory increment = 47596

requests = 500; memory increment = 41406

requests = 200; memory increment = 40853

requests = 1000 repeat 6x:

memory increment = 84436

memory increment = 41967

memory increment = 40166

memory increment = 45353

memory increment = 41281

memory increment = 34740

rmccorm4 · 2025-02-07T07:21:57Z

python/tritonserver/_api/_response.py

+        if self._iterator_queue is None:
+            response = self._get_next_response()
+        else:
+            response = self._iterator_queue.get()


Why is iterator_queue only used when a user_queue is provided? Would it make sense to always use iterator_queue for no if/else here?

One major difference currentlly seems to be only getting next response on demand from iteration calls to __next__ vs. pre-fetching the responses to iterator queue in a background thread, even if the response iterator isn't iterated on.

Can you elaborate on your thoughts/decision here?

Why is iterator_queue only used when a user_queue is provided? Would it make sense to always use iterator_queue for no if/else here?

The code would be a bit simpler if the iterator_queue is always used, but it is not needed if the user did not provide an additional response queue to be feed, so the iterator can pull directly from the binding without the intermediate iterator queue, which in theory will allow for a higher maximum output throughput / lower latency compared to w/ the intermediate iterator queue (and its thread/coroutine).

In other words, we can bypass the intermediate queue when there is only a single consumer for the producer at the binding. If there is more than one consumer for the producer, an additional queue for each consumer is needed to allow those consumer to consume at a different rate without blocking each other. A key part of the test is to make sure the user provided queue is feed with responses without having to iterate over the iterator.

rmccorm4

Majority of changes LGTM - but left a few questions to wrap up tomorrow

rmccorm4 · 2025-02-08T01:18:34Z

Let's prioritize the follow-ups for consistent tests (grpc fork errors), gpu memory support, etc

rmccorm4 · 2025-02-08T01:18:50Z

Nice work Jacky! Thanks for pushing this through!

[PoC] Implement allocators and callbacks entirely in C++ binding

ad0c007

* fix request tensor lifecycle * remove prints for benchmarking * remove one h2h response copy * add allocator delete * schedule future.set_result to be called in event loop

kthui added the PR: perf A code change that improves performance label Jan 11, 2025

kthui self-assigned this Jan 11, 2025

Implement the PoC into production code

131078a

* Use tensor object in output * Enable infer() to use the improved binding * Pass the C pointer directly to Python * Move output device if differ from request setting * Copyright and pre-commit

kthui force-pushed the jacky-py-res-callback branch from 300b0b6 to 131078a Compare January 11, 2025 02:52

kthui added 4 commits January 11, 2025 17:23

Some optimizations and fixes

83f78f3

* Fix py_future object lifecycle * Fix request released after complete final

Fix resposne allocator lifecycle

279af74

Update test_binding.py

5e805da

Update test_api.py and fix from_dlpack with bytes type

2faf303

kthui requested review from GuanLuo, nnshah1, tanmayv25 and rmccorm4 January 15, 2025 19:42

kthui marked this pull request as ready for review January 15, 2025 19:43

nnshah1 reviewed Jan 15, 2025

View reviewed changes

python/test/test_binding.py Show resolved Hide resolved

nnshah1 reviewed Jan 15, 2025

View reviewed changes

python/test/test_binding.py Show resolved Hide resolved

nnshah1 reviewed Jan 15, 2025

View reviewed changes

python/tritonserver/_c/tritonserver_pybind.cc Show resolved Hide resolved

Updates to test_api.py

e1907bb

* Indicate allocator is internal on test_memory_fallback_to_cpu() * Remove test_memory_allocator_exception() as infer no longer use custom allocators * Update reason for skipping test_unsupported_memory_type() * Mark test_stop() with xfail instead of skip

nnshah1 reviewed Jan 16, 2025

View reviewed changes

python/tritonserver/_c/tritonserver_pybind.cc Show resolved Hide resolved

nnshah1 reviewed Jan 16, 2025

View reviewed changes

python/tritonserver/_c/tritonserver_pybind.cc Show resolved Hide resolved

nnshah1 reviewed Jan 16, 2025

View reviewed changes

GuanLuo reviewed Jan 29, 2025

View reviewed changes

kthui added 10 commits February 3, 2025 16:42

Clean up notes on inference_request and request relationship

7df58d5

Unify definition of DeviceOrMemoryType

bc6e7ac

Update inference response conversion interface

023a145

Add ability to cancel a request from response iterator

dbbd68e

Add raise on error to response iterators

898ef02

Clean up base pointer deallocation comment

9575a00

Remove unique_ptr wrapper on release callback

141483f

Return if error when setting buffer attributes

c534d3a

Assert memory type and id on release function

a8f7bda

Update comments related to GetNextResponse() and AddNextResponse()

1d4f98d

kthui added 4 commits February 4, 2025 11:03

Reorder _memorybuffer.py import

bfa74b4

Add test for infer memory increase

adf98e6

Add user custom response queue support

cfab4cf

Adjust mem leak test min and max mem size

f37d9a9

kthui requested review from GuanLuo, rmccorm4 and nnshah1 February 6, 2025 18:00

rmccorm4 reviewed Feb 7, 2025

View reviewed changes

kthui requested a review from rmccorm4 February 7, 2025 21:57

Add _bytes suffix to mem size vars

536cf98

rmccorm4 approved these changes Feb 8, 2025

View reviewed changes

kthui merged commit d53c1f7 into main Feb 8, 2025
2 checks passed

kthui deleted the jacky-py-res-callback branch February 8, 2025 01:28

		@@ -137,6 +137,7 @@ def test_memory_fallback_to_cpu(self, server_options):

		tritonserver.default_memory_allocators[tritonserver.MemoryType.GPU] = allocator

		# max_infer_repeat / min_infer_repeat >> max_mem_inc / min_mem_inc, so the memory
		# increase is not proportional to the increase in inference requests.

perf: Python binding inference performance improvement #426

perf: Python binding inference performance improvement #426

Conversation

kthui commented Jan 11, 2025 • edited Loading

What does the PR do?

Checklist

Commit Type:

Related PRs:

Where should the reviewer start?

Test plan:

Caveats:

Background

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

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

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

kthui Jan 17, 2025 • 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

kthui Jan 17, 2025 • edited Loading

Choose a reason for hiding this comment

kthui Jan 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kthui Jan 28, 2025 • edited Loading

Choose a reason for hiding this comment

rmccorm4 Jan 28, 2025 • 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

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

kthui commented Feb 4, 2025

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

Choose a reason for hiding this comment

rmccorm4 left a comment

Choose a reason for hiding this comment

rmccorm4 commented Feb 8, 2025

rmccorm4 commented Feb 8, 2025

kthui commented Jan 11, 2025 •

edited

Loading

kthui Jan 17, 2025 •

edited

Loading

kthui Jan 17, 2025 •

edited

Loading

kthui Jan 27, 2025 •

edited

Loading

kthui Jan 28, 2025 •

edited

Loading

rmccorm4 Jan 28, 2025 •

edited

Loading