gpu-provider/src/scheduler.py

51 lines
2.0 KiB
Python

from web3 import Web3
from src.utils import load_abi
class Scheduler():
def __init__(self, w3: Web3, address: str, abi_file: str):
self.w3 = w3
abi = load_abi(abi_file)
self.contract = self.w3.eth.contract(address=address, abi=abi)
print(self.contract.functions)
def signed_and_send_transaction(self, account, function):
tx = function.build_transaction({
'from': account.address,
'nonce': self.w3.eth.get_transaction_count(account.address)
})
signed_tx = self.w3.eth.account.sign_transaction(tx, private_key=account.key)
tx_hash = self.w3.eth.send_raw_transaction(signed_tx.rawTransaction)
tx_receipt = self.w3.eth.wait_for_transaction_receipt(tx_hash)
return tx_receipt
def register_cluster(self, account, gpu_id: int, gpu_num: int):
function = self.contract.functions.registerCluster(gpu_id, gpu_num)
return self.signed_and_send_transaction(account, function)
def register_task_with_conditions(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.signed_and_send_transaction(account, function)
def get_clusters(self):
return self.contract.functions.getClusters().call()
def get_tasks(self):
return self.contract.functions.getTasks().call()
'''
def get_start_run_event(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)
'''
def get_cluster_event(self, receipt):
event = self.contract.events.RegisterCluster().process_receipt(receipt)
return event