Merge pull request #34 from FrontierDevelopmentLab/feat/append-data

Feat/append data

Author: frandorr

conf/config.yaml

@@ -4,9 +4,11 @@ output: ./output/${dataset_name}
 log_path: ${output}/main.log
 credentials: ${output}/token.json
 gpd_input: ${output}/aoi.geojson
-item_collection: ${output}/item_collection.geojson
+item_collection: ???
 tiles: ${output}/tiles.pkl
-extraction_tasks: ${output}/extraction_tasks.pkl
+extraction_tasks: ???
+
+overwrite: false
 
 start_date: 2020-01-01
 end_date: 2020-02-01

src/satextractor/cli.py

@@ -14,6 +14,9 @@
 
   Also see (1) from http://click.pocoo.org/5/setuptools/#setuptools-integration
 """
+import base64
+import datetime
+import hashlib
 import os
 import pickle
 
@@ -127,6 +130,7 @@ def preparer(cfg):
         cfg.constellations,
         f"{cfg.cloud.storage_prefix}/{cfg.cloud.storage_root}/{cfg.dataset_name}",
         cfg.tiler.bbox_size,
+        cfg.overwrite,
     )
 
 
@@ -136,9 +140,13 @@ def deployer(cfg):
     extraction_tasks = pickle.load(open(cfg.extraction_tasks, "rb"))
 
     topic = f"projects/{cfg.cloud.project}/topics/{'-'.join([cfg.cloud.user_id, 'stacextractor'])}"
+    job_id = (
+        f"{cfg.cloud.user_id}_{datetime.datetime.now()}_{cfg.extraction_tasks[0:10]}"
+    )
 
     hydra.utils.call(
         cfg.deployer,
+        job_id,
         cfg.credentials,
         extraction_tasks,
         f"{cfg.cloud.storage_prefix}/{cfg.cloud.storage_root}/{cfg.dataset_name}",
@@ -184,7 +192,23 @@ def main(cfg: DictConfig):
             "plugins",
         ], "valid tasks are [build, stac, tile, schedule, prepare, deploy, plugins]"
 
-    logger.info(f"Running tasks {cfg.tasks}")
+    # prepare a hashed representation of key config values (start time, end time, constellations)
+    hash_vals = (
+        "_".join(
+            [cfg.start_date, cfg.end_date]
+            + [constellation for constellation in cfg.constellations],
+        )
+    ).encode("utf-8")
+
+    hash_str_b = hashlib.md5(hash_vals).digest()
+    hash_str = base64.urlsafe_b64encode(hash_str_b).decode("utf-8").rstrip("=")
+
+    cfg.item_collection = os.path.join(cfg.output, f"{hash_str}_item_collection.pkl")
+    cfg.extraction_tasks = os.path.join(cfg.output, f"{hash_str}_extraction_tasks.pkl")
+
+    pickle.dump(cfg, open(os.path.join(cfg.output, f"{hash_str}_cfg.pkl"), "wb"))
+
+    logger.info(f"Running tasks {cfg.tasks} for cfg {hash_str}")
 
     if "build" in cfg.tasks:
         build(cfg)

src/satextractor/deployer/gcp_deployer.py

@@ -1,21 +1,20 @@
 import concurrent
-import hashlib
 import json
-from datetime import datetime
 
 from google.cloud import pubsub_v1
 from loguru import logger
 from satextractor.models.constellation_info import BAND_INFO
 from tqdm import tqdm
 
 
-def deploy_tasks(credentials, extraction_tasks, storage_path, chunk_size, topic):
-
-    user_id = topic.split("/")[-1].split("-")[0]
-
-    job_id = hashlib.sha224(
-        (user_id + str(datetime.now())).encode(),
-    ).hexdigest()[:10]
+def deploy_tasks(
+    job_id,
+    credentials,
+    extraction_tasks,
+    storage_path,
+    chunk_size,
+    topic,
+):
 
     logger.info(f"Deploying {len(extraction_tasks)} tasks with job_id: {job_id}")
 

src/satextractor/preparer/gcp_preparer.py

@@ -23,6 +23,7 @@ def gcp_prepare_archive(
     constellations: List[str],
     storage_root: str,
     patch_size: int,
+    overwrite: bool,
     chunk_size: int,
     n_jobs: int = -1,
     verbose: int = 0,
@@ -69,7 +70,7 @@ def gcp_prepare_archive(
     ):
         Parallel(n_jobs=n_jobs, verbose=verbose, prefer="threads")(
             [
-                delayed(zarr.open)(fs.get_mapper(f"{storage_root}/{tile_id}"))
+                delayed(zarr.open)(fs.get_mapper(f"{storage_root}/{tile_id}"), "a")
                 for tile_id, _ in items
             ],
         )
@@ -88,6 +89,7 @@ def gcp_prepare_archive(
                     sensing_times,
                     constellation,
                     BAND_INFO[constellation],
+                    overwrite,
                 ),
             )
 

src/satextractor/preparer/preparer.py

@@ -1,5 +1,11 @@
+import datetime
+
 import numpy as np
 import zarr
+from zarr.errors import ArrayNotFoundError
+from zarr.errors import ContainsArrayError
+from zarr.errors import ContainsGroupError
+from zarr.errors import PathNotFoundError
 
 
 def create_zarr_patch_structure(
@@ -11,34 +17,129 @@ def create_zarr_patch_structure(
     sensing_times,
     constellation,
     bands,
+    overwrite,
 ):
     if not sensing_times.size == 0:
         patch_size_pixels = patch_size // min(b["gsd"] for _, b in bands.items())
 
         patch_constellation_path = f"{storage_path}/{tile_id}/{constellation}"
-        zarr.open(fs_mapper(patch_constellation_path), mode="a")
+        zarr.open(
+            fs_mapper(patch_constellation_path),
+            mode="a",
+        )  # make sure the path exists
 
         patch_path = f"{patch_constellation_path}/data"
-        zarr.open_array(
-            fs_mapper(patch_path),
-            "w",
-            shape=(
-                len(sensing_times),
-                len(bands),
-                int(patch_size_pixels),
-                int(patch_size_pixels),
-            ),
-            chunks=(1, 1, int(chunk_size), int(chunk_size)),
-            dtype=np.uint16,
-        )
-
-        # Create timestamps array
-        timestamps_path = f"{patch_constellation_path}/timestamps"
-        z_dates = zarr.open_array(
-            fs_mapper(f"{timestamps_path}"),
-            mode="w",
-            shape=(len(sensing_times)),
-            chunks=(len(sensing_times)),
-            dtype="<U27",
-        )
-        z_dates[:] = sensing_times
+
+        if overwrite:
+            zarr.open_array(
+                fs_mapper(patch_path),
+                "w",
+                shape=(
+                    len(sensing_times),
+                    len(bands),
+                    int(patch_size_pixels),
+                    int(patch_size_pixels),
+                ),
+                chunks=(1, 1, int(chunk_size), int(chunk_size)),
+                dtype=np.uint16,
+            )
+
+            # Create timestamps array
+            timestamps_path = f"{patch_constellation_path}/timestamps"
+            z_dates = zarr.open_array(
+                fs_mapper(f"{timestamps_path}"),
+                mode="w",
+                shape=(len(sensing_times)),
+                chunks=(len(sensing_times)),
+                dtype="<U27",
+            )
+            z_dates[:] = sensing_times
+
+        else:
+
+            # read current timestamps
+            timestamps_path = f"{patch_constellation_path}/timestamps"
+
+            try:
+                existing_timestamps = zarr.open_array(fs_mapper(timestamps_path), "r")[
+                    :
+                ]
+                existing_timestamps = np.array(
+                    [
+                        np.datetime64(datetime.datetime.fromisoformat(el))
+                        for el in existing_timestamps
+                    ],
+                )
+
+                new_timesteps = np.array(sensing_times)[
+                    ~np.isin(sensing_times, existing_timestamps)
+                ]
+
+                if new_timesteps.shape[0] > 0:
+                    assert max(existing_timestamps) <= min(
+                        new_timesteps,
+                    ), "Sat-Extractor can only append more recent data or overwrite existing data. "
+
+                # get union of sensing times
+                timestamps_union = np.union1d(existing_timestamps, sensing_times)
+
+            except (PathNotFoundError, ContainsGroupError, ArrayNotFoundError):
+                timestamps_union = sensing_times
+            except Exception as e:
+                raise e
+
+            # write sensing times fresh
+            z_dates = zarr.open_array(
+                fs_mapper(f"{timestamps_path}"),
+                mode="w",
+                shape=(len(timestamps_union)),
+                chunks=(len(timestamps_union)),
+                dtype="<U27",
+            )
+            z_dates[:] = timestamps_union
+
+            # resize any existing array based thereon
+
+            # data
+            try:
+                # if it doesn't exist, create it.
+                z_data = zarr.open_array(
+                    fs_mapper(patch_path),
+                    "w-",
+                    shape=(
+                        len(timestamps_union),
+                        len(bands),
+                        int(patch_size_pixels),
+                        int(patch_size_pixels),
+                    ),
+                    chunks=(1, 1, int(chunk_size), int(chunk_size)),
+                    dtype=np.uint16,
+                )
+            except ContainsArrayError:
+                z_data = zarr.open_array(fs_mapper(patch_path), "r+")
+
+                data_shape = z_data.shape
+                z_data.resize(len(timestamps_union), *data_shape[1:])
+
+            except Exception as e:
+                raise e
+
+            # masks
+            mask_root_path = f"{patch_constellation_path}/mask"
+            z_mask_dir = zarr.open_group(
+                fs_mapper(mask_root_path),
+                "a",
+            )
+
+            for mask_key in z_mask_dir.keys():
+
+                mask_path = f"{mask_root_path}/{mask_key}"
+
+                z_mask = zarr.open_array(
+                    fs_mapper(mask_path),
+                    "r+",
+                )
+
+                mask_shape = z_mask.shape
+
+                z_mask.resize(len(timestamps_union), *mask_shape[1:])