Merge pull request #258 from macrocosm-os/staging

Staging

Author: mccrindlebrian

environment.yml

@@ -13,6 +13,7 @@ dependencies:
   - gzip==1.13
   - pip:
       - bittensor==6.9.4
+      - tenacity==9.0.0
       - python-dotenv==1.0.1
       - wandb==0.17.2
       - eth-utils==2.1.1

folding/__init__.py

@@ -1,7 +1,7 @@
 from .protocol import JobSubmissionSynapse
 from .validators.protein import Protein
 
-__version__ = "1.1.0"
+__version__ = "1.1.1"
 version_split = __version__.split(".")
 __spec_version__ = (
     (10000 * int(version_split[0]))

folding/base/neuron.py

@@ -21,6 +21,7 @@
 import os
 
 import openmm
+from tenacity import RetryError
 
 # Sync calls set weights and also resyncs the metagraph.
 from folding.utils.config import check_config, add_args, config
@@ -154,7 +155,11 @@ def sync(self):
             self.resync_metagraph()
 
         if self.should_set_weights():
-            self.set_weights()
+            try:
+                self.set_weights()
+            except RetryError as e:
+                bt.logging.error(f"Failed to set weights after retry attempts. Skipping for {self.config.neuron.epoch_length} blocks.")
+                
 
         # Always save state.
         self.save_state()

folding/base/validator.py

@@ -19,21 +19,21 @@
 import os
 import json
 import copy
-import time
 import torch
 import asyncio
 import argparse
 import threading
 import bittensor as bt
 
 from typing import List
-from traceback import print_exception
+from pathlib import Path
 
 from folding.base.neuron import BaseNeuron
 from folding.mock import MockDendrite
 from folding.utils.config import add_validator_args
+from folding.utils.ops import print_on_retry
 
-from pathlib import Path
+from tenacity import retry, stop_after_attempt, wait_fixed, retry_if_result
 
 ROOT_DIR = Path(__file__).resolve().parents[2]
 
@@ -67,9 +67,6 @@ def __init__(self, config=None):
             self.metagraph.n, dtype=torch.float32, device=self.device
         )
 
-        # Init sync with the network. Updates the metagraph.
-        self.sync()
-
         # Serve axon to enable external connections.
         if not self.config.neuron.axon_off:
             self.serve_axon()
@@ -110,6 +107,14 @@ def serve_axon(self):
             bt.logging.error(f"Failed to create Axon initialize with exception: {e}")
             pass
 
+    @retry(
+        stop=stop_after_attempt(3),  # Retry up to 3 times
+        wait=wait_fixed(1),  # Wait 1 second between retries
+        retry=retry_if_result(
+            lambda result: result is False
+        ),  # Retry if the result is False
+        after=print_on_retry
+    )
     def set_weights(self):
         """
         Sets the validator weights to the metagraph hotkeys based on the scores it has received from the miners. The weights determine the trust and incentive level the validator assigns to miner nodes on the network.
@@ -161,10 +166,9 @@ def set_weights(self):
             wait_for_inclusion=False,
             version_key=self.spec_version,
         )
-        if result is True:
-            bt.logging.info("set_weights on chain successfully!")
-        else:
-            bt.logging.error("set_weights failed")
+
+            
+        return result
 
     def resync_metagraph(self):
         """Resyncs the metagraph and updates the hotkeys and moving averages based on the new metagraph."""

folding/utils/config.py

@@ -82,7 +82,7 @@ def add_args(cls, parser):
         "--neuron.epoch_length",
         type=int,
         help="The default epoch length (how often we set weights, measured in 12 second blocks).",
-        default=100,
+        default=250,
     )
 
     parser.add_argument(

folding/utils/logging.py

@@ -3,9 +3,7 @@
 from typing import List
 from loguru import logger
 from dataclasses import asdict, dataclass
-import datetime as dt
 import os
-import asyncio
 
 import folding
 import bittensor as bt
@@ -36,7 +34,7 @@ def should_reinit_wandb(self):
     )
 
 
-async def init_wandb(self, pdb_id: str, reinit=True, failed=False):
+def init_wandb(self, pdb_id: str, reinit=True, failed=False):
     """Starts a new wandb run."""
 
     tags = [
@@ -53,7 +51,6 @@ async def init_wandb(self, pdb_id: str, reinit=True, failed=False):
         tags.append("mock")
     if self.config.neuron.disable_set_weights:
         tags.append("disable_set_weights")
-    tags.append("async")
     wandb_config = {
         key: copy.deepcopy(self.config.get(key, None))
         for key in ("neuron", "reward", "netuid", "wandb")
@@ -62,7 +59,7 @@ async def init_wandb(self, pdb_id: str, reinit=True, failed=False):
 
     id = None if pdb_id not in self.wandb_ids.keys() else self.wandb_ids[pdb_id]
 
-    run = await asyncio.to_thread(wandb.init,
+    run = wandb.init(
         anonymous="allow",
         name=pdb_id,
         reinit=reinit,
@@ -93,27 +90,27 @@ async def init_wandb(self, pdb_id: str, reinit=True, failed=False):
     return run
 
 
-async def log_protein(run, pdb_id_path: str):
+def log_protein(run, pdb_id_path: str):
     """Logs the protein visualization to wandb.
     pdb_id_path: str: path to the pdb file on disk.
     """
     try:
-        await asyncio.to_thread(run.log,{"protein_vis": wandb.Molecule(pdb_id_path)})
+        run.log({"protein_vis": wandb.Molecule(pdb_id_path)})
     except:
         bt.logging.warning("Failed to log protein visualization")
 
 
-async def log_folded_protein(run, pdb_id_path: str):
+def log_folded_protein(run, pdb_id_path: str):
     """Logs the folded protein visualization to wandb.
     pdb_id_path: str: path to the pdb file on disk.
     """
     try:
-        await asyncio.to_thread(run.log,{"folded_protein_vis": wandb.Molecule(pdb_id_path)})
+        run.log({"folded_protein_vis": wandb.Molecule(pdb_id_path)})
     except:
         bt.logging.warning("Failed to log folded protein visualization")
 
 
-async def log_event(
+def log_event(
     self,
     event,
     failed=False,
@@ -127,16 +124,16 @@ async def log_event(
         return
     pdb_id = event["pdb_id"]
 
-    run = await init_wandb(self, pdb_id=pdb_id, failed=failed)
+    run = init_wandb(self, pdb_id=pdb_id, failed=failed)
 
     # Log the event to wandb.
     run.log(event)
     wandb.save(os.path.join(self.config.neuron.full_path, f"events.log"))
 
     if pdb_location is not None:
-        await log_protein(run, pdb_id_path=pdb_location)
+        log_protein(run, pdb_id_path=pdb_location)
     if folded_protein_location is not None:
-        await log_folded_protein(run, pdb_id_path=folded_protein_location)
+        log_folded_protein(run, pdb_id_path=folded_protein_location)
         wandb.save(folded_protein_location)
 
     run.finish()

folding/utils/ops.py

@@ -71,6 +71,12 @@ def wrapper(*args, **kwargs):
 
     return decorator
 
+def print_on_retry(retry_state):
+    function_name = retry_state.fn.__name__
+    max_retries = retry_state.retry_object.stop.max_attempt_number
+    bt.logging.warning(f"Retrying {function_name}: retry #{retry_state.attempt_number} out of {max_retries}")
+    
+
 
 def delete_directory(directory: str):
     """We create a lot of files in the process of tracking pdb files.

