expose items in job_metadata.json #1111

Author: bossie

openeogeotrellis/deploy/batch_job.py

@@ -80,6 +80,7 @@
     log_memory,
     to_jsonable,
     wait_till_path_available,
+    unzip,
 )
 
 logger = logging.getLogger('openeogeotrellis.deploy.batch_job')
@@ -367,6 +368,7 @@ def run_job(
         )
         # perform a first metadata write _before_ actually computing the result. This provides a bit more info, even if the job fails.
         write_metadata({**result_metadata, **_get_tracker_metadata("")}, metadata_file)
+        items = {}
 
         for result in results:
             result.options["batch_mode"] = True
@@ -387,17 +389,16 @@ def run_job(
                 ml_model_metadata = result.get_model_metadata(str(output_file))
                 logger.info("Extracted ml model metadata from %s" % output_file)
 
-        def result_write_assets(result_arg) -> dict:
+        def result_write_assets(result_arg) -> (dict, dict):
             items = result_arg.write_assets(str(output_file))
             assets = {
                 asset_key: asset for item in items.values() for asset_key, asset in item.get("assets", {}).items()
             }
-            # TODO: write items to job_metadata.json
-            return assets
+            return assets, items
 
         concurrent_save_results = int(job_options.get("concurrent-save-results", 1))
         if concurrent_save_results == 1:
-            assets_metadata = list(map(result_write_assets, results))
+            assets_metadata, results_items = unzip(*map(result_write_assets, results))
         elif concurrent_save_results > 1:
             with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_save_results) as executor:
                 futures = []
@@ -406,9 +407,13 @@ def result_write_assets(result_arg) -> dict:
 
                 for _ in concurrent.futures.as_completed(futures):
                     continue
-            assets_metadata = list(map(lambda f: f.result(), futures))
+            assets_metadata, results_items = unzip(*map(lambda f: f.result(), futures))
         else:
             raise ValueError(f"Invalid concurrent_save_results: {concurrent_save_results}")
+        assets_metadata = list(assets_metadata)
+
+        # flattens items for each results into one list
+        items = [item for result in results_items for item in result.values()]
 
         for the_assets_metadata in assets_metadata:
             for name, asset in the_assets_metadata.items():
@@ -491,7 +496,7 @@ def result_write_assets(result_arg) -> dict:
             ml_model_metadata=ml_model_metadata,
         )
 
-        write_metadata({**result_metadata, **_get_tracker_metadata("")}, metadata_file)
+        write_metadata({**result_metadata, **_get_tracker_metadata(""), **{"items": items}}, metadata_file)
         logger.debug("Starting GDAL-based retrieval of asset metadata")
         result_metadata = _assemble_result_metadata(
             tracer=tracer,
@@ -519,7 +524,7 @@ def result_write_assets(result_arg) -> dict:
                 enable_merge=job_options.get("export-workspace-enable-merge", False),
             )
     finally:
-        write_metadata({**result_metadata, **_get_tracker_metadata("")}, metadata_file)
+        write_metadata({**result_metadata, **_get_tracker_metadata(""), **{"items": items}}, metadata_file)
 
 
 def write_metadata(metadata: dict, metadata_file: Path):

tests/test_batch_result.py

@@ -2674,12 +2674,12 @@ def test_unified_asset_keys(tmp_path, separate_asset_per_band, expected_tiff_fil
     assert tiff_files == expected_tiff_files
 
     with open(metadata_file) as f:
-        items = json.load(f).get("items")
+        items = json.load(f)["items"]
 
     print(f"items={json.dumps(items, indent=2)}")
-    return
 
     assert len(items) == 2
+    return
 
     item1_assets = items[0]["assets"]
     assert len(item1_assets) == 1