Use async with when using using subtensor

api/src/endpoints/upload.py

@@ -1,4 +1,4 @@
-from bittensor import Subtensor
+from bittensor import AsyncSubtensor
 
 import asyncio
 import os
@@ -24,9 +24,6 @@
 from utils.coingecko import get_tao_price
 from api import config 
 
-# TODO STEPHEN: we should have a global singleton
-subtensor = Subtensor(network=config.SUBTENSOR_NETWORK)
-
 # We use a lock per hotkey to prevent multiple agents being uploaded at the same time for the same hotkey
 hotkey_locks: dict[str, asyncio.Lock] = {}
 hotkey_locks_lock = asyncio.Lock()
@@ -73,8 +70,9 @@ async def check_agent_post(
     await check_hotkey_registered(miner_hotkey)
     await check_agent_banned(miner_hotkey=miner_hotkey) 
     check_if_python_file(agent_file.filename)
-    coldkey = subtensor.get_hotkey_owner(hotkey_ss58=miner_hotkey)
-    miner_balance = subtensor.get_balance(address=coldkey).rao
+    async with AsyncSubtensor(network=config.SUBTENSOR_NETWORK) as subtensor:
+        coldkey = await subtensor.get_hotkey_owner(hotkey_ss58=miner_hotkey)
+        miner_balance = (await subtensor.get_balance(address=coldkey)).rao
     payment_cost = await get_upload_price(cache_time=payment_time)
     if payment_cost.amount_rao > miner_balance:
         raise HTTPException(
@@ -168,7 +166,8 @@ async def post_agent(
 
         # Retrieve payment details from the chain
         try:
-            payment_block = subtensor.substrate.get_block(block_hash=payment_block_hash)
+            async with AsyncSubtensor(network=config.SUBTENSOR_NETWORK) as subtensor:
+                payment_block = await subtensor.substrate.get_block(block_hash=payment_block_hash)
         except Exception as e:
             logger.error(f"Error retrieving payment block: {e}")
             raise HTTPException(
@@ -190,7 +189,8 @@ async def post_agent(
             'stateRoot': '0x301a04303fb97143649e44ca9c1d674606c8004082d11973c816ff67f2a13998'}}
         """
         block_number = payment_block['header']['number']
-        coldkey = subtensor.get_hotkey_owner(hotkey_ss58=miner_hotkey, block=int(block_number))
+        async with AsyncSubtensor(network=config.SUBTENSOR_NETWORK) as subtensor:
+            coldkey = await subtensor.get_hotkey_owner(hotkey_ss58=miner_hotkey, block=int(block_number))
         payment_extrinsic = payment_block['extrinsics'][int(payment_extrinsic_index)]
 
         payment_cost = await get_upload_price(cache_time=payment_time)
@@ -205,7 +205,7 @@ async def post_agent(
                 payment_value = arg['value']
                 break
 
-        if payment_value is None or check_if_extrinsic_failed(payment_block_hash, int(payment_extrinsic_index)):
+        if payment_value is None or await check_if_extrinsic_failed(payment_block_hash, int(payment_extrinsic_index)):
             raise HTTPException(
                 status_code=402,
                 detail="Payment value not found"
@@ -355,8 +355,9 @@ async def get_upload_price() -> UploadPriceResponse:
         send_address=config.UPLOAD_SEND_ADDRESS
     )
 
-def check_if_extrinsic_failed(block_hash: str, extrinsic_index: int) -> bool:
-    events = subtensor.substrate.get_events(block_hash=block_hash)
+async def check_if_extrinsic_failed(block_hash: str, extrinsic_index: int) -> bool:
+    async with AsyncSubtensor(network=config.SUBTENSOR_NETWORK) as subtensor:
+        events = await subtensor.substrate.get_events(block_hash=block_hash)
 
     for event in events:
         if event.get("extrinsic_idx") != extrinsic_index:

ridges.py

@@ -143,20 +143,19 @@ def upload(ctx, file: Optional[str], coldkey_name: Optional[str], hotkey_name: O
                 console.print("[bold red]Payment cancelled by user. Upload aborted.[/bold red]")
                 return
 
-            subtensor = Subtensor(network=os.environ.get('SUBTENSOR_NETWORK', 'finney'))
-
-            # Transfer
-            payment_payload = subtensor.substrate.compose_call(
-                call_module="Balances",
-                call_function="transfer_keep_alive",
-                call_params={
-                    'dest': payment_method_details['send_address'], 
-                    'value': payment_method_details['amount_rao'],
-                }
-            )
-
-            payment_extrinsic = subtensor.substrate.create_signed_extrinsic(call=payment_payload, keypair=wallet.coldkey)
-            receipt = subtensor.substrate.submit_extrinsic(payment_extrinsic, wait_for_inclusion=True)
+            with Subtensor(network=os.environ.get('SUBTENSOR_NETWORK', 'finney')) as subtensor:
+                # Transfer
+                payment_payload = subtensor.substrate.compose_call(
+                    call_module="Balances",
+                    call_function="transfer_keep_alive",
+                    call_params={
+                        'dest': payment_method_details['send_address'], 
+                        'value': payment_method_details['amount_rao'],
+                    }
+                )
+
+                payment_extrinsic = subtensor.substrate.create_signed_extrinsic(call=payment_payload, keypair=wallet.coldkey)
+                receipt = subtensor.substrate.submit_extrinsic(payment_extrinsic, wait_for_inclusion=True)
 
             file_info = f"{wallet.hotkey.ss58_address}:{content_hash}:{version_num}"
             signature = wallet.hotkey.sign(file_info).hex()

utils/bittensor.py

@@ -7,12 +7,9 @@
 
 
 
-subtensor = AsyncSubtensor(network=config.SUBTENSOR_NETWORK)
-
-
-
 async def check_if_hotkey_is_registered(hotkey: str) -> bool:
-    return await subtensor.is_hotkey_registered(hotkey_ss58=hotkey, netuid=config.NETUID)
+    async with AsyncSubtensor(network=config.SUBTENSOR_NETWORK) as subtensor:
+        return await subtensor.is_hotkey_registered(hotkey_ss58=hotkey, netuid=config.NETUID)
 
 
 

validator/main.py

@@ -81,6 +81,9 @@ async def set_weights_loop():
             await asyncio.wait_for(set_weights_from_mapping(weights_mapping), timeout=config.SET_WEIGHTS_TIMEOUT_SECONDS)
         except asyncio.TimeoutError as e:
             logger.error(f"asyncio.TimeoutError in set_weights_from_mapping(): {e}")
+        except Exception as e:
+            logger.error(f"Error in set_weights_from_mapping(): {type(e).__name__}: {e}")
+            logger.error(traceback.format_exc())
 
         await asyncio.sleep(config.SET_WEIGHTS_INTERVAL_SECONDS)
 

validator/set_weights.py

@@ -8,10 +8,6 @@
 
 
 
-subtensor = AsyncSubtensor(network=config.SUBTENSOR_NETWORK, fallback_endpoints=[config.SUBTENSOR_ADDRESS])
-
-
-
 async def set_weights_from_mapping(weights_mapping: Dict[str, float]) -> None:
     if len(weights_mapping.keys()) != 1:
         logger.error("Expected one hotkey")
@@ -21,22 +17,22 @@ async def set_weights_from_mapping(weights_mapping: Dict[str, float]) -> None:
     if weights_mapping[weight_receiving_hotkey] != 1:
         logger.error("Expected weight of 1")
         return
-
-    weight_receiving_uid = await subtensor.get_uid_for_hotkey_on_subnet(hotkey_ss58=weight_receiving_hotkey, netuid=config.NETUID)
-    if weight_receiving_uid is None:
-        logger.error(f"Weight receiving hotkey {weight_receiving_hotkey} not found")
-        return
-
-    logger.info(f"Setting weight of {weight_receiving_hotkey} to 1...")
-
-    success, message = await subtensor.set_weights(
-        wallet=config.VALIDATOR_WALLET,
-        netuid=config.NETUID,
-        uids=[weight_receiving_uid],
-        weights=[1],
-        wait_for_inclusion=True,
-        wait_for_finalization=True
-    )
+    async with AsyncSubtensor(network=config.SUBTENSOR_NETWORK, fallback_endpoints=[config.SUBTENSOR_ADDRESS]) as subtensor:
+        weight_receiving_uid = await subtensor.get_uid_for_hotkey_on_subnet(hotkey_ss58=weight_receiving_hotkey, netuid=config.NETUID)
+        if weight_receiving_uid is None:
+            logger.error(f"Weight receiving hotkey {weight_receiving_hotkey} not found")
+            return
+
+        logger.info(f"Setting weight of {weight_receiving_hotkey} to 1...")
+
+        success, message = await subtensor.set_weights(
+            wallet=config.VALIDATOR_WALLET,
+            netuid=config.NETUID,
+            uids=[weight_receiving_uid],
+            weights=[1],
+            wait_for_inclusion=True,
+            wait_for_finalization=True
+        )
 
     if success:
         logger.info(f"Set weight of hotkey {weight_receiving_hotkey} to 1")