From cdef94135959c8dce2398692600b170080270bf9 Mon Sep 17 00:00:00 2001 From: Ruge Li <91452427+rugeli@users.noreply.github.com> Date: Tue, 12 Nov 2024 14:36:11 -0800 Subject: [PATCH] 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 | 9 ++++-- cellpack/autopack/__init__.py | 4 ++- .../autopack/interface_objects/partners.py | 8 +++-- cellpack/autopack/loaders/recipe_loader.py | 4 ++- cellpack/autopack/upy/colors.py | 1 + cellpack/autopack/upy/simularium/__init__.py | 1 + .../upy/simularium/simularium_helper.py | 32 +++++++++++++------ 7 files changed, 42 insertions(+), 17 deletions(-) diff --git a/cellpack/autopack/DBRecipeHandler.py b/cellpack/autopack/DBRecipeHandler.py index 08c12cc36..61c867459 100644 --- a/cellpack/autopack/DBRecipeHandler.py +++ b/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, + }, ) diff --git a/cellpack/autopack/__init__.py b/cellpack/autopack/__init__.py index dc8467c8e..f2ed6bb07 100755 --- a/cellpack/autopack/__init__.py +++ b/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]` diff --git a/cellpack/autopack/interface_objects/partners.py b/cellpack/autopack/interface_objects/partners.py index 72cac5acf..514b82e11 100644 --- a/cellpack/autopack/interface_objects/partners.py +++ b/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) diff --git a/cellpack/autopack/loaders/recipe_loader.py b/cellpack/autopack/loaders/recipe_loader.py index d289a7037..a9e32c0a6 100644 --- a/cellpack/autopack/loaders/recipe_loader.py +++ b/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"] diff --git a/cellpack/autopack/upy/colors.py b/cellpack/autopack/upy/colors.py index 69d60b144..3fc2c5b31 100644 --- a/cellpack/autopack/upy/colors.py +++ b/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 . """ + import numpy from math import floor diff --git a/cellpack/autopack/upy/simularium/__init__.py b/cellpack/autopack/upy/simularium/__init__.py index acafa855e..46a2d6008 100644 --- a/cellpack/autopack/upy/simularium/__init__.py +++ b/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 . """ + # CRITICAL_DEPENDENCIES = ['blender','c4d'] __revision__ = "01" diff --git a/cellpack/autopack/upy/simularium/simularium_helper.py b/cellpack/autopack/upy/simularium/simularium_helper.py index bd09210b8..845bec2c8 100644 --- a/cellpack/autopack/upy/simularium/simularium_helper.py +++ b/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,7 +1436,7 @@ 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( @@ -1432,7 +1444,7 @@ def store_metadata(file_name, url, db=None): ) # 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(