folding/validators/forward.py

@@ -183,7 +183,7 @@ async def create_new_challenge(self, exclude: List) -> Dict:
             event["hp_search_time"] = time.time() - forward_start_time
 
             # only log the event if the simulation was not successful
-            await log_event(self, event, failed=True)
+            log_event(self, event, failed=True)
             bt.logging.debug(
                 f"❌❌ All hyperparameter combinations failed for pdb_id {pdb_id}.. Skipping! ❌❌"
             )

folding/validators/protein.py

@@ -587,7 +587,7 @@ def is_run_valid(self):
     def get_ns_computed(self):
         """Calculate the number of nanoseconds computed by the miner."""
 
-        return (self.cpt_step * self.system_config.time_step_size) / 1e6
+        return (self.cpt_step * self.system_config.time_step_size) / 1e3
 
     def save_pdb(self, output_path: str):
         """Save the pdb file to the output path."""

neurons/validator.py

@@ -50,6 +50,9 @@ def __init__(self, config=None):
 
         self.load_state()
 
+        # Init sync with the network. Updates the metagraph.
+        self.sync()
+
         # TODO: Change the store to SQLiteJobStore if you want to use SQLite
         self.store = SQLiteJobStore()
         self.mdrun_args = self.parse_mdrun_args()
@@ -322,7 +325,7 @@ async def prepare_event_for_logging(event: Dict):
             bt.logging.error(f"Protein.from_job returns NONE for protein {job.pdb}")
 
         # Remove these keys from the log because they polute the terminal.
-        await log_event(
+        log_event(
             self,
             event=event,
             pdb_location=pdb_location,
@@ -366,20 +369,28 @@ async def create_jobs(self):
                     # Here is where we select, download and preprocess a pdb
                     # We also assign the pdb to a group of workers (miners), based on their workloads
                     await self.add_jobs(k=self.config.neuron.queue_size - queue.qsize())
-                bt.logging.info(
-                    f"Sleeping {self.config.neuron.update_interval} seconds before next job creation loop."
-                )
+                    bt.logging.info(
+                        f"Sleeping 60 seconds before next job creation loop."
+                    )
+                else:
+                    bt.logging.info(
+                        "Job queue is full. Sleeping 60 seconds before next job creation loop."
+                    )
+
             except Exception as e:
                 bt.logging.error(f"Error in create_jobs: {e}")
-            await asyncio.sleep(self.config.neuron.update_interval)
+
+            await asyncio.sleep(60)
 
     async def update_jobs(self):
         while True:
             try:
-                bt.logging.info(f"step({self.step}) block({self.block})")
-
+                # Wait at the beginning of update_jobs since we want to avoid attemping to update jobs before we get data back.
                 await asyncio.sleep(self.config.neuron.update_interval)
+
                 bt.logging.info("Updating jobs.")
+                bt.logging.info(f"step({self.step}) block({self.block})")
+
                 for job in self.store.get_queue(ready=False).queue:
                     # Remove any deregistered hotkeys from current job. This will update the store when the job is updated.
                     if not job.check_for_available_hotkeys(self.metagraph.hotkeys):
@@ -404,7 +415,9 @@ async def update_jobs(self):
                     await self.update_job(job=job)
             except Exception as e:
                 bt.logging.error(f"Error in update_jobs: {e}")
+
             self.step += 1
+
             bt.logging.info(
                 f"Sleeping {self.config.neuron.update_interval} seconds before next job update loop."
             )