fix(processing): scale TCD timeout for large orthos

.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=

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)

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

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: