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