update

commune/subspace/subspace.py

@@ -15,13 +15,9 @@ class Subspace(c.Module):
     Handles interactions with the subspace chain.
     """
 
+
     whitelist = ['query', 
                  'score',
-                 'keys',
-                 'incentives',
-                 'stakes',
-                 'netuids',
-                 'subnet2netuid',
                  'query_map', 
                  'get_module', 
                  'get_balance', 
@@ -97,10 +93,6 @@ def __init__(
         **kwargs,
     ):
         self.set_config(kwargs=kwargs)
-        c.print(self.config)
-        if self.config.run_loop == True:
-            c.print('Running loop')
-            c.thread(self.run_loop)
 
     network_mode = 'ws'
 
@@ -2330,11 +2322,7 @@ def state_dict(self ,
 
     def sync(self,*args, **kwargs):
 
-        self.stake_from(update=1)
-        self.global_params(update=1)
-        self.subnet_params(update=1)
         self.get_balances(update=1)
-
 
 
     def check_storage(self, block_hash = None, network=network):
@@ -2864,7 +2852,6 @@ def unstake(
             key : 'c.Key' = None,  # defaults to first key
             netuid : Union[str, int] = 0, # defaults to module.netuid
             network: str= None,
-            buffer = 0,
             **kwargs
         ) -> dict:
         """
@@ -2906,12 +2893,10 @@ def unstake(
             raise Exception('Invalid input')
 
         if amount == None:
-            stake_to = self.get_stake_to(netuid=netuid, names = False, fmt='nano', key=key.ss58_address)
-            amount = stake_to[module_key]
+            stake_to = self.get_stake_to(netuid=netuid, names = False, fmt='nano', key=module_key)
+            amount = stake_to[module_key] - 100000
         else:
             amount = int(self.to_nanos(amount))
-
-        amount = max(0, amount - buffer)
         # convert to nanos
         params={
             'amount': amount ,
@@ -3305,10 +3290,35 @@ def set_weights(
 
     vote = set_weights
 
+    def register_servers(self,  search=None, netuid = 0, timeout=60, max_age=None,  key=None, update=False, **kwargs):
+        '''
+        key2address : dict
+            A dictionary of module names to their keys
+        timeout : int 
+            The timeout for each registration
+        netuid : int
+            The netuid of the modules
+        
+        '''
+        infos = c.infos(search=search, **kwargs)
+        keys = c.submit(self.keys, dict(netuid=netuid, update=update, max_age=max_age))
+        names = c.submit(self.names, dict(netuid=netuid, update=update, max_age=max_age))
+        keys, names = c.wait([keys, names], timeout=timeout)
+        should_register_fn = lambda x: x['ss58_address'] not in keys and x['name'] not in names
+        infos = [i for i in infos if should_register_fn(i)]
+        c.print(f'Found {infos} modules to register')
+        for i, info in enumerate(infos):
+            r = c.register(name=info['name'], 
+                         address= info['address'],
+                         module_key=info['ss58_address'], 
+                          key=key)
+            c.print(r, color='green')
+
+        return {'success': True, 'message': 'All modules registered'}
 
 
     @classmethod
-    def register_servers(cls, servers = None , key2address = None, timeout=60, netuid = 0):
+    def register_servers_parallel(cls, servers = None , key2address = None, timeout=60, netuid = 0):
         '''
         key2address : dict
             A dictionary of module names to their keys
@@ -3321,12 +3331,11 @@ def register_servers(cls, servers = None , key2address = None, timeout=60, netui
         key2address = key2address or c.key2address()
         if servers == None:
             servers = list(key2address.keys())
-        key2address = {s['name']: s['ss58_address'] for s in servers}
+        key2address = {k:v for k,v in key2address.items() if k in servers}
         futures = []
         launcher_keys = c.launcher_keys()
         future2launcher = {}
         future2module = {}
-        netuid = c.resolve_netuid(netuid)
         registered_keys = c.m('subspace')().keys(netuid=netuid)
         progress = c.tqdm(total=len(key2address))
 
@@ -3340,6 +3349,9 @@ def register_servers(cls, servers = None , key2address = None, timeout=60, netui
                     progress.update(1)
                     continue
                 c.print(f"Registering {module} with key {module}")
+                if len(launcher_keys) == 0:
+                    c.print(f"No launchers found")
+                    launcher_keys = [c.root_key()]
                 launcher_key = launcher_keys[i % len(launcher_keys)]
                 kwargs=dict(name=module, module_key=module_key, serve=True, key=launcher_key)
                 future = c.submit(c.register, kwargs=kwargs, timeout=timeout)
@@ -3354,7 +3366,8 @@ def register_servers(cls, servers = None , key2address = None, timeout=60, netui
                 module_key = key2address.pop(module)
                 c.print(f"Registered {module} module_key:{module_key} launcher_key:{launcher_key}")
                 r = f.result()
-                if c.is_error(r):
+                c.print(r)
+                if not c.is_error(r):
                     progress.update(1)
         return {'success': True, 'message': 'All modules registered'}
 
@@ -3371,19 +3384,6 @@ def unregistered_servers(self, search=None, netuid = 0, network = network,  time
             if  key2address[s] not in keys:
                 unregister_servers += [s]
         return unregister_servers
-
-
-    def rm_min_value_keys(self, min_value=0.1, **kwargs):
-        key2balance = self.key2value(**kwargs)
-        key2address = c.address2key()
-        for k,v in key2balance.items():
-            if v < min_value and k in key2address:
-                c.print(f'Removing {k} with value {v}')
-                # try:
-                #     c.rm_key(k)
-                # except Exception as e:
-                #     c.print(f'Error removing {k} with value {v} {e}')
-        return {'success': True, 'message': f'Removed keys with value < {min_value}'}
 
     def get_balances(self, 
                     keys=None,
@@ -3404,7 +3404,6 @@ def get_balances(self,
         key2address = c.key2address(search=search)
         if keys == None:
             keys = list(key2address.keys())
-
         if len(keys) > n:
             c.print(f'Getting balances for {len(keys)} keys > {n} keys, using batch_size {batch_size}')
             balances = self.balances(network=network, **kwargs)
@@ -3709,6 +3708,9 @@ def test_global_storage(self):
     def test_module_storage(self):
         modules = self.get_modules(netuid=0)
         return modules 
+
+
+