Merge pull request #432 from macrocosm-os/staging

Staging

Author: mccrindlebrian

.gitignore

@@ -194,3 +194,4 @@ db
 local-gjp
 *.ipynb
 miner-data
+api_keys.json

folding/__init__.py

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

folding/base/simulation.py

@@ -26,7 +26,7 @@ def timed(*args, **kwargs):
             start_time = time.time()
             result = method(*args, **kwargs)
             end_time = time.time()
-            logger.info(
+            logger.debug(
                 f"Method {method.__name__} took {end_time - start_time:.4f} seconds"
             )
             return result

folding/miners/folding_miner.py

@@ -10,7 +10,8 @@
 import traceback
 import concurrent.futures
 import asyncio
-import pandas as pd
+import aiohttp
+from aiohttp import FormData
 from collections import defaultdict
 from typing import Dict, List, Tuple, Any
 
@@ -28,8 +29,6 @@
     ExitFileReporter,
     LastTwoCheckpointsReporter,
     SequentialCheckpointReporter,
-    ProteinStructureReporter
-
 )
 from folding.utils.ops import (
     check_if_directory_exists,
@@ -59,14 +58,39 @@ def attach_files(
     return synapse
 
 
+async def upload_to_s3(session: aiohttp.ClientSession, 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()}")
+    except Exception as e:
+        logger.error(f"Error uploading to S3: {e}")
+        get_tracebacks()
+    finally:
+        end_time = time.time()
+        logger.info(f"Upload finished in {end_time - start_time} seconds")
+
+
 def attach_files_to_synapse(
     synapse: JobSubmissionSynapse,
     data_directory: str,
     state: str,
     seed: int,
 ) -> JobSubmissionSynapse:
     """Load output files and attach them to synapse.md_output.
-    
+
     This function handles attaching simulation output files to the synapse object
     for communication with the validator. It combines log files and attaches state files.
 
@@ -78,64 +102,55 @@ def attach_files_to_synapse(
 
     Returns:
         JobSubmissionSynapse: The synapse object with attached files in md_output
-    
+
     Note:
         If state is 'finished', it will be treated as 'md_0_1' for file collection purposes.
     """
     # Initialize empty md_output
     synapse.md_output = {}
 
     try:
+        # Start async S3 upload
+        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
+            ))
         # Normalize state for file collection
         file_collection_state = "md_0_1" if state == "finished" else state
-        
-        # Get state files (excluding logs)
-        state_files_pattern = os.path.join(data_directory, f"{file_collection_state}*")
-        state_files = [f for f in glob.glob(state_files_pattern) if not f.endswith('.log')]
-        
-        if not state_files:
-            raise FileNotFoundError(f"No state files found for {state} in {data_directory}")
 
+        # Get cpt files (excluding logs)
+        cpt_files_pattern = os.path.join(data_directory, f"{file_collection_state}*.cpt")
+        cpt_files = [
+            f for f in glob.glob(cpt_files_pattern)
+        ]
+
+        if not cpt_files:
+            raise FileNotFoundError(
+                f"No cpt files found for {state} in {data_directory}"
+            )
+
+        files_to_attach = cpt_files.copy()
         # Get log files if they exist
-        log_files = sorted(glob.glob(os.path.join(data_directory, "*.log")))
-        files_to_attach = state_files.copy()
-        
-        if log_files:
-            # Read and combine log files using pandas
-            dfs = []
-            
-            for log_file in log_files:
-                try:
-                    # Try reading the file with pandas
-                    df = pd.read_csv(log_file)
-                    if not df.empty:
-                        dfs.append(df)
-                except Exception as e:
-                    logger.warning(f"Could not read log file {log_file}: {e}")
-                    continue
-            
-            if dfs:
-                # Combine all dataframes and sort by step
-                combined_df = pd.concat(dfs, ignore_index=True)  
-                combined_df = combined_df.sort_values('#"Step"')
-                
-                # Write combined log file
-                combined_log_path = os.path.join(data_directory, f"{state}_combined.log")
-                try:
-                    combined_df.to_csv(combined_log_path, index=False)
-                    files_to_attach.append(combined_log_path)
-                except Exception as e:
-                    logger.error(f"Failed to write combined log file: {e}")
-                
-                # Attach all files
-                for filename in files_to_attach:
-                    try:
-                        with open(filename, "rb") as f:
-                            base_filename = os.path.basename(filename)
-                            synapse.md_output[base_filename] = base64.b64encode(f.read())
-                    except Exception as e:
-                        logger.error(f"Failed to read file {filename!r}: {e}")
-                        get_tracebacks()
+        log_file = os.path.join(data_directory, "simulation.log")
+
+        if os.path.exists(log_file):
+            files_to_attach.append(log_file)
+       
+
+        # Attach all files
+        for filename in files_to_attach:
+            try:
+                with open(filename, "rb") as f:
+                    base_filename = os.path.basename(filename)
+                    synapse.md_output[base_filename] = base64.b64encode(
+                        f.read()
+                    )
+            except Exception as e:
+                logger.error(f"Failed to read file {filename!r}: {e}")
+                get_tracebacks()
         else:
             # Just attach state files if no logs exist
             for filename in files_to_attach:
