support vsifile IO support

Author: vincentsarago

tilebench/__init__.py

@@ -19,7 +19,7 @@
 log.add(sys.stderr, format=fmt)
 
 
-def parse_logs(logs: List[str]) -> Dict[str, Any]:
+def parse_rasterio_io_logs(logs: List[str]) -> Dict[str, Any]:
     """Parse Rasterio and CURL logs."""
     # HEAD
     head_requests = len([line for line in logs if "CURL_INFO_HEADER_OUT: HEAD" in line])
@@ -53,25 +53,67 @@ def parse_logs(logs: List[str]) -> Dict[str, Any]:
     }
 
 
+def parse_vsifile_io_logs(logs: List[str]) -> Dict[str, Any]:
+    """Parse VSIFILE IO logs."""
+    # HEAD
+    head_requests = len(
+        [line for line in logs if "VSIFILE_INFO_HEADER_OUT: HEAD" in line]
+    )
+    head_summary = {
+        "count": head_requests,
+    }
+
+    # GET
+    all_get_requests = len(
+        [line for line in logs if "VSIFILE_INFO_HEADER_OUT: GET" in line]
+    )
+
+    get_requests = [line for line in logs if "VSIFILE: Downloading: " in line]
+
+    get_values_str = []
+    for get in get_requests:
+        get_values_str.extend(get.split("VSIFILE: Downloading: ")[1].split(", "))
+
+    get_values = [list(map(int, r.split("-"))) for r in get_values_str]
+    data_transfer = sum([j - i + 1 for i, j in get_values])
+
+    get_summary = {
+        "count": all_get_requests,
+        "bytes": data_transfer,
+        "ranges": get_values_str,
+    }
+
+    warp_kernel = [line.split(" ")[-2:] for line in logs if "GDALWarpKernel" in line]
+
+    return {
+        "HEAD": head_summary,
+        "GET": get_summary,
+        "WarpKernels": warp_kernel,
+    }
+
+
 def profile(
     kernels: bool = False,
     add_to_return: bool = False,
     quiet: bool = False,
     raw: bool = False,
     cprofile: bool = False,
     config: Optional[Dict] = None,
+    io="rasterio",
 ):
     """Profiling."""
+    if io not in ["rasterio", "vsifile"]:
+        raise ValueError(f"Unsupported {io} IO backend")
 
     def wrapper(func: Callable):
         """Wrap a function."""
 
         def wrapped_f(*args, **kwargs):
             """Wrapped function."""
-            rio_stream = StringIO()
-            logger = logging.getLogger("rasterio")
+            io_stream = StringIO()
+            logger = logging.getLogger(io)
             logger.setLevel(logging.DEBUG)
-            handler = logging.StreamHandler(rio_stream)
+            handler = logging.StreamHandler(io_stream)
             logger.addHandler(handler)
 
             gdal_config = config or {}
@@ -88,10 +130,15 @@ def wrapped_f(*args, **kwargs):
             logger.removeHandler(handler)
             handler.close()
 
-            logs = rio_stream.getvalue().splitlines()
+            logs = io_stream.getvalue().splitlines()
             profile_lines = [p for p in profile_stream.getvalue().splitlines() if p]
 
-            results = parse_logs(logs)
+            results = {}
+            if io == "vsifile":
+                results.update(parse_vsifile_io_logs(logs))
+            else:
+                results.update(parse_rasterio_io_logs(logs))
+
             results["Timing"] = t.elapsed
 
             if cprofile:

tilebench/middleware.py

@@ -10,7 +10,7 @@
 from starlette.requests import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
-from tilebench import parse_logs
+from tilebench import parse_rasterio_io_logs
 
 
 class VSIStatsMiddleware(BaseHTTPMiddleware):
@@ -21,22 +21,24 @@ def __init__(
         app: ASGIApp,
         config: Optional[Dict] = None,
         exclude_paths: Optional[List] = None,
+        io: str = "rasterio",
     ) -> None:
         """Init Middleware."""
         super().__init__(app)
         self.config: Dict = config or {}
         self.exclude_paths: List = exclude_paths or []
+        self.io_backend = io
 
     async def dispatch(self, request: Request, call_next):
         """Add VSI stats in headers."""
 
         if request.scope["path"] in self.exclude_paths:
             return await call_next(request)
 
-        rio_stream = StringIO()
-        logger = logging.getLogger("rasterio")
+        io_stream = StringIO()
+        logger = logging.getLogger(self.io_backend)
         logger.setLevel(logging.DEBUG)
-        handler = logging.StreamHandler(rio_stream)
+        handler = logging.StreamHandler(io_stream)
         logger.addHandler(handler)
 
         gdal_config = {"CPL_DEBUG": "ON", "CPL_CURL_VERBOSE": "TRUE"}
@@ -46,10 +48,10 @@ async def dispatch(self, request: Request, call_next):
         logger.removeHandler(handler)
         handler.close()
 
-        if rio_stream:
-            rio_lines = rio_stream.getvalue().splitlines()
+        if io_stream:
+            io_lines = io_stream.getvalue().splitlines()
 
-            results = parse_logs(rio_lines)
+            results = parse_rasterio_io_logs(io_lines)
             head_results = "head;count={count}".format(**results["HEAD"])
             get_results = "get;count={count};size={bytes}".format(**results["GET"])
             ranges_results = "ranges; values={}".format(

tilebench/scripts/cli.py

@@ -103,6 +103,13 @@ def cli():
     callback=options_to_dict,
     help="Reader Options.",
 )
+@click.option(
+    "--io",
+    "io_backend",
+    type=click.Choice(["vsifile", "rasterio"], case_sensitive=True),
+    help="IO Backend Options.",
+    default="rasterio",
+)
 def profile(
     input,
     tile,
@@ -115,6 +122,7 @@ def profile(
     tms,
     config,
     reader_params,
+    io_backend,
 ):
     """Profile Reader Tile read."""
     tilematrixset = default_tms
@@ -166,6 +174,7 @@ def profile(
         raw=add_stdout,
         cprofile=add_cprofile,
         config=config,
+        io=io_backend,
     )
     def _read_tile(src_path: str, x: int, y: int, z: int, tilesize: int = 256):
         with DstReader(src_path, tms=tilematrixset, **reader_params) as cog: