You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by GitBox <gi...@apache.org> on 2022/11/08 19:32:42 UTC
[GitHub] [beam] jrmccluskey opened a new pull request, #24039: [Draft] TensorRT Custom Inference Function Implementation
jrmccluskey opened a new pull request, #24039:
URL: https://github.com/apache/beam/pull/24039
Adds support for specifying a custom inference function to the TensorRT model handler, enabling custom logic to be run during RunInference calls.
Part of #22572
------------------------
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
- [ ] [**Choose reviewer(s)**](https://beam.apache.org/contribute/#make-your-change) and mention them in a comment (`R: @username`).
- [ ] Mention the appropriate issue in your description (for example: `addresses #123`), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment `fixes #<ISSUE NUMBER>` instead.
- [ ] Update `CHANGES.md` with noteworthy changes.
- [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://beam.apache.org/contribute/get-started-contributing/#make-the-reviewers-job-easier).
To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
GitHub Actions Tests Status (on master branch)
------------------------------------------------------------------------------------------------
[](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24039:
URL: https://github.com/apache/beam/pull/24039#issuecomment-1316914959
Reminder, please take a look at this pr: @ryanthompson591
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] jrmccluskey commented on pull request #24039: TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
jrmccluskey commented on PR #24039:
URL: https://github.com/apache/beam/pull/24039#issuecomment-1322574427
Only test failure is apache_beam/transforms/stats_test.py::ApproximateQuantilesTest::test_lots_of_duplicates which is unrelated to this change
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] codecov[bot] commented on pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
codecov[bot] commented on PR #24039:
URL: https://github.com/apache/beam/pull/24039#issuecomment-1307772514
# [Codecov](https://codecov.io/gh/apache/beam/pull/24039?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) Report
> Merging [#24039](https://codecov.io/gh/apache/beam/pull/24039?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (60e289e) into [master](https://codecov.io/gh/apache/beam/commit/62d8bdc604dfeccae08d00225bce8c9b9c3ffb28?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (62d8bdc) will **decrease** coverage by `0.01%`.
> The diff coverage is `0.00%`.
```diff
@@ Coverage Diff @@
## master #24039 +/- ##
==========================================
- Coverage 73.39% 73.38% -0.02%
==========================================
Files 710 710
Lines 96122 96127 +5
==========================================
- Hits 70552 70539 -13
- Misses 24253 24271 +18
Partials 1317 1317
```
| Flag | Coverage Δ | |
|---|---|---|
| python | `83.18% <0.00%> (-0.03%)` | :arrow_down: |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#carryforward-flags-in-the-pull-request-comment) to find out more.
| [Impacted Files](https://codecov.io/gh/apache/beam/pull/24039?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | Coverage Δ | |
|---|---|---|
| [...hon/apache\_beam/ml/inference/tensorrt\_inference.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vbWwvaW5mZXJlbmNlL3RlbnNvcnJ0X2luZmVyZW5jZS5weQ==) | `0.00% <0.00%> (ø)` | |
| [.../python/apache\_beam/testing/test\_stream\_service.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vdGVzdGluZy90ZXN0X3N0cmVhbV9zZXJ2aWNlLnB5) | `88.09% <0.00%> (-4.77%)` | :arrow_down: |
| [.../apache\_beam/runners/interactive/dataproc/types.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy9pbnRlcmFjdGl2ZS9kYXRhcHJvYy90eXBlcy5weQ==) | `93.10% <0.00%> (-3.45%)` | :arrow_down: |
| [sdks/python/apache\_beam/utils/interactive\_utils.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vdXRpbHMvaW50ZXJhY3RpdmVfdXRpbHMucHk=) | `95.12% <0.00%> (-2.44%)` | :arrow_down: |
| [.../python/apache\_beam/transforms/periodicsequence.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vdHJhbnNmb3Jtcy9wZXJpb2RpY3NlcXVlbmNlLnB5) | `97.01% <0.00%> (-1.50%)` | :arrow_down: |
| [...python/apache\_beam/runners/worker/worker\_status.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy93b3JrZXIvd29ya2VyX3N0YXR1cy5weQ==) | `75.33% <0.00%> (-1.34%)` | :arrow_down: |
| [...che\_beam/runners/interactive/interactive\_runner.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy9pbnRlcmFjdGl2ZS9pbnRlcmFjdGl2ZV9ydW5uZXIucHk=) | `90.50% <0.00%> (-1.27%)` | :arrow_down: |
| [...ython/apache\_beam/io/gcp/bigquery\_read\_internal.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vaW8vZ2NwL2JpZ3F1ZXJ5X3JlYWRfaW50ZXJuYWwucHk=) | `57.21% <0.00%> (-0.49%)` | :arrow_down: |
| [sdks/python/apache\_beam/io/gcp/bigquery.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vaW8vZ2NwL2JpZ3F1ZXJ5LnB5) | `71.11% <0.00%> (-0.24%)` | :arrow_down: |
| [...ks/python/apache\_beam/runners/worker/sdk\_worker.py](https://codecov.io/gh/apache/beam/pull/24039/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy93b3JrZXIvc2RrX3dvcmtlci5weQ==) | `89.08% <0.00%> (-0.17%)` | :arrow_down: |
:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] jrmccluskey commented on a diff in pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
jrmccluskey commented on code in PR #24039:
URL: https://github.com/apache/beam/pull/24039#discussion_r1028358991
##########
sdks/python/apache_beam/ml/inference/tensorrt_inference.py:
##########
@@ -164,11 +165,63 @@ def get_engine_attrs(self):
self.stream)
+TensorRTInferenceFn = Callable[
+ [Sequence[np.ndarray], TensorRTEngine, Optional[Dict[str, Any]]],
+ Iterable[PredictionResult]]
+
+
+def _default_tensorRT_inference_fn(
+ batch: Sequence[np.ndarray],
+ engine: TensorRTEngine,
+ inference_args: Optional[Dict[str,
+ Any]] = None) -> Iterable[PredictionResult]:
+ from cuda import cuda
+ (
+ engine,
+ context,
+ context_lock,
+ inputs,
+ outputs,
+ gpu_allocations,
+ cpu_allocations,
+ stream) = engine.get_engine_attrs()
+
+ # Process I/O and execute the network
+ with context_lock:
+ _assign_or_fail(
+ cuda.cuMemcpyHtoDAsync(
+ inputs[0]['allocation'],
+ np.ascontiguousarray(batch),
+ inputs[0]['size'],
+ stream))
+ context.execute_async_v2(gpu_allocations, stream)
+ for output in range(len(cpu_allocations)):
+ _assign_or_fail(
+ cuda.cuMemcpyDtoHAsync(
+ cpu_allocations[output],
+ outputs[output]['allocation'],
+ outputs[output]['size'],
+ stream))
+ _assign_or_fail(cuda.cuStreamSynchronize(stream))
+
+ return [
+ PredictionResult(
+ x, [prediction[idx] for prediction in cpu_allocations]) for idx,
+ x in enumerate(batch)
+ ]
+
+
@experimental(extra_message="No backwards-compatibility guarantees.")
class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
PredictionResult,
TensorRTEngine]):
- def __init__(self, min_batch_size: int, max_batch_size: int, **kwargs):
+ def __init__(
+ self,
+ min_batch_size: int,
+ max_batch_size: int,
+ *,
Review Comment:
That establishes `inference_fn` as a keyword-only parameter. Anand suggested it as the convention in #23642 to avoid it being handled as a positional argument. More information on this is at https://peps.python.org/pep-3102/
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] jrmccluskey commented on a diff in pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
jrmccluskey commented on code in PR #24039:
URL: https://github.com/apache/beam/pull/24039#discussion_r1028381004
##########
sdks/python/apache_beam/ml/inference/tensorrt_inference.py:
##########
@@ -241,40 +297,7 @@ def run_inference(
Returns:
An Iterable of type PredictionResult.
"""
- from cuda import cuda
- (
- engine,
- context,
- context_lock,
- inputs,
- outputs,
- gpu_allocations,
- cpu_allocations,
- stream) = engine.get_engine_attrs()
-
- # Process I/O and execute the network
- with context_lock:
- _assign_or_fail(
- cuda.cuMemcpyHtoDAsync(
- inputs[0]['allocation'],
- np.ascontiguousarray(batch),
- inputs[0]['size'],
- stream))
- context.execute_async_v2(gpu_allocations, stream)
- for output in range(len(cpu_allocations)):
- _assign_or_fail(
- cuda.cuMemcpyDtoHAsync(
- cpu_allocations[output],
- outputs[output]['allocation'],
- outputs[output]['size'],
- stream))
- _assign_or_fail(cuda.cuStreamSynchronize(stream))
-
- return [
- PredictionResult(
- x, [prediction[idx] for prediction in cpu_allocations]) for idx,
- x in enumerate(batch)
- ]
+ return self.inference_fn(batch, engine, inference_args)
Review Comment:
Yeah the format is relatively simple, add a multiple to the original prediction output and make sure we get the modified output. I've made a run at it but my dev environment isn't meshing with the tests so it may take a few iterations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24039:
URL: https://github.com/apache/beam/pull/24039#issuecomment-1307858867
Assigning reviewers. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @ryanthompson591 for label python.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
The PR bot will only process comments in the main thread (not review comments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24039:
URL: https://github.com/apache/beam/pull/24039#issuecomment-1321964799
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @pabloem for label python.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] jrmccluskey commented on pull request #24039: TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
jrmccluskey commented on PR #24039:
URL: https://github.com/apache/beam/pull/24039#issuecomment-1323844585
Run Python PreCommit
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] jrmccluskey merged pull request #24039: TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
jrmccluskey merged PR #24039:
URL: https://github.com/apache/beam/pull/24039
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] ryanthompson591 commented on a diff in pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
ryanthompson591 commented on code in PR #24039:
URL: https://github.com/apache/beam/pull/24039#discussion_r1028101807
##########
sdks/python/apache_beam/ml/inference/tensorrt_inference.py:
##########
@@ -164,11 +165,63 @@ def get_engine_attrs(self):
self.stream)
+TensorRTInferenceFn = Callable[
+ [Sequence[np.ndarray], TensorRTEngine, Optional[Dict[str, Any]]],
+ Iterable[PredictionResult]]
+
+
+def _default_tensorRT_inference_fn(
+ batch: Sequence[np.ndarray],
+ engine: TensorRTEngine,
+ inference_args: Optional[Dict[str,
+ Any]] = None) -> Iterable[PredictionResult]:
+ from cuda import cuda
+ (
+ engine,
+ context,
+ context_lock,
+ inputs,
+ outputs,
+ gpu_allocations,
+ cpu_allocations,
+ stream) = engine.get_engine_attrs()
+
+ # Process I/O and execute the network
+ with context_lock:
+ _assign_or_fail(
+ cuda.cuMemcpyHtoDAsync(
+ inputs[0]['allocation'],
+ np.ascontiguousarray(batch),
+ inputs[0]['size'],
+ stream))
+ context.execute_async_v2(gpu_allocations, stream)
+ for output in range(len(cpu_allocations)):
+ _assign_or_fail(
+ cuda.cuMemcpyDtoHAsync(
+ cpu_allocations[output],
+ outputs[output]['allocation'],
+ outputs[output]['size'],
+ stream))
+ _assign_or_fail(cuda.cuStreamSynchronize(stream))
+
+ return [
+ PredictionResult(
+ x, [prediction[idx] for prediction in cpu_allocations]) for idx,
+ x in enumerate(batch)
+ ]
+
+
@experimental(extra_message="No backwards-compatibility guarantees.")
class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
PredictionResult,
TensorRTEngine]):
- def __init__(self, min_batch_size: int, max_batch_size: int, **kwargs):
+ def __init__(
+ self,
+ min_batch_size: int,
+ max_batch_size: int,
+ *,
Review Comment:
what is this * parameter for?
##########
sdks/python/apache_beam/ml/inference/tensorrt_inference.py:
##########
@@ -241,40 +297,7 @@ def run_inference(
Returns:
An Iterable of type PredictionResult.
"""
- from cuda import cuda
- (
- engine,
- context,
- context_lock,
- inputs,
- outputs,
- gpu_allocations,
- cpu_allocations,
- stream) = engine.get_engine_attrs()
-
- # Process I/O and execute the network
- with context_lock:
- _assign_or_fail(
- cuda.cuMemcpyHtoDAsync(
- inputs[0]['allocation'],
- np.ascontiguousarray(batch),
- inputs[0]['size'],
- stream))
- context.execute_async_v2(gpu_allocations, stream)
- for output in range(len(cpu_allocations)):
- _assign_or_fail(
- cuda.cuMemcpyDtoHAsync(
- cpu_allocations[output],
- outputs[output]['allocation'],
- outputs[output]['size'],
- stream))
- _assign_or_fail(cuda.cuStreamSynchronize(stream))
-
- return [
- PredictionResult(
- x, [prediction[idx] for prediction in cpu_allocations]) for idx,
- x in enumerate(batch)
- ]
+ return self.inference_fn(batch, engine, inference_args)
Review Comment:
is there any way to add a unit test for using a custom function?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] jrmccluskey commented on a diff in pull request #24039: [Draft] TensorRT Custom Inference Function Implementation
Posted by GitBox <gi...@apache.org>.
jrmccluskey commented on code in PR #24039:
URL: https://github.com/apache/beam/pull/24039#discussion_r1017054307
##########
sdks/python/apache_beam/ml/inference/tensorrt_inference.py:
##########
@@ -164,11 +165,63 @@ def get_engine_attrs(self):
self.stream)
+TensorRTInferenceFn = Callable[
+ [Sequence[np.ndarray], TensorRTEngine, Optional[Dict[str, Any]]],
+ Iterable[PredictionResult]]
+
+
+def _default_tensorRT_inference_fn(
+ batch: Sequence[np.ndarray],
+ engine: TensorRTEngine,
+ inference_args: Optional[Dict[str,
+ Any]] = None) -> Iterable[PredictionResult]:
+ from cuda import cuda
+ (
+ engine,
+ context,
+ context_lock,
+ inputs,
+ outputs,
+ gpu_allocations,
+ cpu_allocations,
+ stream) = engine.get_engine_attrs()
+
+ # Process I/O and execute the network
+ with context_lock:
+ _assign_or_fail(
+ cuda.cuMemcpyHtoDAsync(
+ inputs[0]['allocation'],
+ np.ascontiguousarray(batch),
+ inputs[0]['size'],
+ stream))
+ context.execute_async_v2(gpu_allocations, stream)
+ for output in range(len(cpu_allocations)):
+ _assign_or_fail(
+ cuda.cuMemcpyDtoHAsync(
+ cpu_allocations[output],
+ outputs[output]['allocation'],
+ outputs[output]['size'],
+ stream))
+ _assign_or_fail(cuda.cuStreamSynchronize(stream))
+
+ return [
+ PredictionResult(
+ x, [prediction[idx] for prediction in cpu_allocations]) for idx,
+ x in enumerate(batch)
+ ]
+
+
@experimental(extra_message="No backwards-compatibility guarantees.")
class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
PredictionResult,
TensorRTEngine]):
- def __init__(self, min_batch_size: int, max_batch_size: int, **kwargs):
+ def __init__(
+ self,
+ min_batch_size: int,
+ max_batch_size: int,
+ *,
+ inference_fn: TensorRTInferenceFn = _default_tensorRT_inference_fn,
+ **kwargs):
Review Comment:
This pattern looks a little quirky here, but the keyword-only arg has to appear after varargs and before kwargs. This shouldn't change anything, fortunately.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org