GTC-3375 Add "copy_solo_tiles" option to optimize creation of int-dist alerts raster (#710)

* GTC-3375 Add "copy_solo_tiles" option to optimize creation of int-dist alerts

When we merge integrated alerts and dist alerts rasters, much of the
globe has only dist alerts. So, we want to optimize the process by
copying the dist alerts tiles directly when to the final raster if there
is no corresponding integrated alerts tile. We provide the
copy_solo_tiles option to do this, which runs the "copy_solo_tiles.sh"
script after the main raster is created with "union_bands = False". The
script also correctly updates extent.geojson and tiles.geojson.

* Validate that auxiliary_asset is not provided if source_uri is provided.

Author: danscales

app/crud/assets.py

@@ -187,6 +187,9 @@ async def get_default_asset(dataset: str, version: str) -> ORMAsset:
 
 
 async def create_asset(dataset, version, **data) -> ORMAsset:
+    '''Write the asset information into the database.  put_asset runs the actual
+    asset pipeline to create the asset data.'''
+
     v: ORMVersion = await versions.get_version(dataset, version)
 
     # default to version.is_downloadable if not set

app/models/pydantic/creation_options.py

@@ -122,6 +122,14 @@ class RasterTileSetAssetCreationOptions(StrictBaseModel):
             "when input files are in different projections from each other."
         )
     )
+    copy_solo_tiles: bool = Field(
+        False,
+        description=(
+            "For raster calculations with multiple inputs, copy last source tile "
+            "directly to the destination if all other source tiles are non-existent,"
+            "even though union_bands is false."
+        )
+    )
     pixel_meaning: str = Field(
         ..., description="Description of what the pixel value in the "
         "raster represents. This is used to clarify the meaning of the raster "
@@ -212,7 +220,7 @@ class RasterTileSetAssetCreationOptions(StrictBaseModel):
     auxiliary_assets: Optional[List[UUID]] = Field(
         None,
         description="Asset IDs of additional rasters you might want to include "
-        "in your calc expression."
+        "in your calc expression. Ignored if source_uri is set."
     )
     photometric: Optional[PhotometricType] = None
     num_processes: Optional[StrictInt] = None
@@ -240,13 +248,15 @@ class PixETLCreationOptions(RasterTileSetAssetCreationOptions):
     source_uri: Optional[List[str]] = Field(
         description="List of input sources. Sources must be the URI of either a "
         "tiles.geojson file on S3 or a folder (prefix) on S3 or GCS. "
-        "Features in tiles.geojson must have path starting with either /vsis3/ or /vsigs/",
+        "Features in tiles.geojson must have path starting with either /vsis3/ or /vsigs/"
+        "auxiliary_assets is ignored if source_uri is set (for creating new versions)",
     )
 
     @validator("source_uri")
     def validate_source_uri(cls, v, values, **kwargs):
         if values.get("source_type") == SourceType.raster:
             assert v, "Raster source types require source_uri"
+            assert not values.get("auxiliary_assets"), "auxiliary_assets should not be specified with source_uri"
         else:
             assert not v, "Only raster source type require source_uri"
         return v

app/tasks/assets.py

@@ -46,7 +46,7 @@ async def put_asset(
     input_data: Dict[str, Any],
     constructor: FrozenSet[Union[AssetType, SourceType]] = ASSET_PIPELINES,
 ) -> None:
-    """Call Asset Pipeline.
+    """Call Asset Pipeline to actually create the data of the asset.
 
     Default assets use source_type for identification. All other assets
     use asset_type directly.

app/tasks/raster_tile_set_assets/raster_tile_set_assets.py

@@ -21,7 +21,7 @@
 from app.settings.globals import DATA_LAKE_BUCKET
 from app.tasks import Callback, callback_constructor
 from app.tasks.batch import execute
-from app.tasks.raster_tile_set_assets.utils import create_pixetl_job, create_unify_projection_job
+from app.tasks.raster_tile_set_assets.utils import create_pixetl_job, create_unify_projection_job, create_copy_solo_tiles_job
 from app.utils.aws import get_s3_client
 from app.utils.path import (
     get_asset_uri,
@@ -43,7 +43,7 @@ async def raster_tile_set_asset(
     # If being created as a source (default) asset, creation_options["source_uri"]
     # will be a list. When being created as an auxiliary asset, it will be None.
     # In the latter case we will generate one for pixETL based on the default asset,
-    # below.
+    # below.  NOTE: auxiliary_assets is ignored if source_uri option is set.
 
     co = deepcopy(input_data["creation_options"])
 
@@ -92,16 +92,24 @@ async def raster_tile_set_asset(
         jobs.append(unify_job)
         creation_options.source_uri = new_src_uris
 
-    jobs.append(
-        await create_pixetl_job(
-            dataset,
-            version,
-            creation_options,
-            "create_raster_tile_set",
-            callback,
-            [unify_job] if unify_job is not None else None,
-        )
+    pixetl_job = await create_pixetl_job(
+        dataset,
+        version,
+        creation_options,
+        "create_raster_tile_set",
+        callback,
+        [unify_job] if unify_job is not None else None,
     )
+    jobs.append(pixetl_job)
+
+    if creation_options.copy_solo_tiles:
+        # Copy the solo tiles from the last source URI after the main raster job
+        # (which should have union_bands = False) has finished.
+        copy_solo_job = await create_copy_solo_tiles_job(
+            dataset, creation_options.source_uri[-1],
+            default_asset.asset_uri, "copy_solo_tiles",
+            callback, [pixetl_job])
+        jobs.append(copy_solo_job)
 
     log: ChangeLog = await execute(jobs)
 

app/tasks/raster_tile_set_assets/utils.py

@@ -260,3 +260,29 @@ async def create_unify_projection_job(
         environment=JOB_ENV,
         callback=callback,
     )
+
+
+async def create_copy_solo_tiles_job(
+    dataset: str,
+    source_uri: str,
+    target_uri: str,
+    job_name: str,
+    callback: Callback,
+    parents: Optional[List[Job]] = None,
+) -> GDAL2TilesJob:
+    command = [
+        "copy_solo_tiles.sh",
+        "--source",
+        source_uri,
+        "--target",
+        target_uri,
+    ]
+
+    return GDAL2TilesJob(
+        dataset=dataset,
+        job_name=job_name,
+        command=command,
+        environment=JOB_ENV,
+        callback=callback,
+        parents=[parent.job_name for parent in parents] if parents else None,
+    )

batch/scripts/copy_solo_tiles.sh

@@ -0,0 +1,84 @@
+#!/bin/bash
+
+set -e
+
+# requires arguments
+# -s | --source
+# -T | --target
+
+# merge_dist --source source_uri --target target_uri
+# where target/source are URIs like
+# "s3://gfw-data-lake/umd_glad_dist_alerts/v20251018/raster/epsg-4326/10/100000/resample10m/geotiff/{tile_id}.tif"
+
+# Copy all tiles in the target raster that don't exist in the source raster, and
+# update tiles.geojson and extent.geojson. Requires that the target raster has every
+# tile that is in the source raster, and fails if this is not true.
+
+ME=$(basename "$0")
+. get_arguments.sh "$@"
+
+# Remove s3:// start of the uris.
+spath="${SOURCE#s3://}"
+tpath="${TARGET#s3://}"
+
+# Separate into path components
+IFS='/' read -r -a scomponents <<< "$spath"
+IFS='/' read -r -a tcomponents <<< "$tpath"
+
+# The uris should end in either geotiff/tiles.geojson or geotiff/{tile_id}.tif
+if [[ ${#scomponents[@]} -ne 10 || ${scomponents[8]} -ne "geotiff" ]]; then
+    echo "Error: bad format for target $source"
+    exit 1
+fi
+if [[ ${#tcomponents[@]} -ne 10 || ${tcomponents[8]} -ne "geotiff" ]]; then
+    echo "Error: bad format for target $target"
+    exit 1
+fi
+
+# dataset/version and pixel_mean for both source and target
+sversion="${scomponents[1]}/${scomponents[2]}"
+smeaning="${scomponents[7]}"
+tversion="${tcomponents[1]}/${tcomponents[2]}"
+tmeaning="${tcomponents[7]}"
+
+# Same as SOURCE and TARGET, but with geotiff/{...} removed
+source="s3://${scomponents[0]}/$sversion/${scomponents[3]}/${scomponents[4]}/${scomponents[5]}/${scomponents[6]}/${scomponents[7]}"
+target="s3://${tcomponents[0]}/$tversion/${tcomponents[3]}/${tcomponents[4]}/${tcomponents[5]}/${tcomponents[6]}/${tcomponents[7]}"
+
+# Get the set of tiles in each raster
+stiles=$(aws s3 ls $source/geotiff/ | grep '.*.tif$' | cut -c32-)
+ttiles=$(aws s3 ls $target/geotiff/ | grep '.*.tif$' | cut -c32-)
+
+uniqSource=$(comm -13 <(echo "$ttiles" | tr ' ' '\n') <(echo "$stiles" | tr ' ' '\n'))
+uniqTarget=$(comm -23 <(echo "$ttiles" | tr ' ' '\n') <(echo "$stiles" | tr ' ' '\n'))
+
+if [ ! -z "$uniqTarget" ]; then 
+   echo "There are tiles in the target that are not in the source, so failing:"
+   echo $uniqTarget
+   exit 1
+fi
+
+if [ -z "$uniqSource" ]; then 
+   echo "There are no unique tiles in source to copy"
+   exit 1
+fi
+
+l=($uniqSource)
+len=${#l[@]}
+echo Copying $len tiles
+
+j=1
+for f in $uniqSource; do
+  echo "Copying $j/$len tile"
+  ((j++))
+  aws s3 cp $source/geotiff/$f $target/geotiff/$f
+  aws s3 cp $source/gdal-geotiff/$f $target/gdal-geotiff/$f
+done
+
+echo Copying extent.geojson
+aws s3 cp $source/geotiff/extent.geojson $target/geotiff/extent.geojson
+aws s3 cp $source/gdal-geotiff/extent.geojson $target/gdal-geotiff/extent.geojson
+
+echo Copying tiles.geojson
+aws s3 cp $source/geotiff/tiles.geojson  - | sed "s#/$sversion/#/$tversion/#g; s#/$smeaning/#/$tmeaning/#g"  | aws s3 cp - $target/geotiff/tiles.geojson
+aws s3 cp $source/gdal-geotiff/tiles.geojson  - | sed "s#/$sversion/#/$tversion/#g; s#/$smeaning/#/$tmeaning/#g"  | aws s3 cp - $target/gdal-geotiff/tiles.geojson