This is an automated email from the ASF dual-hosted git repository. Caideyipi pushed a commit to branch patch-2094 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit f43aaefc98642453267250c008facba615b2e6b9 Author: Yongzao <[email protected]> AuthorDate: Wed Jun 3 16:32:20 2026 +0800 [AINode] Remove Chronos2 DataLoader pin_memory option (#17822) (cherry picked from commit e2322e9e9e5d0ef5a4b982f27fa682c066e94cca) --- .../request_scheduler/basic_request_scheduler.py | 19 ++++++++----------- .../ainode/core/model/chronos2/pipeline_chronos2.py | 1 - 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py index ef5d37a18a1..2ee15fabca4 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py @@ -16,8 +16,6 @@ # under the License. # -import os - import psutil import torch @@ -53,23 +51,22 @@ class BasicRequestScheduler(AbstractRequestScheduler): def memory_is_available(self): if "cuda" in self.device.type: - used = torch.cuda.memory_allocated(self.device) - reserved = torch.cuda.memory_reserved(self.device) + available, total = torch.cuda.mem_get_info(self.device) elif "cpu" in self.device.type: - process = psutil.Process(os.getpid()) - used = process.memory_info().rss - reserved = used + memory = psutil.virtual_memory() + available = memory.available + total = memory.total else: - used = 0 - reserved = 0 logger.warning( f"[Inference] Unsupported device type: {self.device.type}. Memory checks will not be performed." ) + return True logger.debug( f"[Inference][Device-{self.device}][Pool-{self.pool_id}] " - f"Memory used: {used/1024**2:.2f} MB, Max memory: {self.max_memory_bytes/1024**2:.2f} MB" + f"Memory available: {available/1024**2:.2f} MB, Total memory: {total/1024**2:.2f} MB, " + f"Required free memory: {self.max_memory_bytes/1024**2:.2f} MB" ) - return used < self.max_memory_bytes + return available > self.max_memory_bytes def schedule_activate(self) -> list: requests = [] diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/chronos2/pipeline_chronos2.py b/iotdb-core/ainode/iotdb/ainode/core/model/chronos2/pipeline_chronos2.py index 01ff78ba48d..083125e4b99 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/chronos2/pipeline_chronos2.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/chronos2/pipeline_chronos2.py @@ -309,7 +309,6 @@ class Chronos2Pipeline(ForecastPipeline): test_loader = DataLoader( test_dataset, batch_size=None, - pin_memory=True, shuffle=False, drop_last=False, )
