tabpfn-time-series 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

tabpfn_time_series/predictor.py

@@ -2,7 +2,11 @@ import logging
 from enum import Enum
 
 from tabpfn_time_series.ts_dataframe import TimeSeriesDataFrame
-from tabpfn_time_series.tabpfn_worker import TabPFNClient, LocalTabPFN, MockTabPFN
+from tabpfn_time_series.tabpfn_worker import (
+    TabPFNClient,
+    LocalTabPFN,
+    MockTabPFN,
+)
 from tabpfn_time_series.defaults import TABPFN_TS_DEFAULT_CONFIG
 
 logger = logging.getLogger(__name__)
@@ -29,6 +33,7 @@ class TabPFNTimeSeriesPredictor:
             TabPFNMode.LOCAL: lambda: LocalTabPFN(config),
             TabPFNMode.MOCK: lambda: MockTabPFN(config),
         }
+        self.tabpfn_mode = tabpfn_mode
         self.tabpfn_worker = worker_mapping[tabpfn_mode]()
 
     def predict(

tabpfn_time_series/tabpfn_worker.py

@@ -1,6 +1,8 @@
+import contextvars
 import logging
 from abc import ABC, abstractmethod
-from joblib import Parallel, delayed
+from joblib import Parallel, delayed, parallel_config
+import backoff
 
 from tqdm import tqdm
 import pandas as pd
@@ -14,6 +16,9 @@ from tabpfn_time_series.defaults import TABPFN_TS_DEFAULT_QUANTILE_CONFIG
 
 logger = logging.getLogger(__name__)
 
+# Per-call attempt counter, isolated per thread & task
+_retry_attempts = contextvars.ContextVar("predict_attempts", default=0)
+
 
 class TabPFNWorker(ABC):
     def __init__(
@@ -29,24 +34,7 @@ class TabPFNWorker(ABC):
         train_tsdf: TimeSeriesDataFrame,
         test_tsdf: TimeSeriesDataFrame,
     ):
-        predictions = Parallel(
-            n_jobs=self.num_workers,
-            backend="loky",
-        )(
-            delayed(self._prediction_routine)(
-                item_id,
-                train_tsdf.loc[item_id],
-                test_tsdf.loc[item_id],
-            )
-            for item_id in tqdm(train_tsdf.item_ids, desc="Predicting time series")
-        )
-
-        predictions = pd.concat(predictions)
-
-        # Sort predictions according to original item_ids order (important for MASE and WQL calculation)
-        predictions = predictions.loc[train_tsdf.item_ids]
-
-        return TimeSeriesDataFrame(predictions)
+        raise NotImplementedError("Predict method must be implemented in subclass")
 
     def _prediction_routine(
         self,
@@ -54,8 +42,6 @@ class TabPFNWorker(ABC):
         single_train_tsdf: TimeSeriesDataFrame,
         single_test_tsdf: TimeSeriesDataFrame,
     ) -> pd.DataFrame:
-        # logger.debug(f"Predicting on item_id: {item_id}")
-
         test_index = single_test_tsdf.index
         train_X, train_y = split_time_series_to_X_y(single_train_tsdf.copy())
         test_X, _ = split_time_series_to_X_y(single_test_tsdf.copy())
@@ -116,6 +102,39 @@ class TabPFNWorker(ABC):
         return result
 
 
+def _reset_attempts(_details=None):
+    """Convenience function to reset the attempt counter."""
+    _retry_attempts.set(0)
+
+
+def _predict_giveup_mixed(exc: Exception) -> bool:
+    """Determine whether to give up on a prediction call or not.
+
+    Returns:
+        True if the prediction call should be given up on, False otherwise.
+    """
+    if _is_tabpfn_gcs_429(exc):
+        return False
+
+    # Stop after first retry for non-429
+    return _retry_attempts.get() >= 2
+
+
+def _is_tabpfn_gcs_429(err: Exception) -> bool:
+    """Determine if an error is a 429 error raised from TabPFN API
+    and relates to GCS 429 errors.
+
+    Returns:
+        True if the error is a 429 error raised from TabPFN API.
+    """
+    markers = (
+        "TooManyRequests: 429",
+        "rateLimitExceeded",
+        "cloud.google.com/storage/docs/gcs429",
+    )
+    return any(m in str(err) for m in markers)
+
+
 class TabPFNClient(TabPFNWorker):
     def __init__(
         self,
@@ -135,6 +154,53 @@ class TabPFNClient(TabPFNWorker):
 
         super().__init__(config, num_workers)
 
+    def predict(
+        self,
+        train_tsdf: TimeSeriesDataFrame,
+        test_tsdf: TimeSeriesDataFrame,
+    ):
+        # Run the predictions in parallel
+        with parallel_config(backend="threading"):
+            results = Parallel(
+                n_jobs=self.num_workers,
+            )(
+                delayed(self._prediction_routine)(
+                    item_id,
+                    train_tsdf.loc[item_id],
+                    test_tsdf.loc[item_id],
+                )
+                for item_id in tqdm(train_tsdf.item_ids, desc="Predicting time series")
+            )
+
+        # Convert list to DataFrame
+        predictions = pd.concat(results)
+
+        # Sort predictions according to original item_ids order (important for MASE and WQL calculation)
+        predictions = predictions.loc[train_tsdf.item_ids]
+
+        return TimeSeriesDataFrame(predictions)
+
+    @backoff.on_exception(
+        backoff.expo,
+        Exception,
+        base=1,
+        factor=2,
+        max_tries=5,
+        jitter=backoff.full_jitter,
+        giveup=_predict_giveup_mixed,
+        on_success=_reset_attempts,
+    )
+    def _prediction_routine(
+        self,
+        item_id: str,
+        single_train_tsdf: TimeSeriesDataFrame,
+        single_test_tsdf: TimeSeriesDataFrame,
+    ) -> pd.DataFrame:
+        # Increment attempt count at start of each try
+        _retry_attempts.set(_retry_attempts.get() + 1)
+
+        return super()._prediction_routine(item_id, single_train_tsdf, single_test_tsdf)
+
     def _get_tabpfn_engine(self):
         from tabpfn_client import TabPFNRegressor
 

{tabpfn_time_series-1.0.2.dist-info → tabpfn_time_series-1.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabpfn_time_series
-Version: 1.0.2
+Version: 1.0.4
 Summary: Zero-shot time series forecasting with TabPFNv2
 Project-URL: Homepage, https://github.com/liam-sbhoo/tabpfn-time-series
 Project-URL: Bug Tracker, https://github.com/liam-sbhoo/tabpfn-time-series/issues
@@ -10,6 +10,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
+Requires-Dist: backoff>=2.2.1
 Requires-Dist: datasets>=4.0
 Requires-Dist: gluonts>=0.16.0
 Requires-Dist: pandas<2.2.0,>=2.1.2

{tabpfn_time_series-1.0.2.dist-info → tabpfn_time_series-1.0.4.dist-info}/RECORD

@@ -2,15 +2,15 @@ tabpfn_time_series/__init__.py,sha256=XJXSKqWp3AF9mAaWi-4KCgHQG7NzNTaBkLOYOMxvhS
 tabpfn_time_series/data_preparation.py,sha256=wWjSaKgV9KqKonMtSuDbYnW59ixflrScKIP_HSJ_MlA,5427
 tabpfn_time_series/defaults.py,sha256=ki1y38FR4zmbHWgRjcryA5T88GzNMwhlZC-sTRjuK2U,248
 tabpfn_time_series/plot.py,sha256=UXgLR2S94vi-vv1ArQKI6uYl_QwSAwAau5jFzGmQ7hw,6582
-tabpfn_time_series/predictor.py,sha256=UDhvH7reB3HAuxOyggE4yl2ntQXivvxgXCN_RhoddHc,1467
-tabpfn_time_series/tabpfn_worker.py,sha256=ZlJrU0O1dxfKh_As5Le4phm0P4RCDeXKpqp9X-h5bQs,9619
+tabpfn_time_series/predictor.py,sha256=2wnBAHfU5sOSJoHUm65Ej8tJjA4jGYP8yHvWeo1MzyA,1523
+tabpfn_time_series/tabpfn_worker.py,sha256=k6td4Ml0E3Xr1gERze-S0kyvBB6q_hbLMzvSurdaSp0,11589
 tabpfn_time_series/ts_dataframe.py,sha256=X94mssw_mSFedjplG55hjwTzKj8mM3VwWynveX3fegA,52834
 tabpfn_time_series/features/__init__.py,sha256=lzdZWkEfntfg3ZHqNNbfbg-3o_VIzju0tebdRu3AzF4,421
 tabpfn_time_series/features/auto_features.py,sha256=3OqqY2h7umcoLjLx4hOXypLTjwzrMtd6cQKTNi83vrU,11561
 tabpfn_time_series/features/basic_features.py,sha256=OV3B__S30-CX88vGjwYQDWqAbJajQw80PxcnvJVUbm4,2955
 tabpfn_time_series/features/feature_generator_base.py,sha256=jtySWLJyX4E31v6CbX44EHa8cdz7OMyauf4ltNEQeAQ,534
 tabpfn_time_series/features/feature_transformer.py,sha256=JzxswTGRGlt00QoYFyvAILlUVD68njdvoU3v-phnyi8,1774
-tabpfn_time_series-1.0.2.dist-info/METADATA,sha256=tnwSlCc7EscY_9utB6ouU5FUlQulml4wIz18rE6QwJA,4917
-tabpfn_time_series-1.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-tabpfn_time_series-1.0.2.dist-info/licenses/LICENSE.txt,sha256=iwhPL7kIWQG6gyLZZwIMDItGrNgxMDIq9itxkUSMapY,11345
-tabpfn_time_series-1.0.2.dist-info/RECORD,,
+tabpfn_time_series-1.0.4.dist-info/METADATA,sha256=n7fvApkQVQYw_N4aFC6wuOaMxNsr5Lqa9l7nEyhuL6g,4947
+tabpfn_time_series-1.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+tabpfn_time_series-1.0.4.dist-info/licenses/LICENSE.txt,sha256=iwhPL7kIWQG6gyLZZwIMDItGrNgxMDIq9itxkUSMapY,11345
+tabpfn_time_series-1.0.4.dist-info/RECORD,,