from __future__ import annotations import asyncio import importlib import logging import os import time from typing import Any, Dict, Optional, Tuple import numpy as np import pandas as pd from fastapi import HTTPException try: import torch TORCH_IMPORT_ERROR: Optional[str] = None except Exception as exc: # pragma: no cover - environment dependent torch = None # type: ignore[assignment] TORCH_IMPORT_ERROR = str(exc) TIMESFM_AVAILABLE = False TIMESFM_IMPORT_ERROR: Optional[str] = None timesfm = None if TORCH_IMPORT_ERROR: TIMESFM_IMPORT_ERROR = TORCH_IMPORT_ERROR else: try: import timesfm as _timesfm # type: ignore[import-not-found] timesfm = _timesfm TIMESFM_AVAILABLE = True except Exception as exc: # pragma: no cover - import environment dependent TIMESFM_IMPORT_ERROR = str(exc) _TIMESFM_RUNTIME_PATCHED = False def _patch_timesfm_runtime_compatibility(logger: logging.Logger) -> None: global _TIMESFM_RUNTIME_PATCHED if _TIMESFM_RUNTIME_PATCHED or timesfm is None: return try: internal_module = importlib.import_module("timesfm.timesfm_2p5.timesfm_2p5_torch") model_module_cls = getattr(internal_module, "TimesFM_2p5_200M_torch_module", None) if model_module_cls is None or getattr(model_module_cls, "_aiforecast_meta_patch", False): _TIMESFM_RUNTIME_PATCHED = True return def _patched_load_checkpoint(self: Any, path: str, **kwargs: Any) -> None: tensors = internal_module.load_file(path) has_meta_parameters = any( getattr(parameter, "is_meta", False) for parameter in self.parameters() ) try: if has_meta_parameters: self.load_state_dict(tensors, strict=True, assign=True) else: self.load_state_dict(tensors, strict=True) except TypeError: if has_meta_parameters: self.to_empty(device=self.device) self.load_state_dict(tensors, strict=True) self.to(self.device) torch_compile = kwargs.get("torch_compile", True) if torch_compile: internal_module.logging.info("Compiling model...") self = torch.compile(self) self.eval() model_module_cls.load_checkpoint = _patched_load_checkpoint model_module_cls._aiforecast_meta_patch = True _TIMESFM_RUNTIME_PATCHED = True except Exception as exc: # pragma: no cover - defensive guard logger.warning("[TimesFM] Runtime compatibility patch skipped: %s", exc) class TimesFMForecaster: key = "timesfm" label = "TimesFM" MODEL_HF_ID = os.getenv("TIMESFM_MODEL_HF_ID", "google/timesfm-2.5-200m-pytorch") MAX_CONTEXT = 8_192 MAX_HORIZON = 300 _Q10 = 1 _Q50 = 5 _Q90 = 9 def __init__(self, logger: logging.Logger | None = None) -> None: self._logger = logger or logging.getLogger("ai-forecast.timesfm") self._model: Optional[Any] = None self._loaded = False self._lock: Optional[asyncio.Lock] = None self._predict_lock: Optional[asyncio.Lock] = None self._compiled_horizon: Optional[int] = None @property def available(self) -> bool: return TIMESFM_AVAILABLE @property def import_error(self) -> Optional[str]: return TIMESFM_IMPORT_ERROR @staticmethod def _prepare_feature_frame(df: pd.DataFrame) -> pd.DataFrame: ohlc4 = ( df[["open", "high", "low", "close"]] .mean(axis=1) .astype(np.float32) ) return pd.DataFrame({"ohlc4": ohlc4}) @classmethod def _extract_ohlc4_series(cls, df: pd.DataFrame) -> np.ndarray: feature_frame = cls._prepare_feature_frame(df) ohlc4_series = feature_frame["ohlc4"].to_numpy(dtype=np.float32, copy=False) cls._validate_input_series(ohlc4_series) return ohlc4_series @staticmethod def _validate_input_series(series: np.ndarray) -> None: if series.ndim != 1: raise HTTPException(status_code=422, detail="TimesFM input series must be 1-D") if len(series) < 32: raise HTTPException(status_code=422, detail="TimesFM requires at least 32 OHLC4 points") if not np.isfinite(series).all(): raise HTTPException(status_code=422, detail="TimesFM input contains non-finite OHLC4 values") @classmethod def _validate_output_tensors( cls, point_forecast: np.ndarray, quantile_forecast: np.ndarray, horizon: int, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]: point_arr = np.asarray(point_forecast) quant_arr = np.asarray(quantile_forecast) if point_arr.ndim != 2 or point_arr.shape[0] != 1 or point_arr.shape[1] < horizon: raise HTTPException( status_code=500, detail=f"TimesFM point forecast has invalid shape: {list(point_arr.shape)}", ) if ( quant_arr.ndim != 3 or quant_arr.shape[0] != 1 or quant_arr.shape[1] < horizon or quant_arr.shape[2] <= cls._Q90 ): raise HTTPException( status_code=500, detail=f"TimesFM quantile forecast has invalid shape: {list(quant_arr.shape)}", ) p10 = quant_arr[0, :horizon, cls._Q10].astype(float) p50 = point_arr[0, :horizon].astype(float) p50_from_quantiles = quant_arr[0, :horizon, cls._Q50].astype(float) p90 = quant_arr[0, :horizon, cls._Q90].astype(float) if not (np.isfinite(p10).all() and np.isfinite(p50).all() and np.isfinite(p90).all()): raise HTTPException(status_code=500, detail="TimesFM output contains non-finite values") quantiles_monotonic = bool( np.all(p10 <= (p50 + 1e-6)) and np.all(p50 <= (p90 + 1e-6)) ) if not quantiles_monotonic: raise HTTPException(status_code=500, detail="TimesFM returned non-monotonic quantiles") median_matches_point_forecast = bool( np.allclose(p50, p50_from_quantiles, atol=1e-4, rtol=1e-4) ) if not median_matches_point_forecast: logging.getLogger("ai-forecast").warning( "[TimesFM] point_forecast differs from q50 quantile output" ) output_validation = { "point_shape": [int(dim) for dim in point_arr.shape], "quantile_shape": [int(dim) for dim in quant_arr.shape], "median_matches_point_forecast": median_matches_point_forecast, "quantiles_monotonic": quantiles_monotonic, } return p10, p50, p90, output_validation async def _get_lock(self) -> asyncio.Lock: if self._lock is None: self._lock = asyncio.Lock() return self._lock async def _get_predict_lock(self) -> asyncio.Lock: if self._predict_lock is None: self._predict_lock = asyncio.Lock() return self._predict_lock @property def is_ready(self) -> bool: return self._loaded @property def device(self) -> str: if self._model is None: return "not_loaded" try: return str(self._model.model.device) except Exception: return "cpu" def _compile(self, horizon: int) -> None: output_patch = 128 max_horizon = int(np.ceil(horizon / output_patch) * output_patch) max_horizon = max(max_horizon, output_patch) max_horizon = min(max_horizon, self.MAX_HORIZON) context_len = self.MAX_CONTEXT assert timesfm is not None self._model.compile( timesfm.ForecastConfig( max_context=context_len, max_horizon=max_horizon, normalize_inputs=True, use_continuous_quantile_head=True, force_flip_invariance=True, infer_is_positive=True, fix_quantile_crossing=True, ) ) compiled_config = getattr(self._model, "forecast_config", None) actual_ctx = int(getattr(compiled_config, "max_context", context_len)) actual_max_horizon = int(getattr(compiled_config, "max_horizon", max_horizon)) self._compiled_horizon = actual_max_horizon self._logger.info("[TimesFM] Compiled: ctx=%d max_horizon=%d", actual_ctx, actual_max_horizon) async def _lazy_load(self) -> None: if self._loaded: return lock = await self._get_lock() async with lock: if self._loaded: return if not TIMESFM_AVAILABLE or timesfm is None: raise HTTPException( status_code=503, detail="TimesFM not installed. Run: pip install timesfm[torch]", ) try: _patch_timesfm_runtime_compatibility(self._logger) self._logger.info("[TimesFM] Loading %s ...", self.MODEL_HF_ID) model = await asyncio.to_thread( timesfm.TimesFM_2p5_200M_torch.from_pretrained, self.MODEL_HF_ID, torch_compile=False, ) self._model = model self._compile(self.MAX_HORIZON) self._loaded = True self._logger.info("[TimesFM] Ready on %s", self.device) except Exception as exc: self._logger.error("[TimesFM] Init failed: %s", exc, exc_info=True) raise HTTPException(status_code=500, detail=f"TimesFM init failed: {exc}") async def forecast( self, df: pd.DataFrame, horizon: int, interval: str = "", step_seconds: int = 0, symbol: str = "", ) -> Dict[str, Any]: del interval, step_seconds, symbol await self._lazy_load() if self._model is None: raise HTTPException(status_code=500, detail="TimesFM model is not initialized") try: ohlc4_series = self._extract_ohlc4_series(df) context_len = min(len(ohlc4_series), self.MAX_CONTEXT) ohlc4_context = ohlc4_series[-context_len:] if self._compiled_horizon is None or horizon > self._compiled_horizon: self._compile(horizon) started_at = time.time() predict_lock = await self._get_predict_lock() async with predict_lock: point_forecast, quantile_forecast = await asyncio.to_thread( self._model.forecast, horizon, [ohlc4_context], ) elapsed = time.time() - started_at self._logger.info( "[TimesFM] %.2fs | horizon=%d ctx=%d device=%s", elapsed, horizon, context_len, self.device, ) if torch is not None and torch.cuda.is_available(): torch.cuda.empty_cache() p10, p50, p90, output_validation = self._validate_output_tensors( point_forecast=point_forecast, quantile_forecast=quantile_forecast, horizon=horizon, ) return { "p10": p10, "p50": p50, "p90": p90, "model_name": self.MODEL_HF_ID, "context_length": context_len, "output_horizon": horizon, "input_semantics": { "feature_channels": ["ohlc4"], "active_forecast_channels": ["ohlc4"], "ignored_channels": ["volume"], "price_mode": "ohlc4_single_channel", "base_signal": "ohlc4", "volume_mode": "omitted", "amount_mode": "omitted", "adapter_mode": "timesfm_native", "normalization": "timesfm_internal_revin", }, "input_validation": { "series_field": "ohlc4", "dtype": str(ohlc4_context.dtype), "is_1d": True, "finite": True, }, "output_semantics": { "forecast_channel": "ohlc4", "forecast_mode": "single_future_ohlc4_line", "quantile_fields": ["p10", "p50", "p90"], "candle_projection": "omitted", "reference_baseline": "last_ohlc4", }, "output_validation": output_validation, } except HTTPException: raise except Exception as exc: self._logger.error("[TimesFM] Forecast failed: %s", exc, exc_info=True) raise HTTPException(status_code=500, detail=f"TimesFM prediction failed: {exc}")