Merge pull request #288 from macrocosm-os/staging

Author: Sarkosos

README.md

@@ -69,7 +69,7 @@ We use a combination of `conda` and `poetry` to manage our environments. It is v
 git clone https://github.com/macrocosm-os/folding.git
 cd folding
 
-conda create --name folding python3.11
+conda create --name folding python=3.11
 bash install.sh
 ```
 

folding/__init__.py

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

folding/store.py

@@ -1,15 +1,17 @@
+import json
 import os
 import random
-import string
-import json
 import sqlite3
-from typing import List
+import string
+from dataclasses import asdict, dataclass
+from datetime import datetime
 from queue import Queue
-from datetime import datetime, timedelta
-from dataclasses import dataclass, asdict
+from typing import List
 
 import numpy as np
 import pandas as pd
+import requests
+from atom.epistula.epistula import Epistula
 
 DB_DIR = os.path.join(os.path.dirname(__file__), "db")
 
@@ -22,6 +24,7 @@ def __init__(self, db_path=DB_DIR, table_name="protein_jobs"):
         self.table_name = table_name
         os.makedirs(self.db_path, exist_ok=True)
         self.db_file = os.path.join(self.db_path, "jobs.db")
+        self.epistula = Epistula()
 
         self.init_db()
 
@@ -218,6 +221,41 @@ def __repr__(self):
             df = pd.read_sql_query(f"SELECT * FROM {self.table_name}", conn)
         return f"{self.__class__.__name__}\n{df.__repr__()}"
 
+    def upload_job(
+        self,
+        pdb: str,
+        ff: str,
+        box: str,
+        water: str,
+        hotkeys: list,
+        system_kwargs: dict,
+        hotkey,
+        gjp_address: str,
+    ):
+        """Upload a job to the database."""
+
+        body = {
+            "pdb_id": pdb,
+            "hotkeys": hotkeys,
+            "system_config": {
+                "ff": ff,
+                "box": box,
+                "water": water,
+                "system_kwargs": system_kwargs,
+            },
+            "em_s3_link": "s3://path/to/em",
+            "priority": 1,
+            "organic": False
+        }
+        body_bytes = self.epistula.create_message_body(body)
+        headers = self.epistula.generate_header(hotkey=hotkey, body=body_bytes)
+
+        response = requests.post(
+            f"http://{gjp_address}/jobs", headers=headers, data=body_bytes
+        )
+        if response.status_code != 200:
+            raise ValueError(f"Failed to upload job: {response.text}")
+
 
 # Keep the Job and MockJob classes as they are, they work well with both implementations
 @dataclass

folding/utils/config.py

@@ -432,6 +432,12 @@ def add_validator_args(cls, parser):
             "5CQ9KNHy9qvRGhLWeV37agEpmLckSgMXzbZWEEXwbupSCTQy",
         ],
     )
+    parser.add_argument(
+        "--neuron.gjp_address",
+        type=str,
+        help="The IP address and port of the global job pool.",
+        default="174.138.3.61:8030",
+    )
 
 
 def config(cls):

neurons/validator.py

@@ -249,6 +249,19 @@ async def add_job(self, job_event: dict[str, Any], uids: List[int] = None) -> bo
                 system_kwargs=job_event["system_kwargs"],
                 event=job_event,
             )
+            try:
+                self.store.upload_job(
+                    pdb=job_event["pdb_id"],
+                    ff=job_event["ff"],
+                    water=job_event["water"],
+                    box=job_event["box"],
+                    hotkeys=selected_hotkeys,
+                    system_kwargs=job_event["system_kwargs"],
+                    hotkey=self.wallet.hotkey,
+                    gjp_address=self.config.neuron.gjp_address,
+                )
+            except Exception as e:
+                logger.warning(f"Error uploading job: {e}")
 
             return True
         else: