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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
