You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ju...@apache.org on 2026/06/11 09:42:57 UTC
(iotdb) 44/48: [AINode] Remove Chronos2 DataLoader pin_memory option (#17822)
This is an automated email from the ASF dual-hosted git repository.
Caideyipi pushed a commit to branch hotfix/2.0.9.4-sjzt
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit f43aaefc98642453267250c008facba615b2e6b9
Author: Yongzao <yo...@apache.org>
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,
)