ProgAI Forecasting Framework 📈

This project is a multi-component forecasting framework that exposes a unified API for various state-of-the-art Foundation- and DL-models. It's designed for both rapid prototyping and robust production use, allowing users to perform zero-shot forecasting, fine-tune FMs or train DL-models on custom data, and evaluate model performance.

This project is an experimental playground for exploring foundation models and deep learning approaches for time series forecasting. The code and methods were originally developed as part of the Master’s Thesis “Foundation Models for Time Series Forecasting with Exogenous Variables: A Survey and Benchmark” (GitHub).

Table of Contents

• ProgAI Forecasting Framework 📈
  • Table of Contents
  • Features 🌟
  • Technology Stack 💻
  • Architecture 🏛️
  • Getting Started 🚀
    • Prerequisites
    • Development
    • Production
  • Usage 🎯
    • Forecasting
    • Fine-Tuning
  • Testing ✅
  • Security 🛡️
  • Data Management 📂
  • Documentation 📚
  • Roadmap 🗺️
  • Contributing 🤝
  • License 📄

Features 🌟

• Unified API: A single, consistent API for multiple forecasting models.
• Asynchronous Processing: A robust job queue system using Redis and Arq to handle long-running forecasting and fine-tuning tasks without blocking.
• Model Specialization: A multi-worker architecture where each worker is specialized for a specific set of models (e.g., Moirai, Chronos, TimesFM) to handle the different dependencies of new DL based models.
• Zero-Shot Forecasting: Generate forecasts immediately without prior model training.
• Model Training: Train DL models or fine-tune pre-trained foundation models on your specific datasets for improved performance.
• Advanced GPU Handling: Efficiently manages single and multi-GPU setups with parallel processing capabilities. See GPU Configuration.

Technology Stack 💻

• Backend: Python, FastAPI, Arq, Uvicorn
• ML/Forecasting: GluonTS, PyTorch, TimesFM, Moirai (uni2ts), AutoGluon (Chronos, TabPFN), transformers...
• Data Handling: Pandas, Pydantic
• Infrastructure: Docker, Redis, NVIDIA GPU Support
• Dependencies: Managed with uv.

Architecture 🏛️

The system is built around a central API gateway that dispatches jobs to a set of specialized workers via a Redis queue.

This multi-worker architecture was chosen because different forecasting models often have very specific and conflicting dependency requirements. By isolating models into specialized worker containers, we can manage these diverse dependencies effectively, reduce the overall worker count by grouping compatible models, and ensure comprehensive support for all integrated models.

• API Gateway (main.py): A FastAPI application that serves as the main entry point. It handles HTTP requests, validates data, and enqueues jobs.
• Redis: Acts as the message broker for the task queue.

This Redis and Arq combination was specifically chosen to efficiently manage diverse GPU setups and task types. It allows for sequential processing on systems with limited VRAM (e.g., a single GPU) and parallel execution on multi-GPU environments. Furthermore, it enables effective prioritization, ensuring that fast zero-shot forecasting tasks are processed with higher priority than longer-running fine-tuning tasks. This approach ensures that the available GPU resources are utilized to their full potential.

• Workers (workers/): Dedicated Python processes that listen to specific Redis queues. They perform the forecasting or fine-tuning and provide the results.
  • Worker 1: Handles Moirai, TimesFM, and NeuralForecast.
  • Worker 2: Handles Chronos and TabPFN-TS.
  • Worker 3: Handles TiRex.
  • Worker 4: Handles TOTO and YingLong.
  • Worker 5: Handles TinyTimeMixers and Sundial.

Getting Started 🚀

Prerequisites

• Docker Engine
• Docker Compose
• Git
• NVIDIA Container Toolkit

Development

1. Clone the repository.

2. Copy the environment template and configure:

   cp .env.example .env
   # Edit .env with your configuration (S3 credentials if using S3 storage)
   

3. If you're using VS Code with Dev Containers, you'll need to specify which service to work on. Open .devcontainer/devcontainer.json and set the "service" property to the container you want to develop (e.g., "service": "worker2"). Then, rebuild the dev container.

4. If you're not using a dev container, run the following command from the project root in your local terminal to start the API server with hot-reloading:

   docker-compose -f docker-compose.yml -f docker-compose.dev.yml build
   docker-compose -f docker-compose.yml -f docker-compose.dev.yml up -d
   

5. The API will be available at http://localhost:8080.

6. If any changes were made to the api or worker code you have to restart the container.

   # restart api
   docker-compose restart api
   # restart worker
   docker-compose restart worker2
   

Production

