CRZbulabula commented on code in PR #16595:
URL: https://github.com/apache/iotdb/pull/16595#discussion_r2435382647
##########
iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py:
##########
@@ -119,39 +120,41 @@ def schedule(self, model_id: str) -> List[ScaleAction]:
return [ScaleAction(ScaleActionType.SCALE_UP, pool_num, model_id)]
def schedule_load_model_to_device(
- self, model_id: str, device_id: str
+ self, model_info: ModelInfo, device_id: str
) -> List[ScaleAction]:
- existing_model_ids = [
- existing_model_id
+ existing_model_infos = [
+ MODEL_MANAGER.get_model_info(existing_model_id)
for existing_model_id, pool_group_map in
self._request_pool_map.items()
- if existing_model_id != model_id and device_id in pool_group_map
+ if existing_model_id != model_info.model_id and device_id in
pool_group_map
]
allocation_result = _estimate_shared_pool_size_by_total_mem(
device=convert_device_id_to_torch_device(device_id),
- existing_model_ids=existing_model_ids,
- new_model_id=model_id,
+ existing_model_infos=existing_model_infos,
+ new_model_info=model_info,
)
return self._convert_allocation_result_to_scale_actions(
allocation_result, device_id
)
def schedule_unload_model_from_device(
- self, model_id: str, device_id: str
+ self, model_info: ModelInfo, device_id: str
) -> List[ScaleAction]:
- existing_model_ids = [
- existing_model_id
+ existing_model_infos = [
+ MODEL_MANAGER.get_model_info(existing_model_id)
for existing_model_id, pool_group_map in
self._request_pool_map.items()
- if existing_model_id != model_id and device_id in pool_group_map
+ if existing_model_id != model_info.model_id and device_id in
pool_group_map
]
allocation_result = (
_estimate_shared_pool_size_by_total_mem(
device=convert_device_id_to_torch_device(device_id),
- existing_model_ids=existing_model_ids,
- new_model_id=None,
+ existing_model_infos=existing_model_infos,
+ new_model_info=None,
)
- if len(existing_model_ids) > 0
- else {model_id: 0}
+ if len(existing_model_infos) > 0
+ else {model_info.model_id: 0}
)
+ if len(existing_model_infos) > 0:
+ allocation_result[model_info.model_id] = 0
Review Comment:
U R wrong
--
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]
