Merge pull request #249 from macrocosm-os/staging

Author: Sarkosos

environment.yml

@@ -10,6 +10,7 @@ dependencies:
   - loguru=0.7.0
   - pip=24.2
   - pdbfixer=1.9
+  - gzip==1.13
   - pip:
       - bittensor==6.9.4
       - python-dotenv==1.0.1
@@ -18,4 +19,5 @@ dependencies:
       - numpy==1.26.4
       - parmed==4.2.2
       - plotly==5.22.0
-      - kaleido==0.2.1
+      - kaleido==0.2.1
+      - async-timeout==4.0.3

folding/__init__.py

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

folding/base/neuron.py

@@ -122,17 +122,19 @@ def check_openmm_version(self):
             raise e
 
         bt.logging.success(f"Running OpenMM version: {self.openmm_version}")
-        
+
     def setup_wandb_logging(self):
         if os.path.isfile(f"{self.config.neuron.full_path}/wandb_ids.pkl"):
-            self.wandb_ids = load_pkl(f"{self.config.neuron.full_path}/wandb_ids.pkl", "rb")
+            self.wandb_ids = load_pkl(
+                f"{self.config.neuron.full_path}/wandb_ids.pkl", "rb"
+            )
         else:
             self.wandb_ids = {}
-            
+
     def add_wandb_id(self, pdb_id: str, wandb_id: str):
         self.wandb_ids[pdb_id] = wandb_id
         write_pkl(self.wandb_ids, f"{self.config.neuron.full_path}/wandb_ids.pkl", "wb")
-        
+
     def remove_wandb_id(self, pdb_id: str):
         self.wandb_ids.pop(pdb_id)
         write_pkl(self.wandb_ids, f"{self.config.neuron.full_path}/wandb_ids.pkl", "wb")
@@ -141,10 +143,6 @@ def remove_wandb_id(self, pdb_id: str):
     async def forward(self, synapse: bt.Synapse) -> bt.Synapse:
         ...
 
-    @abstractmethod
-    def run(self):
-        ...
-
     def sync(self):
         """
         Wrapper for synchronizing the state of the network for the given miner or validator.

folding/base/validator.py

@@ -110,158 +110,6 @@ def serve_axon(self):
             bt.logging.error(f"Failed to create Axon initialize with exception: {e}")
             pass
 
-    async def concurrent_forward(self):
-        coroutines = [
-            self.forward() for _ in range(self.config.neuron.num_concurrent_forwards)
-        ]
-        await asyncio.gather(*coroutines)
-
-    def run(self):
-        """
-        Initiates and manages the main loop for the miner on the Bittensor network. The main loop handles graceful shutdown on keyboard interrupts and logs unforeseen errors.
-
-        This function performs the following primary tasks:
-        1. Check for registration on the Bittensor network.
-        2. Continuously forwards queries to the miners on the network, rewarding their responses and updating the scores accordingly.
-        3. Periodically resynchronizes with the chain; updating the metagraph with the latest network state and setting weights.
-
-        The essence of the validator's operations is in the forward function, which is called every step. The forward function is responsible for querying the network and scoring the responses.
-
-        Note:
-            - The function leverages the global configurations set during the initialization of the miner.
-            - The miner's axon serves as its interface to the Bittensor network, handling incoming and outgoing requests.
-
-        Raises:
-            KeyboardInterrupt: If the miner is stopped by a manual interruption.
-            Exception: For unforeseen errors during the miner's operation, which are logged for diagnosis.
-        """
-
-        # Check that validator is registered on the network.
-        self.sync()
-
-        bt.logging.info(f"Validator starting at block: {self.block}")
-
-        # This loop maintains the validator's operations until intentionally stopped.
-        try:
-            while True:
-                # Our BaseValidator logic is intentionally as generic as possible so that the Validator neuron can apply problem-specific logic
-                bt.logging.info(f"step({self.step}) block({self.block})")
-
-                # Check if we need to add more jobs to the queue
-                queue = self.store.get_queue(ready=False)
-                if queue.qsize() < self.config.neuron.queue_size:
-                    # Potential situation where (sample_size * queue_size) > available uids on the metagraph.
-                    # Therefore, this product must be less than the number of uids on the metagraph.
-                    if (
-                        self.config.neuron.sample_size * self.config.neuron.queue_size
-                    ) > self.metagraph.n:
-                        raise ValueError(
-                            f"sample_size * queue_size must be less than the number of uids on the metagraph ({self.metagraph.n})."
-                        )
-
-                    bt.logging.debug(f"✅ Creating jobs! ✅")
-                    # 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
-                    self.add_jobs(k=self.config.neuron.queue_size - queue.qsize())
-
-                # TODO: maybe concurrency for the loop below
-                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):
-                        self.store.update(job=job)
-                        continue
-
-                    # Here we straightforwardly query the workers associated with each job and update the jobs accordingly
-                    job_event = self.forward(job=job)
-
-                    # If we don't have any miners reply to the query, we will make it inactive.
-                    if len(job_event["energies"]) == 0:
-                        job.active = False
-                        self.store.update(job=job)
-                        continue
-
-                    if isinstance(job.event, str):
-                        job.event = eval(job.event)  # if str, convert to dict.
-
-                    job.event.update(job_event)
-                    # Determine the status of the job based on the current energy and the previous values (early stopping)
-                    # Update the DB with the current status
-                    self.update_job(job)
-
-                # Check if we should exit.
-                if self.should_exit:
-                    break
-
-                # Sync metagraph and potentially set weights.
-                self.sync()
-
-                self.step += 1
-                bt.logging.warning(
-                    f"Sleeping for {self.config.neuron.update_interval} before resampling..."
-                )
-                time.sleep(self.config.neuron.update_interval)
-
-        # If someone intentionally stops the validator, it'll safely terminate operations.
-        except KeyboardInterrupt:
-            self.axon.stop()
-            bt.logging.debug("Validator killed by keyboard interrupt.")
-            exit()
-
-        # In case of unforeseen errors, the validator will log the error and continue operations.
-        except Exception as err:
-            bt.logging.error("Error during validation", str(err))
-            bt.logging.debug(print_exception(type(err), err, err.__traceback__))
-            self.should_exit = True
-
-    def run_in_background_thread(self):
-        """
-        Starts the validator's operations in a background thread upon entering the context.
-        This method facilitates the use of the validator in a 'with' statement.
-        """
-        if not self.is_running:
-            bt.logging.debug("Starting validator in background thread.")
-            self.should_exit = False
-            self.thread = threading.Thread(target=self.run, daemon=True)
-            self.thread.start()
-            self.is_running = True
-            bt.logging.debug("Started")
-
-    def stop_run_thread(self):
-        """
-        Stops the validator's operations that are running in the background thread.
-        """
-        if self.is_running:
-            bt.logging.debug("Stopping validator in background thread.")
-            self.should_exit = True
-            self.thread.join(5)
-            self.is_running = False
-            bt.logging.debug("Stopped")
-
-    def __enter__(self):
-        # self.run_in_background_thread()
-        self.run()
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        """
-        Stops the validator's background operations upon exiting the context.
-        This method facilitates the use of the validator in a 'with' statement.
-
-        Args:
-            exc_type: The type of the exception that caused the context to be exited.
-                      None if the context was exited without an exception.
-            exc_value: The instance of the exception that caused the context to be exited.
-                       None if the context was exited without an exception.
-            traceback: A traceback object encoding the stack trace.
-                       None if the context was exited without an exception.
-        """
-        if self.is_running:
-            bt.logging.debug("Stopping validator in background thread.")
-            self.should_exit = True
-            self.thread.join(5)
-            self.is_running = False
-            bt.logging.debug("Stopped")
-
     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.
@@ -352,7 +200,7 @@ def resync_metagraph(self):
         # Update the hotkeys.
         self.hotkeys = copy.deepcopy(self.metagraph.hotkeys)
 
-    def update_scores(self, rewards: torch.FloatTensor, uids: List[int]):
+    async def update_scores(self, rewards: torch.FloatTensor, uids: List[int]):
         """Performs exponential moving average on the scores based on the rewards received from the miners."""
 
         # Check if rewards contains NaN values.

folding/rewards/reward_pipeline.py

@@ -4,7 +4,7 @@
 from folding.rewards.linear_reward import divide_decreasing
 
 
-def reward_pipeline(
+async def reward_pipeline(
     energies: torch.Tensor, rewards: torch.Tensor, top_reward: float, job: Job
 ):
     """A reward pipeline that determines how to place rewards onto the miners sampled within the batch.