@@ -146,13 +161,6 @@ def attach_files_to_synapse(
                 except Exception as e:
                     logger.error(f"Failed to read file {filename!r}: {e}")
                     get_tracebacks()
-                    
-        try:
-            # Clean up combined log file
-            if os.path.exists(combined_log_path):
-                os.remove(combined_log_path)
-        except Exception as e:
-            logger.warning(f"Failed to remove temporary combined log: {e}")
 
         # Set synapse metadata
         synapse.miner_seed = seed
@@ -167,6 +175,7 @@ def attach_files_to_synapse(
 
     return synapse
 
+
 def check_synapse(
     synapse: JobSubmissionSynapse, event: Dict = None
 ) -> JobSubmissionSynapse:
@@ -488,9 +497,7 @@ def intermediate_submission_forward(self, synapse: IntermediateSubmissionSynapse
         output_dir = event["output_dir"]
         cpt_files = {}
         for checkpoint_number in checkpoint_numbers:
-            cpt_file = os.path.join(
-                output_dir, f"{checkpoint_number}.cpt"
-            )
+            cpt_file = os.path.join(output_dir, f"{checkpoint_number}.cpt")
             if os.path.exists(cpt_file):
                 with open(cpt_file, "rb") as f:
                     cpt_files[checkpoint_number] = base64.b64encode(f.read())
@@ -524,6 +531,7 @@ def forward(self, synapse: JobSubmissionSynapse) -> JobSubmissionSynapse:
             new jobs to avoid duplicate work. State files are only attached if valid
             simulation data is found.
         """
+        start_time = time.time()
         job_id = synapse.job_id
 
         has_worked_on_job, condition, event = self.check_if_job_was_worked_on(
@@ -576,18 +584,20 @@ def forward(self, synapse: JobSubmissionSynapse) -> JobSubmissionSynapse:
                             state=state,
                             seed=seed,
                         )
+                        event["condition"] = "found_existing_data"
+                        event["state"] = state
+                        logger.info(f"Time taken for forward function: {time.time() - start_time} seconds")
+                        return check_synapse(synapse=synapse, event=event)
+
                     except Exception as e:
                         logger.error(
                             f"Failed to read state file for protein {event['pdb_id']} with error: {e}"
                         )
                         state = None
-
-                    finally:
-                        event["condition"] = "found_existing_data"
-                        event["state"] = state
-
+                        logger.info(f"Time taken for forward function: {time.time() - start_time} seconds")
                         return check_synapse(synapse=synapse, event=event)
 
+
             # The set of RUNNING simulations.
             elif condition == "running_simulation":
                 self.simulations[event["pdb_hash"]]["queried_at"] = time.time()
@@ -605,6 +615,7 @@ def forward(self, synapse: JobSubmissionSynapse) -> JobSubmissionSynapse:
                 event["condition"] = "running_simulation"
                 event["state"] = current_executor_state
                 event["queried_at"] = simulation["queried_at"]
+                logger.info(f"Time taken for forward function: {time.time() - start_time} seconds")
 
                 return check_synapse(synapse=synapse, event=event)
 
@@ -911,6 +922,7 @@ def __init__(
 
         self.STATES = ["nvt", "npt", "md_0_1"]
         self.CHECKPOINT_INTERVAL = 10000
+        self.TRAJECTORY_INTERVAL = 10000
         self.STATE_DATA_REPORTER_INTERVAL = 10
         self.EXIT_REPORTER_INTERVAL = 10
 
@@ -1072,17 +1084,50 @@ def configure_commands(
                 )
             )
 
-            simulation.reporters.append(
-                ProteinStructureReporter(
-                    file=f"{self.output_dir}/{state}.log",
-                    reportInterval=self.STATE_DATA_REPORTER_INTERVAL,
-                    step=True,
-                    potentialEnergy=True,
-                    reference_pdb=os.path.join(self.output_dir, f"{self.pdb_id}.pdb"),
-                    speed=True,
-
+            if state =="nvt":
+                simulation.reporters.append(
+                    app.DCDReporter(
+                        file=os.path.join(self.output_dir, "trajectory.dcd"),
+                        reportInterval=self.TRAJECTORY_INTERVAL,
+                    )
+                )
+                simulation.reporters.append(
+                    app.StateDataReporter(
+                        file=os.path.join(self.output_dir, f"simulation.log"),
+                        reportInterval=self.STATE_DATA_REPORTER_INTERVAL,
+                        step=True,
+                        potentialEnergy=True,
+                        time=True,
+                        temperature=True,
+                        volume=True,
+                        density=True,
+                        speed=True,
+                        elapsedTime=True,
+                    )
+                )
+            else:
+                simulation.reporters.append(
+                    app.StateDataReporter(
+                        file=os.path.join(self.output_dir, f"simulation.log"),
+                        reportInterval=self.STATE_DATA_REPORTER_INTERVAL,
+                        step=True,
+                        potentialEnergy=True,
+                        time=True,
+                        temperature=True,
+                        volume=True,
+                        density=True,
+                        speed=True,
+                        elapsedTime=True,
+                        append=True,
+                    )
+                )
+                simulation.reporters.append(
+                    app.DCDReporter(
+                        file=os.path.join(self.output_dir, "trajectory.dcd"),
+                        reportInterval=self.TRAJECTORY_INTERVAL,
+                        append=True,
+                    )
                 )
-            )
             state_commands[state] = simulation
 
         return state_commands

folding/protocol.py

@@ -46,6 +46,7 @@ class JobSubmissionSynapse(bt.Synapse):
 
     pdb_id: str
     job_id: str
+    presigned_url: dict
 
     # Optional request output, filled by receiving axon.
     md_output: typing.Optional[dict] = None