add batch_job_id to result metadata (#295)

* retrieve job id

* formatting

* add a bucket for batch jobs

* formatting

* remove print statement and add comments

cellpack/autopack/DBRecipeHandler.py

@@ -644,7 +644,7 @@ def upload_recipe(self, recipe_meta_data, recipe_data):
         recipe_to_save["recipe_path"] = self.db.create_path("recipes", recipe_id)
         self.upload_data("recipes", recipe_to_save, recipe_id)
 
-    def upload_result_metadata(self, file_name, url):
+    def upload_result_metadata(self, file_name, url, job_id=None):
         """
         Upload the metadata of the result file to the database.
         """
@@ -654,7 +654,12 @@ def upload_result_metadata(self, file_name, url):
             self.db.update_or_create(
                 "results",
                 file_name,
-                {"user": username, "timestamp": timestamp, "url": url},
+                {
+                    "user": username,
+                    "timestamp": timestamp,
+                    "url": url,
+                    "batch_job_id": job_id,
+                },
             )
 
 

cellpack/autopack/__init__.py

@@ -392,7 +392,9 @@ def read_text_file(filename, destination="", cache="collisionTrees", force=None)
     return sphere_data
 
 
-def load_file(filename, destination="", cache="geometries", force=None, use_docker=False):
+def load_file(
+    filename, destination="", cache="geometries", force=None, use_docker=False
+):
     if is_remote_path(filename):
         database_name, file_path = convert_db_shortname_to_url(filename)
         # command example: `pack -r firebase:recipes/[FIREBASE-RECIPE-ID] -c [CONFIG-FILE-PATH]`

cellpack/autopack/interface_objects/partners.py

@@ -60,9 +60,11 @@ def __init__(self, partners):
                 partner["name"],
                 partner["position"] if "position" in partner else [0, 0, 0],
                 partner["weight"] if "weight" in partner else weight,
-                partner["binding_probability"]
-                if "binding_probability" in partner
-                else 1.0,
+                (
+                    partner["binding_probability"]
+                    if "binding_probability" in partner
+                    else 1.0
+                ),
             )
             self.all_partners.append(partner)
 

cellpack/autopack/loaders/recipe_loader.py

@@ -157,7 +157,9 @@ def _migrate_version(self, old_recipe):
             )
 
     def _read(self, resolve_inheritance=True, use_docker=False):
-        new_values, database_name = autopack.load_file(self.file_path, cache="recipes", use_docker=use_docker)
+        new_values, database_name = autopack.load_file(
+            self.file_path, cache="recipes", use_docker=use_docker
+        )
         if database_name == "firebase":
             objects, gradients, composition = DBRecipeLoader.collect_and_sort_data(
                 new_values["composition"]

cellpack/autopack/upy/colors.py

@@ -16,6 +16,7 @@
     You should have received a copy of the GNU General Public License
     along with upy.  If not, see <http://www.gnu.org/licenses/gpl-3.0.html>.
 """
+
 import numpy
 from math import floor
 

cellpack/autopack/upy/simularium/__init__.py

@@ -16,5 +16,6 @@
     You should have received a copy of the GNU General Public License
     along with upy.  If not, see <http://www.gnu.org/licenses/gpl-3.0.html>.
 """
+
 # CRITICAL_DEPENDENCIES = ['blender','c4d']
 __revision__ = "01"

cellpack/autopack/upy/simularium/simularium_helper.py

@@ -1398,23 +1398,35 @@ def raycast_test(self, obj, start, end, length, **kw):
     def post_and_open_file(self, file_name, open_results_in_browser=True):
         simularium_file = Path(f"{file_name}.simularium")
         url = None
+        job_id = os.environ.get("AWS_BATCH_JOB_ID", None)
         file_name, url = simulariumHelper.store_result_file(
-            simularium_file, storage="aws"
+            simularium_file, storage="aws", batch_job_id=job_id
         )
         if file_name and url:
-            simulariumHelper.store_metadata(file_name, url, db="firebase")
+            simulariumHelper.store_metadata(
+                file_name, url, db="firebase", job_id=job_id
+            )
             if open_results_in_browser:
                 simulariumHelper.open_in_simularium(url)
 
     @staticmethod
-    def store_result_file(file_path, storage=None):
+    def store_result_file(file_path, storage=None, batch_job_id=None):
         if storage == "aws":
             handler = DATABASE_IDS.handlers().get(storage)
-            initialized_handler = handler(
-                bucket_name="cellpack-results",
-                sub_folder_name="simularium",
-                region_name="us-west-2",
-            )
+            # if batch_job_id is not None, then we are in a batch job and should use the temp bucket
+            # TODO: use cellpack-results bucket for batch jobs once we have the correct permissions
+            if batch_job_id:
+                initialized_handler = handler(
+                    bucket_name="cellpack-demo",
+                    sub_folder_name="simularium",
+                    region_name="us-west-2",
+                )
+            else:
+                initialized_handler = handler(
+                    bucket_name="cellpack-results",
+                    sub_folder_name="simularium",
+                    region_name="us-west-2",
+                )
             file_name, url = initialized_handler.save_file_and_get_url(file_path)
             if not file_name or not url:
                 db_maintainer = DBMaintenance(initialized_handler)
@@ -1424,15 +1436,15 @@ def store_result_file(file_path, storage=None):
         return file_name, url
 
     @staticmethod
-    def store_metadata(file_name, url, db=None):
+    def store_metadata(file_name, url, db=None, job_id=None):
         if db == "firebase":
             handler = DATABASE_IDS.handlers().get(db)
             initialized_db = handler(
                 default_db="staging"
             )  # default to staging for metadata uploads
             if initialized_db._initialized:
                 db_uploader = DBUploader(initialized_db)
-                db_uploader.upload_result_metadata(file_name, url)
+                db_uploader.upload_result_metadata(file_name, url, job_id)
             else:
                 db_maintainer = DBMaintenance(initialized_db)
                 print(