You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@iotdb.apache.org by "jtmer (via GitHub)" <gi...@apache.org> on 2025/09/15 01:35:12 UTC
[PR] [AINode] Add a batcher for inference [iotdb]
jtmer opened a new pull request, #16411:
URL: https://github.com/apache/iotdb/pull/16411
## Description
This PR adds a Batcher for handling inference request batching.
### Integrate Batching into InferenceRequestPool
In `_step`, requests scheduled by BasicRequestScheduler are grouped by (seq_len, max_new_tokens) to guarantee compatibility.
Each group of requests is batched using BasicBatcher.batch_request and passed to the model’s generate method.
Finished requests trigger post_inference and are pushed to the finished_queue.
<!--
In each section, please describe design decisions made, including:
- Choice of algorithms
- Behavioral aspects. What configuration values are acceptable? How are corner cases and error
conditions handled, such as when there are insufficient resources?
- Class organization and design (how the logic is split between classes, inheritance, composition,
design patterns)
- Method organization and design (how the logic is split between methods, parameters and return types)
- Naming (class, method, API, configuration, HTTP endpoint, names of emitted metrics)
-->
<!-- It's good to describe an alternative design (or mention an alternative name) for every design
(or naming) decision point and compare the alternatives with the designs that you've implemented
(or the names you've chosen) to highlight the advantages of the chosen designs and names. -->
<!-- If there was a discussion of the design of the feature implemented in this PR elsewhere
(e. g. a "Proposal" issue, any other issue, or a thread in the development mailing list),
link to that discussion from this PR description and explain what have changed in your final design
compared to your original proposal or the consensus version in the end of the discussion.
If something hasn't changed since the original discussion, you can omit a detailed discussion of
those aspects of the design here, perhaps apart from brief mentioning for the sake of readability
of this PR description. -->
<!-- Some of the aspects mentioned above may be omitted for simple and small changes. -->
<hr>
This PR has:
- [x] been self-reviewed.
<!-- Check the items by putting "x" in the brackets for the done things. Not all of these items
apply to every PR. Remove the items which are not done or not relevant to the PR. None of the items
from the checklist above are strictly necessary, but it would be very helpful if you at least
self-review the PR. -->
<hr>
##### Key changed/added classes (or packages if there are too many classes) in this PR
`iotdb.ainode.core.inference.batcher.AbstractBatcher`
`iotdb.ainode.core.inference.batcher.BasicBatcher`
`iotdb.ainode.core.inference.inference_request_pool`
--
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: reviews-unsubscribe@iotdb.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [AINode] Add a batcher for inference [iotdb]
Posted by "CRZbulabula (via GitHub)" <gi...@apache.org>.
CRZbulabula merged PR #16411:
URL: https://github.com/apache/iotdb/pull/16411
--
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: reviews-unsubscribe@iotdb.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [AINode] Add a batcher for inference [iotdb]
Posted by "CRZbulabula (via GitHub)" <gi...@apache.org>.
CRZbulabula commented on code in PR #16411:
URL: https://github.com/apache/iotdb/pull/16411#discussion_r2348748434
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py:
##########
@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from ainode.core.inference.inference_request import InferenceRequest
+
+
+class AbstractBatcher(ABC):
+ """
+ Abstract base class for batchers that batch inference requests.
+ """
+
+ def __init__(self):
+ """
+ Args:
+
+ """
+
+ @abstractmethod
+ def batch_request(self, reqs: List[InferenceRequest]):
+ """
+ batch given requests.
Review Comment:
```suggestion
batch given requests, such that they can be delivered to the model and be executed concurrently.
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py:
##########
@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from ainode.core.inference.inference_request import InferenceRequest
+
+
+class AbstractBatcher(ABC):
+ """
+ Abstract base class for batchers that batch inference requests.
+ """
+
+ def __init__(self):
+ """
+ Args:
+
+ """
Review Comment:
```suggestion
pass
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py:
##########
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List
+
+import torch
+
+from ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
+from ainode.core.inference.inference_request import InferenceRequest
+
+
+class BasicBatcher(AbstractBatcher):
+ """
+ Basic batcher for inference requests.
+ """
+
+ def __init__(self):
+ """
+ Args:
+
+ """
+
+ def batch_request(self, reqs: List[InferenceRequest]) -> torch.Tensor:
+ """
+ Batch given requests by concatenating their inputs.
+
+ - Considering the current implementation of AINode, we might merely be piecing together the input for now.
+
+ Args:
+ reqs (List[InferenceRequest]): List of inference requests.
+
+ Returns:
+ torch.Tensor: Concatenated input tensor of shape
+ [sum(num_var), length].
Review Comment:
```suggestion
torch.Tensor: Concatenated input tensor of shape
[sum(req.batch_size), length].
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py:
##########
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List
+
+import torch
+
+from ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
+from ainode.core.inference.inference_request import InferenceRequest
+
+
+class BasicBatcher(AbstractBatcher):
+ """
+ Basic batcher for inference requests.
+ """
+
+ def __init__(self):
+ """
+ Args:
+
+ """
Review Comment:
pass
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
num_samples=10,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0].mean(dim=0))
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
+ offset += b
+ # ! Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here
Review Comment:
```suggestion
# TODO: Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
num_samples=10,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0].mean(dim=0))
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
Review Comment:
Rename them for more direct expression.
```suggestion
cur_batch_size = request.batch_size
cur_output = batch_output[offset : offset + b]
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
num_samples=10,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0].mean(dim=0))
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
+ offset += b
+ # ! Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here
+ request.write_step_output(output_i[0].mean(dim=0))
+
+ request.inference_pipeline.post_decode()
+ if request.is_finished():
+ request.inference_pipeline.post_inference()
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
+ )
+ # ensure the output tensor is on CPU before sending to result queue
+ request.output_tensor = request.output_tensor.cpu()
+ self._finished_queue.put(request)
+ else:
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
+ )
+ self._waiting_queue.put(request)
+
elif self.model_id == "timer_xl":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0])
- request.inference_pipeline.post_decode()
- if request.is_finished():
- request.inference_pipeline.post_inference()
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][Req-{request.req_id}] Request is finished"
- )
- # ensure the output tensor is on CPU before sending to result queue
- request.output_tensor = request.output_tensor.cpu()
- self._finished_queue.put(request)
- else:
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][Req-{request.req_id}] Request is not finished, re-queueing"
- )
- self._waiting_queue.put(request)
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
Review Comment:
The same as above.
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py:
##########
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List
+
+import torch
+
+from ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
+from ainode.core.inference.inference_request import InferenceRequest
+
+
+class BasicBatcher(AbstractBatcher):
+ """
+ Basic batcher for inference requests.
+ """
+
+ def __init__(self):
+ """
+ Args:
+
+ """
+
+ def batch_request(self, reqs: List[InferenceRequest]) -> torch.Tensor:
+ """
+ Batch given requests by concatenating their inputs.
Review Comment:
```suggestion
Batch given requests by simply concatenating their inputs, only requests with uniformed output length can be batched.
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
num_samples=10,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0].mean(dim=0))
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
+ offset += b
+ # ! Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here
+ request.write_step_output(output_i[0].mean(dim=0))
+
+ request.inference_pipeline.post_decode()
+ if request.is_finished():
+ request.inference_pipeline.post_inference()
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
+ )
+ # ensure the output tensor is on CPU before sending to result queue
+ request.output_tensor = request.output_tensor.cpu()
+ self._finished_queue.put(request)
+ else:
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
+ )
+ self._waiting_queue.put(request)
+
elif self.model_id == "timer_xl":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
Review Comment:
Plz append concurrent inference test for `timer_xl` in `AINodeConcurrentInferenceIT`.
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
num_samples=10,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0].mean(dim=0))
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
+ offset += b
+ # ! Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here
+ request.write_step_output(output_i[0].mean(dim=0))
Review Comment:
Seems strange in the current implementation. Could u plz explain the shape of `sundial`'s output in the batched case?
--
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: reviews-unsubscribe@iotdb.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [AINode] Add a batcher for inference [iotdb]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #16411:
URL: https://github.com/apache/iotdb/pull/16411#discussion_r2347963099
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py:
##########
@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from ainode.core.inference.inference_request import InferenceRequest
Review Comment:
Import path is missing the 'iotdb.' prefix. This should be 'from iotdb.ainode.core.inference.inference_request import InferenceRequest' to match the package structure used elsewhere in the codebase.
```suggestion
from iotdb.ainode.core.inference.inference_request import InferenceRequest
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
Review Comment:
Should use 'self._model' instead of 'self.model' to match the class attribute definition and maintain consistency with the rest of the method.
```suggestion
batch_output = self._model.generate(
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py:
##########
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List
+
+import torch
+
+from ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
+from ainode.core.inference.inference_request import InferenceRequest
Review Comment:
Import paths are missing the 'iotdb.' prefix. These should be 'from iotdb.ainode.core.inference.batcher.abstract_batcher import AbstractBatcher' and 'from iotdb.ainode.core.inference.inference_request import InferenceRequest' to match the package structure used elsewhere in the codebase.
```suggestion
from iotdb.ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
from iotdb.ainode.core.inference.inference_request import InferenceRequest
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
num_samples=10,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0].mean(dim=0))
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
+ offset += b
+ # ! Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here
Review Comment:
This comment indicates a significant limitation that should be documented more formally. Consider adding a TODO or FIXME tag and documenting this limitation in the method's docstring to ensure future developers are aware of this constraint.
##########
iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py:
##########
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List
+
+import torch
+
+from ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
+from ainode.core.inference.inference_request import InferenceRequest
+
+
+class BasicBatcher(AbstractBatcher):
+ """
+ Basic batcher for inference requests.
+ """
+
+ def __init__(self):
+ """
+ Args:
+
+ """
+
+ def batch_request(self, reqs: List[InferenceRequest]) -> torch.Tensor:
+ """
+ Batch given requests by concatenating their inputs.
+
+ - Considering the current implementation of AINode, we might merely be piecing together the input for now.
+
+ Args:
+ reqs (List[InferenceRequest]): List of inference requests.
+
+ Returns:
+ torch.Tensor: Concatenated input tensor of shape
+ [sum(num_var), length].
+ """
+ if not reqs:
+ raise ValueError("No requests provided to batch_request.")
+
+ # 确保 length 一致
Review Comment:
Comment is in Chinese. Should be translated to English: '# Ensure length consistency' to maintain code consistency with the rest of the codebase.
```suggestion
# Ensure length consistency
```
##########
iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py:
##########
@@ -107,42 +111,77 @@ def _requests_activate_loop(self):
self._activate_requests()
def _step(self):
- requests = self._request_scheduler.schedule_step()
- # TODO: We need a batcher to accelerate the concurrent inference
- for request in requests:
+ all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+
+ grouped_requests = defaultdict(list)
+ for req in all_requests:
+ key = (req.inputs.shape[1], req.max_new_tokens)
+ grouped_requests[key].append(req)
+ grouped_requests = list(grouped_requests.values())
+
+ for requests in grouped_requests:
+ batch_inputs = self._batcher.batch_request(requests).to(self.device)
if self.model_id == "sundial":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
+ batch_inputs,
+ max_new_tokens=requests[0].max_new_tokens,
num_samples=10,
revin=True,
)
- request.output_tensor = request.output_tensor.to(self.device)
- request.write_step_output(output[0].mean(dim=0))
+
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ b = request.batch_size
+ output_i = batch_output[offset : offset + b]
+ offset += b
+ # ! Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here
+ request.write_step_output(output_i[0].mean(dim=0))
+
+ request.inference_pipeline.post_decode()
+ if request.is_finished():
+ request.inference_pipeline.post_inference()
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
+ )
+ # ensure the output tensor is on CPU before sending to result queue
+ request.output_tensor = request.output_tensor.cpu()
+ self._finished_queue.put(request)
+ else:
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
+ )
+ self._waiting_queue.put(request)
+
elif self.model_id == "timer_xl":
- request.inputs = request.inputs.to(self.device)
- output = self._model.generate(
- request.inputs,
- max_new_tokens=request.max_new_tokens,
+ batch_output = self.model.generate(
Review Comment:
Should use 'self._model' instead of 'self.model' to match the class attribute definition and maintain consistency with the rest of the method.
```suggestion
batch_output = self._model.generate(
```
--
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: reviews-unsubscribe@iotdb.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org