From c0e7fea3bbba3dd3473f3689b2cfdbf4e95f00cf Mon Sep 17 00:00:00 2001
From: snsd0805 <levi900227@gmail.com>
Date: Sat, 11 May 2024 21:16:35 +0800
Subject: [PATCH] feat: register tasks & get tasks

---
 src/scheduler.py   | 25 +++++++++++++++++++++----
 main.py => test.py | 10 +++++++---
 2 files changed, 28 insertions(+), 7 deletions(-)
 rename main.py => test.py (71%)

diff --git a/src/scheduler.py b/src/scheduler.py
index 7c385c2..9185869 100644
--- a/src/scheduler.py
+++ b/src/scheduler.py
@@ -6,9 +6,10 @@ class Scheduler():
         self.w3 = w3
         abi = load_abi(abi_file)
         self.contract = self.w3.eth.contract(address=address, abi=abi)
+        print(self.contract.functions)
 
-    def registerCluster(self, account, gpu_id: int, gpu_num: int):
-        tx = self.contract.functions.registerCluster(gpu_id, gpu_num).build_transaction({
+    def signedAndSendTransaction(self, account, function):
+        tx = function.build_transaction({
             'from': account.address,
             'nonce': self.w3.eth.get_transaction_count(account.address)
         })
@@ -16,14 +17,30 @@ class Scheduler():
 
         tx_hash = self.w3.eth.send_raw_transaction(signed_tx.rawTransaction)
         tx_receipt = self.w3.eth.wait_for_transaction_receipt(tx_hash)
-        print(tx_receipt)
         _hash = tx_receipt['transactionHash']
-        print(self.w3.eth.get_transaction(_hash))
+        return _hash
+
+    def registerCluster(self, account, gpu_id: int, gpu_num: int):
+        function = self.contract.functions.registerCluster(gpu_id, gpu_num)
+        return self.signedAndSendTransaction(account, function)
+
+    def registerTaskWithConditions(self, account, data_image: str, train_image: str, gpu_id: int, cluster_size: int):
+        function = self.contract.functions.registerTaskWithConditions(data_image, train_image, gpu_id, cluster_size)
+        return self.signedAndSendTransaction(account, function)
 
     def getClusters(self):
         return self.contract.functions.getClusters().call()
 
+    def getTasks(self):
+        return self.contract.functions.getTasks().call()
+
     def getStartRunEvent(self):
         event_filter = self.contract.events.StartRun.create_filter(fromBlock=0)
         events = event_filter.get_new_entries()
         print(events)
+
+    def getRegisterClusterEvent(self):
+        event_filter = self.contract.events.RegisterCluster.create_filter(fromBlock=0)
+        events = event_filter.get_new_entries()
+        print(events)
+
diff --git a/main.py b/test.py
similarity index 71%
rename from main.py
rename to test.py
index e7ac84f..6d16cab 100644
--- a/main.py
+++ b/test.py
@@ -1,8 +1,7 @@
-import json
 from web3 import Web3
 from src.scheduler import Scheduler
 
-SCHEDULER_ADDR = "0x13D69Cf7d6CE4218F646B759Dcf334D82c023d8e"
+SCHEDULER_ADDR = "0x544eAe853EA3774A8857573C6423E6Db95b79258"
 SCHEDULER_ABI_FILE = "../gpu-contract/abi/Scheduler.abi"
 
 PROVIDER1_KEY = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"
@@ -13,7 +12,6 @@ CLIENT_KEY = "0x7c852118294e51e653712a81e05800f419141751be58f605c371e15141b007a6
 
 w3 = Web3(Web3.HTTPProvider('http://127.0.0.1:8545'))
 
-
 if __name__ == '__main__':
     if w3.is_connected():
         scheduler = Scheduler(w3, SCHEDULER_ADDR, SCHEDULER_ABI_FILE)
@@ -24,6 +22,12 @@ if __name__ == '__main__':
         client = w3.eth.account.from_key(CLIENT_KEY)
 
         print(scheduler.getClusters())
+        scheduler.registerCluster(provider1, 1, 4)
+        scheduler.registerCluster(provider2, 2, 2)
+        scheduler.registerCluster(provider3, 3, 1)
+        scheduler.registerTaskWithConditions(client, "https://data.com", "http://train.tw", 3, 1)
+        print(scheduler.getClusters())
+        print(scheduler.getTasks())
 
     else:
         print("cannot connected to the chain")