From c75296dda02df172a15d72eb037960b24b4fc058 Mon Sep 17 00:00:00 2001
From: leavesster <11785335+leavesster@users.noreply.github.com>
Date: Tue, 23 Sep 2025 15:01:23 +0800
Subject: [PATCH 1/2] feat: support update node weight

---
 oocana/oocana/context.py  |  6 ++----
 oocana/oocana/internal.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 oocana/oocana/internal.py

diff --git a/oocana/oocana/context.py b/oocana/oocana/context.py
index 697e98e..07606f6 100644
--- a/oocana/oocana/context.py
+++ b/oocana/oocana/context.py
@@ -10,16 +10,13 @@
 from .throttler import throttle
 from .preview import PreviewPayload, DataFrame, PreviewPayloadInternal, ShapeDataFrame
 from .data import EXECUTOR_NAME
+from .internal import random_string, InternalAPI
 import os.path
 import logging
-import random
-import string
 import hashlib
 
 __all__ = ["Context", "HandleDefDict", "BlockJob", "BlockExecuteException"]
 
-def random_string(length=8):
-    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))
 
 def string_hash(text: str) -> str:
     """
@@ -194,6 +191,7 @@ def __init__(
     ) -> None:
 
         self.__block_info = blockInfo
+        self.internal: InternalAPI = InternalAPI(mainframe, blockInfo.job_info())
 
         self.__mainframe = mainframe
         self.__store = store
diff --git a/oocana/oocana/internal.py b/oocana/oocana/internal.py
new file mode 100644
index 0000000..ef40c75
--- /dev/null
+++ b/oocana/oocana/internal.py
@@ -0,0 +1,31 @@
+from .mainframe import Mainframe
+from .data import JobDict
+import random
+import string
+
+class InternalAPI:
+
+    def __init__(self, client: Mainframe, job_id: JobDict) -> None:
+        self._client = client
+        self._job_id = job_id
+
+    async def update_node_weight(self, node_id: str, weight: int | float) -> None:
+        """
+        Update the weight of a node.
+
+        :param node_id: The ID of the node to update.
+        :param weight: The new weight for the node.
+        """
+        if not isinstance(weight, (int, float)) or weight < 0:
+            raise ValueError("Weight must be a non-negative number.")
+
+        self._client.send(self._job_id, {
+            "type": "BlockRequest",
+            "action": "UpdateNodeWeight",
+            "node_id": node_id,
+            "weight": weight,
+            "request_id": random_string(16),
+        })
+
+def random_string(length=8):
+    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

From dc7d4ca519fb5a830b258cdb1eb247a950463328 Mon Sep 17 00:00:00 2001
From: leavesster <11785335+leavesster@users.noreply.github.com>
Date: Tue, 23 Sep 2025 15:17:58 +0800
Subject: [PATCH 2/2] update check

---
 oocana/oocana/internal.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/oocana/oocana/internal.py b/oocana/oocana/internal.py
index ef40c75..cbecf74 100644
--- a/oocana/oocana/internal.py
+++ b/oocana/oocana/internal.py
@@ -2,6 +2,7 @@
 from .data import JobDict
 import random
 import string
+import math
 
 class InternalAPI:
 
@@ -9,15 +10,15 @@ def __init__(self, client: Mainframe, job_id: JobDict) -> None:
         self._client = client
         self._job_id = job_id
 
+    # keep this method async for future use
     async def update_node_weight(self, node_id: str, weight: int | float) -> None:
         """
         Update the weight of a node.
-
         :param node_id: The ID of the node to update.
         :param weight: The new weight for the node.
         """
-        if not isinstance(weight, (int, float)) or weight < 0:
-            raise ValueError("Weight must be a non-negative number.")
+        if not isinstance(weight, (int, float)) or not math.isfinite(weight) or weight < 0:
+            raise ValueError("Weight must be a non-negative finite number.")
 
         self._client.send(self._job_id, {
             "type": "BlockRequest",