Merge pull request #189 from mapbox/nickgerancher/cu-estimation-testing

Nickgerancher/cu estimation testing

Author: nicklvsa

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 =======
 
+# 1.12.0 (2024-10-23)
+- Added command `tilesets estimate-cu` that returns an estimated compute unit value for a user's tileset.
+
 # 1.11.1 (2024-08-01)
 - Added command `tilesets upload-raster-source` to upload raster files as sources
 

README.md

@@ -45,6 +45,30 @@ Note, Windows is not officially supported at this time.
 Windows users need to install [GDAL](http://www.lfd.uci.edu/~gohlke/pythonlibs/#gdal) and [rasterio](http://www.lfd.uci.edu/~gohlke/pythonlibs/#rasterio).
 Then `pip install 'mapbox-tilesets[estimate-area]'`
 
+## Installing optional `estimate-cu` command
+
+If you are using an x86 Mac or Linux machine, run:
+`pip install 'mapbox-tilesets[estimate-cu]'`
+
+Otherwise, you will need to install some dependencies.
+
+### arm64 MacOS
+
+If you're on an arm64 Mac (e.g., with an M1 chip), you'll need to install [GDAL](https://gdal.org/) first. On Mac, a simple way is to use [Homebrew](https://brew.sh/):
+
+```sh
+$ brew install gdal
+...
+$ pip install 'mapbox-tilesets[estimate-cu]'
+```
+
+### Windows
+
+Note, Windows is not officially supported at this time.
+
+Windows users need to install [GDAL](http://www.lfd.uci.edu/~gohlke/pythonlibs/#gdal) and [rasterio](http://www.lfd.uci.edu/~gohlke/pythonlibs/#rasterio).
+Then `pip install 'mapbox-tilesets[estimate-cu]'`
+
 ## Mapbox Access Tokens
 
 In order to use the tilesets endpoints, you need a Mapbox Access Token with `tilesets:write`, `tilesets:read`, and `tilesets:list` scopes. This is a secret token, so do not share it publicly!
@@ -71,6 +95,7 @@ export MAPBOX_ACCESS_TOKEN=my.token
   - [`list-sources`](#list-sources)
   - [`delete-source`](#delete-source)
   - [`estimate-area`](#estimate-area)
+  - [`estimate-cu`](#estimate-cu)
 - Recipes
   - [`view-recipe`](#view-recipe)
   - [`validate-recipe`](#validate-recipe)
@@ -222,6 +247,33 @@ Usage
 tilesets delete-source user source_id
 ```
 
+### estimate-cu
+
+```shell
+tilesets estimate-cu <tileset> -s/--sources <sources> -b/--num-bands <number> --raw
+```
+
+Estimates the CU value of a tileset before publishing it. This is useful to understand the estimated cost a given tileset before you start processing the data. Note: This is currently only available to tileset recipes with type `raster` or `rasterarray`.
+
+See https://docs.mapbox.com/help/glossary/compute-unit/ for more information.
+
+Flags:
+- `-s` or `--sources` [optional]: Local path to the sources that your recipe points at. This is highly recommeneded.
+- `-b` or `--num-bands` [optional]: The number of bands you expect your recipe to select across all layers. This is recommended.
+- `--minzoom` [optional]: Use this flag if your recipe does not contain a minzoom value.
+- `--maxzoom` [optional]: Use this flag if your recipe does not contain a maxzoom value.
+- `--raw` [optional]: This will toggle the pretty print output.
+
+Usage
+
+```shell
+# Estimate the CUs for 'account.tileset' with sources located in 'path/to/sources/' and a band count of 20.
+tilesets estimate-cu account.tileset -s 'path/to/sources/*.grib2' -b 20
+
+# Estimate the CUs for 'account.tileset' for a single source and a band count of 10 (pretty print the results)
+tilesets estimate-cu account.tileset -s 'path/to/sources/helloworld.grib2' -b 10 --raw
+```
+
 ### estimate-area
 
 ```shell

mapbox_tilesets/__init__.py

@@ -1,3 +1,3 @@
 """mapbox_tilesets package"""
 
-__version__ = "1.11.1"
+__version__ = "1.12.0"

mapbox_tilesets/scripts/cli.py

@@ -1,10 +1,12 @@
 """Tilesets command line interface"""
 
+import os
 import base64
 import builtins
 import json
 import re
 import tempfile
+from glob import glob
 from urllib.parse import parse_qs, urlencode, urlparse
 
 import click
@@ -840,6 +842,112 @@ def validate_stream(features):
         yield feature
 
 
+@cli.command("estimate-cu")
+@click.argument("tileset", required=True, type=str)
+@click.option(
+    "--sources",
+    "-s",
+    required=False,
+    type=click.Path(exists=False),
+    help="Local sources represented in your tileset's recipe",
+)
+@click.option(
+    "--num-bands",
+    "-b",
+    required=False,
+    type=int,
+    default=15,
+    help="The number of bands your recipe is selecting",
+)
+@click.option("--minzoom", required=False, type=int, help="The minzoom value override")
+@click.option("--maxzoom", required=False, type=int, help="The maxzoom value override")
+@click.option(
+    "--raw", required=False, type=bool, default=False, is_flag=True, help="Raw CU value"
+)
+@click.option("--token", "-t", required=False, type=str, help="Mapbox access token")
+def estimate_cu(
+    tileset,
+    num_bands=15,
+    minzoom=None,
+    maxzoom=None,
+    sources=None,
+    raw=False,
+    token=None,
+):
+    """
+    Estimates the CUs that will be consumed when processing your recipe into a tileset.
+    Requires extra installation steps: see https://github.com/mapbox/tilesets-cli/blob/master/README.md
+    """
+
+    rio = utils.load_module("rasterio")
+
+    if sources is None:
+        click.echo(f"[warning] estimating '{tileset}' with a default global bounds")
+        sources = ""
+
+    total_size = 0
+    overall_bounds = None
+    src_list = glob(sources)
+
+    if len(src_list) > 0:
+        with rio.open(src_list[0], mode="r") as ds:
+            overall_bounds = ds.bounds
+
+        total_size += os.path.getsize(src_list[0])
+
+    if len(src_list) > 1:
+        for source in src_list:
+            try:
+                with rio.open(source, mode="r") as ds:
+                    if ds.bounds.left < overall_bounds.left:
+                        overall_bounds.left = ds.bounds.left
+                    if ds.bounds.right > overall_bounds.right:
+                        overall_bounds.right = ds.bounds.right
+                    if ds.bounds.top > overall_bounds.top:
+                        overall_bounds.top = ds.bounds.top
+                    if ds.bounds.bottom < overall_bounds.bottom:
+                        overall_bounds.bottom = ds.bounds.bottom
+
+                total_size += os.path.getsize(source)
+            except Exception:
+                click.echo(f"[warning] skipping invalid source '{source}'")
+
+    s = utils._get_session()
+    mapbox_api = utils._get_api()
+    mapbox_token = utils._get_token(token)
+    url = "{0}/tilesets/v1/{1}/estimate".format(mapbox_api, tileset)
+
+    query_params = {
+        "filesize": total_size,
+        "band_count": num_bands,
+        "access_token": mapbox_token,
+    }
+
+    if overall_bounds is not None:
+        query_params["bounds"] = json.dumps([*overall_bounds])
+
+    if minzoom is not None:
+        query_params["minzoom"] = minzoom
+
+    if maxzoom is not None:
+        query_params["maxzoom"] = maxzoom
+
+    response = s.get(url, params=query_params)
+
+    if not response.ok:
+        raise errors.TilesetsError(response.text)
+
+    parsed = json.loads(response.text)
+    if "cu" not in parsed:
+        raise errors.TilesetsError(response.text)
+
+    click.echo(
+        response.text
+        if raw
+        else f"\nEstimated CUs for '{tileset}': {click.style(parsed['cu'], bold=True, fg=155)}. To publish your tileset, run 'tilesets publish'."
+    )
+
+
 @cli.command("estimate-area")
 @cligj.features_in_arg
 @click.option(

setup.py

@@ -44,6 +44,9 @@ def read(fname):
         "estimate-area": [
             "supermercado~=0.2.0",
         ],
+        "estimate-cu": [
+            "rasterio~=1.4.1",
+        ],
         "test": [
             "codecov",
             "pytest==6.2.5",

tests/conftest.py

@@ -11,11 +11,17 @@ def token_environ(monkeypatch):
     monkeypatch.setenv("MapboxAccessToken", "test-token")
 
 
+@pytest.fixture(scope="function")
+def api_environ(monkeypatch):
+    monkeypatch.setenv("MAPBOX_API", "https://api.mapbox.com")
+
+
 class _MockResponse:
     def __init__(self, mock_json, status_code=200):
         self.text = json.dumps(mock_json)
         self._json = mock_json
         self.status_code = status_code
+        self.ok = status_code < 400
 
     def MockResponse(self):
         return self