1. Clone the repository on your production server.
2. Build the production images:

   docker-compose -f docker-compose.yml -f docker-compose.prod.yml build
   

3. Start the services:

   docker-compose -f docker-compose.yml -f docker-compose.prod.yml up -d
   

Usage 🎯

The primary way to interact with the framework is through its REST API.

Forecasting

Submit a POST request to the /forecast endpoint. The time series data should be in a "long format".

Example curl request:

curl -X 'POST' 
  'http://localhost:8080/forecast' 
  -H 'accept: application/json' 
  -H 'Content-Type: application/json' 
  -d 
  '''
  {
  "model_name": "timesfm",
  "model_path": "google/timesfm-1.0-200m-pytorch",
  "prediction_length": 7,
  "time_series": {
    "unique_id": ["series1", "series1", "series1", "series1", "series1"],
    "ds": ["2023-01-01", "2023-01-02", "2023-01-03", "2023-01-04", "2023-01-05"],
    "y": [10, 12, 15, 11, 13]
  }
}
'''

This will return a job_id. You can use this job_id as query parameter for the /status and /results endpoints to track and retrieve your forecast (e.g. /status?job_id=f40b6b52b4154897a24a115cbc7ee470). For more details on request and response schemas, refer to the API Reference.

Fine-Tuning

Submit a POST request to the /train endpoint. You must first upload your training data (as a CSV) to the S3 bucket's datasets/ directory.

Example curl request:

curl -X 'POST' 
  'http://localhost:8080/train' 
  -H 'accept: application/json' 
  -H 'Content-Type: application/json' 
  -d 
  '''
  {
  "base_model_name": "moirai",
  "base_model_path": "salesforce/moirai-1.1-R-base",
  "finetuned_model_name": "my_custom_chronos_model",
  "training_data": "training_data.csv",
  "prediction_length": 14
}
'''

This will queue a fine-tuning job. The resulting model will be saved to the persistent data volume and can be used for forecasting by referencing my_custom_chronos_model.

For more detailed examples, see the Example Usage Notebook.

Testing ✅

Tests are run using pytest. The test for the API are located in root (./tests) and for the workers they are in the corresponding worker directory e.g. ./workers/worker2/tests.

• Using a Dev Container:

  • You can run pytest directly in the Dev Container terminal.
  • Note: You can only run the tests for the service you've specified in your .devcontainer/devcontainer.json file (e.g., if "service": "worker2", only worker2's dependencies are installed, so only its tests will pass).
  • E.g.:

    # api
    pytest tests
    # worker1
    pytest workers/worker1/tests
    

• Without a Dev Container:

  • To run tests for a specific service, use docker-compose exec from your host machine.
  • E.g.:

    # api
    docker-compose exec api pytest tests
    # worker 2
    docker-compose exec worker2 pytest workers/worker2/tests
    

Security 🛡️

This project uses Grype to scan for vulnerabilities in the Docker images. The process is managed via a dedicated Docker Compose file.

For more detailed information, see the Security Documentation.

Data Management 📂

The project supports flexible data storage with two backend options:

• S3-compatible Storage (Default for Production):

  • Secure, scalable storage for datasets and models
  • Managed through a dedicated S3Handler
  • Full validation and error handling

• Local Storage (Development-friendly):

  • Store data directly in local filesystem
  • Perfect for development and testing
  • Same interface as S3 storage

The storage backend is configurable via environment variables, and the application uses a unified StorageHandler interface that works identically regardless of the chosen backend.

For detailed instructions on configuring and using either storage backend, see the Storage Documentation.

Documentation 📚

For more detailed information, please refer to the documents in the /docs directory:

• API Reference: Detailed information about API endpoints, request/response schemas, and data formats.
• GPU Configuration: Guide for configuring single and multi-GPU setups with parallel processing.
• Data Storage: Guide to managing datasets and models with S3 or local storage.
• Supported Models: Overview of supported forecasting models and their features.
• Model References: Comprehensive links to official repositories, documentation, and implementation guides for all supported models.
• Project Roadmap: The development plan for the project.
• Security: Information about security practices and vulnerability scanning.
• Adding a new Model: Step by step guide on how to add new models.

Roadmap 🗺️

The project is under active development. Key future milestones include:

• Implementing zero-shot forecasting for all remaining workers.
• Developing a comprehensive /evaluate endpoint for model comparison.
• Implementing fine-tuning for all remaining models that support this functionality.

See the full Project Roadmap for more details.

Contributing 🤝

We welcome contributions! Please see our Contributing Guidelines for details on how to get started.

License 📄

This project is licensed under the MIT License - see the LICENSE file for details.