# Training Pipeline Project

## Overview

This project was scaffolded using the ACC CLI — a developer tool that generates a standardised Databricks project structure.

By selecting the **Training** template during `acc init`, the CLI automatically provisions a step-based ML training architecture (Data Loading → Validation → Train → Save Model) built on PySpark and MLflow, so you can focus on writing model logic instead of setting up boilerplate.

## Project Structure
```
{{cookiecutter.project_slug}}/
│
├── configs/
│   ├── dev.yml
│   └── prod.yml
│
├── jobs/
│   └── training_job.json
│
├── src/
│   ├── training/
│   │   ├── __init__.py
│   │   ├── data_loader.py
│   │   ├── validator.py
│   │   ├── trainer.py
│   │   └── model_saver.py
│   ├── utils/
│   │   ├── __init__.py
│   │   └── helpers.py
│   ├── __init__.py
│   └── main.py
│
├── tests/
│   ├── __init__.py
│   ├── test_data_loader.py
│   ├── test_validator.py
│   ├── test_trainer.py
│   └── test_model_saver.py
│
├── .gitignore
├── __init__.py
├── Makefile
├── pyproject.toml
├── README.md
├── requirements.txt
└── requirements-dev.txt
```

## How It Works

### `configs/dev.yml` & `configs/prod.yml`
Environment-specific configuration files. Hold source data paths, feature settings, model parameters, and MLflow tracking URIs for dev and prod environments.

### `jobs/training_job.json`
Databricks job definition. Defines the cluster config, task dependencies, and entry point for running the training pipeline on Databricks.

### `src/training/data_loader.py`
Entry point for the training pipeline. Reads raw training data from the configured source (Delta table, parquet, CSV, etc.) and returns a clean DataFrame ready for validation.

### `src/training/validator.py`
Validates the loaded dataset before training begins. Checks schema expectations, null thresholds, class distributions, and feature completeness. Raises early if data does not meet quality standards.

### `src/training/trainer.py`
Core model training logic. Loads hyperparameters from config, fits the model on validated data, and logs parameters, metrics, and artefacts to MLflow for experiment tracking.

### `src/training/model_saver.py`
Saves the trained model to the MLflow Model Registry or Unity Catalog. Handles model versioning, stage transitions (Staging → Production), and writes model metadata back to config.

### `src/utils/helpers.py`
Shared utility functions used across the pipeline. Includes helpers for Spark session creation, config loading, MLflow setup, logging, and common DataFrame operations.

### `src/main.py`
Main entry point for the pipeline. Loads config, initialises the Spark session, and orchestrates the Data Loading → Validation → Training → Save Model execution sequence.

### `tests/test_data_loader.py`
Unit tests for `data_loader.py`. Validates raw ingestion logic, schema expectations, and that the returned DataFrame meets downstream requirements.

### `tests/test_validator.py`
Unit tests for `validator.py`. Validates null handling, schema checks, and that invalid datasets are correctly caught and rejected.

### `tests/test_trainer.py`
Unit tests for `trainer.py`. Validates model fitting logic, hyperparameter loading, and MLflow logging behaviour against expected outputs.

### `tests/test_model_saver.py`
Unit tests for `model_saver.py`. Validates model registration logic, version handling, and stage transition behaviour in the MLflow registry.

### `Makefile`
Convenience commands for common tasks — `make build`, `make test`, `make lint`, `make format`.

### `pyproject.toml`
Project metadata and dependency management. Defines the package name, version, and build system configuration.

#### Optional Cloud & Databricks Dependencies
If your pipeline needs to upload or interact with Amazon S3, Azure Data Lake Storage Gen2, Azure Blob Storage, or Databricks Unity Catalog, the required dependencies can be added directly to the generated pyproject.toml file.

### `requirements.txt` & `requirements-dev.txt`
Runtime and development dependencies respectively. `requirements-dev.txt` includes testing and linting libraries on top of the base requirements.

## Building a `.whl` File

To package this project as a Python wheel (`.whl`) for deployment to Databricks, AWS or Azure, run the following command from the project root:

\`\`\`
acc build
\`\`\`

This will compile the project into a distributable `.whl` file.