ML Project Template

A modular template for building machine learning training pipelines and inference APIs with PyTorch and FastAPI.

This template provides a structured starting point for ML projects, including:

• Model training utilities
• Evaluation metrics
• Early stopping
• Configuration validation
• FastAPI inference endpoints
• Schema-based request/response validation
• Logging setup
• Structured project architecture

The goal is to provide a clean foundation for building production-ready ML services.

---

General Workflow Diagram

---

Features

• PyTorch Training Framework

  • Modular trainer class
  • Device-aware training
  • Support for binary and multiclass classification
  • Extendable to regression

• Metrics System

  • Built with torchmetrics
  • Accuracy
  • Precision
  • Recall
  • F1 Score
  • Regression metrics

• MLflow Integration

  • Centralized experiment tracking and model versioning.
  • Used alias production and stage to signify production and staging models

• ONNX Serving Strategy

  • High-performance inference decoupled from the training framework.

• Strategy Pattern Architecture

  • Clean separation of Binary, Multiclass, and Regression logic (Easily extendable for any task type)

• Fail-Safe Preprocessing

  • Automatic recovery of normalization pipelines from JSON backups. (As a mockup, but easily extendable to any preprocessing specific attributes you need)

• Early Stopping

  • Patience-based training stopping
  • Automatic best model restoration

• FastAPI Inference API

  • /api/predict endpoint
  • /health health check
  • Pydantic request and response schemas

• Configuration System

  • Environment validation
  • Model configuration management

• Structured Logging

---

Project Structure

.
├── README.md
├── configs                             # Where the YAML configurations live
│   └── training_configs.yaml           # Example YAML configurations for training.
├── data                                # Where the data is stored.
│   ├── features                        # Where you store the features ready for inference or training (might expand to use an actual feature store, still exploring this path).
│   ├── preprocessed                    # Preprocessed data.
│   └── raw                             # Raw data (CSV, Images, etc.)
├── docker                              # All Docker related files live here.
│   ├── Dockerfile
│   ├── docker-compose.dev.yaml
│   ├── docker-compose.test.yaml
│   └── docker-compose.yaml
├── justfile                            # A Justfile to save and run project-specific commands. 
├── pyproject.toml                      # TOML file generated by UV.
├── pytest.ini                          # INI file to configure PyTest.
├── src                                 # Where the source code lives.
│   ├── __init__.py
│   └── ml_project_template
│       ├── __init__.py
│       ├── core                        # Core functionality shared between serving and training.
│       │   ├── configs                 # Configuration loading, environment variable validation, and model specific configurations.
│       │   │   ├── __init__.py
│       │   │   ├── loader.py           # Loads app configurations from a YAML file and matches them to a PyDantic model of your choice.
│       │   │   ├── model_configs.py    # Define model specific configurations here.
│       │   │   └── validator.py        # Validate and load environment variables
│       │   ├── errors                  # You can define all of them in the `exceptions.py` file, or be even more specific and have an exceptions module for each type of erros.
│       │   │   ├── __init__.py
│       │   │   └── exceptions.py       # Where custom exceptions are defined.
│       │   ├── logging
│       │   │   ├── __init__.py
│       │   │   └── setup.py            # Logging setup function defined here.
│       │   ├── schemas                 # Specific Schema for different cases are defined here.
│       │   │   ├── __init__.py
│       │   │   ├── app_configs.py      # App specific configuration schema
│       │   │   ├── request_schemas.py
│       │   │   ├── response_schemas.py
│       │   │   └── training_configs.py # Training configurations schema.
│       │   └── utils
│       │       ├── __init__.py
│       │       ├── postprocessing.py   # Post-processing logic and classes.
│       │       ├── preprocessing.py    # Preprocessing logic and classes.
│       │       └── utils.py            # Shared Utilities
│       ├── notebooks                   # Exploration notebooks (defined some templates for convenience but go wild.)
│       │   ├── binary_classifier_model_training.ipynb
│       │   └── multiclass_classifier_model_training.ipynb
│       ├── scripts                     # Project Specific Scripts
│       │   └── export_model_to_onnx_and_save_to_mlflow.py
│       ├── serving
│       │   ├── __init__.py
│       │   ├── api                     # FastAPI specific configurations and logic lives here.
│       │   │   ├── app.py              # FastAPI Entry Point.
│       │   │   └── routes
│       │   │       └── api.py          # Routes for your app.
│       │   ├── model
│       │   │   ├── __init__.py
│       │   │   └── loader.py           # Where the `Model` class resides, which is a model wrapper.
│       │   └── services
│       │       ├── __init__.py
│       │       └── inference.py        # Inference interface decoupled from the API itself, but this is what the API uses to perform inference.
│       └── training
│           ├── __init__.py
│           ├── architecture            # Where project specific architectures live, like model architecture and dataset setup.
│           │   ├── __init__.py
│           │   ├── dataset.py
│           │   └── model.py
│           ├── early_stopping          # Early Stopping implementation.
│           │   ├── __init__.py
│           │   └── early_stopping.py
│           ├── metrics
│           │   ├── __init__.py
│           │   ├── classification
│           │   │   └── metrics.py      # Classification metrics.
│           │   └── regression
│           │       └── metrics.py      # Regression metrics.
│           ├── model
│           │   ├── __init__.py
│           │   └── trainer.py          # Where the model `Trainer` and all adaptors live.
│           ├── run.py
│           └── utils                   # Training specific utilities.
│               ├── __init__.py
│               ├── count_model_params.py
│               └── create_dataloaders.py
├── tests                               # Unit tests reside here (it mocks the same directory structure as src/ml_project_template).
│   ├── __init__.py
│   ├── core
│   │   ├── configs
│   │   │   ├── test_configs.yaml       # Test configurations.
│   │   │   ├── test_configs_loader.py  # Tests the configurations loader.
│   │   │   └── test_validator.py       # Tests the environment variable validator (modify this depending on the env vars your project has).
│   │   └── logging
│   │       └── test_setup.py           # Tests the logger setup
│   ├── inference
│   │   ├── model
│   │   │   └── test_model_loader.py    # Tests the functionality of the `Model` class.
│   │   └── services
│   │       └── test_inference.py       # Tests the inference interface used by the API.
│   └── training
│       ├── early_stopping
│       │   └── test_early_stopping.py
│       ├── metrics
│       │   ├── classification
│       │   │   └── test_classification_metrics.py
│       │   └── regression
│       │       └── test_regression_metrics.py
│       ├── model
│       │   └── test_trainer.py
│       └── utils
│           ├── test_count_model_params.py
│           └── test_create_dataloaders.py
└── uv.lock                             # Project's Lockfile

---

Justfile

There were too many commands to keep up with, and honestly it was inconvenient, so I created a Justfile to have an easy way to collect and execute all necessary commands for this project

To check all available commands and their description, run:

just --list

And to run a command:

just command          # if the command takes no arguments
just command arg_val  # if the command takes some argument(s)

Just installation guide

Most common commands I found myself working with

• just train: Starts the training pipeline using the default config.
• just serve: Launches the FastAPI server with hot-reload.
• just test: Runs the full suite of tests with MLflow mocking.
• just clean: Cleans the directory from all pycache directories, .ruff_cache, and .pytest_cache

---

Installation

You can clonse the repository:

git clone https://github.com/MarwanMohammed2500/ml_project_template
cd ml_project_template

and then rename it to what you want

Or you can create a new repo using this as a template

---

Dependancies

All dependancies are outlined in pyproject.toml, and can be installed with:

uv sync

---

Core Workflow

• Configure: Define your task and hyperparameters in configs/training_configs.yaml.
• Train: Run the training pipeline. The Trainer automatically logs metrics, params, and the state_dict to MLflow.
• Export: Use the scripts/export_model_to_onnx... to convert your best PyTorch model to ONNX. This script doesn't bundles your preproc_pipeline.pkl as an MLflow extra_file, this is done when you are training the model, manual exports don't (for now, Will add it later).
• Register: Assign the @production alias to your ONNX model in the MLflow Registry.
• Serve: The FastAPI app uses the Model loader to pull the @production model, its ONNX weights, and its preprocessing assets automatically.

---

Adding New Task Types

To extend the template for a new task:

• Define Strategy: Create a new subclass in serving/model/loader.py (e.g., _TimeSeriesModel).
• Add Metrics: Implement the corresponding torchmetrics logic in training/metrics/.
• Update Config: Add the new TASK_TYPE to the validator.py allowed list.

---

Configurations

You can define environment variables and then validate them using validate_env_vars, you can also define model-specific configs in src/ml_project_template/configs/model_configs.py, which can then be used for model loading, inference, training, etc.

Example Config:

TASK_TYPE: Literal["binary", "multiclass", "regression"] = "binary"
CLASS_MAP: dict[int, str] = {0: "your", 1: "class", 2: "map"}
DECISION_THRESHOLD: float = 0.5
PRODUCTION_MODEL_URI: str = "models:/SimpleModel_ONNX@production"

---

Training a Model

Use the Trainer class.

Example:

trainer = Trainer(
        task_type=training_configs.trainer_configs.task_type,
        num_epochs=training_configs.trainer_configs.num_epochs,
        loss_fn=loss_fn,
        optimizer_class=optim.Adam,
        learning_rate=training_configs.trainer_configs.learning_rate,
        train_dataloader=train_dataloader,
        test_dataloader=test_dataloader,
        num_classes=num_classes,
        model_uri=training_configs.trainer_configs.model_uri,
        lr_scheduler_class=torch.optim.lr_scheduler.StepLR,
        early_stopper=early_stopper,
        binary_decision_threshold=training_configs.trainer_configs.binary_decision_threshold,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )

trainer.train()

---

Early Stopping

Example usage:

early_stopper = EarlyStopping(patience=5)

trainer = Trainer(
    ...
    early_stopper=early_stopper
)

Training will stop immediately if validation loss stop improving by a certain value delta

---

Metrics

Classification metrics include:

• Accuracy
• Precision
• Recall
• F1 Score

Regression metrics include:

• MSE
• RMSE
• Normalized RMSE

Metrics are implemented using torchmetrics.

---

Resilience & Fail-safes

This template is built for production stability:

• Pickle vs. JSON Handshake: If the serialized preprocessing pipeline (.pkl) fails to load (e.g., due to version mismatch), the Model loader automatically attempts to reconstruct the pipeline using the logged normalization_constants.json.
• Mocked CI: Every pull request is validated against a mocked MLflow environment, ensuring that changes to the core logic don't break the inference "contract."

---

Performance & Efficiency

• Zero-Disk Export: Torch-to-ONNX conversion happens entirely in-memory using BytesIO buffers before being streamed to MLflow.
• Warm-Worker Pattern: The Model manager supports pre-loading, ensuring that workers are "warm" and ready for sub-millisecond inference the moment they receive a request.

---

Future Plans

For me, this template is something I actually use, and so, improving it is an iterative process. I will end up adding more unit tests, cleaning up all around, implementing best practices, etc. But again, this is an iterative process, and thus, collaborations and contributions are more than welcomed.

---

Contributions

Feel free to fork the repo and add to it what you feel fits and then open a PR, we'll discuss it and if we reach a conclusion that the additinos are improvements, they'll be added and you'll be listed as a contributor

Contributions must pass the automated GitHub Actions suite, which includes mocked MLflow/ONNX tests to ensure architectural integrity.