diff --git a/.env.example b/.env.example index d5a8daee..08d01a4f 100644 --- a/.env.example +++ b/.env.example @@ -15,6 +15,8 @@ STORAGE_SERVER_DATA_PATH=/data SSH_PRIVATE_KEY_PASSPHRASE= TCD_CONTAINER_IMAGE=deadtrees-tcd:latest TCD_CONTAINER_TIMEOUT_SECONDS=14400 +TCD_CONTAINER_TIMEOUT_MAX_SECONDS=43200 +TCD_CONTAINER_TIMEOUT_BASE_PIXELS=2000000000 # API/reference export runtime. DEADTREES_USER= diff --git a/processor/src/treecover_segmentation_oam_tcd/predict_treecover.py b/processor/src/treecover_segmentation_oam_tcd/predict_treecover.py index edefd19a..19d7bdf0 100644 --- a/processor/src/treecover_segmentation_oam_tcd/predict_treecover.py +++ b/processor/src/treecover_segmentation_oam_tcd/predict_treecover.py @@ -3,6 +3,8 @@ import uuid import json import time +import math +from dataclasses import dataclass from pathlib import Path import numpy as np import rasterio @@ -49,6 +51,16 @@ CHECKPOINT_NAME = TCD_MODEL +@dataclass(frozen=True) +class _TCDTimeoutPolicy: + timeout_seconds: int + base_timeout_seconds: int + max_timeout_seconds: int + base_pixels: int + input_pixels: int | None + capped: bool + + class _TCDContainerTimeout(Exception): """Raised when the TCD Docker wait call reaches its allowed wall time.""" @@ -80,6 +92,38 @@ def _is_docker_wait_timeout(exc: BaseException) -> bool: return isinstance(exc, requests.exceptions.ConnectionError) and 'Read timed out' in str(exc) +def _compute_tcd_timeout_policy(input_width: int | None = None, input_height: int | None = None) -> _TCDTimeoutPolicy: + """ + Choose a TCD wall-time limit from configured base timeout and input size. + + TCD runtime scales with the number of pixels passed to the model. The configured + timeout remains the floor for ordinary inputs; very large orthos get more time + up to a separate cap so a single stuck container cannot block the processor + indefinitely. + """ + base_timeout_seconds = max(1, settings.TCD_CONTAINER_TIMEOUT_SECONDS) + max_timeout_seconds = max(base_timeout_seconds, settings.TCD_CONTAINER_TIMEOUT_MAX_SECONDS) + base_pixels = max(1, settings.TCD_CONTAINER_TIMEOUT_BASE_PIXELS) + input_pixels = None + timeout_seconds = base_timeout_seconds + + if input_width is not None and input_height is not None and input_width > 0 and input_height > 0: + input_pixels = input_width * input_height + if input_pixels > base_pixels: + timeout_seconds = math.ceil(base_timeout_seconds * input_pixels / base_pixels) + + capped = timeout_seconds > max_timeout_seconds + timeout_seconds = min(timeout_seconds, max_timeout_seconds) + return _TCDTimeoutPolicy( + timeout_seconds=timeout_seconds, + base_timeout_seconds=base_timeout_seconds, + max_timeout_seconds=max_timeout_seconds, + base_pixels=base_pixels, + input_pixels=input_pixels, + capped=capped, + ) + + class _GeneratorStream(io.RawIOBase): """ Wraps a generator to provide a file-like interface for tarfile. @@ -275,7 +319,13 @@ def _copy_files_to_tcd_volume(ortho_path: str, volume_name: str, dataset_id: int ) -def _run_tcd_pipeline_container(volume_name: str, dataset_id: int, token: str) -> str: +def _run_tcd_pipeline_container( + volume_name: str, + dataset_id: int, + token: str, + input_width: int | None = None, + input_height: int | None = None, +) -> str: """ Execute TCD container using Pipeline class via Python script for complete confidence map output. @@ -299,7 +349,18 @@ def _run_tcd_pipeline_container(volume_name: str, dataset_id: int, token: str) - LogContext(category=LogCategory.TREECOVER, token=token, dataset_id=dataset_id), ) - tcd_timeout_seconds = max(1, settings.TCD_CONTAINER_TIMEOUT_SECONDS) + timeout_policy = _compute_tcd_timeout_policy(input_width=input_width, input_height=input_height) + tcd_timeout_seconds = timeout_policy.timeout_seconds + logger.info( + 'TCD container timeout policy: ' + f'timeout_seconds={timeout_policy.timeout_seconds}, ' + f'base_timeout_seconds={timeout_policy.base_timeout_seconds}, ' + f'max_timeout_seconds={timeout_policy.max_timeout_seconds}, ' + f'base_pixels={timeout_policy.base_pixels}, ' + f'input_pixels={timeout_policy.input_pixels}, ' + f'capped={timeout_policy.capped}', + LogContext(category=LogCategory.TREECOVER, token=token, dataset_id=dataset_id), + ) try: # Preflight: ensure image exists @@ -322,7 +383,9 @@ def _run(device_requests=None): environment={ 'NVIDIA_VISIBLE_DEVICES': 'all', 'NVIDIA_DRIVER_CAPABILITIES': 'compute,utility', - } if use_gpu else {}, + } + if use_gpu + else {}, labels={ **resource_labels, 'dt_role': 'tcd_pipeline', @@ -614,6 +677,9 @@ def predict_treecover(dataset_id: int, file_path: Path, user_id: str, token: str LogContext(category=LogCategory.TREECOVER, token=token, dataset_id=dataset_id), ) reprojected_path = Path(_reproject_orthomosaic_for_tcd(str(file_path), str(reprojected_temp_path))) + with rasterio.open(str(reprojected_path)) as reprojected_src: + reprojected_width = reprojected_src.width + reprojected_height = reprojected_src.height # Step 2: Container Setup - Create shared volume and copy reprojected ortho volume_name = f'tcd_volume_{dataset_id}_{uuid.uuid4().hex[:8]}' @@ -636,7 +702,13 @@ def predict_treecover(dataset_id: int, file_path: Path, user_id: str, token: str LogContext(category=LogCategory.TREECOVER, token=token, dataset_id=dataset_id), ) - _run_tcd_pipeline_container(volume_name, dataset_id, token) + _run_tcd_pipeline_container( + volume_name, + dataset_id, + token, + input_width=reprojected_width, + input_height=reprojected_height, + ) # Refresh token before extraction - TCD containers can run for hours and token may have expired token = login(settings.PROCESSOR_USERNAME, settings.PROCESSOR_PASSWORD) diff --git a/processor/tests/test_tcd_pipeline_container_timeout.py b/processor/tests/test_tcd_pipeline_container_timeout.py index 68bedede..ffc59496 100644 --- a/processor/tests/test_tcd_pipeline_container_timeout.py +++ b/processor/tests/test_tcd_pipeline_container_timeout.py @@ -1,8 +1,63 @@ +import os +import sys +import types +from enum import Enum + import requests import pytest from urllib3.exceptions import ReadTimeoutError -from processor.src.treecover_segmentation_oam_tcd import predict_treecover +os.environ.setdefault('SUPABASE_URL', 'http://localhost') +os.environ.setdefault('SUPABASE_KEY', 'test') + + +class _NoopLogger: + def info(self, *_args, **_kwargs): + pass + + def warning(self, *_args, **_kwargs): + pass + + def error(self, *_args, **_kwargs): + pass + + +class _TestLogCategory(Enum): + TREECOVER = 'treecover' + + +class _TestLogContext: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +logger_module = types.ModuleType('shared.logger') +logger_module.logger = _NoopLogger() +sys.modules.setdefault('shared.logger', logger_module) + +logging_module = types.ModuleType('shared.logging') +logging_module.LogCategory = _TestLogCategory +logging_module.LogContext = _TestLogContext +sys.modules.setdefault('shared.logging', logging_module) + +db_module = types.ModuleType('shared.db') +db_module.login = lambda *_args, **_kwargs: 'token' +db_module.verify_token = lambda *_args, **_kwargs: True +sys.modules.setdefault('shared.db', db_module) + +labels_module = types.ModuleType('shared.labels') +labels_module.create_label_with_geometries = lambda *_args, **_kwargs: None +labels_module.delete_model_prediction_labels = lambda *_args, **_kwargs: 0 +sys.modules.setdefault('shared.labels', labels_module) + +segmentation_module = types.ModuleType('processor.src.utils.segmentation') +segmentation_module.mask_to_polygons = lambda *_args, **_kwargs: [] +segmentation_module.reproject_polygons = lambda polygons, *_args, **_kwargs: polygons +segmentation_module.filter_polygons_by_area = lambda polygons, *_args, **_kwargs: polygons +segmentation_module.get_utm_string_from_latlon = lambda *_args, **_kwargs: 'EPSG:32632' +sys.modules.setdefault('processor.src.utils.segmentation', segmentation_module) + +from processor.src.treecover_segmentation_oam_tcd import predict_treecover # noqa: E402 class _FakeImages: @@ -15,11 +70,11 @@ def __init__(self): self.killed = False self.removed = False self.name = 'fake-tcd-container' + self.wait_timeout = None def wait(self, timeout): - raise requests.exceptions.ConnectionError( - ReadTimeoutError(None, None, 'Read timed out.') - ) + self.wait_timeout = timeout + raise requests.exceptions.ConnectionError(ReadTimeoutError(None, None, 'Read timed out.')) def kill(self): self.killed = True @@ -60,6 +115,8 @@ def test_tcd_wait_connection_timeout_is_controlled_and_not_retried(monkeypatch): lambda **kwargs: bundles.append(kwargs), ) monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_SECONDS', 14400) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_MAX_SECONDS', 43200) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_BASE_PIXELS', 2_000_000_000) with pytest.raises( predict_treecover._TCDContainerTimeout, @@ -70,4 +127,62 @@ def test_tcd_wait_connection_timeout_is_controlled_and_not_retried(monkeypatch): assert len(client.containers.runs) == 1 assert client.containers.runs[0].killed is True assert client.containers.runs[0].removed is True + assert client.containers.runs[0].wait_timeout == 14400 assert bundles + + +def test_tcd_timeout_policy_uses_base_timeout_without_input_size(monkeypatch): + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_SECONDS', 14400) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_MAX_SECONDS', 43200) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_BASE_PIXELS', 2_000_000_000) + + policy = predict_treecover._compute_tcd_timeout_policy() + + assert policy.timeout_seconds == 14400 + assert policy.input_pixels is None + assert policy.capped is False + + +def test_tcd_timeout_policy_scales_large_ortho_to_cap(monkeypatch): + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_SECONDS', 14400) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_MAX_SECONDS', 43200) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_BASE_PIXELS', 2_000_000_000) + + policy = predict_treecover._compute_tcd_timeout_policy( + input_width=90377, + input_height=104977, + ) + + assert policy.input_pixels == 9_487_506_329 + assert policy.timeout_seconds == 43200 + assert policy.capped is True + + +def test_tcd_container_wait_uses_adaptive_timeout_for_large_input(monkeypatch): + client = _FakeDockerClient() + + monkeypatch.setattr(predict_treecover.docker, 'from_env', lambda: client) + monkeypatch.setattr( + predict_treecover, + 'build_container_forensics', + lambda *args, **kwargs: {'dataset_id': kwargs['dataset_id'], 'stage': kwargs['stage']}, + ) + monkeypatch.setattr(predict_treecover, 'write_debug_bundle', lambda **_kwargs: None) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_SECONDS', 14400) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_MAX_SECONDS', 43200) + monkeypatch.setattr(predict_treecover.settings, 'TCD_CONTAINER_TIMEOUT_BASE_PIXELS', 2_000_000_000) + + with pytest.raises( + predict_treecover._TCDContainerTimeout, + match='TCD container timed out after 12 hours', + ): + predict_treecover._run_tcd_pipeline_container( + 'tcd_volume_10183_test', + 10183, + 'token', + input_width=90377, + input_height=104977, + ) + + assert len(client.containers.runs) == 1 + assert client.containers.runs[0].wait_timeout == 43200 diff --git a/shared/settings.py b/shared/settings.py index 499c2b2f..3c23186c 100644 --- a/shared/settings.py +++ b/shared/settings.py @@ -44,6 +44,8 @@ class Settings(BaseSettings): # Containers TCD_CONTAINER_IMAGE: str = 'deadtrees-tcd:latest' TCD_CONTAINER_TIMEOUT_SECONDS: int = 14400 + TCD_CONTAINER_TIMEOUT_MAX_SECONDS: int = 43200 + TCD_CONTAINER_TIMEOUT_BASE_PIXELS: int = 2_000_000_000 # Base paths and directories BASE_DIR: str = str(BASE) @@ -160,7 +162,9 @@ class Settings(BaseSettings): def model_post_init(self, __context): if 'API_ENDPOINT' not in self.model_fields_set: - self.API_ENDPOINT = 'http://localhost:8080/api/v1/' if self.DEV_MODE else 'https://data2.deadtrees.earth/api/v1/' + self.API_ENDPOINT = ( + 'http://localhost:8080/api/v1/' if self.DEV_MODE else 'https://data2.deadtrees.earth/api/v1/' + ) if 'API_ENTPOINT_DATASETS' not in self.model_fields_set: self.API_ENTPOINT_DATASETS = self.API_ENDPOINT + 'datasets/chunk' if 'PREPACKAGED_DOWNLOAD_BASE_URL' not in self.model_fields_set: