Merge pull request #443 from macrocosm-os/staging

Staging

Author: Sarkosos

folding/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.5.2"
+__version__ = "2.6.0"
 version_split = __version__.split(".")
 __spec_version__ = (
     (10000 * int(version_split[0]))

folding/miners/folding_miner.py

@@ -58,23 +58,25 @@ def attach_files(
     return synapse
 
 
-async def upload_to_s3(session: aiohttp.ClientSession, presigned_url: dict, file_path: str) -> None:
+async def upload_to_s3(presigned_url: dict, file_path: str) -> None:
     """Asynchronously upload a file to S3 using presigned URL"""
     try:
         start_time = time.time()
         data = FormData()
         for key, value in presigned_url["fields"].items():
             data.add_field(key, value)
-        
+
         with open(file_path, "rb") as f:
             data.add_field("file", f, filename="trajectory.dcd")
-            
-            async with session.post(
-                presigned_url["url"],
-                data=data
-            ) as response:
-                if response.status != 204:
-                    logger.error(f"Failed to upload trajectory to s3: {await response.text()}")
+
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    presigned_url["url"],
+                    data=data
+                ) as response:
+                    if response.status != 204:
+                        logger.error(f"Failed to upload trajectory to s3: {await response.text()}")
+
     except Exception as e:
         logger.error(f"Error uploading to S3: {e}")
         get_tracebacks()
@@ -114,7 +116,6 @@ def attach_files_to_synapse(
         trajectory_path = os.path.join(data_directory, "trajectory.dcd")
         if os.path.exists(trajectory_path):
             asyncio.create_task(upload_to_s3(
-                session=aiohttp.ClientSession(),
                 presigned_url=synapse.presigned_url,
                 file_path=trajectory_path
             ))
@@ -921,9 +922,9 @@ def __init__(
         }
 
         self.STATES = ["nvt", "npt", "md_0_1"]
-        self.CHECKPOINT_INTERVAL = 10000
-        self.TRAJECTORY_INTERVAL = 10000
-        self.STATE_DATA_REPORTER_INTERVAL = 10
+        self.CHECKPOINT_INTERVAL = self.system_config.save_interval_checkpoint
+        self.TRAJECTORY_INTERVAL = self.system_config.save_interval_trajectory
+        self.STATE_DATA_REPORTER_INTERVAL = self.system_config.save_interval_log
         self.EXIT_REPORTER_INTERVAL = 10
 
     def create_empty_file(self, file_path: str):

folding/organic/api.py

@@ -3,6 +3,10 @@
 from folding.base import validator
 from folding.organic.organic import router as organic_router
 from folding.utils.logging import logger
+import multiprocessing
+from multiprocessing.connection import Connection
+import pickle
+from typing import Optional
 
 app = FastAPI()
 
@@ -25,3 +29,96 @@ async def start_organic_api(organic_validator, config):
     )
     server = uvicorn.Server(config)
     await server.serve()
+
+
+def api_process_main(pipe_connection: Connection, config):
+    """
+    Main function to run in the separate API process.
+    Receives jobs from the API and sends them back to the main process.
+
+    Args:
+        pipe_connection: Connection to communicate with the main process
+        config: Configuration for the API
+    """
+    from folding.organic.api import app
+    import uvicorn
+    from atom.organic_scoring.organic_queue import OrganicQueue
+    import asyncio
+    from asyncio import Task
+
+    # Create a dummy validator object that will send jobs through the pipe
+    class PipeOrganicValidator:
+        def __init__(self, pipe_connection):
+            self._organic_queue = OrganicQueue()
+            self._pipe_connection = pipe_connection
+            self._check_queue_task: Optional[Task] = None
+
+        async def check_queue(self):
+            """Periodically check the queue and send jobs to the main process"""
+            while True:
+                try:
+                    if not self._organic_queue.is_empty():
+                        # Get all items from the queue
+                        items = []
+                        while not self._organic_queue.is_empty():
+                            item = self._organic_queue.sample()
+                            if item:
+                                items.append(item)
+
+                        # Send items through the pipe
+                        if items:
+                            logger.info(f"Sending {len(items)} jobs to main process")
+                            self._pipe_connection.send(pickle.dumps(items))
+                except Exception as e:
+                    logger.error(f"Error checking queue: {e}")
+                await asyncio.sleep(
+                    1
+                )  # Check more frequently than the main process reads
+
+    # Set up the API
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    # Create the validator
+    organic_validator = PipeOrganicValidator(pipe_connection)
+
+    # Start the queue checking task
+    organic_validator._check_queue_task = loop.create_task(
+        organic_validator.check_queue()
+    )
+
+    # Set up the app state
+    app.state.validator = organic_validator
+    app.state.config = config
+
+    # Start the API
+    uvicorn_config = uvicorn.Config(
+        "folding.organic.api:app",
+        host="0.0.0.0",
+        port=config.neuron.organic_api.port,
+        loop="asyncio",
+        reload=False,
+    )
+
+    server = uvicorn.Server(uvicorn_config)
+    loop.run_until_complete(server.serve())
+
+
+def start_organic_api_in_process(config):
+    """
+    Start the organic API in a separate process and return a pipe connection
+    to receive jobs from it.
+
+    Args:
+        config: Configuration for the API
+
+    Returns:
+        Connection: Pipe connection to receive jobs from the API process
+    """
+    parent_conn, child_conn = multiprocessing.Pipe()
+    process = multiprocessing.Process(
+        target=api_process_main, args=(child_conn, config), daemon=True
+    )
+    process.start()
+    logger.info(f"Started organic API in separate process (PID: {process.pid})")
+    return parent_conn, process

folding/registries/evaluation_registry.py

@@ -7,7 +7,8 @@
 import numpy as np
 import pandas as pd
 from openmm import app
-
+import MDAnalysis as mda
+from MDAnalysis.analysis.rms import rmsd
 from folding.base.evaluation import BaseEvaluator
 from folding.base.simulation import OpenMMSimulation
 from folding.utils import constants as c
@@ -106,10 +107,24 @@ def process_md_output(self) -> bool:
         )
 
         try:
+
+            self.trajectory_path = os.path.join(
+                self.miner_data_directory, "trajectory.dcd"
+            )
+
+            # download the trajectory from s3
+            self.s3_handler.get(self.trajectory_s3_path, self.trajectory_path)
+
+            # check if file exists
+            if not os.path.exists(self.trajectory_path):
+                logger.error(
+                    f"Trajectory file {self.trajectory_path} does not exist... Skipping!"
+                )
+                return False
+
             # NOTE: The seed written in the self.system_config is not used here
             # because the miner could have used something different and we want to
             # make sure that we are using the correct seed.
-
             logger.info(
                 f"Recreating miner {self.hotkey_alias} simulation in state: {self.current_state}"
             )
@@ -129,12 +144,6 @@ def process_md_output(self) -> bool:
             self.log_file_path = os.path.join(
                 self.miner_data_directory, self.md_outputs_exts["log"]
             )
-            self.trajectory_path = os.path.join(
-                self.miner_data_directory, "trajectory.dcd"
-            )
-
-            # download the trajectory from s3
-            self.s3_handler.get(self.trajectory_s3_path, self.trajectory_path)
 
             simulation.loadCheckpoint(checkpoint_path)
 
@@ -210,7 +219,7 @@ def process_md_output(self) -> bool:
                     f"Miner {self.hotkey_alias} has modified the system in unintended ways... Skipping!"
                 )
             self.number_of_checkpoints = (
-                int(self.log_file['#"Step"'].iloc[-1] / 10000) - 1
+                int(self.cpt_step // self.system_config.save_interval_checkpoint) - 1
             )
             if self.number_of_checkpoints < c.MAX_CHECKPOINTS_TO_VALIDATE:
                 raise ValidationError(
@@ -338,6 +347,7 @@ async def is_run_valid(
                 result,
             ) = self.is_checkpoint_valid(
                 checkpoint_path=self.checkpoint_path,
+                current_cpt_step=self.cpt_step,
                 steps_to_run=c.MAX_SIMULATION_STEPS_FOR_EVALUATION,
                 checkpoint_num="final",
             )
@@ -395,13 +405,18 @@ async def is_run_valid(
                     with open(temp_checkpoint_path, "wb") as f:
                         f.write(checkpoint_data)
 
+                    cpt_step = (
+                        int(checkpoint_num) + 1
+                    ) * self.system_config.save_interval_checkpoint
+
                     (
                         is_valid,
                         checked_energies,
                         miner_energies,
                         result,
                     ) = self.is_checkpoint_valid(
                         checkpoint_path=temp_checkpoint_path,
+                        current_cpt_step=cpt_step,
                         steps_to_run=c.INTERMEDIATE_CHECKPOINT_STEPS,
                         checkpoint_num=checkpoint_num,
                     )
@@ -544,9 +559,45 @@ def get_miner_log_file_energies(
 
         return miner_energies
 
+    def calculate_rmsd(
+        self, miner_trajectory, validator_trajectory, start_frame, end_frame
+    ):
+        """Calculate the RMSD between the miner and validator trajectories for every frame.
+
+        Args:
+            miner_trajectory: MDAnalysis Universe for miner trajectory
+            validator_trajectory: MDAnalysis Universe for validator trajectory
+
+        Returns:
+            list[float]: List of RMSD values for each frame
+        """
+        # Get backbone atoms for both trajectories
+        miner_bb = miner_trajectory.select_atoms("backbone")
+        validator_bb = validator_trajectory.select_atoms("backbone")
+
+        rmsds = []
+
+        # Iterate through frames in both trajectories
+        for miner_ts, validator_ts in zip(
+            miner_trajectory.trajectory[start_frame:end_frame],
+            validator_trajectory.trajectory,
+        ):
+            # Get coordinates for current frame
+            miner_positions = miner_bb.positions.copy()
+            validator_positions = validator_bb.positions.copy()
+
+            # Calculate RMSD for current frame
+            rmsd_value = rmsd(
+                miner_positions, validator_positions, center=True, superposition=True
+            )
+            rmsds.append(rmsd_value)
+
+        return rmsds
+
     def is_checkpoint_valid(
         self,
         checkpoint_path: str,
+        current_cpt_step: int,
         steps_to_run: int = c.MIN_SIMULATION_STEPS,
         checkpoint_num: str = "final",
     ):
@@ -580,7 +631,6 @@ def is_checkpoint_valid(
 
         # Load checkpoint
         simulation.loadCheckpoint(checkpoint_path)
-        current_cpt_step = simulation.currentStep
 
         if current_cpt_step + steps_to_run > self.log_step:
             raise ValidationError(message="simulation-step-out-of-range")
@@ -620,6 +670,10 @@ def is_checkpoint_valid(
                 self.miner_data_directory, f"check_{checkpoint_num}.log"
             )
 
+            current_state_trajectory = os.path.join(
+                self.miner_data_directory, f"check_{checkpoint_num}.dcd"
+            )
+
             simulation, _ = self.md_simulator.create_simulation(
                 pdb=load_pdb_file(pdb_file=self.pdb_location),
                 system_config=self.system_config.get_config(),
@@ -639,6 +693,13 @@ def is_checkpoint_valid(
                 )
             )
 
+            simulation.reporters.append(
+                app.DCDReporter(
+                    current_state_trajectory,
+                    self.system_config.save_interval_trajectory,
+                )
+            )
+
             logger.info(
                 f"Running {steps_to_run} steps. log_step: {self.log_step}, cpt_step: {current_cpt_step}"
             )
@@ -695,6 +756,29 @@ def is_checkpoint_valid(
                 )
                 raise ValidationError(message="anomaly")
 
+            # convert steps to frames
+            start_frame = (
+                current_cpt_step // self.system_config.save_interval_trajectory
+            )
+            end_frame = max_step // self.system_config.save_interval_trajectory
+
+            check_universe = mda.Universe(self.pdb_location, current_state_trajectory)
+            miner_universe = mda.Universe(self.pdb_location, self.trajectory_path)
+
+            # Calculate RMSD between the trajectories for each frame
+            rmsds = self.calculate_rmsd(
+                miner_universe, check_universe, start_frame, end_frame
+            )
+
+            # Get median RMSD value
+            median_rmsd = np.median(rmsds)
+
+            if median_rmsd > 1:
+                logger.warning(
+                    f"hotkey {self.hotkey_alias} failed trajectory RMSD check for {self.pdb_id}, checkpoint_num: {checkpoint_num}, with median RMSD: {median_rmsd} ... Skipping!"
+                )
+                raise ValidationError(message="trajectory-rmsd")
+
             # Save the intermediate or final pdb file if the run is valid
             positions = simulation.context.getState(getPositions=True).getPositions()
             topology = simulation.topology

folding/utils/opemm_simulation_config.py

@@ -26,6 +26,7 @@ class SimulationConfig(BaseModel):
     time_step_size: float = 0.002
     time_units: str = "picosecond"
     save_interval_checkpoint: int = 10000
+    save_interval_trajectory: int = 100
     save_interval_log: int = 10
     box_padding: float = 1.0
     friction: float = 1.0