Repository: jeroenvdhoven/mlpype (language: python)

Issue:
Add conversion of experiment to inferencer



Discussion / hints:

Fix patch:
diff --git a/.coveragerc b/.coveragerc
index 003c131..873dff2 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -4,3 +4,8 @@ show_missing=True
 exclude_lines =
     pragma: no cover
     raise NotImplementedError
+omit = 
+    # all python setup.py files
+    */setup.py
+    # cmd_tools, since it's only used for dev now
+    pype/cmd_tools/*
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index ab12548..9e2d2ed 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,10 +14,22 @@ __pycache__/
 .xlsx
 
 # coverage
-.coverage
+.coverage*
+.coveragerc!
 
 # outputs generated by example scripts
 outputs/
 
 # notebooks
-*.ipynb
\ No newline at end of file
+*.ipynb
+
+# mlflow
+mlruns*/
+
+# build artifacts
+dist/
+packages/
+build/
+
+# personal todo list
+todo.txt
\ No newline at end of file
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 791a1f4..e1ae761 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -30,7 +30,6 @@ repos:
   hooks:
   -   id: mypy
       # This setting is not that strict: see https://mypy.readthedocs.io/en/stable/config_file.html for more config options.
-      args: [--no-strict-optional, --ignore-missing-imports, --disallow-untyped-defs, --python-version, "3.10"]
       exclude: ^(tests)|(docs)/
       additional_dependencies:
         - "types-PyYAML"
@@ -56,7 +55,7 @@ repos:
   - id: unittest
   # todo: reuse test environment!
     name: unittest
-    entry: 'make test'
+    entry: 'make test-without-spark'
     language: conda
     types: [python]
     verbose: true
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 0000000..3a4c12f
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,3 @@
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
\ No newline at end of file
diff --git a/Makefile b/Makefile
index a2c4254..2c77080 100644
--- a/Makefile
+++ b/Makefile
@@ -1,13 +1,31 @@
 
+build:
+	./scripts/build.sh
+
+clean:
+	rm -rf dist/ packages/
 
 dev-install:
 	pip install -e .
-	pype dev-install --editable
+	./scripts/dev_install.sh -e 1
+
+host-pypi-local:
+	mkdir -p packages
+	pypi-server run -p 8080 packages -a . -P . &
+	twine upload --repository-url http://0.0.0.0:8080 dist/* -u '' -p ''
+
+build-and-host-local: clean build host-pypi-local
 
-test:
+test-without-spark:
+	python -m pytest -m "not spark"
+
+test-all:
 	python -m pytest
 
-coverage:
+coverage-without-spark:
+	python -m pytest --cov-report term-missing --cov pype -ra
+
+coverage-all:
 	python -m pytest --cov-report term-missing --cov pype -ra
 
 # Pre-commit defaults
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..0255b78
--- /dev/null
+++ b/README.md
@@ -0,0 +1,101 @@
+# Pype: standardise model training across libraries
+
+## Why yet another ML package?
+
+## Usage and examples
+There are a couple of examples available of training scripts using pype and various subpackages in the `examples` folder.
+
+In general, you'll need the following components:
+- DataSetSource's: a collection of DataSource's, that can be read to produce a DataSet. This will specify which data is used by your
+model's input and how to collect it. This can range from cloud-storage files, SQL databases, etc.
+- Pipeline: Another trainable pipeline format? Yes. Those familiar with kedro will see some similarities, but we specify a difference
+between the fitting and transforming phase, much like sklearn's pipelines. The framework should be flexible enough to allow a variety
+of packages to integrate with the base Operator class, or at least comply with its interface.
+- Model: An interface that should detail well what functions are needed to make a Model work. We have intergrations and examples for sklearn, pyspark,
+and keras models, but you can also make your own integration.
+- Evaluator:
+- ExperimentLogger: 
+- input/output type checkers:
+
+Optionally, you can specify the following:
+- serialiser
+- output_folder
+- additional_files_to_store
+
+You can check the documentation for `Experiment` for more details.
+
+## Advantages
+- Modular: `pype` is setup as a namespace package. This allows you to install parts of the library to suit your needs.
+    Development is done by default on a full installation of `pype`, so all packages *should* work together.
+- Standardised: Don't code the same old steps every single time, just re-use them. `pype` handles a lot of good standard
+    steps in ML training, like experiment logging and keeping track of artifacts.
+- Extensive: Currently `pype` has a couple of integrations with other popular packages like `sklearn` and `tensorflow`, but
+    also `mlflow`. Ideally we'd like to expand this to other common ML libraries and also platforms. We'd like to see this
+    become a simple, universal starting point for machine learning with proper standardisation
+- Flexibility: Bring your own models, data sources, and preprocessing steps. Bring your own loggers and serialisers. Just
+    follow the interfaces we've set up, and you should be good to go!
+
+# Features
+## Model training
+
+The goal is to help standardise the following steps:
+- Fetching data
+- Pipeline fitting
+- Model training
+- Serialisation
+- Logging
+
+We aim to do this by providing basic interfaces that should be easy to implement for other packages.
+This should help keep the base package lightweight but flexible, while allowing subpackages to implement
+specific functionality for other packages. You can also directly extend the base interfaces to work with
+`pype`.
+
+## Model loading
+With most steps standardised, it becomes a lot easier to standardise model deployment as well. Currently we offer:
+
+- Loading your model in a standardised way back into memory from an output folder
+- Deploying your model using a FastAPI server
+- Packaging your model as a wheel package, with dependencies linked from the experiment.
+    - It is possible to extend this wheel package with your own functions as well, such as the FastAPI server
+
+## Subpackages
+Pype is setup in a modular way. The base package provides the interfaces (e.g. `Model`, `DataSource` classes) and interactions between them (e.g. `Experiment` class).
+Its dependency footprint is quite low to make sure no excess packages are installed when working with `pype` in different ways.
+
+Subpackages like `pype.sklearn` provide implementations, such as sklearn, tensorflow, or mlflow integrations. Not all subpackages need to be installed,
+making deployments lighter and less prone to dependency problems. This also reduces dependencies between subpackages, reducing complicated internal dependency issues.
+
+
+Currently the following subpackages are available:
+
+- `pype.base`: The base package. Required for other packages. Provides basics of training Pype-compliant models.
+    - Fetching data
+    - Pipeline fitting
+    - Model training
+    - Serialisation
+    - Logging
+    - Deployment
+- `pype.mlflow`: MLflow integration package. Allows easy use of logging through mlflow and downloading/loading models trained through Pype and logged on MLflow.
+- `pype.sklearn`: Sklearn integration package. Simplifies training sklearn models on Pype.
+- `pype.xgboost`: Minor extension to the `pype.sklearn` package, including XGBoost models.
+- `pype.tensorflow`: Incorporates tensorflow/keras into Pype.
+- `pype.hyperopt`: Use hyperopt to do hyperparameter tuning in Pype.
+- `pype.spark`: In progress, will provide pyspark integrations for Pype.
+
+# Contributing
+
+## Basics
+
+Besides regular good coding practices, we'd like to ask you to use type hints where possible and realistic.
+This greatly helps the users and other developers develop using your new code quicker. We have mypy enabled
+to help you out and find missing type hints quicker.
+## Setting up development environment
+
+If you're interested in helping out, you can install all packages into your environment using `make dev-install`.
+Please make sure to also install the pre-commit steps using `make pre-commit-install`. Check out the `.pre-commit-config.yaml`
+file for the exact steps we use, but it's summarised as:
+
+- sorting imports
+- remove unused imports
+- black
+- mypy for type hitns
diff --git a/examples/custom_model_example.py b/examples/custom_model_example.py
new file mode 100644
index 0000000..d94b73c
--- /dev/null
+++ b/examples/custom_model_example.py
@@ -0,0 +1,120 @@
+"""Please run this file using `python -m examples.custom_model_example`.
+
+We do not guarantee results if you use `python examples/custom_model_example.py`
+"""
+
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+import pandas as pd
+from sklearn.datasets import load_iris
+from sklearn.metrics import accuracy_score
+from sklearn.model_selection import train_test_split
+
+from pype.base.data import DataSetSource
+from pype.base.deploy.inference import Inferencer
+from pype.base.evaluate.evaluator import Evaluator
+from pype.base.experiment.experiment import Experiment
+from pype.base.logger.local_logger import LocalLogger
+from pype.base.pipeline.pipe import Pipe
+from pype.base.pipeline.pipeline import Pipeline
+from pype.base.pipeline.type_checker import TypeCheckerPipe
+from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
+from pype.sklearn.data.data_frame_source import DataFrameSource
+from pype.sklearn.model.logistic_regression_model import LogisticRegressionModel
+
+from .custom_processer import CustomStandardScaler, tcc
+
+
+# classes / functions defined in the core file can be imported
+class CustomModel(LogisticRegressionModel):
+    pass
+
+
+if __name__ == "__main__":
+
+    def _make_data() -> Iterable[np.ndarray]:
+        iris = load_iris(as_frame=True)
+        x = pd.DataFrame(iris["data"])
+        y = pd.DataFrame(iris["target"])
+
+        kept_rows = y["target"] < 2
+        x = x.loc[kept_rows, :]
+        y = y.loc[kept_rows, :]
+
+        return train_test_split(x, y, test_size=0.2)
+
+    train_x, test_x, train_y, test_y = _make_data()
+
+    model = CustomModel(
+        model=None,
+        inputs=["x"],
+        outputs=["y"],
+    )
+
+    ds = {
+        "train": DataSetSource(
+            x=DataFrameSource(train_x),
+            y=DataFrameSource(train_y),
+        ),
+        "test": DataSetSource(
+            x=DataFrameSource(test_x),
+            y=DataFrameSource(test_y),
+        ),
+    }
+
+    evaluator = Evaluator(
+        {
+            "accuracy": accuracy_score,
+        }
+    )
+
+    input_ds_type_checker = TypeCheckerPipe(
+        "type_checker-in",
+        input_names=["x"],
+        type_checker_classes=tcc,
+    )
+
+    output_ds_type_checker = TypeCheckerPipe(
+        "type_checker-out",
+        input_names=["y"],
+        type_checker_classes=tcc,
+    )
+
+    pipeline = Pipeline([Pipe("scale", CustomStandardScaler, inputs=["x"], outputs=["x"])])
+    of = Path("outputs")
+
+    this_file = Path(__file__)
+
+    experiment = Experiment(
+        data_sources=ds,
+        model=model,
+        pipeline=pipeline,
+        evaluator=evaluator,
+        logger=LocalLogger(),
+        input_type_checker=input_ds_type_checker,
+        output_type_checker=output_ds_type_checker,
+        serialiser=JoblibSerialiser(),
+        output_folder=of,
+        # Need to add this file to output to make sure we can import CustomModel and CustomStandardScaler/tcc
+        additional_files_to_store=[this_file, this_file.parent / "custom_processer.py"],
+    )
+
+    metrics = experiment.run()
+
+    print(metrics)
+
+    # Try loading results again
+
+    folder = Path("outputs")
+
+    inferencer = Inferencer.from_folder(folder)
+
+    train_x, test_x, train_y, test_y = _make_data()
+    test_data = DataSetSource(
+        x=DataFrameSource(test_x),
+        y=DataFrameSource(test_y),
+    )
+    result = inferencer.predict(test_data)
+    print(result)
diff --git a/examples/custom_processer.py b/examples/custom_processer.py
new file mode 100644
index 0000000..674651a
--- /dev/null
+++ b/examples/custom_processer.py
@@ -0,0 +1,15 @@
+import numpy as np
+import pandas as pd
+from sklearn.preprocessing import StandardScaler
+
+from pype.sklearn.pipeline.numpy_type_checker import NumpyTypeChecker
+from pype.sklearn.pipeline.pandas_type_checker import PandasTypeChecker
+
+tcc = [
+    (np.ndarray, NumpyTypeChecker),
+    (pd.DataFrame, PandasTypeChecker),
+]
+
+
+class CustomStandardScaler(StandardScaler):
+    pass
diff --git a/examples/mlflow/mlflow_sklearn_example.py b/examples/mlflow/mlflow_sklearn_example.py
new file mode 100644
index 0000000..addfb88
--- /dev/null
+++ b/examples/mlflow/mlflow_sklearn_example.py
@@ -0,0 +1,151 @@
+"""Please run this file using `python -m examples.mlflow.mlflow_sklearn_example`.
+
+We do not guarantee results if you use `python examples/mlflow/mlflow_sklearn_example.py`
+You can use command line arguments in this example, such as:
+python -m examples.mlflow.mlflow_sklearn_example --model__fit_intercept=False
+
+This requires the pype.sklearn package to also be installed.
+
+Please make sure mlflow is running locally before starting this script, e.g. by using `mlflow ui`
+If you are running this example locally, make sure you run this from the top level directory.
+This will make sure the artifacts will show up.
+"""
+# %%
+
+from pathlib import Path
+from typing import Iterable
+from unittest.mock import MagicMock
+
+import numpy as np
+import pandas as pd
+from mlflow.tracking.fluent import get_experiment_by_name
+from sklearn.datasets import load_iris
+from sklearn.metrics import accuracy_score
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+
+from pype.base.data import DataSetSource
+from pype.base.deploy.inference import Inferencer
+from pype.base.evaluate.evaluator import Evaluator
+from pype.base.experiment.experiment import Experiment
+from pype.base.pipeline.pipe import Pipe
+from pype.base.pipeline.pipeline import Pipeline
+from pype.base.pipeline.type_checker import TypeCheckerPipe
+from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
+from pype.mlflow.logger.mlflow_logger import MlflowLogger
+from pype.sklearn.data.data_frame_source import DataFrameSource
+from pype.sklearn.model.linear_regression_model import LinearRegressionModel
+from pype.sklearn.model.logistic_regression_model import LogisticRegressionModel
+from pype.sklearn.pipeline.numpy_type_checker import NumpyTypeChecker
+from pype.sklearn.pipeline.pandas_type_checker import PandasTypeChecker
+
+# %%
+
+parser = MagicMock()
+# parser = ArgumentParser()
+model = LinearRegressionModel.get_parameters(parser)
+
+print(parser.add_argument.call_args_list)
+
+# %% [markdown]
+# Try a run with sklearn
+
+# %%
+# Make a logger using Mlflow. Ensure it starts with 'http://', or you will get connection issues.
+# Make sure you start mlflow before running this!
+experiment_name = "jeroen-example-experiment"
+logger = MlflowLogger(experiment_name, "http://localhost:5000")
+
+# for on-databricks logging. This will also log artifacts.
+# experiment_name = "/Users/<user name>/<experiment name>"
+# logger = MlflowLogger(experiment_name, "databricks")
+
+# %%
+#  Try a run with sklearn and argument reading
+
+
+def _make_data() -> Iterable[np.ndarray]:
+    iris = load_iris()
+    x = iris["data"]
+    y = iris["target"]
+
+    kept_rows = y < 2
+    x = x[kept_rows, :]
+    y = y[kept_rows]
+
+    return train_test_split(x, y, test_size=0.2)
+
+
+train_x, test_x, train_y, test_y = _make_data()
+
+ds = {
+    "train": DataSetSource(
+        x=DataFrameSource(train_x),
+        y=DataFrameSource(train_y),
+    ),
+    "test": DataSetSource(
+        x=DataFrameSource(test_x),
+        y=DataFrameSource(test_y),
+    ),
+}
+
+evaluator = Evaluator(
+    {
+        "accuracy": accuracy_score,
+    }
+)
+
+tcc = [
+    (np.ndarray, NumpyTypeChecker),
+    (pd.DataFrame, PandasTypeChecker),
+]
+
+input_ds_type_checker = TypeCheckerPipe(
+    "type_checker-in",
+    input_names=["x"],
+    type_checker_classes=tcc,
+)
+
+output_ds_type_checker = TypeCheckerPipe(
+    "type_checker-out",
+    input_names=["y"],
+    type_checker_classes=tcc,
+)
+
+pipeline = Pipeline([Pipe("scale", StandardScaler, inputs=["x"], outputs=["x"])])
+of = Path("outputs")
+
+experiment = Experiment.from_command_line(
+    data_sources=ds,
+    model_class=LogisticRegressionModel,
+    model_inputs=["x"],
+    model_outputs=["y"],
+    pipeline=pipeline,
+    evaluator=evaluator,
+    logger=logger,
+    input_type_checker=input_ds_type_checker,
+    output_type_checker=output_ds_type_checker,
+    serialiser=JoblibSerialiser(),
+    output_folder=of,
+)
+
+metrics = experiment.run()
+print("Metrics:", metrics)
+
+# Try loading results again
+folder = Path("outputs")
+inferencer = Inferencer.from_folder(folder)
+
+train_x, test_x, train_y, test_y = _make_data()
+test_data = DataSetSource(
+    x=DataFrameSource(test_x),
+    y=DataFrameSource(test_y),
+)
+result = inferencer.predict(test_data)
+print(result)
+
+# %%
+# Get artifact path from mlflow
+
+exp = get_experiment_by_name(experiment_name)
+print("artifact_location:", exp.artifact_location)
diff --git a/examples/sklearn_example.py b/examples/sklearn/sklearn_example.py
similarity index 93%
rename from examples/sklearn_example.py
rename to examples/sklearn/sklearn_example.py
index 2e3d306..dbde03e 100644
--- a/examples/sklearn_example.py
+++ b/examples/sklearn/sklearn_example.py
@@ -1,3 +1,7 @@
+"""Please run this file using `python -m examples.sklearn.sklearn_example`.
+
+We do not guarantee results if you use `python examples/sklearn/sklearn_example.py`
+"""
 # %%
 
 from pathlib import Path
@@ -96,6 +100,7 @@ def _make_data() -> Iterable[np.ndarray]:
 )
 
 pipeline = Pipeline([Pipe("scale", StandardScaler, inputs=["x"], outputs=["x"])])
+of = Path("outputs")
 
 experiment = Experiment(
     data_sources=ds,
@@ -106,7 +111,7 @@ def _make_data() -> Iterable[np.ndarray]:
     input_type_checker=input_ds_type_checker,
     output_type_checker=output_ds_type_checker,
     serialiser=JoblibSerialiser(),
-    output_folder="outputs",
+    output_folder=of,
 )
 
 metrics = experiment.run()
diff --git a/examples/sklearn_example_from_cmd_args.py b/examples/sklearn/sklearn_example_from_cmd_args.py
similarity index 73%
rename from examples/sklearn_example_from_cmd_args.py
rename to examples/sklearn/sklearn_example_from_cmd_args.py
index 4335d84..8a2daf5 100644
--- a/examples/sklearn_example_from_cmd_args.py
+++ b/examples/sklearn/sklearn_example_from_cmd_args.py
@@ -1,12 +1,13 @@
 """Call this using some form of command line.
 
-e.g. python -m examples.sklearn_example_from_cmd_args --model__fit_intercept=False
+e.g. python -m examples.sklearn.sklearn_example_from_cmd_args --model__fit_intercept=False
 """
 
 from pathlib import Path
 from typing import Iterable
 
 import numpy as np
+import pandas as pd
 from sklearn.datasets import load_iris
 from sklearn.metrics import accuracy_score
 from sklearn.model_selection import train_test_split
@@ -19,9 +20,12 @@
 from pype.base.logger.local_logger import LocalLogger
 from pype.base.pipeline.pipe import Pipe
 from pype.base.pipeline.pipeline import Pipeline
+from pype.base.pipeline.type_checker import TypeCheckerPipe
 from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
 from pype.sklearn.data.data_frame_source import DataFrameSource
 from pype.sklearn.model.logistic_regression_model import LogisticRegressionModel
+from pype.sklearn.pipeline.numpy_type_checker import NumpyTypeChecker
+from pype.sklearn.pipeline.pandas_type_checker import PandasTypeChecker
 
 
 #  Try a run with sklearn and argument reading
@@ -56,7 +60,26 @@ def _make_data() -> Iterable[np.ndarray]:
     }
 )
 
+tcc = [
+    (np.ndarray, NumpyTypeChecker),
+    (pd.DataFrame, PandasTypeChecker),
+]
+
+input_ds_type_checker = TypeCheckerPipe(
+    "type_checker-in",
+    input_names=["x"],
+    type_checker_classes=tcc,
+)
+
+output_ds_type_checker = TypeCheckerPipe(
+    "type_checker-out",
+    input_names=["y"],
+    type_checker_classes=tcc,
+)
+
 pipeline = Pipeline([Pipe("scale", StandardScaler, inputs=["x"], outputs=["x"])])
+of = Path("outputs")
+
 experiment = Experiment.from_command_line(
     data_sources=ds,
     model_class=LogisticRegressionModel,
@@ -65,8 +88,10 @@ def _make_data() -> Iterable[np.ndarray]:
     pipeline=pipeline,
     evaluator=evaluator,
     logger=LocalLogger(),
+    input_type_checker=input_ds_type_checker,
+    output_type_checker=output_ds_type_checker,
     serialiser=JoblibSerialiser(),
-    output_folder="outputs",
+    output_folder=of,
 )
 
 metrics = experiment.run()
diff --git a/examples/spark/example.py b/examples/spark/example.py
new file mode 100644
index 0000000..b75fcab
--- /dev/null
+++ b/examples/spark/example.py
@@ -0,0 +1,56 @@
+"""Please run this file using `python -m examples.spark.example`.
+
+We do not guarantee results if you use `python examples/spark/example.py`
+"""
+
+
+import pandas as pd
+from pipeline.type_checker import TypeCheckerPipe
+from pyspark.ml.feature import StandardScaler, VectorAssembler
+from pyspark.sql import DataFrame as SparkDataFrame
+from pyspark.sql import SparkSession
+
+from pype.base.data import DataSetSource
+from pype.base.evaluate import Evaluator
+from pype.base.experiment import Experiment
+from pype.base.logger import LocalLogger
+from pype.base.pipeline import Pipeline
+from pype.spark.data.spark_data_frame_source import SparkDataFrameSource
+from pype.spark.model import LinearSparkModel
+from pype.spark.pipeline import SparkTypeChecker
+from pype.spark.pipeline.spark_pipe import SparkPipe
+from pype.spark.serialisation.spark_serialiser import SparkSerialiser
+
+ss = SparkSession.builder.getOrCreate()
+
+data = {
+    "train": DataSetSource(
+        x=SparkDataFrameSource(ss.createDataFrame(pd.DataFrame({"a": [1, 2, 3.0], "target": [2, 3, 4.0]}))),
+        helper=SparkDataFrameSource(ss.createDataFrame(pd.DataFrame({"a": [2, 3, 4]}))),
+    )
+}
+
+pipeline = Pipeline(
+    [
+        SparkPipe("assemble", VectorAssembler, ["x"], ["x"], kw_args={"inputCols": ["a"], "outputCol": "a_gathered"}),
+        SparkPipe("scale", StandardScaler, ["x"], ["x"], kw_args={"inputCol": "a_gathered", "outputCol": "a_scaled"}),
+    ]
+)
+
+model = LinearSparkModel(["x"], ["x"], featuresCol="a_scaled", labelCol="target", predictionCol="prediction")
+
+logger = LocalLogger()
+
+experiment = Experiment(
+    data_sources=data,
+    model=model,
+    pipeline=pipeline,
+    evaluator=Evaluator({}),
+    logger=logger,
+    serialiser=SparkSerialiser(),
+    output_folder="outputs",
+    input_type_checker=TypeCheckerPipe("inputs", ["x"], [(SparkDataFrame, SparkTypeChecker)]),
+    output_type_checker=TypeCheckerPipe("outputs", ["x"], [(SparkDataFrame, SparkTypeChecker)]),
+)
+
+experiment.run()
diff --git a/pype/base/constants.py b/pype/base/constants.py
index 05b87eb..886fd61 100644
--- a/pype/base/constants.py
+++ b/pype/base/constants.py
@@ -5,3 +5,7 @@ class Constants:
     PIPELINE_FILE: str = "pipeline.pkl"
     INPUT_TYPE_CHECKER_FILE: str = "input_type_checker.pkl"
     OUTPUT_TYPE_CHECKER_FILE: str = "output_type_checker.pkl"
+    EXTRA_FILES: str = "extra_files.json"
+
+    PYTHON_VERSION_FILE: str = "python_version.json"
+    REQUIREMENTS_FILE: str = "requirements.txt"
diff --git a/pype/base/data/__init__.py b/pype/base/data/__init__.py
index 5a7589d..adb3d60 100644
--- a/pype/base/data/__init__.py
+++ b/pype/base/data/__init__.py
@@ -1,3 +1,4 @@
+from .data_sink import DataSink
 from .data_source import DataSource
 from .dataset import DataSet
 from .dataset_source import DataSetSource
diff --git a/pype/base/data/data_sink.py b/pype/base/data/data_sink.py
new file mode 100644
index 0000000..450ebaa
--- /dev/null
+++ b/pype/base/data/data_sink.py
@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+Data = TypeVar("Data")
+
+
+class DataSink(ABC, Generic[Data]):
+    @abstractmethod
+    def write(self, data: Data) -> None:
+        """Writes data to a given source."""
+        raise NotImplementedError
diff --git a/pype/base/data/dataset_source.py b/pype/base/data/dataset_source.py
index 5138e19..dfbfb3d 100644
--- a/pype/base/data/dataset_source.py
+++ b/pype/base/data/dataset_source.py
@@ -18,3 +18,15 @@ def read(self) -> DataSet[Data]:
             DataSet[Data]: The DataSet constructed from the DataSources.
         """
         return DataSet.from_dict({name: data.read() for name, data in self.items()})
+
+    def __str__(self, indents: int = 0) -> str:
+        """Create string representation of this DataSetSource.
+
+        Args:
+            indents (int, optional): The number of preceding tabs. Defaults to 0.
+
+        Returns:
+            str: A string representation of this DataSetSource.
+        """
+        tab = "\t" * indents
+        return "\n".join([f"{tab}{name}: {source}" for name, source in self.items()])
diff --git a/pype/base/deploy/__init__.py b/pype/base/deploy/__init__.py
index e69de29..c2cbbb5 100644
--- a/pype/base/deploy/__init__.py
+++ b/pype/base/deploy/__init__.py
@@ -0,0 +1,1 @@
+from .inference import Inferencer
diff --git a/pype/base/deploy/app.py b/pype/base/deploy/app.py
deleted file mode 100644
index 610c942..0000000
--- a/pype/base/deploy/app.py
+++ /dev/null
@@ -1,47 +0,0 @@
-from dataclasses import dataclass
-from pathlib import Path
-
-from fastapi import FastAPI
-
-from pype.base.data.dataset import DataSet
-from pype.base.deploy.inference import Inferencer
-from pype.base.pipeline.type_checker import DataSetModel
-
-
-@dataclass
-class PypeApp:
-    name: str
-    folder: str | Path
-
-    def __post_init__(self):  # type: ignore
-        """Makes sure the folder is an actual Path."""
-        if isinstance(self.folder, str):
-            self.folder = Path(self.folder)
-
-    def create_app(self) -> FastAPI:
-        """Creates a FastAPI app for a saved Pype experiment.
-
-        Returns:
-            FastAPI: The FastAPI server that can be served.
-        """
-        app = FastAPI()
-        inferencer = self._load_model()
-
-        InputType: type[DataSetModel] = inferencer.input_type_checker.get_pydantic_types()
-        OutputType: type[DataSetModel] = inferencer.output_type_checker.get_pydantic_types()
-
-        @app.get("/")
-        async def home_page() -> str:
-            return f"Welcome to the Pype FastAPI app for {self.name}"
-
-        @app.post("/predict")
-        async def predict(inputs: InputType) -> OutputType:  # type: ignore
-            converted: DataSet = inputs.convert()  # type: ignore
-            prediction = inferencer.predict(converted)
-            return OutputType.to_model(prediction)
-
-        return app
-
-    def _load_model(self) -> Inferencer:
-        assert isinstance(self.folder, Path)
-        return Inferencer.from_folder(self.folder)
diff --git a/pype/base/deploy/inference.py b/pype/base/deploy/inference.py
index ad7eab5..2eed853 100644
--- a/pype/base/deploy/inference.py
+++ b/pype/base/deploy/inference.py
@@ -1,12 +1,16 @@
+import json
 from pathlib import Path
 from typing import Type
 
 from pype.base.constants import Constants
 from pype.base.data import DataSet, DataSetSource
+from pype.base.experiment.experiment import Experiment
 from pype.base.model import Model
 from pype.base.pipeline import Pipeline
 from pype.base.pipeline.type_checker import TypeCheckerPipe
 from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
+from pype.base.serialiser.serialiser import Serialiser
+from pype.base.utils.workspace import switch_workspace
 
 
 class Inferencer:
@@ -50,21 +54,32 @@ def predict(self, data: DataSet | DataSetSource) -> DataSet:
         return predicted
 
     @classmethod
-    def from_folder(cls: Type["Inferencer"], folder: Path) -> "Inferencer":
+    def from_folder(cls: Type["Inferencer"], folder: Path, serialiser: Serialiser | None = None) -> "Inferencer":
         """Loads a Inferencer from the results of an Experiment.
 
+        We use the absolute version of the path to try and prevent loading issues.
+
         Args:
             folder (Path): The output folder from an Experiment, from which we load
                 all required elements to make a inference pipeline.
+            serialiser (Serialiser | None): The Serialiser used to deserialise the pipeline and input/output
+                type checkers. Defaults to a JoblibSerialiser.
 
         Returns:
             Inferencer: The inference pipeline that can predict for new data.
         """
-        serialiser = JoblibSerialiser()
-        model = Model.load(folder / Constants.MODEL_FOLDER)
-        pipeline = serialiser.deserialise(folder / Constants.PIPELINE_FILE)
-        input_type_checker = serialiser.deserialise(folder / Constants.INPUT_TYPE_CHECKER_FILE)
-        output_type_checker = serialiser.deserialise(folder / Constants.OUTPUT_TYPE_CHECKER_FILE)
+        if serialiser is None:
+            serialiser = JoblibSerialiser()
+
+        folder = folder.absolute()
+        with open(folder / Constants.EXTRA_FILES, "r") as f:
+            extra_files = json.load(f)["paths"]
+
+        with switch_workspace(folder, extra_files):
+            model = Model.load(folder / Constants.MODEL_FOLDER)
+            pipeline = serialiser.deserialise(folder / Constants.PIPELINE_FILE)
+            input_type_checker = serialiser.deserialise(folder / Constants.INPUT_TYPE_CHECKER_FILE)
+            output_type_checker = serialiser.deserialise(folder / Constants.OUTPUT_TYPE_CHECKER_FILE)
 
         return cls(
             model=model,
@@ -72,3 +87,37 @@ def from_folder(cls: Type["Inferencer"], folder: Path) -> "Inferencer":
             input_type_checker=input_type_checker,
             output_type_checker=output_type_checker,
         )
+
+    @classmethod
+    def from_experiment(cls, experiment: Experiment) -> "Inferencer":
+        """Generates an Inferencer from an experiment.
+
+        We currently cannot check if the experiment has been trained already.
+
+        Args:
+            experiment (Experiment): The experiment to convert to a inferencing
+                pipeline.
+
+        Returns:
+            Inferencer: An inferencing pipeline obtained from the given experiment.
+        """
+        return cls(
+            model=experiment.model,
+            pipeline=experiment.pipeline,
+            input_type_checker=experiment.input_type_checker,
+            output_type_checker=experiment.output_type_checker,
+        )
+
+    def __str__(self) -> str:
+        """Create string representation of this Inferencer.
+
+        Returns:
+            str: A string representation of this Inferencer.
+        """
+        pipeline_part = "pipeline:\n" + self.pipeline.__str__(indents=1)
+        model_part = f"model:\n\t{str(self.model)}"
+
+        inputs = f"inputs:\n\t{str(self.input_type_checker)}"
+        outputs = f"outputs:\n\t{str(self.output_type_checker)}"
+
+        return "\n\n".join([inputs, outputs, pipeline_part, model_part])
diff --git a/pype/base/deploy/wheel/__init__.py b/pype/base/deploy/wheel/__init__.py
new file mode 100644
index 0000000..fdcfc00
--- /dev/null
+++ b/pype/base/deploy/wheel/__init__.py
@@ -0,0 +1,2 @@
+from .builder import WheelBuilder
+from .extensions import WheelExtension
diff --git a/pype/base/deploy/wheel/builder.py b/pype/base/deploy/wheel/builder.py
new file mode 100644
index 0000000..df32d09
--- /dev/null
+++ b/pype/base/deploy/wheel/builder.py
@@ -0,0 +1,126 @@
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from shutil import copytree
+from tempfile import TemporaryDirectory
+
+from pype.base.constants import Constants
+from pype.base.utils.workspace import switch_directory
+
+from .extensions import WheelExtension
+
+BaseExtension = WheelExtension(
+    "base",
+    functionality=[(Path(__file__).parent / "helpers" / "load_model.py", ["load_model"])],
+    libraries=["pype.base"],
+)
+
+
+@dataclass
+class WheelBuilder:
+    """Class to help turn an output folder from an Experiment into a wheel file.
+
+    This should allow you to save models into shippable packages
+    quite easily.
+    # TODO: add proper tests
+
+    The functionality can be extended quite easily using the `extensions` argument.
+    See `WheelExtension` for how to make your own. The base functionality is also
+    included as an extension, feel free to check out how that was done as well.
+
+    Args:
+        model_folder (Path): Path to the output folder generated by an Experiment
+        model_name (str): The name of the model. Will be used as package name. Should follow
+            proper package naming conventions.
+        version (str): Model/package version.
+        output_wheel_file (Path | str | None, optional): The output directory. Defaults to "wheel_output"
+            in the current directory.
+        libraries (list[str] | None, optional): A list of (versioned) libraries that are used as
+            dependencies. By default we use the requirements.txt file from the experiment run.
+        extensions (list[WheelExtension] | None, optional): A list of extension to add to this
+            wheel file. Extensions add extra functionality besides the base "import pipeline" functionality.
+            This can include hosting a FastAPI app for instance.
+    """
+
+    model_folder: Path
+    model_name: str
+    version: str
+    output_wheel_file: Path | str | None = None
+    libraries: list[str] | None = None
+    extensions: list[WheelExtension] | None = None
+
+    def __post_init__(self) -> None:
+        """Performs checks and default initialisations."""
+        if self.output_wheel_file is None:
+            self.output_wheel_file = Path("").parent.absolute() / "wheel_output"
+        if self.extensions is None:
+            self.extensions = []
+        if BaseExtension not in self.extensions:
+            self.extensions.append(BaseExtension)
+        self._validate_extensions()
+
+    def _validate_extensions(self) -> None:
+        names = set()
+        function_names = set()
+
+        assert isinstance(self.extensions, list)
+        for extension in self.extensions:
+            assert extension.name not in names, f"{extension.name} is a duplicated extension name, this is not allowed."
+            names.add(extension.name)
+
+            for _, imports in extension.functionality:
+                for imp in imports:
+                    assert (
+                        imp not in function_names
+                    ), f"{imp} is a duplicated function name from {extension.name}, this is not allowed."
+                    function_names.add(imp)
+
+    def build(self) -> None:
+        """Builds the actual package."""
+        local_dir = Path(__file__).parent / "helpers"
+        template_file = local_dir / "setup_template.py"
+
+        if self.libraries is None:
+            with open(self.model_folder / Constants.REQUIREMENTS_FILE, "r") as f:
+                requirements_raw = f.readlines()
+                libraries = list(map(lambda x: x.replace("\n", ""), filter(lambda x: "==" in x, requirements_raw)))
+        else:
+            libraries = self.libraries
+        assert len(libraries) > 0, "Using pype, it is assumend there is at least 1 dependency"
+
+        with TemporaryDirectory() as tmp_directory:
+            tmp_directory_path = Path(tmp_directory)
+            tmp_package_dir = tmp_directory_path / self.model_name
+            tmp_model_dir = tmp_package_dir / "outputs"
+            tmp_setup_file = tmp_directory_path / "setup.py"
+            tmp_init_file = tmp_package_dir / "__init__.py"
+
+            tmp_package_dir.mkdir()
+            with open(tmp_init_file, "w") as f:
+                f.write("")
+
+            assert isinstance(self.extensions, list)
+            for extension in self.extensions:
+                print(extension, extension.name)
+                extension.extend(self.model_name, libraries, tmp_package_dir)
+
+            requirements = '", "'.join(libraries)
+
+            with open(template_file, "r") as f:
+                setup_file_raw = f.read()
+                setup_file_formatted = (
+                    setup_file_raw.replace("{install_requires}", requirements)
+                    .replace("{package_name}", self.model_name)
+                    .replace("{version}", self.version)
+                )
+
+            copytree(self.model_folder, tmp_model_dir)
+            with open(tmp_setup_file, "w") as f:
+                f.write(setup_file_formatted)
+
+            output_wheel_file = str(self.output_wheel_file)
+            with switch_directory(tmp_directory):
+                subprocess.check_output(
+                    [sys.executable, "setup.py", "bdist_wheel", "--dist-dir", output_wheel_file]
+                ).decode()
diff --git a/pype/base/deploy/wheel/extensions.py b/pype/base/deploy/wheel/extensions.py
new file mode 100644
index 0000000..4fdb330
--- /dev/null
+++ b/pype/base/deploy/wheel/extensions.py
@@ -0,0 +1,100 @@
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from shutil import copyfile
+
+
+@dataclass
+class WheelExtension:
+    """An extension to Pypes deployment process in wheel files.
+
+    This class provides standardised configuration to add plugins to Pype wheel files.
+
+    Extensions require 3 parameters:
+        - name: A unique name for the extension. This will be used to split imports and files.
+        - functionality: A list of (path, functions) pairs. These are files to be copied over to any
+            wheel package being created. Please make sure any imports in these files are NOT relative
+            and do not otherwise rely on file structure/ordering. Ideally you should make sure each
+            file is self-contained, only relying on external dependencies.
+        - libraries that should be installed as extra dependencies to make the plugin work.
+
+    All extensions will be placed in a similar fashion inside the wheel file:
+        - <name of the model/package>
+        |
+        --  - <plugin name 1>
+                - files for plugin 1
+            - <plugin name 2>
+            - <plugin name 3>
+    Files from a plugin will be place in a flat manner inside the folder. Duplicate file names will cause issues.
+    """
+
+    name: str
+    functionality: list[tuple[Path, list[str]]]
+    libraries: list[str]
+
+    def __post_init__(self) -> None:
+        """Perform checks on inputs."""
+        self._check_file_name(self.name)
+        for path, _ in self.functionality:
+            self._check_file_name(path.name.replace(".py", ""))
+
+        files = set()
+        function_names = set()
+
+        for path, functions in self.functionality:
+            assert path not in files, f"{path} is a duplicated file name, this is not allowed."
+            files.add(path)
+
+            for func in functions:
+                assert (
+                    func not in function_names
+                ), f"{func} is a duplicated function name from {path}, this is not allowed."
+                function_names.add(func)
+
+    @staticmethod
+    def _check_file_name(name: str) -> None:
+        """Check that the given name is an okay filename for imports.
+
+        Args:
+            name (str): The name to check
+
+        Raises:
+            AssertionError: if the name starts with a non-letter or contains a space or dot.
+        """
+        assert (
+            re.search(r"^[a-zA-Z]", name) is not None
+        ), f"Wheel extension file names need to start with a letter: {name}"
+
+        for character in [" ", "."]:
+            assert character not in name, f"{character} is not supported in wheel extension names: {name}"
+
+    def extend(
+        self,
+        package_name: str,
+        libraries: list[str],
+        wheel_package_dir: Path,
+    ) -> None:
+        """Extends the given wheel file with this extension.
+
+        Args:
+            package_name (str): The name of the main package.
+            libraries (list[str]): The current set of libraries to be installed.
+            wheel_package_dir (Path): The Path to the main package.
+        """
+        libraries.extend(self.libraries)
+
+        init_file_path = wheel_package_dir / "__init__.py"
+        subfolder = wheel_package_dir / self.name
+        subfolder.mkdir(exist_ok=False)
+
+        with open(subfolder / "__init__.py", "w") as f:
+            f.write("")
+
+        with open(init_file_path, "a") as init_file:
+            for file, imports in self.functionality:
+                copyfile(file, subfolder / file.name)
+                import_file_name = file.name.replace(".py", "")
+                imports_str = ", ".join(imports)
+
+                init_line = f"from {package_name}.{self.name}.{import_file_name} import {imports_str}\n"
+                init_file.write(init_line)
diff --git a/pype/base/deploy/wheel/helpers/setup_template.py b/pype/base/deploy/wheel/helpers/setup_template.py
new file mode 100644
index 0000000..4d0a7a5
--- /dev/null
+++ b/pype/base/deploy/wheel/helpers/setup_template.py
@@ -0,0 +1,31 @@
+import os
+from pathlib import Path
+
+from setuptools import setup
+
+
+def package_files(directory: Path | str) -> list[str]:
+    """Find all files from the output directory of an Experiment run to include.
+
+    Args:
+        directory (Path | str): The main output directory to search.
+
+    Returns:
+        list[str]: The files to include.
+    """
+    paths = []
+    for (path, _, filenames) in os.walk(directory):
+        for filename in filenames:
+            paths.append(os.path.join("..", path, filename))
+    return paths
+
+
+extra_files = package_files(Path().parent.absolute())
+
+setup(
+    name="{package_name}",
+    install_requires=["{install_requires}"],
+    version="{version}",
+    include_package_data=True,
+    package_data={"": extra_files},
+)
diff --git a/pype/base/deploy/wheel/helpers/wheel_main.py b/pype/base/deploy/wheel/helpers/wheel_main.py
new file mode 100644
index 0000000..f1d9dc2
--- /dev/null
+++ b/pype/base/deploy/wheel/helpers/wheel_main.py
@@ -0,0 +1,14 @@
+from pathlib import Path
+
+from pype.base.deploy import Inferencer
+
+
+def load_model() -> Inferencer:
+    """Wrapper function to load an Inferencer from a fixed location in a package.
+
+    Returns:
+        Inferencer: The Inferencer contained in this package.
+    """
+    folder = Path(__file__).parent.parent / "outputs"
+    print(f"Loading inferencer from {str(folder)}")
+    return Inferencer.from_folder(folder)
diff --git a/pype/base/deploy/wheel/wheel.py b/pype/base/deploy/wheel/wheel.py
new file mode 100644
index 0000000..344bd4b
--- /dev/null
+++ b/pype/base/deploy/wheel/wheel.py
@@ -0,0 +1,78 @@
+import subprocess
+import sys
+from pathlib import Path
+from shutil import copyfile, copytree
+from tempfile import TemporaryDirectory
+
+from pype.base.constants import Constants
+from pype.base.utils.workspace import switch_directory
+
+
+def make_wheel(
+    model_folder: Path,
+    model_name: str,
+    version: str,
+    output_wheel_file: Path | str | None = None,
+    libraries: list[str] | None = None,
+) -> None:
+    """Turns an output folder from an Experiment into a wheel file.
+
+    This should allow you to save models into shippable packages
+    quite easily.
+    # TODO: add proper tests
+
+    Args:
+        model_folder (Path): Path to the output folder generated by an Experiment
+        model_name (str): The name of the model. Will be used as package name. Should follow
+            proper package naming conventions.
+        version (str): Model/package version.
+        output_wheel_file (Path | str | None, optional): The output directory. Defaults to "wheel_output"
+            in the current directory.
+        libraries (list[str] | None, optional): A list of (versioned) libraries that are used as
+            dependencies. By default we use the requirements.txt file from the experiment run.
+    """
+    if output_wheel_file is None:
+        output_wheel_file = Path("").parent.absolute() / "wheel_output"
+
+    local_dir = Path(__file__).parent / "helpers"
+    template_file = local_dir / "setup_template.py"
+    main_file = local_dir / "wheel_main.py"
+
+    if libraries is None:
+        with open(model_folder / Constants.REQUIREMENTS_FILE, "r") as f:
+            requirements_raw = f.readlines()
+            assert len(requirements_raw) > 0, "Using pype, it is assumend there is at least 1 dependency"
+            requirements = '", "'.join(filter(lambda x: "==" in x, requirements_raw)).replace("\n", "")
+    else:
+        assert len(libraries) > 0, "Using pype, it is assumend there is at least 1 dependency"
+        requirements = '", "'.join(libraries)
+
+    with open(template_file, "r") as f:
+        setup_file_raw = f.read()
+        setup_file_formatted = (
+            setup_file_raw.replace("{install_requires}", requirements)
+            .replace("{package_name}", model_name)
+            .replace("{version}", version)
+        )
+
+    with TemporaryDirectory() as tmp_directory:
+        tmp_directory_path = Path(tmp_directory)
+        tmp_model_dir = tmp_directory_path / model_name / "outputs"
+        tmp_setup_file = tmp_directory_path / "setup.py"
+        tmp_main_file = tmp_directory_path / model_name / "main.py"
+
+        copytree(model_folder, tmp_model_dir)
+        copyfile(main_file, tmp_main_file)
+
+        with open(tmp_setup_file, "w") as f:
+            f.write(setup_file_formatted)
+        with open(tmp_directory_path / "__init__.py", "w") as f:
+            f.write("")
+        with open(tmp_directory_path / model_name / "__init__.py", "w") as f:
+            f.write("from .main import load_model, load_app")
+
+        output_wheel_file = str(output_wheel_file)
+        with switch_directory(tmp_directory):
+            subprocess.check_output(
+                [sys.executable, "setup.py", "bdist_wheel", "--dist-dir", output_wheel_file]
+            ).decode()
diff --git a/pype/base/evaluate/evaluator.py b/pype/base/evaluate/evaluator.py
index 4ffa8f0..bc3f173 100644
--- a/pype/base/evaluate/evaluator.py
+++ b/pype/base/evaluate/evaluator.py
@@ -46,3 +46,15 @@ def evaluate(
         output_data = data.get_all(model.outputs)
 
         return {name: func(*output_data, *predictions) for name, func in self.functions.items()}
+
+    def __str__(self, indents: int = 0) -> str:
+        """Create string representation of this Evaluator.
+
+        Args:
+            indents (int, optional): The number of preceding tabs. Defaults to 0.
+
+        Returns:
+            str: A string representation of this Evaluator.
+        """
+        tabs = "\t" * indents
+        return "\n".join([f"{tabs}{name}: {func.__name__}" for name, func in self.functions.items()])
diff --git a/pype/base/experiment/experiment.py b/pype/base/experiment/experiment.py
index 25104a5..396c528 100644
--- a/pype/base/experiment/experiment.py
+++ b/pype/base/experiment/experiment.py
@@ -1,4 +1,7 @@
-import warnings
+import json
+import os
+import subprocess
+import sys
 from argparse import ArgumentParser
 from logging import getLogger
 from pathlib import Path
@@ -12,7 +15,7 @@
 from pype.base.model import Model
 from pype.base.pipeline import Pipeline
 from pype.base.pipeline.type_checker import TypeCheckerPipe
-from pype.base.serialiser.serialiser import Serialiser
+from pype.base.serialiser import JoblibSerialiser, Serialiser
 from pype.base.utils.parsing import get_args_for_prefix
 
 
@@ -24,11 +27,11 @@ def __init__(
         pipeline: Pipeline,
         evaluator: Evaluator,
         logger: ExperimentLogger,
-        serialiser: Serialiser,
-        output_folder: Path | str,
         input_type_checker: TypeCheckerPipe,
         output_type_checker: TypeCheckerPipe,
-        additional_files_to_store: list[str] | None = None,
+        serialiser: Serialiser | None = None,
+        output_folder: Path | str = "outputs",
+        additional_files_to_store: list[str | Path] | None = None,
         parameters: dict[str, Any] | None = None,
     ):
         """The core of the pype library: run a standardised ML experiment with the given parameters.
@@ -46,28 +49,31 @@ def __init__(
             evaluator (Evaluator): The evaluator to test how good your Model performs.
             logger (ExperimentLogger): The experiment logger to make sure you record how well your experiment worked,
                 and log any artifacts such as the trained model.
-            serialiser (Serialiser): The serialiser to serialise any Python objects (expect the Model).
-            output_folder: (Path | str): The output folder to log artifacts to.
+            serialiser (Serialiser | None, optional): The serialiser to serialise any Python objects (expect the Model).
+                Defaults to a joblib serialiser.
+            output_folder: (Path | str): The output folder to log artifacts to. Defaults to "outputs".
             input_type_checker: (TypeCheckerPipe): A type checker for all input data. Will be used to verify incoming
                 data and standardise the order of data. Will be used later to help serialise/deserialise data.
             output_type_checker: (TypeCheckerPipe): A type checker for all output data. Will be used to verify outgoing
                 data and standardise the order of data. Will be used later to help serialise/deserialise data.
-            additional_files_to_store (list[str] | None, optional): Extra files to store, such as python files.
+            additional_files_to_store (list[str | Path] | None, optional): Extra files to store, such as python files.
                 Defaults to no extra files (None).
             parameters (dict[str, Any] | None, optional): Any parameters to log as part of this experiment.
                 Defaults to None.
         """
         assert "train" in data_sources, "Must provide a 'train' entry in the data_sources dictionary."
+        self.logger = getLogger(__name__)
+        if serialiser is None:
+            serialiser = JoblibSerialiser()
+
         if additional_files_to_store is None:
             additional_files_to_store = []
         if parameters is None:
             parameters = {}
-            warnings.warn(
-                """
-It is highly recommended to provide the parameters used to initialise your
+            self.logger.warning(
+                """It is highly recommended to provide the parameters used to initialise your
 run here for logging purposes. Consider using the `from_command_line` or
-`from_dictionary` initialisation methods
-                """
+`from_dictionary` initialisation methods"""
             )
         if isinstance(output_folder, str):
             output_folder = Path(output_folder)
@@ -79,7 +85,6 @@ def __init__(
         self.input_type_checker = input_type_checker
         self.output_type_checker = output_type_checker
 
-        self.logger = getLogger(__name__)
         self.experiment_logger = logger
         self.parameters = parameters
         self.serialiser = serialiser
@@ -111,7 +116,7 @@ def run(self) -> dict[str, dict[str, str | float | int | bool]]:
             self.model.fit(transformed["train"])
 
             self.logger.info("Evaluate model")
-            metrics = {name: self.evaluator.evaluate(self.model, data) for name, data in datasets.items()}
+            metrics = {name: self.evaluator.evaluate(self.model, data) for name, data in transformed.items()}
 
             self.logger.info("Create output type checker")
             predicted_train = self.model.transform(transformed["train"])
@@ -134,17 +139,92 @@ def run(self) -> dict[str, dict[str, str | float | int | bool]]:
             )
             self.experiment_logger.log_parameters(self.parameters)
 
-            for extra_file in self.additional_files_to_store:
-                self.experiment_logger.log_file(extra_file)
+            # log requirements.txt
+            self._log_requirements()
 
+            # extra py files
+            self._log_extra_files()
             self.logger.info("Done")
         return metrics
 
+    def _log_requirements(self) -> None:
+        python_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+        version_info = {
+            "python_version": python_version,
+            "major": sys.version_info.major,
+            "minor": sys.version_info.minor,
+            "micro": sys.version_info.micro,
+        }
+        version_file = self.output_folder / Constants.PYTHON_VERSION_FILE
+        with open(version_file, "w") as f:
+            json.dump(version_info, f)
+        self.experiment_logger.log_file(version_file)
+
+        reqs = subprocess.check_output([sys.executable, "-m", "pip", "freeze"]).decode()
+        requirements_file = self.output_folder / Constants.REQUIREMENTS_FILE
+        with open(requirements_file, "w") as f:
+            f.write(reqs)
+        self.experiment_logger.log_file(requirements_file)
+
+    def _log_extra_files(self) -> None:
+        """Logs the extra files for an experiment, as specified in the constructor."""
+        paths_to_log = []
+        self.logger.info("Log extra files")
+        for extra_file in self.additional_files_to_store:
+            relative_path = Path(extra_file).relative_to(os.getcwd())
+            self.experiment_logger.log_local_file(relative_path, self.output_folder / relative_path)
+            paths_to_log.append(str(relative_path))
+
+        self.logger.info("Log `extra files`-file")
+        extra_files_file = self.output_folder / Constants.EXTRA_FILES
+        with open(extra_files_file, "w") as f:
+            data = {"paths": paths_to_log}
+            json.dump(data, f)
+        # Make sure the file path is closed before logging anything, to make sure all writes have flushed.
+        self.experiment_logger.log_file(extra_files_file)
+
     def _create_output_folders(self) -> None:
         of = self.output_folder
         of.mkdir(exist_ok=True, parents=True)
         (of / Constants.MODEL_FOLDER).mkdir(exist_ok=True)
 
+    def copy(self, parameters: dict[str, Any], seed: int = 1) -> "Experiment":
+        """Create a fresh copy of this Experiment.
+
+        The Model & Pipeline will be recreated, so any trained versions will be
+        re-initialised.
+
+        Args:
+            parameters (dict[str, Any]): New parameters for the Model and Pipeline
+                to be re-initialised with.
+            seed (int, optional): Training seed. Defaults to 1.
+
+        Returns:
+            Experiment: _description_
+        """
+        model_class = self.model.__class__
+
+        model_args = get_args_for_prefix("model__", parameters)
+        model = model_class(seed=seed, inputs=self.model.inputs, outputs=self.model.outputs, **model_args)
+
+        pipeline_args = get_args_for_prefix("pipeline__", parameters)
+        new_pipeline = self.pipeline.copy(pipeline_args)
+
+        return Experiment(
+            # We might want to also deep copy data_sources, evaluator, type checkers, and experiment_logger.
+            data_sources=self.data_sources,
+            model=model,
+            pipeline=new_pipeline,
+            evaluator=self.evaluator,
+            logger=self.experiment_logger,
+            serialiser=self.serialiser,
+            output_folder=self.output_folder,
+            input_type_checker=self.input_type_checker,
+            output_type_checker=self.output_type_checker,
+            additional_files_to_store=self.additional_files_to_store,
+            parameters=parameters,
+        )
+
     @classmethod
     def from_dictionary(
         cls,
@@ -154,14 +234,14 @@ def from_dictionary(
         evaluator: Evaluator,
         logger: ExperimentLogger,
         serialiser: Serialiser,
-        output_folder: Path | str,
         input_type_checker: TypeCheckerPipe,
         output_type_checker: TypeCheckerPipe,
         model_inputs: list[str],
         model_outputs: list[str],
         parameters: dict[str, Any],
+        output_folder: Path | str = "outputs",
         seed: int = 1,
-        additional_files_to_store: list[str] | None = None,
+        additional_files_to_store: list[str | Path] | None = None,
     ) -> "Experiment":
         """Creates an Experiment from a dictionary with parameters.
 
@@ -184,7 +264,7 @@ def from_dictionary(
             seed (int): The RNG seed to ensure reproducability.
             parameters (dict[str, Any] | None, optional): Any parameters to log as part of this experiment.
                 Defaults to None.
-            additional_files_to_store (list[str] | None, optional): Extra files to store, such as python files.
+            additional_files_to_store (list[str | Path] | None, optional): Extra files to store, such as python files.
                 Defaults to no extra files (None).
 
         Returns:
@@ -197,13 +277,13 @@ def from_dictionary(
         pipeline.reinitialise(pipeline_args)
 
         return Experiment(
-            data_sources,
-            model,
-            pipeline,
-            evaluator,
-            logger,
-            serialiser,
-            output_folder,
+            data_sources=data_sources,
+            model=model,
+            pipeline=pipeline,
+            evaluator=evaluator,
+            logger=logger,
+            serialiser=serialiser,
+            output_folder=output_folder,
             input_type_checker=input_type_checker,
             output_type_checker=output_type_checker,
             additional_files_to_store=additional_files_to_store,
@@ -219,13 +299,14 @@ def from_command_line(
         evaluator: Evaluator,
         logger: ExperimentLogger,
         serialiser: Serialiser,
-        output_folder: Path | str,
         input_type_checker: TypeCheckerPipe,
         output_type_checker: TypeCheckerPipe,
         model_inputs: list[str],
         model_outputs: list[str],
+        output_folder: Path | str = "outputs",
         seed: int = 1,
-        additional_files_to_store: list[str] | None = None,
+        additional_files_to_store: list[str | Path] | None = None,
+        fixed_arguments: dict[str, Any] | None = None,
     ) -> "Experiment":
         """Automatically initialises an Experiment from command line arguments.
 
@@ -248,8 +329,13 @@ def from_command_line(
             model_inputs: (list[str]): Input dataset names to the model.
             model_outputs: (list[str]): Output dataset names to the model.
             seed (int): The RNG seed to ensure reproducability.
-            additional_files_to_store (list[str] | None, optional): Extra files to store, such as python files.
+            additional_files_to_store (list[str | Path] | None, optional): Extra files to store, such as python files.
                 Defaults to no extra files (None).
+            fixed_arguments (dict[str, Any] | None, optional): Arguments that won't be read from command line.
+                Useful to pass complex objects that you don't want to optimize. Think of classes, loss functions,
+                metrics, etc. Any value in this dict will overwrite values from the command line. These need to be
+                presented in the same `model__<arg>` or `pipeline__<step>__<arg>` format as the command line arguments.
+
 
         Returns:
             Experiment: An Experiment created with the given parameters from the command line.
@@ -257,13 +343,17 @@ def from_command_line(
         arg_parser = cls._get_cmd_args(model_class, pipeline)
         parsed_args, _ = arg_parser.parse_known_args()
 
+        arguments = parsed_args.__dict__
+        if fixed_arguments is not None:
+            arguments.update(fixed_arguments)
+
         return cls.from_dictionary(
-            data_sources,
-            model_class,
-            pipeline,
-            evaluator,
-            logger,
-            serialiser,
+            data_sources=data_sources,
+            model_class=model_class,
+            pipeline=pipeline,
+            evaluator=evaluator,
+            logger=logger,
+            serialiser=serialiser,
             output_folder=output_folder,
             input_type_checker=input_type_checker,
             output_type_checker=output_type_checker,
@@ -281,3 +371,24 @@ def _get_cmd_args(cls, model_class: Type[Model], pipeline: Pipeline) -> Argument
         model_class.get_parameters(parser)
         add_args_to_parser_for_pipeline(parser, pipeline)
         return parser
+
+    def __str__(self) -> str:
+        """Create string representation of this Experiment.
+
+        Returns:
+            str: A string representation of this Experiment.
+        """
+        data_part = "datasets:\n" + "\n\t".join(
+            [
+                f"""\t{name}:
+{ds_source.__str__(indents=2)}"""
+                for name, ds_source in self.data_sources.items()
+            ]
+        )
+
+        pipeline_part = "pipeline:\n" + self.pipeline.__str__(indents=1)
+        model_part = f"model:\n\t{str(self.model)}"
+        logger_part = f"logger:\n\t{str(self.experiment_logger)}"
+        evaluator_part = f"evaluator:\n{self.evaluator.__str__(indents=1)}"
+
+        return "\n\n".join([data_part, pipeline_part, model_part, logger_part, evaluator_part])
diff --git a/pype/base/logger/experiment_logger.py b/pype/base/logger/experiment_logger.py
index 5fa3161..ee8f885 100644
--- a/pype/base/logger/experiment_logger.py
+++ b/pype/base/logger/experiment_logger.py
@@ -1,19 +1,20 @@
+import shutil
 from abc import ABC, abstractmethod
 from pathlib import Path
 from types import TracebackType
-from typing import Any, Dict, Union
+from typing import Any
 
 from pype.base.model.model import Model
 from pype.base.serialiser import Serialiser
 
 
 class ExperimentLogger(ABC):
-    # TODO: make a LocalLogger.
-
+    @abstractmethod
     def __enter__(self) -> None:
         """Start the experiment."""
         return
 
+    @abstractmethod
     def __exit__(
         self,
         exc_type: type[BaseException] | None,
@@ -23,7 +24,7 @@ def __exit__(
         """Stop the experiment."""
         return
 
-    def log_metrics(self, dataset_name: str, metrics: Dict[str, Union[float, int, str | bool]]) -> None:
+    def log_metrics(self, dataset_name: str, metrics: dict[str, float | int | str | bool]) -> None:
         """Logs the metrics for a given dataset.
 
         Metrics will be stored by default as `<dataset_name>_<metric_name>`.
@@ -74,14 +75,28 @@ def log_artifact(self, file: str | Path, serialiser: Serialiser, object: Any) ->
 
         Args:
             file (str | Path): The file to serialise to and subsequently log.
-            serialiser (Serialiser): Serialiser used to serialise the given object, if any.
+            serialiser (Serialiser): Serialiser used to serialise the given object.
             object (Any): The object to log. This object will be serialised to the given file
                 using the serialiser and logged using `log_file`
         """
-        if object is not None and serialiser is not None:
-            serialiser.serialise(object, file)
+        serialiser.serialise(object, file)
         self.log_file(file)
 
+    def log_local_file(self, file: str | Path, output_target: str | Path) -> None:
+        """Logs a given local file as part of an experiment.
+
+        First the file is copied to the output directory, then
+        log_file is called to make it part of the experiment.
+
+        Args:
+            file (str | Path): The file to log.
+            output_target (str | Path): The target location of the file.
+        """
+        output_target = Path(output_target)
+        output_target.parent.mkdir(exist_ok=True)
+        shutil.copy(file, output_target)
+        self.log_file(output_target)
+
     @abstractmethod
     def log_file(self, file: str | Path) -> None:
         """Logs a given file as part of an experiment.
diff --git a/pype/base/logger/local_logger.py b/pype/base/logger/local_logger.py
index e689a7c..4771add 100644
--- a/pype/base/logger/local_logger.py
+++ b/pype/base/logger/local_logger.py
@@ -1,10 +1,19 @@
 from pathlib import Path
+from types import TracebackType
 from typing import Any
 
 from .experiment_logger import ExperimentLogger
 
 
 class LocalLogger(ExperimentLogger):
+    def __enter__(self) -> None:
+        """Start the experiment."""
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> None:
+        """Stop the experiment."""
+
     def __init__(self, min_whitespace: int = 30) -> None:
         """A basic logger which only saves local copies of the data and prints results."""
         super().__init__()
@@ -34,5 +43,5 @@ def log_file(self, file: str | Path) -> None:
         """This method does no additional logging.
 
         Args:
-            file (str | Path): The file to log. No actual logging is done.
+            file (str | Path): The file to log.
         """
diff --git a/pype/base/model/model.py b/pype/base/model/model.py
index 2c4a335..3aecbaa 100644
--- a/pype/base/model/model.py
+++ b/pype/base/model/model.py
@@ -1,5 +1,5 @@
 import json
-from abc import ABC, abstractmethod
+from abc import ABC, abstractclassmethod, abstractmethod
 from argparse import ArgumentParser
 from pathlib import Path
 from typing import Generic, Tuple, TypeVar
@@ -97,8 +97,7 @@ def load(cls, folder: str | Path) -> "Model":
         model_class = joblib_serialiser.deserialise(folder / Constants.MODEL_CLASS_FILE)
         return model_class._load(folder, lists["inputs"], lists["outputs"])
 
-    @classmethod
-    @abstractmethod
+    @abstractclassmethod
     def _load(cls, folder: Path, inputs: list[str], outputs: list[str]) -> "Model":
         """Loads a model from file into this Model.
 
@@ -143,3 +142,11 @@ def transform(self, data: DataSet) -> DataSet:
     @abstractmethod
     def _transform(self, *data: Data) -> Tuple[Data] | Data:
         raise NotImplementedError
+
+    def __str__(self) -> str:
+        """Create string representation of this Model.
+
+        Returns:
+            str: A string representation of this Model.
+        """
+        return f"{type(self).__name__}: {self.inputs} -> {self.outputs}"
diff --git a/pype/base/pipeline/__init__.py b/pype/base/pipeline/__init__.py
index 784db8a..0538d30 100644
--- a/pype/base/pipeline/__init__.py
+++ b/pype/base/pipeline/__init__.py
@@ -1,3 +1,4 @@
 from .operator import Operator
 from .pipe import Pipe
 from .pipeline import Pipeline
+from .type_checker import DataModel, DataSetModel, TypeChecker, TypeCheckerPipe
diff --git a/pype/base/pipeline/pipe.py b/pype/base/pipeline/pipe.py
index 8186217..37f0a7c 100644
--- a/pype/base/pipeline/pipe.py
+++ b/pype/base/pipeline/pipe.py
@@ -27,6 +27,7 @@ def __init__(
             fit_inputs: (list[str] | None): optional additional arguments to fit().
                 Will not be used in transform().
         """
+        assert "__" not in name, "Pipe names cannot contain the string `__`"
         if fit_inputs is None:
             fit_inputs = []
         if kw_args is None:
@@ -106,3 +107,36 @@ def reinitialise(self, args: Dict[str, Any]) -> None:
             args (Dict[str, Any]): The dictionary of arguments to use in re-initialisation.
         """
         self.operator = self.operator_class(**args)
+
+    def copy(self, args: Dict[str, Any] | None = None) -> "Pipe":
+        """Create a copy of this Pipe's Operator given a dict of arguments.
+
+        Args:
+            args (Dict[str, Any] | None): The dictionary of arguments to use in re-initialisation.
+                If set to None, we'll use the same arguments as before.
+
+        Returns:
+            A new copy of this Pipe, with a new Operator object.
+        """
+        if args is None:
+            args = self.args
+
+        return Pipe(
+            self.name,
+            self.operator_class,
+            inputs=self.inputs,
+            outputs=self.outputs,
+            kw_args=args,
+            fit_inputs=self.fit_inputs,
+        )
+
+    def __str__(self) -> str:
+        """Create string representation of this Pipe.
+
+        Returns:
+            str: A string representation of this Pipe.
+        """
+        fit_section = f" (+ {self.fit_inputs})" if len(self.fit_inputs) > 0 else ""
+        input_output_section = f"{self.inputs}{fit_section} -> {self.outputs}"
+
+        return f"Pipe `{self.name}`, {input_output_section}: {self.operator}"
diff --git a/pype/base/pipeline/pipeline.py b/pype/base/pipeline/pipeline.py
index 518b24e..a778754 100644
--- a/pype/base/pipeline/pipeline.py
+++ b/pype/base/pipeline/pipeline.py
@@ -90,6 +90,30 @@ def reinitialise(self, args: Dict[str, Any]) -> None:
                 if len(pipe_args) > 0:
                     pipe.reinitialise(pipe_args)
 
+    def copy(self, args: Dict[str, Any] | None = None) -> "Pipeline":
+        """Create a copy of this Pipeline using the given dictionary.
+
+        Args:
+            args (Dict[str, Any]): Dictionary containing new arguments.
+                The keys of the `args` dict should be of the form `<name>__<arg_name>`:
+                - name: the name of the Pipe.
+                - arg_name: the argument name to use.
+                The values will be used as argument parameters.
+        """
+        result = Pipeline([])
+        for pipe in self.pipes:
+            if isinstance(pipe, Pipeline):
+                extra_pipeline = pipe.copy(args)
+                result += extra_pipeline
+            else:
+                if args is None:
+                    pipe_args = {}
+                else:
+                    pipe_args = get_args_for_prefix(f"{pipe.name}__", args)
+                new_pipe = pipe.copy(pipe_args)
+                result += new_pipe
+        return result
+
     def __iter__(self) -> "Pipeline":
         """Prepares this object for iteration using next().
 
@@ -193,3 +217,22 @@ def __len__(self) -> int:
             else:
                 result += 1
         return result
+
+    def __str__(self, indents: int = 0) -> str:
+        """Create string representation of this Pipeline.
+
+        Args:
+            indents (int, optional): The number of preceding tabs. Defaults to 0.
+
+        Returns:
+            str: A string representation of this Pipeline.
+        """
+        tabs = "\t" * indents
+        output = [tabs + "Pipeline:"]
+        indents += 1
+        for pipe in self.pipes:
+            if isinstance(pipe, Pipe):
+                output.append(tabs + str(pipe))
+            else:
+                output.append(pipe.__str__(indents))
+        return "\n".join(output)
diff --git a/pype/base/pipeline/type_checker.py b/pype/base/pipeline/type_checker.py
index 14f8b75..787519b 100644
--- a/pype/base/pipeline/type_checker.py
+++ b/pype/base/pipeline/type_checker.py
@@ -1,5 +1,5 @@
-import warnings
-from abc import ABC, abstractmethod
+import logging
+from abc import ABC, abstractclassmethod, abstractmethod
 from typing import Any
 
 from pydantic import BaseModel, create_model
@@ -19,8 +19,7 @@ def convert(self) -> Any:
             Any: The converted data.
         """
 
-    @classmethod
-    @abstractmethod
+    @abstractclassmethod
     def to_model(cls, data: Any) -> "DataModel":
         """Convert the Model to actual data (e.g. numpy or pandas).
 
@@ -153,9 +152,7 @@ def __init__(self, input_names: list[str], type_checker_classes: list[tuple[type
         """
         super().__init__()
         self.input_names = input_names
-
         self.type_checker_classes = type_checker_classes
-
         self.type_checkers: dict[str, TypeChecker] = {}
 
     def fit(self, *data: Data) -> "Operator":
@@ -169,7 +166,8 @@ def fit(self, *data: Data) -> "Operator":
             type_checker_class = self._get_type_checker(dataset)
 
             if type_checker_class is None:
-                warnings.warn(f"{ds_name} has no supported type checker!")
+                logger = logging.getLogger(__name__)
+                logger.warning(f"{ds_name} has no supported type checker!")
             else:
                 checker = type_checker_class()
                 checker.fit(dataset)
diff --git a/pype/base/setup.py b/pype/base/setup.py
index 37287f2..7d54352 100644
--- a/pype/base/setup.py
+++ b/pype/base/setup.py
@@ -1,12 +1,14 @@
 from setuptools import find_namespace_packages, setup
 
-dev_deps = ["pre-commit"]
-test_deps = ["pytest"]
-deps: list[str] = ["docstring_parser==0.14.1", "pydantic==1.9.1", "joblib==1.1.0", "uvicorn==0.18.2", "fastapi==0.79.0"]
+dev_deps = ["pre-commit", "build==0.8.0", "pypiserver==1.5.1", "twine==4.0.1"]
+test_deps = ["pytest", "pytest-cov"]
+deps = ["docstring_parser>=0.14.1", "pydantic>=1.9.1", "joblib>=1.1.0"]
+strict_deps = [s.replace(">=", "==") for s in deps]
 
 setup(
     name="pype-base",
     install_requires=deps,
-    extras_require={"dev": deps + dev_deps + test_deps, "test": deps + test_deps},
+    extras_require={"dev": strict_deps + dev_deps + test_deps, "test": strict_deps + test_deps, "strict": strict_deps},
     packages=find_namespace_packages(include=["pype.*"]),
+    version="0.1.0",
 )
diff --git a/pype/base/utils/__init__.py b/pype/base/utils/__init__.py
index 9a43f58..0657a42 100644
--- a/pype/base/utils/__init__.py
+++ b/pype/base/utils/__init__.py
@@ -1,1 +1,2 @@
 from .parsing import get_args_for_prefix
+from .workspace import switch_directory, switch_workspace
diff --git a/pype/base/utils/workspace.py b/pype/base/utils/workspace.py
new file mode 100644
index 0000000..8a031ab
--- /dev/null
+++ b/pype/base/utils/workspace.py
@@ -0,0 +1,130 @@
+import importlib
+import logging
+import os
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+from types import ModuleType
+from typing import Generator, Set
+
+
+@contextmanager
+def switch_directory(directory: str | Path) -> Generator:
+    """Allows you to temporarily switch Python over to a different directory.
+
+    This changes the following:
+        - The current working directory will change to `target_workspace`
+
+    These changes will be undone once the context manager ends.
+
+    Args:
+        target_workspace (str | Path): The directory to switch to.
+    """
+    # create backups of the workspace and path.
+    target_directory = Path(directory)
+    old_workspace = os.getcwd()
+    os.chdir(target_directory)
+    # Yield so the normal code in the managed context can run.
+    yield
+
+    os.chdir(old_workspace)
+
+
+@contextmanager
+def switch_workspace(target_workspace: str | Path, extra_files: list[str | Path] | None = None) -> Generator:
+    """Allows you to temporarily switch Python over to using a different working directory.
+
+    This changes the following:
+        - The current working directory will change to `target_workspace`
+        - `target_workspace` is added to the system path
+        - Any files provided in `extra_files` will be assumed to follow a relative path
+            from `target_workspace` to python files that will be imported. Variables
+            (incl. classes and functions) defined in those files will be imported into
+            the __main__ module, and the files will be added to the path. We do not overwrite variables
+            already existing in __main__.
+
+    These changes will be undone once the context manager ends.
+
+    Args:
+        target_workspace (str | Path): The directory to switch to.
+        extra_files (list[str | Path] | None, optional): Any additional files whose contents should
+            be added to the current python program. Defaults to no extra files.
+    """
+    if extra_files is None:
+        extra_files = []
+
+    # create backups of the workspace and path.
+    target_workspace = Path(target_workspace)
+    old_workspace = os.getcwd()
+    old_sys_path = sys.path.copy()
+
+    # Add each file to  the path, and import variables specific to that file into globals.
+    main_lib = importlib.import_module("__main__")
+    main_to_remove = _create_temporary_workspace(target_workspace, extra_files, main_lib)
+    # Yield so the normal code in the managed context can run.
+    yield
+
+    _reset_workspace(old_workspace, old_sys_path, main_lib, main_to_remove)
+
+
+def _reset_workspace(
+    old_workspace: str, old_sys_path: list[str], main_lib: ModuleType, main_to_remove: Set[str]
+) -> None:
+    os.chdir(old_workspace)
+    sys.path = old_sys_path
+    for var in main_to_remove:
+        del main_lib.__dict__[var]
+
+
+def _create_temporary_workspace(
+    target_workspace: Path, extra_files: list[str | Path], main_lib: ModuleType
+) -> Set[str]:
+    """Perform the actual workspace switching.
+
+    Args:
+        target_workspace (Path): The directory to switch to.
+        extra_files (list[str  |  Path]): Any additional files whose contents should
+            be added to the current python program. Empty lists are accepted.
+        main_lib (ModuleType): The __main__ module.
+
+    Returns:
+        Set[str]: The set of variable names that should be removed from __main__ again
+            after the context manager quits.
+    """
+    logger = logging.getLogger(__name__)
+    main_to_remove = set()
+
+    # change directory and path
+    os.chdir(target_workspace)
+    sys.path.append(str(target_workspace))
+    directory_name = target_workspace.name
+
+    for file in extra_files:
+        file = str(file)
+        if not file.endswith(".py"):
+            logger.warning(f"Not adding {file} to path since it is not a python file!")
+            continue
+
+        # add the python file to the path. This makes sure we can import variables from other files.
+        relative_file_name = str(os.path.join(".", file))
+        sys.path.append(relative_file_name)
+
+        # import the file as a module.
+        module_name = f'{directory_name}.{file.replace("/", ".").replace(".py", "")}'
+        lib = importlib.import_module(module_name)
+
+        # find all variables defined in this file, avoiding things like numpy/pandas imports.
+        sub_vars = {
+            name: value
+            for name, value in lib.__dict__.items()
+            if hasattr(value, "__module__") and value.__module__ == module_name
+        }
+
+        # add the variable to the main module. This makes sure we can import variables from the
+        # same 'main' file that was used to run the experiment.
+        # out of safety reasons we only add a variable if it doesn't exist in 'main' yet.
+        for var_name, value in sub_vars.items():
+            if var_name not in main_lib.__dict__:
+                main_to_remove.add(var_name)
+                main_lib.__dict__[var_name] = value
+    return main_to_remove
diff --git a/pype/cmd_tools/__init__.py b/pype/cmd_tools/__init__.py
deleted file mode 100644
index d5d45a6..0000000
--- a/pype/cmd_tools/__init__.py
+++ /dev/null
@@ -1,26 +0,0 @@
-"""Only intended for command line tools."""
-from argparse import ArgumentParser
-
-from .dev_tools import run_dev_install
-
-
-def main() -> None:
-    """Main console entrypoint for pype.
-
-    Raises:
-        ValueError: If the subcommand is not allowed. Current values are:
-            - dev-install
-    """
-    allowed_values = ["dev-install"]
-
-    parser = ArgumentParser()
-    parser.add_argument("command", type=str, help=f"The command to run. Allowed values: {allowed_values}")
-
-    parsed, _ = parser.parse_known_args()
-    command = parsed.command
-
-    if command not in allowed_values:
-        raise ValueError(f"`{command}` is not a valid pype command")
-
-    if command == "dev-install":
-        run_dev_install(parser)
diff --git a/pype/cmd_tools/dev_tools.py b/pype/cmd_tools/dev_tools.py
deleted file mode 100644
index 2b687fc..0000000
--- a/pype/cmd_tools/dev_tools.py
+++ /dev/null
@@ -1,61 +0,0 @@
-"""Only intended for development command line tools."""
-
-import re
-import subprocess
-from argparse import ArgumentParser, BooleanOptionalAction
-from pathlib import Path
-
-from setuptools import find_namespace_packages
-
-
-def run_dev_install(parser: ArgumentParser) -> None:
-    """Runs the development installation command.
-
-    This command will install all sub-namespace packages in
-    editable mode, with dependencies.
-
-    The base package will be installed first. Other packages
-    requiring prioritisation can be configured in the code
-    of this function.
-
-    Args:
-        parser (ArgumentParser): Argument parser from the top-level entrypoint.
-    """
-    # please ensure this is properly sorted!
-    prioritised_packages = ["pype.base"]
-
-    parser.add_argument("--editable", default=True, action=BooleanOptionalAction)
-    parsed = parser.parse_args()
-    use_editable = parsed.editable
-
-    base_path = Path(__file__).parent.parent.parent.absolute()
-    packages = find_namespace_packages(str(base_path), include=["pype.*"], exclude=["pype.cmd_tools", "pype.*.*"])
-
-    print("PREPARING TO INSTALL PYPE IN DEV MODE\nInstalling the following packages:\n\t", packages)
-
-    # Installing prioritised packages first.
-    for package in prioritised_packages:
-        assert package in packages, f"Prioritised package was not found: {package}"
-        _install_package(package, use_editable, base_path)
-
-    for package in packages:
-        if package != "pype.base":
-            _install_package(package, use_editable, base_path)
-
-
-def _install_package(package: str, use_editable: bool, base_path: Path) -> None:
-    """Runs an installation command for the given sub-namespace package.
-
-    Args:
-        package (str): The sub-namespace package to be installed.
-        use_editable (bool): If an editable version of the package should be installed.
-        base_path (Path): Top-level directory of this reposity.
-    """
-    assert re.match(r"^pype\.[a-z_0-9]+$", package) is not None, f"`{package}` is not a valid package name. Exiting"
-
-    editable_cmd = "-e" if use_editable else ""
-    path = base_path.joinpath(*package.split("."))
-    command = ["pip", "install", editable_cmd, f"{path}[dev]"]
-
-    print(f"\nInstalling:\n\t{package} from {path}\n")
-    subprocess.run(command)
diff --git a/pype/fastapi/__init__.py b/pype/fastapi/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pype/fastapi/deploy/__init__.py b/pype/fastapi/deploy/__init__.py
new file mode 100644
index 0000000..3381cb2
--- /dev/null
+++ b/pype/fastapi/deploy/__init__.py
@@ -0,0 +1,2 @@
+from .app import PypeApp
+from .wheel_extension import FastApiExtension
diff --git a/pype/fastapi/deploy/app.py b/pype/fastapi/deploy/app.py
new file mode 100644
index 0000000..c454cfc
--- /dev/null
+++ b/pype/fastapi/deploy/app.py
@@ -0,0 +1,105 @@
+from dataclasses import dataclass
+from logging import Logger, getLogger
+from pathlib import Path
+from typing import Any
+
+from fastapi.applications import FastAPI
+from fastapi.background import BackgroundTasks
+
+from pype.base.data import DataSink
+from pype.base.data.dataset import DataSet
+from pype.base.deploy.inference import Inferencer
+from pype.base.pipeline.type_checker import DataModel, DataSetModel
+
+
+@dataclass
+class PypeApp:
+    name: str
+    folder: str | Path
+    tracking_servers: dict[str, DataSink] | None = None
+
+    def __post_init__(self):  # type: ignore
+        """Makes sure the folder is an actual Path."""
+        if isinstance(self.folder, str):
+            self.folder = Path(self.folder)
+
+    def create_app(self) -> FastAPI:
+        """Creates a FastAPI app for a saved Pype experiment.
+
+        Returns:
+            FastAPI: The FastAPI server that can be served.
+        """
+        app = FastAPI()
+        inferencer = self._load_model()
+        logger = getLogger(self.name)
+
+        InputType: type[DataSetModel] = inferencer.input_type_checker.get_pydantic_types()
+        OutputType: type[DataSetModel] = inferencer.output_type_checker.get_pydantic_types()
+        self._verify_tracking_servers(InputType, OutputType, logger)
+
+        @app.get("/")
+        async def home_page() -> str:
+            logger.info("Homepage request")
+            return f"Welcome to the Pype FastAPI app for {self.name}"
+
+        @app.post("/predict")
+        async def predict(inputs: InputType, background: BackgroundTasks) -> OutputType:  # type: ignore
+            logger.info("Prediction request")
+            converted: DataSet = inputs.convert()  # type: ignore
+            prediction = inferencer.predict(converted)
+            result = OutputType.to_model(prediction)
+
+            self._handle_tracking(converted, prediction, logger, background)
+            logger.info("Finishing prediction request")
+            return result
+
+        return app
+
+    def _handle_tracking(self, inputs: DataSet, outputs: DataSet, logger: Logger, background: BackgroundTasks) -> None:
+        if self.tracking_servers is not None:
+            logger.info("Logging records request")
+            for name, sink in self.tracking_servers.items():
+                try:
+                    # Try to find the name of the dataset in the tracking servers.
+                    ds: DataModel | None = None
+                    if name in outputs:
+                        ds = outputs[name]
+                    elif name in inputs:
+                        ds = inputs[name]
+
+                    if ds is None:
+                        logger.warning(f"Not storing data to sink `{name}` since no data by that name was found.")
+                    else:
+                        logger.info(f"Storing data to sink `{name}`.")
+                        background.add_task(write_in_background, sink=sink, data=ds)
+
+                except Exception as e:
+                    logger.error(f"Encountered error while sending data to {name}: {str(e)}")
+
+    def _verify_tracking_servers(
+        self,
+        inputs: type[DataSetModel],
+        outputs: type[DataSetModel],
+        logger: Logger,
+    ) -> None:
+        # Since it could be possible someone wants to log an intermediate dataset,
+        # we can't put a hard restriction on logging only input or output data.
+        # We will only log the latest version of a dataset though.
+        if self.tracking_servers is not None:
+            for name in self.tracking_servers:
+                if not (name in inputs.__fields__ or name in outputs.__fields__):
+                    logger.warning(f"No dataset named `{name}` found in the fields of input or output DataSetModels")
+
+    def _load_model(self) -> Inferencer:
+        assert isinstance(self.folder, Path)
+        return Inferencer.from_folder(self.folder)
+
+
+def write_in_background(sink: DataSink, data: Any) -> None:
+    """Background function to write data to a given DataSink.
+
+    Args:
+        sink (DataSink): The Sink to write to.
+        data (Any): The Data to write.
+    """
+    sink.write(data)
diff --git a/pype/fastapi/deploy/wheel/helpers.py b/pype/fastapi/deploy/wheel/helpers.py
new file mode 100644
index 0000000..b3d8556
--- /dev/null
+++ b/pype/fastapi/deploy/wheel/helpers.py
@@ -0,0 +1,26 @@
+from pathlib import Path
+
+from fastapi.applications import FastAPI
+
+from pype.fastapi.deploy import PypeApp
+
+
+def load_app() -> FastAPI:
+    """Wrapper function to create a FastAPI server from a fixed location in this package.
+
+    Based on PypeApp.
+
+    If you create a python file F where this App is loaded:
+    ```
+    from <model name> import load_app
+    app = load_app()
+    ```
+
+    You can host this model using a command like:
+    uvicorn <python package path to file F>:app
+
+    Returns:
+        FastAPI: The FastAPI/PypeApp server used to host the model.
+    """
+    folder = Path(__file__).parent.parent / "outputs"
+    return PypeApp("model", folder=folder).create_app()
diff --git a/pype/fastapi/deploy/wheel_extension.py b/pype/fastapi/deploy/wheel_extension.py
new file mode 100644
index 0000000..6f9ad93
--- /dev/null
+++ b/pype/fastapi/deploy/wheel_extension.py
@@ -0,0 +1,7 @@
+from pathlib import Path
+
+from pype.base.deploy.wheel import WheelExtension
+
+FastApiExtension = WheelExtension(
+    "fastapi", [(Path(__file__).parent / "wheel" / "helpers.py", ["load_app"])], ["pype.fastapi"]
+)
diff --git a/pype/fastapi/setup.py b/pype/fastapi/setup.py
new file mode 100644
index 0000000..19fdcd4
--- /dev/null
+++ b/pype/fastapi/setup.py
@@ -0,0 +1,14 @@
+from setuptools import find_namespace_packages, setup
+
+version = "0.1.0"
+deps: list[str] = [f"pype-base=={version}", "fastapi>=0.79.0"]
+dev_deps = ["uvicorn==0.18.2"]
+strict_deps = [s.replace(">=", "==") for s in deps]
+
+setup(
+    name="pype-fastapi",
+    install_requires=deps,
+    extras_require={"dev": strict_deps + dev_deps, "strict": strict_deps},
+    packages=find_namespace_packages(include=["pype.*"]),
+    version=version,
+)
diff --git a/pype/hyperopt/__init__.py b/pype/hyperopt/__init__.py
new file mode 100644
index 0000000..9d7e05b
--- /dev/null
+++ b/pype/hyperopt/__init__.py
@@ -0,0 +1,1 @@
+from .optimise import optimise_experiment
diff --git a/pype/hyperopt/optimise.py b/pype/hyperopt/optimise.py
new file mode 100644
index 0000000..49edc4d
--- /dev/null
+++ b/pype/hyperopt/optimise.py
@@ -0,0 +1,161 @@
+from typing import Any, Callable
+
+from hyperopt.base import Trials
+from hyperopt.fmin import fmin, space_eval
+from hyperopt.mongoexp import MongoTrials
+
+from pype.base.experiment import Experiment
+
+
+def create_optimisation_function(
+    experiment_template: Experiment,
+    target_metric: tuple[str, str],
+    minimise_target: bool,
+    seed: int,
+) -> Callable[[dict], float | int]:
+    """Create an optimisation function that hyperopt can use to find the best model.
+
+    Args:
+        experiment_template (Experiment): An example Experiment that will be run
+            multiple times to find the best hyperparameters. We will create clean copies
+            of this Experiment for every run.
+        target_metric (tuple[str, str]): The target metric to optimise. Should be of the form
+            (dataset name, metric name). The dataset name needs to match one in the DataSet from the Experiment.
+            The metric name should match a metric from the Evaluator in the Experiment.
+        minimise_target (bool): A boolean indicating if the target metric should be minimised or not.
+        seed (int): The seed used for training.
+
+    Returns:
+        Callable[[dict], float | int]: The optimisation function to be used by hyperopt's fmin
+    """
+
+    def optimise(hyper_params: dict[str, Any]) -> float | int:
+        """Optimise a given Experiment at the give hyperparameters.
+
+        Args:
+            hyper_params (dict[str, Any]): The hyperparameters to initialise
+                a new Experiment with. Make sure these are properly prefixed with
+                model__ and pipeline__<pipe name>__, as you would initialise a
+                regular Experiment from arguments.
+
+        Returns:
+            float | int: The performance of this Experiment on the given parameters
+                on one metric and dataset.
+        """
+        experiment = experiment_template.copy(hyper_params, seed=seed)
+        metrics = experiment.run()
+        dataset, metric = target_metric
+
+        result = metrics[dataset][metric]
+        assert isinstance(result, (float, int)), "The target metric should be a float or integer, please fix this."
+
+        # Since we use fmin, we'll need to multiply our target with -1 if we want to maxisime our target instead.
+        if minimise_target:
+            return result
+        else:
+            return -result
+
+    return optimise
+
+
+def optimise_experiment(
+    experiment_template: Experiment,
+    search_space: dict[str, tuple[str, Callable]],
+    target_metric: tuple[str, str],
+    max_evals: int,
+    minimise_target: bool,
+    trial_type: str = "normal",
+    mongo_url: str | None = None,
+    mongo_exp_key: str = "experiment",
+    trials: Trials | None = None,
+    training_seed: int = 1,
+    **kwargs: Any,
+) -> tuple[int | None, float | None, dict[str, Any], Trials]:
+    """Optimise a pype Experiment given a search space using Hyperopt.
+
+    We attempt to create a clean copy of Experiments between runs to preserve run independency.
+    To help this, please ensure the input/output checkers can be refitted without issue, and
+    the DataSources provide the same data even if called multiple times.
+
+    Args:
+        experiment_template (Experiment): An Experiment that will serve as the template for each run.
+            We try to ensure a new, clean Experiment is created for each run. Please allow this by
+            not sharing important, trained variables between Models or Operators of the same class.
+            Each call to any DataSource provided should also return the same dataset.
+        search_space (dict[str, tuple[str, Callable]]): A search space defined in the standard hyperopt way.
+            We suggest providing this as a dictionary in the form:
+                {<argument name>: (<name>, <hyperopt callable like hp.choice(...)>)}
+            Make sure these are properly prefixed with model__ and pipeline__<pipe name>__, as you would
+            initialise a regular Experiment from arguments.
+        target_metric (tuple[str, str]): The target metric to optimise. Should be of the form
+            (dataset name, metric name). The dataset name needs to match one in the DataSet from the Experiment.
+            The metric name should match a metric from the Evaluator in the Experiment.
+        max_evals (int): The maximum number of runs to perform for this hyperoptimisation search. Note that
+            hyperopt will not create more runs if you provide a trials object with more trials than max_evals.
+        minimise_target (bool): A boolean indicating if the target metric should be minimised or not.
+        trial_type (str): If `trials` is not provided, we use this to create a new Trials object. There are currently 3
+            options available:
+                - 'normal': A regular Trials object, stored locally.
+                - 'mongo': A MongoTrials object, storing trials in a mongodb. You will need to provide `mongo_url` and
+                    `mongo_exp_key` to choose this type of trial.
+                # TODO: SparkTrials integration.
+                - 'spark': A SparkTrials object will be created. Currently not supported yet, WIP. Providing your own
+                    SparkTrials object should still work.
+        mongo_url (str | None): The link to the mongodb that should be used. Only used if `trial_type` == 'mongo', and
+            is required in that case.
+        mongo_exp_key (str): The experiment name in the mongodb that should be used.
+            Only used if `trial_type` == 'mongo'. We encourage you to set this to distinguish between runs.
+            The default is "experiment"
+        trials (Trials | None): An existing Trials object with previous runs of the optimisation. Useful if you want to
+            continue an old optimisation run without losing information. By default a new Trials object will be created.
+        training_seed (int): Seed used to initialise a new Experiment (and run).
+        **kwargs: Additional arguments to be provided to hyperopt's `fmin`.
+
+    Returns:
+        A tuple consisting of:
+            - The best trial index (or None if no good trial was found)
+            - The performance of the best trial (or None if no good trial was found)
+            - The arguments for the best trial.
+            - The Trials object created by this run.
+    """
+    optimise = create_optimisation_function(experiment_template, target_metric, minimise_target, seed=training_seed)
+
+    if not minimise_target:
+        print("Heads up: printed targets will be maximised. In print output, they will be multiplied with -1.")
+
+    # initialise a trials object, if needed.
+    if trials is None:
+        if trial_type == "mongo":
+            assert mongo_url is not None, "mongo_url can not be None if you want to run hyperopt with Mongo."
+            assert mongo_url.endswith("/jobs"), "hyperopt currently required the DB name to end in /jobs."
+            trials = MongoTrials(mongo_url, exp_key=mongo_exp_key)
+            mongo_url_sub = "://".join(mongo_url.split("://")[1:])
+            print(
+                f"""
+    Please make sure you start hyperopt mongo workers in a separate terminal!
+    Currently, you can use `hyperopt-mongo-worker --mongo={mongo_url_sub} --poll-interval=0.1` If this fails,
+    check https://github.com/hyperopt/hyperopt/wiki/Parallelizing-Evaluations-During-Search-via-MongoDB"""
+            )
+        elif trial_type == "spark":
+            raise ValueError(
+                "Spark trials have not been implemented yet with auto-setup."
+                + "Provide a SparkTrials object to the trials argument instead."
+            )
+        else:
+            trials = Trials()
+
+    # Run the optimisation
+    opt_result = fmin(optimise, search_space, max_evals=max_evals, trials=trials, **kwargs)
+    best_arguments = space_eval(search_space, opt_result)
+
+    # Get the best trial ID and corresponding performance.
+    if trials.best_trial is None:
+        best_trial_id = None
+        performance = None
+    else:
+        best_trial_id = trials.best_trial["tid"]
+        performance = trials.best_trial["result"]["loss"]
+        if not minimise_target:
+            performance = -performance
+
+    return best_trial_id, performance, best_arguments, trials
diff --git a/pype/hyperopt/setup.py b/pype/hyperopt/setup.py
new file mode 100644
index 0000000..a6170e2
--- /dev/null
+++ b/pype/hyperopt/setup.py
@@ -0,0 +1,13 @@
+from setuptools import find_namespace_packages, setup
+
+version = "0.1.0"
+deps: list[str] = [f"pype-base=={version}", "hyperopt>=0.2.7"]
+strict_deps = [s.replace(">=", "==") for s in deps]
+
+setup(
+    name="pype-hyperopt",
+    install_requires=deps,
+    extras_require={"dev": strict_deps, "strict": strict_deps},
+    packages=find_namespace_packages(include=["pype.*"]),
+    version=version,
+)
diff --git a/pype/mlflow/__init__.py b/pype/mlflow/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pype/mlflow/deploy/__init__.py b/pype/mlflow/deploy/__init__.py
new file mode 100644
index 0000000..e458229
--- /dev/null
+++ b/pype/mlflow/deploy/__init__.py
@@ -0,0 +1,1 @@
+from .load_experiment import load_experiment_from_mlflow
diff --git a/pype/mlflow/deploy/load_experiment.py b/pype/mlflow/deploy/load_experiment.py
new file mode 100644
index 0000000..dff92a8
--- /dev/null
+++ b/pype/mlflow/deploy/load_experiment.py
@@ -0,0 +1,53 @@
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from mlflow.artifacts import download_artifacts
+from mlflow.tracking import set_tracking_uri
+from mlflow.tracking.fluent import get_experiment_by_name, list_run_infos
+
+from pype.base.deploy import Inferencer
+
+
+def load_experiment_from_mlflow(
+    url: str,
+    experiment_name: str,
+    run_id: str,
+    directory: Path | str | None = None,
+) -> Inferencer:
+    """Download and import a trained Pype Model from mlflow's artifact store.
+
+    Args:
+        url (str): The tracking url.
+        experiment_name (str): The name of the experiment. Used to verify that the
+            run_id is the correct one.
+        run_id (str): The run id. Used to pull the actual experiment.
+        directory (Path | str | None, optional): Optional directory to store the results
+            in. By default, a temporary directory is used.
+
+    Returns:
+        Inferencer: _description_
+    """
+    set_tracking_uri(url)
+
+    # verify the run ID provided is part of the given experiment.
+    exp_id = get_experiment_by_name(experiment_name)
+    assert exp_id is not None, f"Experiment {experiment_name} does not exist in {url}."
+    assert run_id in [
+        run.run_id for run in list_run_infos(exp_id.experiment_id)
+    ], f"Run ID {run_id} is not present in the given experiment."
+
+    if directory is None:
+        with TemporaryDirectory(prefix="mlflow_model") as tmp_dir:
+            return _download_and_load(run_id, Path(tmp_dir))
+    else:
+        if isinstance(directory, str):
+            directory = Path(directory)
+        return _download_and_load(run_id, directory)
+
+
+def _download_and_load(
+    run_id: str,
+    directory: Path,
+) -> Inferencer:
+    download_artifacts(f"runs:/{run_id}/", dst_path=str(directory))
+    return Inferencer.from_folder(directory)
diff --git a/pype/mlflow/logger/__init__.py b/pype/mlflow/logger/__init__.py
new file mode 100644
index 0000000..cde4a7c
--- /dev/null
+++ b/pype/mlflow/logger/__init__.py
@@ -0,0 +1,1 @@
+from .mlflow_logger import MlflowLogger
diff --git a/pype/mlflow/logger/mlflow_logger.py b/pype/mlflow/logger/mlflow_logger.py
new file mode 100644
index 0000000..fce7dc6
--- /dev/null
+++ b/pype/mlflow/logger/mlflow_logger.py
@@ -0,0 +1,139 @@
+import logging
+from pathlib import Path
+from types import TracebackType
+from typing import Any
+
+from git import InvalidGitRepositoryError
+from git.repo import Repo
+from mlflow import create_experiment  # type: ignore
+from mlflow import get_experiment_by_name  # type: ignore
+from mlflow import log_artifact  # type: ignore
+from mlflow import log_metrics  # type: ignore
+from mlflow import log_params  # type: ignore
+from mlflow import set_experiment  # type: ignore
+from mlflow import set_tag  # type: ignore
+from mlflow import set_tracking_uri  # type: ignore
+from mlflow import start_run  # type: ignore
+
+from pype.base.logger import ExperimentLogger
+
+
+class MlflowLogger(ExperimentLogger):
+    def __init__(self, name: str, uri: str, artifact_location: str | None = None) -> None:
+        """A logger using mlflow for pype.
+
+        Args:
+            name (str): The name of the experiment. For runs done on databricks, make sure you preface
+                this with a workspace, like /Users/<your databricks user name>/<experiment name>.
+            uri (str): The tracking uri. Should most likely start with `http://` or `databricks`.
+                It can also equal `databricks`.
+            artifact_location (str | None, optional): The artificat uri location. This is needed if you set up
+                your own artifact location, to make sure files are copied over to any remote tracking server.
+                This is not required when using hosted databricks. Defaults to None, like in
+        """
+        super().__init__()
+        logger = logging.getLogger(__name__)
+        self.name = name
+
+        if not uri.startswith("http://") and not uri.startswith("databricks"):
+            logger.warning(f"Most often `uri` has to start with `http://` or `databricks`. Got: {uri}")
+        self.uri = uri
+        self.artifact_location = artifact_location
+        self.run = None
+
+    def __enter__(self) -> None:
+        """Setup the mlflow experiment and call start_run."""
+        set_tracking_uri(self.uri)
+        experiment = get_experiment_by_name(self.name)
+
+        if experiment is None:
+            # this is a new experiment: create it explicitly with the artifact uri
+            # TODO: make sure this works on S3 / databricks!
+            # Note: this works on databricks.
+            create_experiment(self.name, self.artifact_location)
+
+        set_experiment(experiment_name=self.name)
+
+        self.run = start_run().__enter__()
+        self.log_branch()
+
+    def log_branch(self) -> None:
+        """Set the current branch as a tag on the experiment."""
+        assert self.run is not None, "Please start the experiment first"
+        repo = self._find_encapsulating_repo(Path("."))
+
+        if repo is None:
+            branch_name = "unknown_branch"
+        else:
+            branch_name = repo.active_branch.name
+
+        set_tag("git_branch", branch_name)
+
+    def _find_encapsulating_repo(self, directory: Path) -> Repo | None:
+        """Finds the current repo we're in, if any.
+
+        We'll recursively move up from `directory` until we find a valid
+        git repo.
+
+        Args:
+            directory (Path): The directory to start looking for. Usually set
+                to the current directory.
+
+        Returns:
+            Repo | None: None if no repo is found, otherwise a GitPython Repo.
+        """
+        repo = None
+        directory = directory.absolute()
+        while str(directory) != "/":
+            try:
+                repo = Repo(directory)
+                break
+            except InvalidGitRepositoryError:
+                directory = directory.parent
+        return repo
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> None:
+        """Exits a run started by __enter__.
+
+        Args:
+            exc_type (type[BaseException] | None): Passed to ActiveRun.__exit__()
+            exc_val (BaseException | None): Passed to ActiveRun.__exit__()
+            exc_tb (TracebackType | None): Passed to ActiveRun.__exit__()
+        """
+        assert self.run is not None, "Run not started, cannot exit!"
+        self.run.__exit__(exc_type, exc_val, exc_tb)
+
+    def _log_metrics(self, metrics: dict[str, float | int | str | bool]) -> None:
+        """Perform the actual logging of metrics in mlflow.
+
+        Args:
+            metrics (dict[str, float | int | str | bool]): A dictionary of metric names and values.
+        """
+        log_metrics(metrics)
+
+    def log_parameters(self, parameters: dict[str, Any]) -> None:
+        """Logs the parameters for a given run.
+
+        These will by default be passed on with prefixes such as `model__`.
+
+        Args:
+            parameters (dict[str, Any]): The parameters of a given run. Ideally these
+                parameters should be no more complicated than string, float, int, bool, or a
+                list of these.
+        """
+        log_params(parameters)
+
+    def log_file(self, file: str | Path) -> None:
+        """Logs a given file as part of an experiment.
+
+        This only works well if:
+            - you work with a databricks-hosted mlflow experiment.
+            - you have set a artifact uri and configured your server properly.
+            - you have run mlflow ui/server at the same folder where mlruns is located.
+
+        Args:
+            file (str | Path): The file to log.
+        """
+        log_artifact(str(file))
diff --git a/pype/mlflow/setup.py b/pype/mlflow/setup.py
new file mode 100644
index 0000000..b6943d9
--- /dev/null
+++ b/pype/mlflow/setup.py
@@ -0,0 +1,18 @@
+from setuptools import find_namespace_packages, setup
+
+version = "0.1.0"
+
+deps = [
+    f"pype-base=={version}",
+    "mlflow>=1.28.0, <2.0",
+    "GitPython>=3.1.27",
+]
+strict_deps = [s.replace(">=", "==") for s in deps]
+
+setup(
+    name="pype-mlflow",
+    install_requires=deps,
+    extras_require={"dev": strict_deps, "strict": strict_deps},
+    packages=find_namespace_packages(include=["pype.*"]),
+    version=version,
+)
diff --git a/pype/sklearn/data/__init__.py b/pype/sklearn/data/__init__.py
index 55208b8..74596bd 100644
--- a/pype/sklearn/data/__init__.py
+++ b/pype/sklearn/data/__init__.py
@@ -1,1 +1,2 @@
+from .data_frame_source import DataFrameSource
 from .sklearn_data import SklearnData
diff --git a/pype/sklearn/data/data_frame_source.py b/pype/sklearn/data/data_frame_source.py
index 98c07e3..a3aa305 100644
--- a/pype/sklearn/data/data_frame_source.py
+++ b/pype/sklearn/data/data_frame_source.py
@@ -8,7 +8,7 @@ def __init__(self, df: pd.DataFrame) -> None:
         """A DataSource based around a created DataFrame.
 
         Good for testing purposes, but please use other sources for actual
-        ML runs.
+        ML runs to prevent any data having to be hard-coded into the script.
 
         Args:
             df (pd.DataFrame): The DataFrame to use as a source.
diff --git a/pype/sklearn/model/__init__.py b/pype/sklearn/model/__init__.py
index e69de29..05a7564 100644
--- a/pype/sklearn/model/__init__.py
+++ b/pype/sklearn/model/__init__.py
@@ -0,0 +1,4 @@
+from .linear_regression_model import LinearRegressionModel
+from .logistic_regression_model import LogisticRegressionModel
+from .sklearn_base_type import SklearnModelBaseType
+from .sklearn_model import SklearnModel
diff --git a/pype/sklearn/model/linear_regression_model.py b/pype/sklearn/model/linear_regression_model.py
index da58aad..5561122 100644
--- a/pype/sklearn/model/linear_regression_model.py
+++ b/pype/sklearn/model/linear_regression_model.py
@@ -1,44 +1,10 @@
+from typing import Any
+
 from sklearn.linear_model import LinearRegression
 
 from pype.sklearn.model.sklearn_model import SklearnModel
 
 
-class LinearRegressionModel(SklearnModel):
-    def __init__(
-        self,
-        inputs: list[str],
-        outputs: list[str],
-        model: LinearRegression = None,
-        fit_intercept: bool = True,
-        normalize: str = "deprecated",
-        copy_X: bool = True,
-        n_jobs: int = None,
-        positive: bool = False,
-        seed: int = 1,
-    ) -> None:
-        """Initialises a sklearn-LinearRegression model ready for pype to use.
-
-        Args:
-            inputs (List[str]): A list of names of input Data. This determines which Data is
-                used to fit the model.
-            outputs (List[str]): A list of names of output Data. This determines the names of
-                output Data.
-            seed (int, optional): The RNG seed to ensure reproducability. Defaults to 1.
-            model (LinearRegression, optional): The LinearRegression to use. Use `from_parameters` to initialise
-                a model from just the parameters. Optional, if not set, we'll create a new model using the parameters.
-            fit_intercept (bool, optional): See docs of `sklearn.linear.LinearRegression`
-            normalize (str, optional): See docs of `sklearn.linear.LinearRegression`
-            copy_X (bool, optional): See docs of `sklearn.linear.LinearRegression`
-            n_jobs (int, optional): See docs of `sklearn.linear.LinearRegression`
-            positive (bool, optional): See docs of `sklearn.linear.LinearRegression`
-        """
-        if model is None:
-            model = LinearRegression(
-                fit_intercept=fit_intercept,
-                normalize=normalize,
-                copy_X=copy_X,
-                n_jobs=n_jobs,
-                positive=positive,
-            )
-
-        super().__init__(inputs, outputs, model, seed=seed)
+class LinearRegressionModel(SklearnModel[LinearRegression]):
+    def _init_model(self, args: dict[str, Any]) -> LinearRegression:
+        return LinearRegression(**args)
diff --git a/pype/sklearn/model/logistic_regression_model.py b/pype/sklearn/model/logistic_regression_model.py
index 09f40ae..e719b0b 100644
--- a/pype/sklearn/model/logistic_regression_model.py
+++ b/pype/sklearn/model/logistic_regression_model.py
@@ -1,61 +1,10 @@
+from typing import Any
+
 from sklearn.linear_model import LogisticRegression
 
 from pype.sklearn.model.sklearn_model import SklearnModel
 
 
-class LogisticRegressionModel(SklearnModel):
-    def __init__(
-        self,
-        inputs: list[str],
-        outputs: list[str],
-        model: LogisticRegression = None,
-        penalty: str = "l2",
-        dual: bool = False,
-        tol: float = 1e-4,
-        C: float = 1.0,
-        fit_intercept: bool = True,
-        intercept_scaling: float = 1,
-        class_weight: dict = None,
-        solver: str = "lbfgs",
-        max_iter: int = 100,
-        multi_class: str = "auto",
-        n_jobs: int = None,
-        seed: int = 1,
-    ) -> None:
-        """Initialises a sklearn-LogisticRegression model ready for pype to use.
-
-        Args:
-            inputs (List[str]): A list of names of input Data. This determines which Data is
-                used to fit the model.
-            outputs (List[str]): A list of names of output Data. This determines the names of
-                output Data.
-            seed (int, optional): The RNG seed to ensure reproducability. Defaults to 1.
-            model (LogisticRegression, optional): The LogisticRegression to use. Use `from_parameters` to initialise
-                a model from just the parameters. Optional, if not set, we'll create a new model using the parameters.
-            penalty (str, optional): See docs of `sklearn.linear.LogisticRegression`
-            dual (bool, optional): See docs of `sklearn.linear.LogisticRegression`
-            tol (float, optional): See docs of `sklearn.linear.LogisticRegression`
-            C (float, optional): See docs of `sklearn.linear.LogisticRegression`
-            fit_intercept (bool, optional): See docs of `sklearn.linear.LogisticRegression`
-            intercept_scaling (float, optional): See docs of `sklearn.linear.LogisticRegression`
-            class_weight (dict, optional): See docs of `sklearn.linear.LogisticRegression`
-            solver (str, optional): See docs of `sklearn.linear.LogisticRegression`
-            max_iter (int, optional): See docs of `sklearn.linear.LogisticRegression`
-            multi_class (str, optional): See docs of `sklearn.linear.LogisticRegression`
-            n_jobs (int, optional): See docs of `sklearn.linear.LogisticRegression`
-        """
-        if model is None:
-            model = LogisticRegression(
-                fit_intercept=fit_intercept,
-                n_jobs=n_jobs,
-                penalty=penalty,
-                dual=dual,
-                tol=tol,
-                C=C,
-                intercept_scaling=intercept_scaling,
-                class_weight=class_weight,
-                max_iter=max_iter,
-                solver=solver,
-                multi_class=multi_class,
-            )
-        super().__init__(inputs, outputs, model, seed=seed)
+class LogisticRegressionModel(SklearnModel[LogisticRegression]):
+    def _init_model(self, args: dict[str, Any]) -> LogisticRegression:
+        return LogisticRegression(**args)
diff --git a/pype/sklearn/model/sklearn_base_type.py b/pype/sklearn/model/sklearn_base_type.py
index 800c9d5..46b03f3 100644
--- a/pype/sklearn/model/sklearn_base_type.py
+++ b/pype/sklearn/model/sklearn_base_type.py
@@ -7,8 +7,8 @@
 class SklearnModelBaseType(ABC):
     @abstractmethod
     def fit(self, *x: SklearnData) -> Any:
-        pass
+        """Fit a model to the given data."""
 
     @abstractmethod
     def predict(self, *x: SklearnData) -> Iterable[SklearnData] | SklearnData:
-        pass
+        """Predict for given data using a trained model."""
diff --git a/pype/sklearn/model/sklearn_model.py b/pype/sklearn/model/sklearn_model.py
index 011f512..62c9a9e 100644
--- a/pype/sklearn/model/sklearn_model.py
+++ b/pype/sklearn/model/sklearn_model.py
@@ -1,7 +1,8 @@
-from abc import ABC
+import typing
+from abc import ABC, abstractmethod
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import Iterable, Type
+from typing import Any, Generic, Iterable, Type, TypeVar
 
 import numpy as np
 
@@ -11,30 +12,55 @@
 from pype.sklearn.data.sklearn_data import SklearnData
 from pype.sklearn.model.sklearn_base_type import SklearnModelBaseType
 
+T = TypeVar("T", bound=SklearnModelBaseType)
 
-class SklearnModel(Model[SklearnData], ABC):
+
+class SklearnModel(Model[SklearnData], ABC, Generic[T]):
     SKLEARN_MODEL_FILE = "model.pkl"
 
     def __init__(
         self,
         inputs: list[str],
         outputs: list[str],
-        model: SklearnModelBaseType,
+        model: T | None = None,
         seed: int = 1,
+        **model_args: Any,
     ) -> None:
         """A generic class for sklearn-like Models.
 
+        You can set a sklearn model as a type hint to this class when defining a new model.
+        This allows us to get the parameters from the documentation of that sklearn model.
+        For an example, see the implementation of LinearModel, especially the `SklearnModel[LinearRegression]` part.
+
+        class LinearRegressionModel(SklearnModel[LinearRegression]):
+            def _init_model(self, **args) -> LinearRegression:
+                return LinearRegression(**args)
+
+
         Args:
             inputs (List[str]): A list of names of input Data. This determines which Data is
                 used to fit the model.
             outputs (List[str]): A list of names of output Data. This determines the names of
                 output variables.
-            model (SklearnModelBaseType): An object that has fit() and predict() methods.
+            model (SklearnModelBaseType | None): An object that has fit() and predict() methods. If none,
+                we will use the model_args to instantiate a new model.
             seed (int, optional): The RNG seed to ensure reproducability.. Defaults to 1.
+            **model_args (Any): Optional keyword arguments passed to the model class to instantiate a new
+                model if `model` is None.
         """
         super().__init__(inputs, outputs, seed)
+        if model is None:
+            model = self._init_model(model_args)
         self.model = model
 
+    @abstractmethod
+    def _init_model(self, args: dict[str, Any]) -> T:
+        raise NotImplementedError
+
+    @classmethod
+    def _get_annotated_class(cls) -> Type[SklearnModelBaseType]:
+        return typing.get_args(cls.__orig_bases__[0])[0]
+
     def set_seed(self) -> None:
         """Sets the RNG seed."""
         np.random.seed(self.seed)
@@ -50,7 +76,7 @@ def _save(self, folder: Path) -> None:
         serialiser.serialise(self.model, folder / self.SKLEARN_MODEL_FILE)
 
     @classmethod
-    def _load(cls: Type["SklearnModel"], folder: Path, inputs: list[str], outputs: list[str]) -> "Model":
+    def _load(cls: Type["SklearnModel"], folder: Path, inputs: list[str], outputs: list[str]) -> "SklearnModel":
         serialiser = JoblibSerialiser()
         model = serialiser.deserialise(folder / cls.SKLEARN_MODEL_FILE)
         return cls(inputs=inputs, outputs=outputs, model=model, seed=1)
@@ -67,6 +93,8 @@ def get_parameters(cls: Type["SklearnModel"], parser: ArgumentParser) -> None:
             parser (ArgumentParser): The ArgumentParser to add arguments to.
         """
         super().get_parameters(parser)
+        BaseModel = cls._get_annotated_class()
+
         add_args_to_parser_for_class(
-            parser, cls, "model", [SklearnModel], excluded_args=["seed", "inputs", "outputs", "model"]
+            parser, BaseModel, "model", [], excluded_args=["seed", "inputs", "outputs", "model"]
         )
diff --git a/pype/sklearn/pipeline/__init__.py b/pype/sklearn/pipeline/__init__.py
index e69de29..d6b8cd3 100644
--- a/pype/sklearn/pipeline/__init__.py
+++ b/pype/sklearn/pipeline/__init__.py
@@ -0,0 +1,2 @@
+from .numpy_type_checker import NumpyTypeChecker
+from .pandas_type_checker import PandasTypeChecker
diff --git a/pype/sklearn/pipeline/numpy_type_checker.py b/pype/sklearn/pipeline/numpy_type_checker.py
index e0ee131..89a48f5 100644
--- a/pype/sklearn/pipeline/numpy_type_checker.py
+++ b/pype/sklearn/pipeline/numpy_type_checker.py
@@ -38,7 +38,7 @@ def __init__(self) -> None:
         """
         super().__init__()
         self.dims: tuple = tuple()
-        self.dtype: type = None
+        self.dtype: type | None = None
 
     def fit(self, data: np.ndarray) -> "NumpyTypeChecker":
         """Fit this Numpy TypeChecker to the given data.
@@ -62,11 +62,12 @@ def transform(self, data: np.ndarray) -> np.ndarray:
         Returns:
             np.ndarray: data, if the data fits the specifications. Otherwise, an assertion error is thrown.
         """
+        assert self.dtype is not None, "Please fit pipeline first"
         assert isinstance(data, np.ndarray), "Please provide a numpy array!"
         assert data.shape[1:] == self.dims, f"Dimensions of numpy arrays do not add up: {data.shape[1:]} vs {self.dims}"
-        assert (
-            self._convert_dtype(data.dtype) == self.dtype
-        ), f"Dtype of data does not add up: {data.dtype} vs {self.dtype}"
+
+        converted_type = self._convert_dtype(data.dtype)
+        assert converted_type == self.dtype, f"Dtype of data does not add up: {converted_type} vs {self.dtype}"
         return data
 
     def _convert_dtype(self, dtype: np.dtype) -> type:
diff --git a/pype/sklearn/setup.py b/pype/sklearn/setup.py
index 0f551d9..4c3e3f8 100644
--- a/pype/sklearn/setup.py
+++ b/pype/sklearn/setup.py
@@ -1,12 +1,19 @@
 from setuptools import find_namespace_packages, setup
 
+version = "0.1.0"
+
+deps = [
+    f"pype-base=={version}",
+    "numpy>=1.23.0",
+    "scikit-learn>=1.1.1",
+    "pandas>=1.4.3",
+]
+strict_deps = [s.replace(">=", "==") for s in deps]
+
 setup(
     name="pype-sklearn",
-    install_requires=[
-        "pype-base",
-        "numpy==1.23.0",
-        "scikit-learn==1.1.1",
-        "pandas==1.4.3",
-    ],
+    install_requires=deps,
+    extras_require={"dev": strict_deps, "strict": strict_deps},
     packages=find_namespace_packages(include=["pype.*"]),
+    version=version,
 )
diff --git a/pype/spark/__init__.py b/pype/spark/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pype/spark/data/__init__.py b/pype/spark/data/__init__.py
new file mode 100644
index 0000000..d5552a0
--- /dev/null
+++ b/pype/spark/data/__init__.py
@@ -0,0 +1,3 @@
+from .spark_data_frame_source import SparkDataFrameSource
+from .spark_read import SparkReadSource
+from .spark_sql_source import SparkSqlSource
diff --git a/pype/spark/data/spark_data_frame_source.py b/pype/spark/data/spark_data_frame_source.py
new file mode 100644
index 0000000..574bc49
--- /dev/null
+++ b/pype/spark/data/spark_data_frame_source.py
@@ -0,0 +1,30 @@
+from pyspark.sql import DataFrame as SparkDataFrame
+
+from pype.base.data.data_source import DataSource
+
+
+class SparkDataFrameSource(DataSource[SparkDataFrame]):
+    # should only be used for testing purposes!
+    def __init__(
+        self,
+        df: SparkDataFrame,
+    ) -> None:
+        """Returns a fixed spark dataframe reference.
+
+        This should only be used for testing/debugging purposes.
+        Consider the SQL or Read classes for proper data reading!
+
+        Args:
+            df (SparkDataFrame): The DataFrame to use as output when
+                read() is called.
+        """
+        super().__init__()
+        self.df = df
+
+    def read(self) -> SparkDataFrame:
+        """Returns the stored DataFrame.
+
+        Returns:
+            SparkDataFrame: The stored DataFrame.
+        """
+        return self.df
diff --git a/pype/spark/data/spark_read.py b/pype/spark/data/spark_read.py
new file mode 100644
index 0000000..da61db0
--- /dev/null
+++ b/pype/spark/data/spark_read.py
@@ -0,0 +1,24 @@
+from typing import Any
+
+from pyspark.sql import DataFrame as SparkDataFrame
+from pyspark.sql import SparkSession
+
+from pype.base.data.data_source import DataSource
+
+
+class SparkReadSource(DataSource[SparkDataFrame]):
+    def __init__(
+        self,
+        spark_session: SparkSession,
+        file: str,
+        format: str,
+        options: dict[str, Any],
+    ) -> None:
+        super().__init__()
+        self.spark_session = spark_session
+        self.file = file
+        self.format = format
+        self.options = options
+
+    def read(self) -> SparkDataFrame:
+        return self.spark_session.read.format(self.format).load(self.file, **self.options)
diff --git a/pype/spark/data/spark_sql_source.py b/pype/spark/data/spark_sql_source.py
new file mode 100644
index 0000000..ace5cc9
--- /dev/null
+++ b/pype/spark/data/spark_sql_source.py
@@ -0,0 +1,18 @@
+from pyspark.sql import DataFrame as SparkDataFrame
+from pyspark.sql import SparkSession
+
+from pype.base.data.data_source import DataSource
+
+
+class SparkSqlSource(DataSource[SparkDataFrame]):
+    def __init__(
+        self,
+        spark_session: SparkSession,
+        query: str,
+    ) -> None:
+        super().__init__()
+        self.query = query
+        self.spark_session = spark_session
+
+    def read(self) -> SparkDataFrame:
+        return self.spark_session.sql(self.query)
diff --git a/pype/spark/model/__init__.py b/pype/spark/model/__init__.py
new file mode 100644
index 0000000..c1b66cf
--- /dev/null
+++ b/pype/spark/model/__init__.py
@@ -0,0 +1,2 @@
+from .linear_spark_model import LinearSparkModel
+from .spark_model import SparkModel
diff --git a/pype/spark/model/linear_spark_model.py b/pype/spark/model/linear_spark_model.py
new file mode 100644
index 0000000..6562112
--- /dev/null
+++ b/pype/spark/model/linear_spark_model.py
@@ -0,0 +1,10 @@
+from typing import Any
+
+from pyspark.ml.regression import LinearRegression
+
+from pype.spark.model.spark_model import SparkModel
+
+
+class LinearSparkModel(SparkModel[LinearRegression]):
+    def _init_model(self, args: dict[str, Any]) -> LinearRegression:
+        return LinearRegression(**args)
diff --git a/pype/spark/model/spark_model.py b/pype/spark/model/spark_model.py
new file mode 100644
index 0000000..a81cc39
--- /dev/null
+++ b/pype/spark/model/spark_model.py
@@ -0,0 +1,143 @@
+import json
+import typing
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import Any, Generic, Iterable, Type, TypeVar
+
+from pyspark.ml import Model as BaseSparkModel
+from pyspark.ml import Predictor
+from pyspark.sql import DataFrame as SparkDataFrame
+
+from pype.base.data import DataSet
+from pype.base.experiment.argument_parsing import add_args_to_parser_for_class
+from pype.base.model import Model
+from pype.base.serialiser import JoblibSerialiser
+
+T = TypeVar("T", bound=Predictor)
+
+
+class SparkModel(Model[SparkDataFrame], ABC, Generic[T]):
+    PYPE_MODEL_CONFIG = "config.json"
+    SPARK_PREDICTOR_PATH = "predictor"
+    SPARK_MODEL_PATH = "spark_model"
+    SPARK_MODEL_CLASS_PATH = "spark_model_class"
+
+    def __init__(
+        self,
+        inputs: list[str],
+        outputs: list[str],
+        output_col: str | None = None,
+        model: BaseSparkModel | None = None,
+        predictor: T | None = None,
+        seed: int = 1,
+        **model_args: Any,
+    ) -> None:
+        """A Pype-compliant framework for using Spark Models.
+
+        Args:
+            inputs (list[str]): The name of the input DataFrame. Should be a list of 1 string.
+            outputs (list[str]): The name of the output DataFrame. Should be a list of 1 string, the same as `inputs`
+            output_col (str | None, optional): The name of the column where the model will put the output.
+                Defaults to None, which means we won't select any columns and instead return the full output
+                of the model.
+            predictor (Predictor, optional): The Spark Predictor. If not set, we try to instantiate it
+                using `model_args`
+            model (BaseSparkModel, optional): The Spark Model. Defaults to None. If set to None,
+                this model can't be serialised or used for inference.
+            seed (int, optional): Spark Seed. Currently ignored, unfortunately. Defaults to 1.
+        """
+        assert len(inputs) == 1, "SparkML only requires a single DataFrame as input and output, the same one."
+        assert len(outputs) == 1, "SparkML only requires a single DataFrame as input and output, the same one."
+        assert inputs == outputs, "SparkML only requires a single DataFrame as input and output, the same one."
+        super().__init__(inputs, outputs, seed)
+        if predictor is None:
+            predictor = self._init_model(model_args)
+
+        self.model = model
+        self.predictor = predictor
+        self.output_col = output_col
+
+    @abstractmethod
+    def _init_model(self, args: dict[str, Any]) -> T:
+        raise NotImplementedError
+
+    @classmethod
+    def _get_annotated_class(cls) -> Type[T]:
+        return typing.get_args(cls.__orig_bases__[0])[0]
+
+    def set_seed(self) -> None:
+        """Sets the RNG seed."""
+        # TODO: set seed properly in pyspark
+        print("Currently, setting a seed is not yet supported.")
+
+    def fit(self, data: DataSet) -> "Model":
+        """Fits the Model to the given DataSet.
+
+        The DataSet should contain all inputs. For Spark models, this means 1 DataFrame
+        that contains a feature vector and the label.
+
+        Args:
+            data (DataSet): The DataSet to fit this Model on.
+        """
+        self._fit(*data.get_all(self.inputs))
+        return self
+
+    def _fit(self, *data: SparkDataFrame) -> None:
+        assert len(data) == 1, f"SparkML needs a single DataFrame as input, got {len(data)}"
+        self.model = self.predictor.fit(data[0])
+
+    def _transform(self, *data: SparkDataFrame) -> Iterable[SparkDataFrame] | SparkDataFrame:
+        assert len(data) == 1, f"SparkML needs a single DataFrame as input, got {len(data)}"
+        assert self.model is not None, "Please fit this model before transforming data."
+
+        result = self.model.transform(data[0])
+        if self.output_col is not None:
+            result = result.select(self.output_col)
+        return result
+
+    def _save(self, folder: Path) -> None:
+        assert self.model is not None, "Please fit this model before transforming data."
+
+        serialiser = JoblibSerialiser()
+        config = {"output_col": self.output_col}
+        with open(folder / self.PYPE_MODEL_CONFIG, "w") as f:
+            json.dump(config, f)
+
+        self.predictor.save(str(folder / self.SPARK_PREDICTOR_PATH))
+        self.model.save(str(folder / self.SPARK_MODEL_PATH))
+        serialiser.serialise(type(self.model), str(folder / self.SPARK_MODEL_CLASS_PATH))
+
+    @classmethod
+    def _load(cls: Type["SparkModel"], folder: Path, inputs: list[str], outputs: list[str]) -> "SparkModel":
+        serialiser = JoblibSerialiser()
+
+        with open(folder / cls.PYPE_MODEL_CONFIG, "r") as f:
+            config = json.load(f)
+        output_col = config["output_col"]
+
+        predictor_class: type[Predictor] = cls._get_annotated_class()
+        model_class: Type[BaseSparkModel] = serialiser.deserialise(str(folder / cls.SPARK_MODEL_CLASS_PATH))
+
+        predictor: Predictor = predictor_class.load(str(folder / cls.SPARK_PREDICTOR_PATH))
+        model: BaseSparkModel = model_class.load(str(folder / cls.SPARK_MODEL_PATH))
+
+        return cls(inputs=inputs, outputs=outputs, predictor=predictor, model=model, output_col=output_col, seed=1)
+
+    @classmethod
+    def get_parameters(cls: Type["SparkModel"], parser: ArgumentParser) -> None:
+        """Get and add parameters to initialise this class.
+
+        SparkModel's will work by requiring 2 ways to instantiate a Model:
+            - through `model`, which is a spark model.
+            - through parameters, which will instantiate the model internally.
+
+        Args:
+            parser (ArgumentParser): The ArgumentParser to add arguments to.
+        """
+        super().get_parameters(parser)
+        BaseModel = cls._get_annotated_class()
+
+        add_args_to_parser_for_class(
+            parser, BaseModel, "model", [], excluded_args=["seed", "inputs", "outputs", "model"]
+        )
diff --git a/pype/spark/pipeline/__init__.py b/pype/spark/pipeline/__init__.py
new file mode 100644
index 0000000..4327d22
--- /dev/null
+++ b/pype/spark/pipeline/__init__.py
@@ -0,0 +1,2 @@
+from .spark_pipe import SparkPipe
+from .spark_type_checker import SparkTypeChecker
diff --git a/pype/spark/pipeline/spark_pipe.py b/pype/spark/pipeline/spark_pipe.py
new file mode 100644
index 0000000..84480a6
--- /dev/null
+++ b/pype/spark/pipeline/spark_pipe.py
@@ -0,0 +1,101 @@
+from typing import Any, Type
+
+from pyspark.ml import Estimator, Transformer
+
+from pype.base.data.dataset import DataSet
+from pype.base.pipeline.operator import Operator
+from pype.base.pipeline.pipe import Pipe
+
+
+class SparkPipe(Pipe):
+    def __init__(
+        self,
+        name: str,
+        operator: Type[Operator],
+        inputs: list[str],
+        outputs: list[str],
+        kw_args: dict[str, Any] | None = None,
+        fit_inputs: list[str] | None = None,
+    ) -> None:
+        """Same init as Pipe."""
+        super().__init__(name, operator, inputs, outputs, kw_args, fit_inputs)
+        self.fitted = None
+
+    def fit(self, data: DataSet) -> "SparkPipe":
+        """Fits the SparkPipe to the given DataSet.
+
+        Args:
+            data (DataSet): The DataSet to use in fitting.
+
+        Returns:
+            SparkPipe: This object.
+        """
+        op: Transformer | Estimator = self.operator
+
+        if isinstance(op, Transformer):
+            self.fitted = op
+        elif isinstance(op, Estimator):
+            self.fitted = op.fit(*data.get_all(self.inputs), *data.get_all(self.fit_inputs))
+        else:
+            raise ValueError(f"In a SparkPipe, the operator must be a Transformer or Estimator. Got: {type(op)}")
+        return self
+
+    def transform(self, data: DataSet) -> DataSet:
+        """Transforms the given data using this Pipe.
+
+        This Pipe should be fitted first using fit().
+        This is version of Pipe is changed to work with Spark's API: transformation is done using the
+        fitted object instead of the Operator directly.
+
+        Args:
+            data (DataSet): The DataSet to use in transforming.
+
+        Returns:
+            DataSet: The transformed Data.
+        """
+        assert self.fitted is not None, "Make sure you fit the pipeline before transforming."
+        transformed = self.fitted.transform(*data.get_all(self.inputs))
+        if len(self.outputs) < 2:
+            transformed = [transformed]
+        result = data.copy()
+        result.set_all(self.outputs, transformed)
+        return result
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Sets the state of this object from a dictionary.
+
+        Used by pickle to properly prepare the "fitted" and "operator" fields as None's.
+
+        Args:
+            state (dict[str, Any]): _description_
+        """
+        state["fitted"] = None
+        state["operator"] = None
+        self.__dict__ = state
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Gets the state of this object, excluding any Spark objects.
+
+        This is to make sure serialisation works as expected.
+
+        Returns:
+            dict[str, Any]: A dict representation of this object.
+        """
+        dct = self.__dict__.copy()
+        if "fitted" in dct:
+            del dct["fitted"]
+        del dct["operator"]
+        return dct
+
+    def __str__(self) -> str:
+        """Create string representation of this Pipe.
+
+        Returns:
+            str: A string representation of this Pipe.
+        """
+        fit_section = f" (+ {self.fit_inputs})" if len(self.fit_inputs) > 0 else ""
+        input_output_section = f"{self.inputs}{fit_section} -> {self.outputs}"
+
+        postfix = "unfitted" if self.fitted is None else str(self.fitted)
+
+        return f"Pipe `{self.name}`, {input_output_section}: {postfix}"
diff --git a/pype/spark/pipeline/spark_type_checker.py b/pype/spark/pipeline/spark_type_checker.py
new file mode 100644
index 0000000..569d7cf
--- /dev/null
+++ b/pype/spark/pipeline/spark_type_checker.py
@@ -0,0 +1,122 @@
+from typing import Any, Dict, List, Type
+
+from pydantic import create_model
+from pyspark.sql import DataFrame as SparkDataFrame
+from pyspark.sql import SparkSession
+
+from pype.base.pipeline.type_checker import DataModel, TypeChecker
+
+
+class SparkData(DataModel):
+    def convert(self) -> SparkDataFrame:
+        """Converts this object to a spark DataFrame.
+
+        Returns:
+            SparkDataFrame: The spark DataFrame contained by this object.
+        """
+        list_format: List[Dict[str, Any]] = []
+        length = None
+        for name, value in self.__dict__.items():
+            if length is None:
+                length = len(value)
+                list_format = [{} for _ in range(length)]
+            else:
+                assert length == len(value), f"Warning: {name} did not have the same length as others"
+            for i, v in enumerate(value):
+                list_format[i][name] = v
+
+        spark_session = SparkSession.builder.getOrCreate()
+
+        return spark_session.createDataFrame(list_format)
+
+    @classmethod
+    def to_model(cls, data: SparkDataFrame) -> "SparkData":
+        """Converts a spark DataFrame to a SparkData model, which can be serialised.
+
+        Args:
+            data (SparkDataFrame): A spark DataFrame to serialise.
+
+        Returns:
+            SparkData: A serialisable version of the DataFrame.
+        """
+        rows = data.collect()
+        columns: Dict[str, List[Any]] = {}
+        for row in rows:
+            for name, value in row.asDict().items():
+                if name in columns:
+                    columns[name].append(value)
+                else:
+                    columns[name] = [value]
+
+        return cls(**columns)
+
+
+class SparkTypeChecker(TypeChecker[SparkDataFrame]):
+    def fit(self, data: SparkDataFrame) -> "SparkTypeChecker":
+        """Fit this SparkTypeChecker to the given data.
+
+        Args:
+            data (SparkDataFrame): The data to fit.
+
+        Returns:
+            SparkTypeChecker: self
+        """
+        self.raw_types = self._convert_dtypes(dict(data.dtypes))
+        return self
+
+    def _convert_dtypes(self, dct: dict[str, str]) -> dict[str, tuple[str, type]]:
+        return {name: self._convert_dtype(string_type) for name, string_type in dct.items()}
+
+    def _convert_dtype(self, string_type: str) -> tuple[str, type]:
+        conv_type: Type | None = None
+        if string_type in ["str", "string"]:
+            conv_type = str
+        elif string_type in ["int", "bigint"]:
+            conv_type = int
+        elif string_type in ["float", "double"]:
+            conv_type = float
+        elif string_type == "bool":
+            conv_type = bool
+        else:
+            raise ValueError(f"{string_type} not supported")
+
+        return (string_type, conv_type)
+
+    def transform(self, data: SparkDataFrame) -> SparkDataFrame:
+        """Checks if the given data fits the specifications this TypeChecker was fitted for.
+
+        Args:
+            data (SparkDataFrame): The data to check.
+
+        Returns:
+            SparkDataFrame: data, if the data fits the specifications. Otherwise, an assertion error is thrown.
+        """
+        assert isinstance(data, SparkDataFrame), "Please provide a spark DataFrame!"
+        colnames = list(self.raw_types.keys())
+        for col in colnames:
+            assert col in data.columns, f"`{col}` is missing from the dataset"
+        dtype_dict = dict(data.dtypes)
+
+        data = data.select(colnames)
+
+        for name, (spark_type, _) in self.raw_types.items():
+            actual_dtype = dtype_dict[name]
+            assert (
+                actual_dtype == spark_type
+            ), f"Dtypes did not match up for col {name}: Expected {spark_type}, got {actual_dtype}"
+        return data
+
+    def get_pydantic_type(self) -> type[SparkData]:
+        """Creates a Pydantic model for this data to handle serialisation/deserialisation.
+
+        Returns:
+            type[SparkData]: A SparkData model that fits the data this wat fitted on.
+        """
+        data_type = {
+            name: (list[dtype] | dict[str | int, dtype], ...)  # type: ignore
+            for name, (_, dtype) in self.raw_types.items()
+        }
+
+        model = create_model("SparkData", **data_type, __base__=SparkData)
+
+        return model
diff --git a/pype/spark/serialisation/__init__.py b/pype/spark/serialisation/__init__.py
new file mode 100644
index 0000000..dfcce47
--- /dev/null
+++ b/pype/spark/serialisation/__init__.py
@@ -0,0 +1,1 @@
+from .spark_serialiser import SparkSerialiser
diff --git a/pype/spark/serialisation/spark_serialiser.py b/pype/spark/serialisation/spark_serialiser.py
new file mode 100644
index 0000000..f24341b
--- /dev/null
+++ b/pype/spark/serialisation/spark_serialiser.py
@@ -0,0 +1,128 @@
+import json
+from pathlib import Path
+from typing import Any
+
+import joblib
+from pyspark.ml import Transformer
+
+from pype.base.pipeline import Pipe, Pipeline
+from pype.base.serialiser import Serialiser
+from pype.spark.pipeline.spark_pipe import SparkPipe
+
+
+class SparkSerialiser(Serialiser):
+    SUB_PIPE_PREFIX = "__sub_pipe_"
+    BASE_PIPE_FILE = "pipe.pkl"
+    SPARK_TRANSFORMER_FILE = "spark_transformer"
+    SPARK_TRANSFORMER_CLASS_FILE = "spark_transformer_class"
+    STEPS_FILE = "steps.json"
+
+    def serialise(self, object: Any, file: str | Path) -> None:
+        """Serialse a given object.
+
+        If the object is:
+            - A Pype Pipeline: we serialise each step separately.
+            - Otherwise, we use joblib.
+
+        Args:
+            object (Any): The object to serialise.
+            file (str | Path): The path to serialise to.
+        """
+        if isinstance(object, Pipeline):
+            self._serialise_pipeline(object, file)
+        else:
+            return self._serialise_joblib(object, file)
+
+    def deserialise(self, file: str | Path) -> Any:
+        """Deserialise the object in the given file.
+
+        This function can handle Pype Pipelines with PySpark elements.
+
+        Args:
+            file (str | Path): The file containing a python object to deserialise.
+
+        Returns:
+            Any: The deserialised object.
+        """
+        file = Path(file)
+        if file.is_dir() and (file / self.STEPS_FILE).is_file():
+            return self._deserialise_pipeline(file)
+        else:
+            return self._deserialise_joblib(file)
+
+    def _serialise_pipeline(self, pipeline: Pipeline, file: str | Path) -> None:
+        file = Path(file)
+
+        file.mkdir(exist_ok=True)
+        steps = []
+        i = 0
+        for step in pipeline:
+            if isinstance(step, SparkPipe):
+                name = step.name
+                pipe_path = file / name
+                pipe_path.mkdir(exist_ok=True)
+                self._serialise_joblib(step, pipe_path / self.BASE_PIPE_FILE)
+
+                # store the class for re-importing
+                transformer: Transformer = step.fitted
+                self._serialise_joblib(type(transformer), pipe_path / self.SPARK_TRANSFORMER_CLASS_FILE)
+                transformer.save(str(pipe_path / self.SPARK_TRANSFORMER_FILE))
+            elif isinstance(step, Pipe):
+                name = step.name
+                self._serialise_joblib(step, file / name)
+            else:
+                name = f"{self.SUB_PIPE_PREFIX}{i}"
+                i += 1
+                self._serialise_pipeline(step, file / name)
+
+            steps.append(name)
+
+        with open(file / self.STEPS_FILE, "w") as f:
+            json.dump({"steps": steps}, f)
+
+    def _serialise_joblib(self, object: Any, file: str | Path) -> None:
+        """Serialise the given object to the given file.
+
+        Args:
+            object (Any): The object to serialise.
+            file (str | Path): The file to serialise to.
+        """
+        joblib.dump(object, file)
+
+    def _deserialise_pipeline(self, file: str | Path) -> Pipeline:
+        file = Path(file)
+        assert file.is_dir()
+
+        with open(file / self.STEPS_FILE, "r") as f:
+            steps: list[str] = json.load(f)["steps"]
+
+        pipes = []
+        for step in steps:
+            step_path = file / step
+            if step_path.is_file():
+                pipes.append(self._deserialise_joblib(step_path))
+            elif step_path.is_dir():
+                if step.startswith(self.SUB_PIPE_PREFIX):
+                    pipes.append(self._deserialise_pipeline(step_path))
+                else:
+                    pipe: SparkPipe = self._deserialise_joblib(step_path / self.BASE_PIPE_FILE)
+                    fitted_class: type[Transformer] = self._deserialise_joblib(
+                        step_path / self.SPARK_TRANSFORMER_CLASS_FILE
+                    )
+                    pipe.fitted = fitted_class.load(str(step_path / self.SPARK_TRANSFORMER_FILE))
+                    pipes.append(pipe)
+            else:
+                raise ValueError(f"{step_path} was not an existing path")
+
+        return Pipeline(pipes)
+
+    def _deserialise_joblib(self, file: str | Path) -> Any:
+        """Deserialise the object in the given file.
+
+        Args:
+            file (str | Path): The file to deserialise.
+
+        Returns:
+            Any: The python object stored in the file.
+        """
+        return joblib.load(file)
diff --git a/pype/spark/setup.py b/pype/spark/setup.py
new file mode 100644
index 0000000..3127621
--- /dev/null
+++ b/pype/spark/setup.py
@@ -0,0 +1,20 @@
+from setuptools import find_namespace_packages, setup
+
+version = "0.1.0"
+
+deps = [
+    f"pype-base=={version}",
+    # We will provide absolute no guarantees that our integration will work with
+    # EVERY version of pyspark. This has been developed under pyspark==3.2.1.
+    "pyspark>=3.2.1",
+]
+strict_deps = [s.replace(">=", "==") for s in deps]
+test_deps = ["pandas"]
+
+setup(
+    name="pype-spark",
+    install_requires=deps,
+    extras_require={"dev": strict_deps + test_deps, "test": strict_deps + test_deps, "strict": strict_deps},
+    packages=find_namespace_packages(include=["pype.*"]),
+    version=version,
+)
diff --git a/pype/tensorflow/__init__.py b/pype/tensorflow/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pype/tensorflow/data/__init__.py b/pype/tensorflow/data/__init__.py
new file mode 100644
index 0000000..193cb9f
--- /dev/null
+++ b/pype/tensorflow/data/__init__.py
@@ -0,0 +1,1 @@
+from .tensor_source import TensorSource
diff --git a/pype/tensorflow/data/tensor_source.py b/pype/tensorflow/data/tensor_source.py
new file mode 100644
index 0000000..433428a
--- /dev/null
+++ b/pype/tensorflow/data/tensor_source.py
@@ -0,0 +1,25 @@
+from tensorflow import Tensor  # type: ignore
+
+from pype.base.data.data_source import DataSource
+
+
+class TensorSource(DataSource[Tensor]):
+    def __init__(self, tensor: Tensor) -> None:
+        """A DataSource based around a created tensorflow Tensor.
+
+        Good for testing purposes, but please use other sources for actual
+        ML runs to prevent any data having to be hard-coded into the script.
+
+        Args:
+            tensor (Tensor): The Tensor to use as a source.
+        """
+        super().__init__()
+        self.tensor = tensor
+
+    def read(self) -> Tensor:
+        """Returns the Tensor given as input before.
+
+        Returns:
+            Tensor: The Tensor used to initialise this object.
+        """
+        return self.tensor
diff --git a/pype/tensorflow/model/__init__.py b/pype/tensorflow/model/__init__.py
new file mode 100644
index 0000000..16bbadf
--- /dev/null
+++ b/pype/tensorflow/model/__init__.py
@@ -0,0 +1,3 @@
+from .keras_pype_model import KerasPypeModel
+from .mlp_keras import MLPKeras
+from .mlp_pype_model import MLPPypeModel
diff --git a/pype/tensorflow/model/keras_pype_model.py b/pype/tensorflow/model/keras_pype_model.py
new file mode 100644
index 0000000..71d7850
--- /dev/null
+++ b/pype/tensorflow/model/keras_pype_model.py
@@ -0,0 +1,133 @@
+import typing
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import Any, Callable, Generic, Iterable, Type, TypeVar
+
+from keras import Model as KerasBaseModel
+from keras.losses import Loss
+from keras.metrics import Metric
+from keras.models import load_model as keras_load_model
+from keras.optimizers import Optimizer
+from tensorflow import Tensor  # type: ignore
+from tensorflow.data import Dataset  # type: ignore
+from tensorflow.random import set_seed  # type: ignore
+
+from pype.base.experiment.argument_parsing import add_args_to_parser_for_class
+from pype.base.model import Model
+from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
+
+T = TypeVar("T", bound=KerasBaseModel)
+
+
+class KerasPypeModel(Model[Tensor], ABC, Generic[T]):
+    KERAS_MODEL_FILE = "keras_model"
+    LOSS_FILE = "loss_func.pkl"
+
+    def __init__(
+        self,
+        inputs: list[str],
+        outputs: list[str],
+        loss: Callable | Loss,
+        optimizer_class: Type[Optimizer],
+        epochs: int = 10,
+        batch_size: int = 32,
+        learning_rate: float = 0.001,
+        metrics: list[Metric] | dict[str, Metric] | None = None,
+        model: T | None = None,
+        seed: int = 1,
+        **model_args: Any,
+    ) -> None:
+        """A Model to integrate Keras models with the Pype framework.
+
+        Args:
+            inputs (list[str]): Input names from the DataSet to use.
+            outputs (list[str]): Output names from the DataSet to use.
+            loss (Callable | Loss): The loss for your Keras Model.
+            optimizer_class (Type[Optimizer]): The optimization class to use.
+            epochs (int, optional): The number of epochs to use for training. Defaults to 10.
+            batch_size (int, optional): The batch size to use for training. Defaults to 32.
+            learning_rate (float, optional): The learning rate to use for training. Defaults to 0.001.
+            metrics (list[Metric] | dict[str, Metric] | None, optional): Additional metrics to
+                use while training th emodel. Defaults to no metrics.
+            model (T | None, optional): The Keras Model to train. By Default we'll use excess arguments
+                to initialise the default model.
+            seed (int, optional): The seed to initialise tensorflow with. Defaults to 1.
+        """
+        super().__init__(inputs, outputs, seed)
+        if metrics is None:
+            metrics = []
+
+        if model_args is None:
+            model_args = {}
+
+        if model is None:
+            model = self._init_model(model_args)
+        self.model = model
+        self.optimizer = optimizer_class(learning_rate=learning_rate)
+        self.loss = loss
+        self.metrics = metrics
+
+        self.epochs = epochs
+        self.batch_size = batch_size
+
+    @abstractmethod
+    def _init_model(self, args: dict[str, Any]) -> T:
+        raise NotImplementedError
+
+    @classmethod
+    def _get_annotated_class(cls) -> Type[KerasBaseModel]:
+        return typing.get_args(cls.__orig_bases__[0])[0]
+
+    def set_seed(self) -> None:
+        """Sets the RNG seed."""
+        set_seed(self.seed)
+
+    def _fit(self, *data: Tensor) -> None:
+        dataset = Dataset.from_tensor_slices(data).batch(self.batch_size)
+
+        self.model.compile(optimizer=self.optimizer, loss=self.loss, metrics=self.metrics)
+        self.model.fit(dataset, epochs=self.epochs)
+
+    def _transform(self, *data: Tensor) -> Iterable[Tensor] | Tensor:
+        return self.model.call(*data)
+
+    def _save(self, folder: Path) -> None:
+        serialiser = JoblibSerialiser()
+        serialiser.serialise(self.loss, folder / self.LOSS_FILE)
+        self.model.save(folder / self.KERAS_MODEL_FILE)
+
+    @classmethod
+    def _load(cls: Type["KerasPypeModel"], folder: Path, inputs: list[str], outputs: list[str]) -> "KerasPypeModel":
+        model: KerasBaseModel = keras_load_model(folder / cls.KERAS_MODEL_FILE)
+
+        serialiser = JoblibSerialiser()
+        loss = serialiser.deserialise(folder / cls.LOSS_FILE)
+        optimizer = model.optimizer
+        return cls(inputs=inputs, outputs=outputs, model=model, seed=1, loss=loss, optimizer_class=optimizer.__class__)
+
+    @classmethod
+    def get_parameters(cls: Type["KerasPypeModel"], parser: ArgumentParser) -> None:
+        """Get and add parameters to initialise this class.
+
+        KerasPypeModel's will work by requiring 2 ways to instantiate a Model:
+            - through `model`, which is a keras model.
+            - through parameters, which will instantiate the model internally.
+
+        Arguments are taken from both the Keras Model class and the Pype Model class.
+
+        Args:
+            parser (ArgumentParser): The ArgumentParser to add arguments to.
+        """
+        super().get_parameters(parser)
+        BaseModel = cls._get_annotated_class()
+
+        add_args_to_parser_for_class(parser, BaseModel, "model", [KerasBaseModel], excluded_args=[])
+
+        add_args_to_parser_for_class(
+            parser,
+            cls,
+            "model",
+            [Model],
+            excluded_args=["seed", "inputs", "outputs", "model", "loss", "optimizer_class"],
+        )
diff --git a/pype/tensorflow/model/mlp_keras.py b/pype/tensorflow/model/mlp_keras.py
new file mode 100644
index 0000000..873032b
--- /dev/null
+++ b/pype/tensorflow/model/mlp_keras.py
@@ -0,0 +1,39 @@
+from typing import Any
+
+from keras import Model
+from keras.layers import Dense
+from tensorflow import Tensor  # type: ignore
+
+
+class MLPKeras(Model):
+    def __init__(
+        self,
+        output_size: int,
+        n_layers: int,
+        layer_size: int,
+        activation: str = None,
+        output_activation: str = None,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        """A simple model showing how to use a MLP from Keras with Pype.
+
+        Args:
+            output_size (int): The dimensions of the output
+            n_layers (int): The number of hidden layers
+            layer_size (int): The dimensions of hidden layers.
+            activation (str, optional): The activation function of hidden layers. Defaults to None.
+            output_activation (str, optional): The output activation function. Defaults to None.
+        """
+        super().__init__(*args, **kwargs)
+        assert n_layers >= 0, "Need at least 1 layer in the model."
+
+        activations = [activation for _ in range(n_layers)] + [output_activation]
+        sizes = [layer_size for _ in range(n_layers)] + [output_size]
+        self._layers = [Dense(size, activation=act) for act, size in zip(activations, sizes)]
+
+    def call(self, inputs: Tensor) -> Tensor:
+        """Standard Keras model API call."""
+        for layer in self._layers:
+            inputs = layer(inputs)
+        return inputs
diff --git a/pype/tensorflow/model/mlp_pype_model.py b/pype/tensorflow/model/mlp_pype_model.py
new file mode 100644
index 0000000..a3322ba
--- /dev/null
+++ b/pype/tensorflow/model/mlp_pype_model.py
@@ -0,0 +1,11 @@
+from typing import Any
+
+from .keras_pype_model import KerasPypeModel
+from .mlp_keras import MLPKeras
+
+
+class MLPPypeModel(KerasPypeModel[MLPKeras]):
+    """A Keras model integrated with Pype's APIs."""
+
+    def _init_model(self, args: dict[str, Any]) -> MLPKeras:
+        return MLPKeras(**args)
diff --git a/pype/tensorflow/pipeline/__init__.py b/pype/tensorflow/pipeline/__init__.py
new file mode 100644
index 0000000..485b3ce
--- /dev/null
+++ b/pype/tensorflow/pipeline/__init__.py
@@ -0,0 +1,1 @@
+from .tensor_checker import TensorflowTypeChecker
diff --git a/pype/tensorflow/pipeline/tensor_checker.py b/pype/tensorflow/pipeline/tensor_checker.py
new file mode 100644
index 0000000..0ae2d4d
--- /dev/null
+++ b/pype/tensorflow/pipeline/tensor_checker.py
@@ -0,0 +1,100 @@
+import tensorflow as tf
+from pydantic import create_model
+from tensorflow import DType, Tensor  # type: ignore
+
+from pype.base.pipeline.type_checker import DataModel, TypeChecker
+
+
+class TensorflowData(DataModel):
+    data: list
+
+    """An object that can be converted to a Tensorflow tensor."""
+
+    def convert(self) -> Tensor:
+        """Converts this object to a Tensorflow tensor.
+
+        Returns:
+            Tensor: The Tensorflow tensor contained by this object.
+        """
+        return tf.convert_to_tensor(self.data)  # type: ignore
+
+    @classmethod
+    def to_model(cls, data: Tensor) -> "TensorflowData":
+        """Converts a Tensorflow tensor to a TensorflowData model, which can be serialised.
+
+        Args:
+            data (Tensor): A Tensorflow tensor to serialise.
+
+        Returns:
+            TensorflowData: A serialisable version of the tensor.
+        """
+        return TensorflowData(data=data.numpy().tolist())
+
+
+class TensorflowTypeChecker(TypeChecker[Tensor]):
+    def __init__(self) -> None:
+        """Type checker for Tensorflow.
+
+        Checks if the incoming data is of the correct non-first dimensions and dtype.
+        """
+        super().__init__()
+        self.dims: tuple = tuple()
+        self.dtype: type | None = None
+
+    def fit(self, data: Tensor) -> "TensorflowTypeChecker":
+        """Fit this Tensorflow TypeChecker to the given data.
+
+        Args:
+            data (Tensor): The data to fit.
+
+        Returns:
+            TensorflowTypeChecker: self.
+        """
+        self.dims = data.shape[1:]
+        self.dtype = self._convert_dtype(data.dtype)
+        return self
+
+    def transform(self, data: Tensor) -> Tensor:
+        """Checks if the given data fits the specifications this TypeChecker was fitted for.
+
+        Args:
+            data (Tensor):  The data to check.
+
+        Returns:
+            Tensor: data, if the data fits the specifications. Otherwise, an assertion error is thrown.
+        """
+        assert self.dtype is not None, "Please fit pipeline first"
+        assert isinstance(data, Tensor), "Please provide a Tensorflow tensor!"
+        assert (
+            data.shape[1:] == self.dims
+        ), f"Dimensions of Tensorflow tensors do not add up: {data.shape[1:]} vs {self.dims}"
+
+        converted_type = self._convert_dtype(data.dtype)
+        assert converted_type == self.dtype, f"Dtype of data does not add up: {converted_type} vs {self.dtype}"
+        return data
+
+    def _convert_dtype(self, dtype: DType) -> type:
+        dtype_name = str(dtype)
+        if "int" in dtype_name:
+            return int
+        elif "float" in dtype_name:
+            return float
+        elif "bool" in dtype_name:
+            return bool
+        else:
+            return str
+
+    def get_pydantic_type(self) -> type[TensorflowData]:
+        """Creates a Pydantic model for this data to handle serialisation/deserialisation.
+
+        Returns:
+            type[TensorflowData]: A TensorflowData model that fits the data this wat fitted on.
+        """
+        base_iter: type = list[self.dtype]  # type: ignore
+
+        for _ in range(len(self.dims)):
+            base_iter = list[base_iter]  # type: ignore
+
+        model = create_model("TensorflowData", data=(base_iter, ...), __base__=TensorflowData)
+
+        return model
diff --git a/pype/tensorflow/setup.py b/pype/tensorflow/setup.py
new file mode 100644
index 0000000..4e7e572
--- /dev/null
+++ b/pype/tensorflow/setup.py
@@ -0,0 +1,20 @@
+from setuptools import find_namespace_packages, setup
+
+version = "0.1.0"
+
+deps = [
+    f"pype-base=={version}",
+    # on mac, it is recommended to use conda/mamba or source to install tensorflow
+    "tensorflow>=2.9.1",
+    "numpy>=1.23.0",
+    "protobuf>=3.20.1",
+]
+strict_deps = [s.replace(">=", "==") for s in deps]
+
+setup(
+    name="pype-tensorflow",
+    install_requires=deps,
+    extras_require={"dev": strict_deps, "strict": strict_deps},
+    packages=find_namespace_packages(include=["pype.*"]),
+    version=version,
+)
diff --git a/pype/xgboost/__init__.py b/pype/xgboost/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pype/xgboost/model/__init__.py b/pype/xgboost/model/__init__.py
new file mode 100644
index 0000000..30ddc30
--- /dev/null
+++ b/pype/xgboost/model/__init__.py
@@ -0,0 +1,2 @@
+from .xgboost_classifier import XGBClassifierModel
+from .xgboost_regressor import XGBRegressorModel
diff --git a/pype/xgboost/model/xgboost_classifier.py b/pype/xgboost/model/xgboost_classifier.py
new file mode 100644
index 0000000..186b908
--- /dev/null
+++ b/pype/xgboost/model/xgboost_classifier.py
@@ -0,0 +1,23 @@
+from pathlib import Path
+from typing import Any, Type
+
+from xgboost.sklearn import XGBClassifier
+
+from pype.sklearn.model import SklearnModel
+
+
+class XGBClassifierModel(SklearnModel[XGBClassifier]):
+    XGB_MODEL_FILE = "model.txt"
+
+    def _init_model(self, args: dict[str, Any]) -> XGBClassifier:
+        return XGBClassifier(**args)
+
+    def _save(self, folder: Path) -> None:
+        assert isinstance(self.model, XGBClassifier)
+        self.model.save_model(folder / self.XGB_MODEL_FILE)
+
+    @classmethod
+    def _load(cls: Type["XGBClassifierModel"], folder: Path, inputs: list[str], outputs: list[str]) -> "SklearnModel":
+        model = XGBClassifier()
+        model.load_model(folder / cls.XGB_MODEL_FILE)
+        return cls(inputs=inputs, outputs=outputs, model=model, seed=1)
diff --git a/pype/xgboost/model/xgboost_regressor.py b/pype/xgboost/model/xgboost_regressor.py
new file mode 100644
index 0000000..0afe900
--- /dev/null
+++ b/pype/xgboost/model/xgboost_regressor.py
@@ -0,0 +1,23 @@
+from pathlib import Path
+from typing import Any, Type
+
+from xgboost.sklearn import XGBRegressor
+
+from pype.sklearn.model import SklearnModel
+
+
+class XGBRegressorModel(SklearnModel[XGBRegressor]):
+    XGB_MODEL_FILE = "model.txt"
+
+    def _init_model(self, args: dict[str, Any]) -> XGBRegressor:
+        return XGBRegressor(**args)
+
+    def _save(self, folder: Path) -> None:
+        assert isinstance(self.model, XGBRegressor)
+        self.model.save_model(folder / self.XGB_MODEL_FILE)
+
+    @classmethod
+    def _load(cls: Type["XGBRegressorModel"], folder: Path, inputs: list[str], outputs: list[str]) -> "SklearnModel":
+        model = XGBRegressor()
+        model.load_model(folder / cls.XGB_MODEL_FILE)
+        return cls(inputs=inputs, outputs=outputs, model=model, seed=1)
diff --git a/pype/xgboost/setup.py b/pype/xgboost/setup.py
new file mode 100644
index 0000000..9c2fdc2
--- /dev/null
+++ b/pype/xgboost/setup.py
@@ -0,0 +1,18 @@
+from setuptools import find_namespace_packages, setup
+
+version = "0.1.0"
+
+deps = [
+    f"pype-base=={version}",
+    "pype-sklearn",
+    "xgboost>=1.6.2",
+]
+strict_deps = [s.replace(">=", "==") for s in deps]
+
+setup(
+    name="pype-xgboost",
+    install_requires=deps,
+    extras_require={"dev": strict_deps, "strict": strict_deps},
+    packages=find_namespace_packages(include=["pype.*"]),
+    version=version,
+)
diff --git a/pyproject.toml b/pyproject.toml
index 7d7bff0..0735012 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,9 @@ markers = [
     "end_to_end: End to end test, tends to be slow.",
     "spark: uses Spark, will take extra time to set up."
 ]
+filterwarnings = [
+    'ignore:Call to deprecated create function',  # Some tensorboard internal stuff
+]
 
 [tool.black]
 line-length = 120
@@ -30,3 +33,11 @@ exclude = '''
   | \.toml
 )/
 '''
+
+[tool.mypy]
+mypy_path = "pype"
+ignore_missing_imports = true
+disallow_untyped_defs = true
+python_version = "3.10"
+implicit_reexport = true
+
diff --git a/scripts/build.sh b/scripts/build.sh
new file mode 100755
index 0000000..1fc423d
--- /dev/null
+++ b/scripts/build.sh
@@ -0,0 +1,24 @@
+#! /bin/bash
+
+# Look for output-dir and packages flags.
+while getopts "o:p:" flag
+do
+    case "${flag}" in
+        h) output_dir=${OPTARG};;
+        e) packages=${OPTARG};;
+    esac
+done
+output_dir="${editable:-dist}"
+if [ -z $packages ]
+then
+    packages=($(ls pype))
+fi
+
+# Build all requested packages.
+for package in "${packages[@]}"
+do
+    package_path="pype/${package}"
+    echo "Building: ${package_path}"
+    # install_package $package_path $host $editable_string
+    python -m build $package_path --outdir $output_dir
+done
\ No newline at end of file
diff --git a/scripts/dev_install.sh b/scripts/dev_install.sh
new file mode 100755
index 0000000..22a7e65
--- /dev/null
+++ b/scripts/dev_install.sh
@@ -0,0 +1,67 @@
+#! /bin/bash
+
+function install_package () {
+    # 0: script name
+    # 1: package name
+    # 2: host
+    # 3: editable. Ignored if host is set.
+    raw_name=$1
+    package="${raw_name/.//}"
+
+    echo "host: $2"
+    
+    if [ $2 != "local" ];
+    then
+        echo "Installing $1 from host $2"
+        ip=`echo $2 | grep -oE [0-9]+\.[0-9]+\.[0-9]+\.[0-9]+`
+        echo "IP: $ip"
+        pip install -i $2 $1 --trusted-host $ip
+    else
+        echo "Installing ${package} from local machine"
+        pip install $3 "${package}[dev]"
+    fi;
+}
+
+# Look for editable and host flags.
+while getopts "e:h:" flag
+do
+    case "${flag}" in
+        h) host=${OPTARG};;
+        e) editable=${OPTARG};;
+    esac
+done
+host="${host:-local}"
+editable="${editable:-1}"
+
+echo "Editable state: $editable, host: $host"
+
+# Set editable
+if [ $editable == 1 ];
+then
+    editable_string=" -e "
+else
+    editable_string=""
+fi;
+
+# Install priority packages first: base and sklearn. These are packages that others depend on.
+# Please note, that especially tensorflow may run into issues. On some Mac machines (M1 versions)
+# the result installation may fail. It is recommended in those cases to first manually install
+# tensorflow, then install all pype packages.
+priority_packages=( "pype.base" "pype.sklearn" )
+for priority_package in "${priority_packages[@]}"
+do
+    echo "Priority installing: ${priority_package}"
+    install_package $priority_package $host $editable_string
+done
+
+# Install all remaining packages.
+packages=($(ls pype))
+for package in "${packages[@]}"
+do
+    package_path="pype.${package}"
+    if [[ ! " ${priority_packages[*]} " =~ " ${package_path} " ]]; then
+        # only install if package wasn't priority package
+        echo "Installing: ${package_path}"
+        install_package $package_path $host $editable_string
+    fi
+done
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
index 268422a..acb1bbe 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -4,4 +4,8 @@ extend-ignore=E203
 
 [tool:pytest]
 # Required to run on BDP, since for some reason the current dir is not added to pythonpath by BDP's python.
-pythonpath="."
\ No newline at end of file
+pythonpath="."
+
+[metadata]
+license_files=LICENSE.txt
+description-file=READM.md
\ No newline at end of file
diff --git a/setup.py b/setup.py
index c5d7c3b..2d218af 100644
--- a/setup.py
+++ b/setup.py
@@ -5,5 +5,8 @@
     packages=find_namespace_packages(include=["pype.*"]),
     namespace_packages=["pype"],
     python_requires=">=3.10",
-    entry_points={"console_scripts": ["pype = pype.cmd_tools:main"]},
+    version="0.1.0",
+    license="MIT",
+    author="Jeroen van den Hoven",
+    url="https://github.com/jeroenvdhoven/pype",
 )


Test patch:
diff --git a/tests/base/deploy/__init__.py b/tests/base/deploy/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/base/deploy/test_app.py b/tests/base/deploy/test_app.py
deleted file mode 100644
index 9e86919..0000000
--- a/tests/base/deploy/test_app.py
+++ /dev/null
@@ -1,79 +0,0 @@
-import json
-import shutil
-from pathlib import Path
-from unittest.mock import patch
-
-from fastapi import FastAPI
-from fastapi.testclient import TestClient
-from pytest import fixture
-
-from pype.base.data.dataset import DataSet
-from pype.base.deploy.app import PypeApp
-from pype.base.deploy.inference import Inferencer
-from tests.test_utils import DummyDataModel, get_dummy_experiment
-
-
-@fixture(scope="module")
-def experiment_path():
-    experiment = get_dummy_experiment()
-    output_folder = experiment.output_folder
-
-    experiment.run()
-
-    yield output_folder
-
-    shutil.rmtree(output_folder)
-
-
-@fixture
-def app(experiment_path: Path):
-    return PypeApp("dummy-app", experiment_path)
-
-
-@fixture()
-def test_client(app: PypeApp) -> TestClient:
-    return TestClient(app.create_app())
-
-
-class Test_create_app:
-    def test_loading(self, app: PypeApp):
-        with patch.object(app, "_load_model") as mock_load:
-            mock_inferencer = mock_load.return_value
-            mock_inferencer.input_type_checker.get_pydantic_types.return_value = DummyDataModel
-            mock_inferencer.output_type_checker.get_pydantic_types.return_value = DummyDataModel
-
-            created_app = app.create_app()
-            assert isinstance(created_app, FastAPI)
-
-            mock_load.assert_called_once_with()
-            mock_inferencer.input_type_checker.get_pydantic_types.assert_called_once_with()
-            mock_inferencer.output_type_checker.get_pydantic_types.assert_called_once_with()
-
-
-class Test_app:
-    def test_home_page(self, test_client: TestClient):
-        response = test_client.get("/")
-        assert response.status_code == 200
-
-        content = json.loads(response.content.decode())
-        assert content == "Welcome to the Pype FastAPI app for dummy-app"
-
-    def test_predict(self, test_client: TestClient, experiment_path: Path):
-        x = [1, 2, 3, 4]
-        response = test_client.post("/predict", json={"x": {"data": x}})
-        assert response.status_code == 200
-
-        content = json.loads(response.content.decode())
-        assert "y" in content
-        assert "data" in content["y"]
-        prediction = content["y"]["data"]
-
-        assert isinstance(prediction, list)
-        for p in prediction:
-            assert isinstance(p, (float, int))
-
-        # compare to local reading
-        inferencer = Inferencer.from_folder(experiment_path)
-        y_true = inferencer.predict(DataSet(x=x))["y"]
-
-        assert y_true == prediction
diff --git a/tests/base/deploy/test_inferencer.py b/tests/base/deploy/test_inferencer.py
index d042fb8..596e531 100644
--- a/tests/base/deploy/test_inferencer.py
+++ b/tests/base/deploy/test_inferencer.py
@@ -1,6 +1,5 @@
-import shutil
 from pathlib import Path
-from unittest.mock import MagicMock, call, patch
+from unittest.mock import MagicMock, call, mock_open, patch
 
 from pytest import fixture
 
@@ -8,42 +7,48 @@
 from pype.base.data.dataset import DataSet
 from pype.base.data.dataset_source import DataSetSource
 from pype.base.deploy.inference import Inferencer
-from tests.test_utils import get_dummy_data, get_dummy_experiment
+from pype.base.experiment import Experiment
+from tests.shared_fixtures import dummy_experiment
+from tests.utils import get_dummy_data
 
+dummy_experiment
 
-@fixture(scope="module")
-def experiment_path():
-    experiment = get_dummy_experiment()
-    output_folder = experiment.output_folder
-
-    experiment.run()
 
-    yield output_folder
-
-    shutil.rmtree(output_folder)
+@fixture(scope="module")
+def experiment_path(dummy_experiment: Experiment):
+    dummy_experiment.run()
+    yield dummy_experiment.output_folder
 
 
 class Test_Inferencer:
     def test_from_folder(self):
         path = Path("something")
+        abs_path = path.absolute()
 
         pipeline = MagicMock()
         inputs = MagicMock()
         outputs = MagicMock()
         deserialised_values = [pipeline, inputs, outputs]
 
+        read_data = '{"paths": ["1", "2"]}'
+        m_open = mock_open(read_data=read_data)
         with patch("pype.base.deploy.inference.Model.load") as mock_model_load, patch(
             "pype.base.deploy.inference.JoblibSerialiser.deserialise", side_effect=deserialised_values
-        ) as mock_deserialise:
+        ) as mock_deserialise, patch("pype.base.deploy.inference.switch_workspace") as mock_switch, patch(
+            "pype.base.deploy.inference.open", m_open
+        ):
             result = Inferencer.from_folder(path)
 
-            mock_model_load.assert_called_once_with(path / Constants.MODEL_FOLDER)
+            m_open.assert_called_once_with(abs_path / Constants.EXTRA_FILES, "r")
+
+            mock_switch.assert_called_once_with(abs_path, ["1", "2"])
+            mock_model_load.assert_called_once_with(abs_path / Constants.MODEL_FOLDER)
 
             mock_deserialise.assert_has_calls(
                 [
-                    call(path / Constants.PIPELINE_FILE),
-                    call(path / Constants.INPUT_TYPE_CHECKER_FILE),
-                    call(path / Constants.OUTPUT_TYPE_CHECKER_FILE),
+                    call(abs_path / Constants.PIPELINE_FILE),
+                    call(abs_path / Constants.INPUT_TYPE_CHECKER_FILE),
+                    call(abs_path / Constants.OUTPUT_TYPE_CHECKER_FILE),
                 ]
             )
 
@@ -52,6 +57,29 @@ def test_from_folder(self):
             assert result.input_type_checker == inputs
             assert result.output_type_checker == outputs
 
+    def test_from_experiment(self):
+        pipeline = MagicMock()
+        model = MagicMock()
+        itc = MagicMock()
+        otc = MagicMock()
+
+        experiment = Experiment(
+            data_sources={"train": MagicMock()},
+            pipeline=pipeline,
+            model=model,
+            evaluator=MagicMock(),
+            input_type_checker=itc,
+            output_type_checker=otc,
+            logger=MagicMock(),
+        )
+
+        result = Inferencer.from_experiment(experiment)
+
+        assert result.model == model
+        assert result.pipeline == pipeline
+        assert result.input_type_checker == itc
+        assert result.output_type_checker == otc
+
     def test_predict(self):
         model = MagicMock()
         pipeline = MagicMock()
diff --git a/tests/base/experiment/__init__.py b/tests/base/experiment/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/base/experiment/a.py b/tests/base/experiment/a.py
new file mode 100644
index 0000000..3769f5b
--- /dev/null
+++ b/tests/base/experiment/a.py
@@ -0,0 +1,1 @@
+# dummy, empty. for tests
diff --git a/tests/base/experiment/test_argument_parsing.py b/tests/base/experiment/test_argument_parsing.py
index 34586e4..554f6cc 100644
--- a/tests/base/experiment/test_argument_parsing.py
+++ b/tests/base/experiment/test_argument_parsing.py
@@ -18,7 +18,7 @@
 )
 from pype.base.pipeline.pipe import Pipe
 from pype.base.pipeline.pipeline import Pipeline
-from tests.test_utils import pytest_assert
+from tests.utils import pytest_assert
 
 
 class Test_convert_bool:
diff --git a/tests/base/experiment/test_experiment.py b/tests/base/experiment/test_experiment.py
new file mode 100644
index 0000000..0fe2211
--- /dev/null
+++ b/tests/base/experiment/test_experiment.py
@@ -0,0 +1,492 @@
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+from unittest.mock import MagicMock, call, mock_open, patch
+
+from pytest import fixture, mark
+
+from pype.base.constants import Constants
+from pype.base.experiment.experiment import Experiment
+from tests.shared_fixtures import dummy_experiment
+from tests.utils import get_dummy_data, pytest_assert
+
+dummy_experiment
+
+
+class Test_run:
+    @fixture(scope="class")
+    def run_experiment(self, dummy_experiment: Experiment) -> Iterable[tuple[Experiment, dict]]:
+        metrics = dummy_experiment.run()
+
+        yield dummy_experiment, metrics
+
+    def test_unit(self):
+        data_sources = {"train": MagicMock(), "test": MagicMock()}
+        model = MagicMock()
+        pipeline = MagicMock()
+        evaluator = MagicMock()
+        logger = MagicMock()
+        serialiser = MagicMock()
+        input_type_checker = MagicMock()
+        output_type_checker = MagicMock()
+        output_folder = Path("tmp")
+        additional_files_to_store = ["a", "b"]
+
+        experiment = Experiment(
+            data_sources=data_sources,
+            model=model,
+            pipeline=pipeline,
+            evaluator=evaluator,
+            logger=logger,
+            serialiser=serialiser,
+            input_type_checker=input_type_checker,
+            output_type_checker=output_type_checker,
+            output_folder=output_folder,
+            additional_files_to_store=additional_files_to_store,
+        )
+
+        train_transform = MagicMock()
+        test_transform = MagicMock()
+        pipeline.transform.side_effect = [train_transform, test_transform]
+
+        train_performance = MagicMock()
+        test_performance = MagicMock()
+        evaluator.evaluate.side_effect = [train_performance, test_performance]
+
+        with patch.object(experiment, "_create_output_folders") as mock_create_folders, patch.object(
+            experiment, "_log_extra_files"
+        ) as mock_log_extra_files, patch.object(experiment, "_log_requirements") as mock_log_requirements:
+            result = experiment.run()
+
+        logger.__enter__.assert_called_once()
+        logger.__exit__.assert_called_once()
+
+        # getting data
+        data_sources["train"].read.assert_called_once_with()
+        data_sources["test"].read.assert_called_once_with()
+
+        # input checker
+        dataset_train = data_sources["train"].read.return_value
+        dataset_test = data_sources["test"].read.return_value
+        input_type_checker.fit.assert_called_once_with(dataset_train)
+        input_type_checker.transform.assert_has_calls([call(dataset_train), call(dataset_test)])
+
+        # pipeline fitting
+        pipeline.fit.assert_called_once_with(dataset_train)
+        pipeline.transform.assert_has_calls([call(dataset_train), call(dataset_test)])
+
+        # model fitting
+        model.fit.assert_called_once_with(train_transform)
+        model.transform.assert_called_once_with(train_transform)
+        output_type_checker.fit.assert_called_once_with(model.transform.return_value)
+
+        # evaluation
+        evaluator.evaluate.assert_has_calls([call(model, train_transform), call(model, test_transform)])
+        assert result == {
+            "train": train_performance,
+            "test": test_performance,
+        }
+
+        # logging
+        logger.log_metrics.assert_has_calls(
+            [
+                call("train", train_performance),
+                call("test", test_performance),
+            ]
+        )
+
+        mock_create_folders.assert_called_once_with()
+        logger.log_model.assert_called_once_with(model, output_folder / Constants.MODEL_FOLDER)
+        logger.log_artifact.assert_has_calls(
+            [
+                call(output_folder / Constants.PIPELINE_FILE, serialiser, object=pipeline),
+                call(output_folder / Constants.INPUT_TYPE_CHECKER_FILE, serialiser, object=input_type_checker),
+                call(output_folder / Constants.OUTPUT_TYPE_CHECKER_FILE, serialiser, object=output_type_checker),
+            ]
+        )
+        logger.log_parameters.assert_called_once_with({})
+
+        # extra files
+        mock_log_extra_files.assert_called_once_with()
+        mock_log_requirements.assert_called_once_with()
+
+    def test_integration(self, run_experiment: tuple[Experiment, dict]):
+        experiment, metrics = run_experiment
+        for ds_name in ["train", "test"]:
+            assert ds_name in metrics
+            assert "diff" in metrics[ds_name]
+
+        test_data = get_dummy_data(10, 2, 3).read()
+        y = test_data["y"]
+        predictions = experiment.model.transform(experiment.pipeline.transform(test_data))["y"]
+
+        assert len(y) == len(predictions)
+
+
+def test_log_extra_files():
+    cwd = Path(__file__).parent
+
+    data_sources = {"train": MagicMock(), "test": MagicMock()}
+    model = MagicMock()
+    pipeline = MagicMock()
+    evaluator = MagicMock()
+    logger = MagicMock()
+    serialiser = MagicMock()
+    input_type_checker = MagicMock()
+    output_type_checker = MagicMock()
+    output_folder = Path("tmp")
+    additional_files_to_store = [cwd / "a.py"]
+
+    experiment = Experiment(
+        data_sources=data_sources,
+        model=model,
+        pipeline=pipeline,
+        evaluator=evaluator,
+        logger=logger,
+        serialiser=serialiser,
+        input_type_checker=input_type_checker,
+        output_type_checker=output_type_checker,
+        output_folder=output_folder,
+        additional_files_to_store=additional_files_to_store,
+    )
+
+    m_open = mock_open()
+    with patch("pype.base.experiment.experiment.os.getcwd", return_value=cwd) as mock_getcwd, patch(
+        "pype.base.experiment.experiment.open", m_open
+    ), patch("pype.base.experiment.experiment.json.dump") as mock_dump:
+        experiment._log_extra_files()
+
+        mock_getcwd.assert_called_once_with()
+        logger.log_local_file.assert_called_once_with(Path("a.py"), output_folder / "a.py")
+        logger.log_file.assert_called_once_with(output_folder / Constants.EXTRA_FILES)
+
+        m_open.assert_called_once_with(output_folder / Constants.EXTRA_FILES, "w")
+        opened_obj = m_open.return_value
+
+        mock_dump.assert_called_once_with({"paths": ["a.py"]}, opened_obj)
+
+
+@dataclass
+class VersionInfo:
+    major: int
+    minor: int
+    micro: int
+
+
+@fixture
+def version_info():
+    old_version = sys.version_info
+
+    sys.version_info = VersionInfo(major=4, minor=10, micro=129)
+
+    yield sys.version_info
+    sys.version_info = old_version
+
+
+def test_log_requirements(version_info: VersionInfo):
+    cwd = Path(__file__).parent
+
+    data_sources = {"train": MagicMock(), "test": MagicMock()}
+    model = MagicMock()
+    pipeline = MagicMock()
+    evaluator = MagicMock()
+    logger = MagicMock()
+    serialiser = MagicMock()
+    input_type_checker = MagicMock()
+    output_type_checker = MagicMock()
+    output_folder = Path("tmp")
+    additional_files_to_store = [cwd / "a.py"]
+
+    experiment = Experiment(
+        data_sources=data_sources,
+        model=model,
+        pipeline=pipeline,
+        evaluator=evaluator,
+        logger=logger,
+        serialiser=serialiser,
+        input_type_checker=input_type_checker,
+        output_type_checker=output_type_checker,
+        output_folder=output_folder,
+        additional_files_to_store=additional_files_to_store,
+    )
+
+    f1 = MagicMock()
+    f2 = MagicMock()
+    req_text = b"pandas==0.1.0\nnumpy=1.0.1"
+    with patch("pype.base.experiment.experiment.open", side_effect=[f1, f2]) as mock_open, patch(
+        "pype.base.experiment.experiment.subprocess.check_output", return_value=req_text
+    ) as mock_check_output, patch("pype.base.experiment.experiment.json.dump") as mock_dump:
+        experiment._log_requirements()
+
+        mock_open.assert_has_calls(
+            [
+                call(output_folder / Constants.PYTHON_VERSION_FILE, "w"),
+                call(output_folder / Constants.REQUIREMENTS_FILE, "w"),
+            ]
+        )
+
+        # logger.log_local_file.assert_called_once_with(Path("a.py"), output_folder / "a.py")
+        # logger.log_file.assert_called_once_with(output_folder / Constants.EXTRA_FILES)
+        # python version
+        mock_dump.assert_called_once_with(
+            {
+                "python_version": "4.10.129",
+                "major": version_info.major,
+                "minor": version_info.minor,
+                "micro": version_info.micro,
+            },
+            f1.__enter__.return_value,
+        )
+
+        # requirements
+        mock_check_output.assert_called_once_with([sys.executable, "-m", "pip", "freeze"])
+        f2.__enter__.return_value.write.assert_called_once_with(req_text.decode())
+
+        # logging
+        logger.log_file.assert_has_calls(
+            [
+                call(output_folder / Constants.PYTHON_VERSION_FILE),
+                call(output_folder / Constants.REQUIREMENTS_FILE),
+            ],
+            any_order=True,
+        )
+
+
+class Test_init:
+    def test_assert_train_name(self):
+        data_sources = {"test": MagicMock()}
+        model = MagicMock()
+        pipeline = MagicMock()
+        evaluator = MagicMock()
+        logger = MagicMock()
+        serialiser = MagicMock()
+        input_type_checker = MagicMock()
+        output_type_checker = MagicMock()
+        output_folder = MagicMock()
+        additional_files_to_store = MagicMock()
+
+        args = {}
+
+        with pytest_assert(AssertionError, "Must provide a 'train' entry in the data_sources dictionary."):
+            Experiment(
+                data_sources=data_sources,
+                model=model,
+                pipeline=pipeline,
+                evaluator=evaluator,
+                logger=logger,
+                serialiser=serialiser,
+                input_type_checker=input_type_checker,
+                output_type_checker=output_type_checker,
+                output_folder=output_folder,
+                additional_files_to_store=additional_files_to_store,
+                parameters=args,
+            )
+
+    def test_warning_on_no_args(self):
+        data_sources = {"train": MagicMock()}
+        model = MagicMock()
+        pipeline = MagicMock()
+        evaluator = MagicMock()
+        logger = MagicMock()
+        serialiser = MagicMock()
+        input_type_checker = MagicMock()
+        output_type_checker = MagicMock()
+        output_folder = MagicMock()
+        additional_files_to_store = MagicMock()
+
+        with patch("pype.base.experiment.experiment.getLogger") as mock_get_logger:
+            Experiment(
+                data_sources=data_sources,
+                model=model,
+                pipeline=pipeline,
+                evaluator=evaluator,
+                logger=logger,
+                serialiser=serialiser,
+                input_type_checker=input_type_checker,
+                output_type_checker=output_type_checker,
+                output_folder=output_folder,
+                additional_files_to_store=additional_files_to_store,
+            )
+
+            mock_get_logger.assert_called_once()
+            logger = mock_get_logger.return_value
+
+            logger.warning.assert_called_once_with(
+                """It is highly recommended to provide the parameters used to initialise your
+run here for logging purposes. Consider using the `from_command_line` or
+`from_dictionary` initialisation methods"""
+            )
+
+
+def test_from_dictionary():
+    data_sources = MagicMock()
+    model_class = MagicMock()
+    pipeline = MagicMock()
+    evaluator = MagicMock()
+    logger = MagicMock()
+    serialiser = MagicMock()
+    input_type_checker = MagicMock()
+    output_type_checker = MagicMock()
+    model_inputs = MagicMock()
+    model_outputs = MagicMock()
+    output_folder = MagicMock()
+    additional_files_to_store = MagicMock()
+    seed = 3
+
+    args = {"model__a": 1, "model__b": 3, "pipeline__step1__c": 3, "pipeline__step2__d": 5}
+
+    with patch.object(Experiment, "__init__", return_value=None) as mock_init:
+        result = Experiment.from_dictionary(
+            data_sources=data_sources,
+            model_class=model_class,
+            pipeline=pipeline,
+            evaluator=evaluator,
+            logger=logger,
+            serialiser=serialiser,
+            input_type_checker=input_type_checker,
+            output_type_checker=output_type_checker,
+            model_inputs=model_inputs,
+            model_outputs=model_outputs,
+            output_folder=output_folder,
+            additional_files_to_store=additional_files_to_store,
+            parameters=args,
+            seed=seed,
+        )
+
+        assert isinstance(result, Experiment)
+        model_class.assert_called_once_with(a=1, b=3, seed=seed, inputs=model_inputs, outputs=model_outputs)
+        pipeline.reinitialise.assert_called_once_with({"step1__c": 3, "step2__d": 5})
+
+        mock_init.assert_called_once_with(
+            data_sources=data_sources,
+            model=model_class.return_value,
+            pipeline=pipeline,
+            evaluator=evaluator,
+            logger=logger,
+            serialiser=serialiser,
+            output_folder=output_folder,
+            input_type_checker=input_type_checker,
+            output_type_checker=output_type_checker,
+            additional_files_to_store=additional_files_to_store,
+            parameters=args,
+        )
+
+
+def test_copy(dummy_experiment: Experiment):
+    exp = dummy_experiment
+    params = {"model__a": 4}
+
+    result = exp.copy(params)
+
+    assert result.model.a == 4
+    assert result.model != exp.model
+    assert result.model.__class__ == exp.model.__class__
+
+    assert result.pipeline != exp.pipeline
+    assert len(result.pipeline) == len(exp.pipeline)
+    assert result.pipeline[0].name == exp.pipeline[0].name
+
+    assert result.data_sources == exp.data_sources
+    assert result.additional_files_to_store == exp.additional_files_to_store
+    assert result.evaluator == exp.evaluator
+    assert result.experiment_logger == exp.experiment_logger
+    assert result.input_type_checker == exp.input_type_checker
+    assert result.output_type_checker == exp.output_type_checker
+    assert result.output_folder == exp.output_folder
+    assert result.serialiser == exp.serialiser
+
+
+@mark.parametrize(
+    ["name", "fixed_args", "expected_dict_extras"],
+    [
+        ["no extras", None, {}],
+        ["extras", {"model__8": 9}, {"model__8": 9}],
+        ["overwrite", {"a": 120}, {"a": 120}],
+    ],
+)
+def test_from_command_line(name, fixed_args, expected_dict_extras):
+    class ParserOutput:
+        def __init__(self, a, b) -> None:
+            self.a = a
+            self.b = b
+
+    data_sources = MagicMock()
+    model_class = MagicMock()
+    pipeline = MagicMock()
+    evaluator = MagicMock()
+    logger = MagicMock()
+    serialiser = MagicMock()
+    input_type_checker = MagicMock()
+    output_type_checker = MagicMock()
+    model_inputs = MagicMock()
+    model_outputs = MagicMock()
+    output_folder = MagicMock()
+    additional_files_to_store = MagicMock()
+    seed = 3
+
+    args = ParserOutput(2, 4)
+
+    with patch.object(Experiment, "_get_cmd_args") as mock_get_cmd, patch.object(
+        Experiment, "from_dictionary"
+    ) as mock_from_dictionary:
+        mock_parser = mock_get_cmd.return_value
+        mock_parser.parse_known_args.return_value = (args, None)
+
+        expected_args = args.__dict__.copy()
+        expected_args.update(expected_dict_extras)
+
+        result = Experiment.from_command_line(
+            data_sources=data_sources,
+            model_class=model_class,
+            pipeline=pipeline,
+            evaluator=evaluator,
+            logger=logger,
+            serialiser=serialiser,
+            input_type_checker=input_type_checker,
+            output_type_checker=output_type_checker,
+            model_inputs=model_inputs,
+            model_outputs=model_outputs,
+            output_folder=output_folder,
+            additional_files_to_store=additional_files_to_store,
+            seed=seed,
+            fixed_arguments=fixed_args,
+        )
+
+        assert result == mock_from_dictionary.return_value
+        mock_get_cmd.assert_called_once_with(model_class, pipeline)
+        mock_parser.parse_known_args.assert_called_once_with()
+
+        mock_from_dictionary.assert_called_once_with(
+            data_sources=data_sources,
+            model_class=model_class,
+            pipeline=pipeline,
+            evaluator=evaluator,
+            logger=logger,
+            serialiser=serialiser,
+            input_type_checker=input_type_checker,
+            output_type_checker=output_type_checker,
+            model_inputs=model_inputs,
+            model_outputs=model_outputs,
+            output_folder=output_folder,
+            additional_files_to_store=additional_files_to_store,
+            parameters=expected_args,
+            seed=seed,
+        )
+
+
+def test_get_cmd_args():
+    with patch("pype.base.experiment.experiment.ArgumentParser") as mock_parser_class, patch(
+        "pype.base.experiment.experiment.add_args_to_parser_for_pipeline"
+    ) as mock_add_args:
+        mock_parser = mock_parser_class.return_value
+        mock_model_class = MagicMock()
+        pipeline = MagicMock()
+
+        result = Experiment._get_cmd_args(mock_model_class, pipeline)
+
+        assert result == mock_parser
+        mock_parser_class.assert_called_once_with()
+
+        mock_model_class.get_parameters.assert_called_once_with(mock_parser)
+        mock_add_args.assert_called_once_with(mock_parser, pipeline)
diff --git a/tests/base/logger/__init__.py b/tests/base/logger/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/base/logger/test_experiment_logger.py b/tests/base/logger/test_experiment_logger.py
new file mode 100644
index 0000000..8f01aa1
--- /dev/null
+++ b/tests/base/logger/test_experiment_logger.py
@@ -0,0 +1,63 @@
+# we test shared functions through LocalLogger
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from pype.base.logger.local_logger import LocalLogger
+from pype.base.serialiser.serialiser import Serialiser
+
+
+class Test_experiment_logger:
+    def test_log_metrics(self):
+        logger = LocalLogger()
+        metrics = {"a": 2, "b4": 9}
+        ds_name = "some name"
+
+        with patch.object(logger, "_log_metrics") as mock_log:
+            logger.log_metrics(ds_name, metrics)
+
+            mock_log.assert_called_once_with(
+                {
+                    f"{ds_name}_a": 2,
+                    f"{ds_name}_b4": 9,
+                }
+            )
+
+    def test_log_model(self):
+        model = MagicMock()
+        folder = Path("some")
+        logger = LocalLogger()
+
+        with patch.object(logger, "log_file") as mock_log:
+            logger.log_model(model, folder)
+        model.save.assert_called_once_with(folder)
+        mock_log.assert_called_once_with(folder)
+
+    def test_log_artifact(self):
+        file = Path("dummy")
+        obj = MagicMock()
+        serialiser = MagicMock(spec=Serialiser)
+        logger = LocalLogger()
+
+        with patch.object(logger, "log_file") as mock_log_file:
+            logger.log_artifact(file, serialiser, obj)
+
+        serialiser.serialise.assert_called_once_with(obj, file)
+        mock_log_file.assert_called_once_with(file)
+
+    def test_log_local_file(self):
+        logger = LocalLogger()
+        file = "a.py"
+        output = "target"
+
+        with patch("pype.base.logger.experiment_logger.shutil") as mock_shutil, patch.object(
+            logger, "log_file"
+        ) as mock_log_file, patch("pype.base.logger.experiment_logger.Path") as mock_path:
+            logger.log_local_file(file, output)
+
+            mock_path.assert_called_once_with(output)
+            real_out_target = mock_path.return_value
+
+            real_out_target.parent.mkdir.assert_called_once_with(exist_ok=True)
+            mock_shutil.copy.assert_called_once_with(file, real_out_target)
+            mock_log_file.assert_called_once_with(real_out_target)
diff --git a/tests/base/logger/test_local_logger.py b/tests/base/logger/test_local_logger.py
new file mode 100644
index 0000000..d58f4a9
--- /dev/null
+++ b/tests/base/logger/test_local_logger.py
@@ -0,0 +1,42 @@
+# we test shared functions through LocalLogger
+
+from unittest.mock import call, patch
+
+from pype.base.logger.local_logger import LocalLogger
+
+
+class Test_local_logger:
+    def test_log_metrics(self):
+        ws = 5
+        logger = LocalLogger(ws)
+        metrics = {"a": 2, "b4": 9}
+
+        with patch("pype.base.logger.local_logger.print") as mock_print:
+            logger._log_metrics(metrics)
+
+            mock_print.assert_has_calls(
+                [
+                    call("a    : 2"),
+                    call("b4   : 9"),
+                ]
+            )
+
+    def test_log_parameters(self):
+        ws = 5
+        logger = LocalLogger(ws)
+        parameters = {"a": 2, "b4": 9}
+
+        with patch("pype.base.logger.local_logger.print") as mock_print:
+            logger.log_parameters(parameters)
+
+            mock_print.assert_has_calls(
+                [
+                    call("a    : 2"),
+                    call("b4   : 9"),
+                ]
+            )
+
+    def test_log_file(self):
+        # just shouldn't break.
+        logger = LocalLogger()
+        logger.log_file("")
diff --git a/tests/base/pipeline/test_pipe.py b/tests/base/pipeline/test_pipe.py
index cda8233..526e57d 100644
--- a/tests/base/pipeline/test_pipe.py
+++ b/tests/base/pipeline/test_pipe.py
@@ -4,6 +4,7 @@
 
 from pype.base.data.dataset import DataSet
 from pype.base.pipeline.pipe import Pipe
+from tests.utils import pytest_assert
 
 
 class TestPipe:
@@ -41,6 +42,17 @@ def test_init(self, operator: MagicMock):
         )
         operator.assert_called_once_with(a=1, b="satr")
 
+    def test_init_name_check(self, operator: MagicMock):
+        with pytest_assert(AssertionError, "Pipe names cannot contain the string `__`"):
+            Pipe(
+                "name__something",
+                operator=operator,
+                inputs=["a", "b"],
+                outputs=["out", "out2"],
+                fit_inputs=["c"],
+                kw_args={"a": 1, "b": "satr"},
+            )
+
     def test_fit(self, data: DataSet[int], pipe: Pipe, operator: MagicMock):
         mock_operator = operator.return_value
         result = pipe.fit(data)
@@ -129,3 +141,36 @@ def test_reinitialise(self, pipe: Pipe, operator: MagicMock):
         assert pipe.operator != operator
         assert pipe.operator == mock_class.return_value
         mock_class.assert_called_once_with(a=3, b=2)
+
+    def test_copy(self, pipe: Pipe, operator: MagicMock):
+        # instantiating the operator generates different objects
+        operator.side_effect = [1, 2]
+        result = pipe.copy()
+
+        assert result.args == pipe.args
+        assert result.operator_class == pipe.operator_class
+        assert result.name == pipe.name
+        assert result.inputs == pipe.inputs
+        assert result.outputs == pipe.outputs
+        assert result.fit_inputs == pipe.fit_inputs
+
+        # the object is not completely the same.
+        assert result.operator != pipe.operator
+
+    def test_copy_with_args(self, pipe: Pipe, operator: MagicMock):
+        # instantiating the operator generates different objects
+        operator.side_effect = [1, 2]
+        args = {"a": 3}
+        result = pipe.copy(args)
+
+        assert result.args != pipe.args
+        assert result.args == args
+        assert result.operator_class == pipe.operator_class
+        assert result.name == pipe.name
+        assert result.inputs == pipe.inputs
+        assert result.outputs == pipe.outputs
+        assert result.fit_inputs == pipe.fit_inputs
+
+        # the object is not completely the same.
+        assert result.operator != pipe.operator
+        result.operator_class.assert_called_with(a=3)
diff --git a/tests/base/pipeline/test_pipeline.py b/tests/base/pipeline/test_pipeline.py
index 24cbd88..688630a 100644
--- a/tests/base/pipeline/test_pipeline.py
+++ b/tests/base/pipeline/test_pipeline.py
@@ -5,7 +5,7 @@
 from pype.base.data.dataset import DataSet
 from pype.base.pipeline.pipe import Pipe
 from pype.base.pipeline.pipeline import Pipeline
-from tests.test_utils import pytest_assert
+from tests.utils import pytest_assert
 
 
 class Test_Pipeline:
@@ -184,6 +184,41 @@ def test_reinitialise_with_pipeline(
             }
         )
 
+    def test_copy(self):
+        sub_pipe_1_op_class = MagicMock()
+        sub_pipe_1 = Pipe("step0", sub_pipe_1_op_class, [], [])
+
+        sub_pipe_2_op_class = MagicMock()
+        sub_pipe_2 = Pipe("step1", sub_pipe_2_op_class, [], [])
+
+        sub_pipe_pipeline_op_class = MagicMock()
+        sub_pipe_pipeline = Pipe("step2", sub_pipe_pipeline_op_class, [], [])
+
+        sub_pipeline = Pipeline([sub_pipe_pipeline])
+        pipeline = Pipeline([sub_pipe_1, sub_pipe_2, sub_pipeline])
+
+        args = {
+            "step0__a": 1,
+            "step0__b": 2,
+            "step1__c": 3,
+            "step2__d": 4,
+        }
+        result = pipeline.copy(args)
+
+        expected_args = [{"a": 1, "b": 2}, {"c": 3}, {"d": 4}]
+        expected_pipes = [sub_pipe_1, sub_pipe_2, sub_pipe_pipeline]
+        for expected, actual, args in zip(expected_pipes, result.pipes, expected_args):
+            assert isinstance(actual, Pipe)
+            assert expected != actual
+
+            assert actual.args == args
+            assert actual.operator_class == expected.operator_class
+            actual.operator_class.assert_called_with(**args)
+            assert actual.name == expected.name
+            assert actual.inputs == expected.inputs
+            assert actual.outputs == expected.outputs
+            assert actual.fit_inputs == expected.fit_inputs
+
     def test_get_item_int(self, pipeline: Pipeline, pipes: list[MagicMock]):
         assert pipeline[0] == pipes[0]
         assert pipeline[1] == pipes[1]
diff --git a/tests/base/pipeline/test_type_checker.py b/tests/base/pipeline/test_type_checker.py
new file mode 100644
index 0000000..bd1245d
--- /dev/null
+++ b/tests/base/pipeline/test_type_checker.py
@@ -0,0 +1,154 @@
+from typing import Type
+from unittest.mock import call, patch
+
+from pydantic import create_model
+from pytest import fixture
+
+from pype.base.data.dataset import DataSet
+from pype.base.pipeline.type_checker import (
+    DataModel,
+    DataSetModel,
+    DataSetTypeChecker,
+    TypeCheckerPipe,
+)
+from tests.utils import DummyDataModel, DummyTypeChecker, get_dummy_data, pytest_assert
+
+
+class DummyDataStr(DataModel):
+    data: list[str]
+
+    def convert(self) -> list[str]:
+        return self.data
+
+    @classmethod
+    def to_model(cls, data: list[str]) -> "DataModel":
+        return cls(data=data)
+
+
+class Test_DataSetModel:
+    @fixture
+    def ModelClass(self) -> Type[DataSetModel]:
+        return create_model(
+            "DummyDataSetModel", floats=(DummyDataModel, ...), strs=(DummyDataStr, ...), __base__=DataSetModel
+        )
+
+    def test_convert(self, ModelClass: Type[DataSetModel]):
+        model = ModelClass(
+            floats=DummyDataModel(data=[1.0, 3, 4]),
+            strs=DummyDataStr(data=["a", "b"]),
+        )
+
+        result = model.convert()
+
+        assert result["floats"] == [1.0, 3, 4]
+        assert result["strs"] == ["a", "b"]
+
+    def test_to_model(self, ModelClass: Type[DataSetModel]):
+        dataset = DataSet(floats=[1.0, 3, 4], strs=["a", "b"])
+        result = ModelClass.to_model(dataset)
+
+        expected = ModelClass(
+            floats=DummyDataModel(data=[1.0, 3, 4]),
+            strs=DummyDataStr(data=["a", "b"]),
+        )
+
+        assert expected == result
+
+
+class Test_TypeCheckerPipe:
+    def test_get_pydantic_types(self):
+        pipe = TypeCheckerPipe(
+            "inputs",
+            ["x"],
+            [tuple()],
+        )
+
+        with patch.object(pipe.operator, "get_pydantic_types") as mock_get:
+            result = pipe.get_pydantic_types()
+            assert result == mock_get.return_value
+
+            mock_get.assert_called_once_with(["x"])
+
+    def test_get_pydantic_types_with_names(self):
+        pipe = TypeCheckerPipe(
+            "inputs",
+            ["x"],
+            [tuple()],
+        )
+
+        with patch.object(pipe.operator, "get_pydantic_types") as mock_get:
+            names = ["a", "b", "c"]
+            result = pipe.get_pydantic_types(names)
+            assert result == mock_get.return_value
+
+            mock_get.assert_called_once_with(names)
+
+
+class Test_DataSetTypeChecker:
+    @fixture
+    def type_checker(self, data: DataSet) -> DataSetTypeChecker:
+        return DataSetTypeChecker(input_names=list(data.keys()), type_checker_classes=[(list, DummyTypeChecker)])
+
+    @fixture
+    def data(self) -> DataSet:
+        return get_dummy_data(10, 2, 3).read()
+
+    def test_fit(self, type_checker: DataSetTypeChecker, data: DataSet):
+        with patch.object(DummyTypeChecker, "fit") as mock_fit:
+            type_checker.fit(data["x"], data["y"])
+
+            mock_fit.assert_has_calls([call(data["x"]), call(data["y"])])
+
+            for name in data.keys():
+                assert name in type_checker.type_checkers
+                assert isinstance(type_checker.type_checkers[name], DummyTypeChecker)
+
+    def test_fit_warning(self, type_checker: DataSetTypeChecker):
+        dummy_data = {"z": 1}
+
+        with patch.object(DummyTypeChecker, "fit") as mock_fit, patch(
+            "pype.base.pipeline.type_checker.logging"
+        ) as mock_logging:
+            mock_logger = mock_logging.getLogger.return_value
+
+            type_checker.fit(dummy_data)
+            mock_fit.assert_not_called()
+            mock_logger.warning.assert_called_once_with("x has no supported type checker!")
+
+    def test_transform(self, type_checker: DataSetTypeChecker, data: DataSet):
+        type_checker.fit(data["x"], data["y"])
+
+        with patch.object(type_checker.type_checkers["x"], "transform") as mock_x_transform, patch.object(
+            type_checker.type_checkers["y"], "transform"
+        ) as mock_y_transform:
+            x, y = type_checker.transform(data["x"], data["y"])
+
+            mock_x_transform.assert_called_once_with(data["x"])
+            mock_y_transform.assert_called_once_with(data["y"])
+
+            assert x == data["x"]
+            assert y == data["y"]
+
+    def test_fit_transform_integration(self, type_checker: DataSetTypeChecker, data: DataSet):
+        type_checker.fit(data["x"], data["y"])
+
+        # good
+        trans_x, trans_y = type_checker.transform(data["x"], data["y"])
+        assert trans_x == data["x"]
+        assert trans_y == data["y"]
+
+        # expected to fail
+        with pytest_assert(AssertionError, "Provide a list!", exact=False):
+            bad_data = {"z": 1}
+            type_checker.transform(bad_data)
+
+    def test_get_pydantic_types(self, type_checker: DataSetTypeChecker, data: DataSet):
+        # done as integration test
+        type_checker.fit(data["x"], data["y"])
+
+        DummyDataSetModel: Type[DataSetModel] = type_checker.get_pydantic_types()
+
+        model = DummyDataSetModel.to_model(data)
+
+        assert model.x.data == data["x"]
+        assert model.y.data == data["y"]
diff --git a/tests/base/serialiser/__init__.py b/tests/base/serialiser/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/base/serialiser/test_joblib_serialiser.py b/tests/base/serialiser/test_joblib_serialiser.py
new file mode 100644
index 0000000..3dba80a
--- /dev/null
+++ b/tests/base/serialiser/test_joblib_serialiser.py
@@ -0,0 +1,23 @@
+from unittest.mock import patch
+
+from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
+
+
+class Test_JoblibSerialiser:
+    def test_save(self):
+        with patch("pype.base.serialiser.joblib_serialiser.dump") as mock_dump:
+            serialiser = JoblibSerialiser()
+            obj = [None]
+            file = "some file"
+            serialiser.serialise(obj, file)
+
+            mock_dump.assert_called_once_with(obj, file)
+
+    def test_load(self):
+        with patch("pype.base.serialiser.joblib_serialiser.load") as mock_load:
+            serialiser = JoblibSerialiser()
+            file = "some file"
+            result = serialiser.deserialise(file)
+
+            mock_load.assert_called_once_with(file)
+            assert mock_load.return_value == result
diff --git a/tests/base/utils/dummy_folder/a.py b/tests/base/utils/dummy_folder/a.py
new file mode 100644
index 0000000..c7d8b58
--- /dev/null
+++ b/tests/base/utils/dummy_folder/a.py
@@ -0,0 +1,2 @@
+def probably_not_known_form_of_hello():
+    return "Hello!"
diff --git a/tests/base/utils/test_switch_workspace.py b/tests/base/utils/test_switch_workspace.py
new file mode 100644
index 0000000..a116655
--- /dev/null
+++ b/tests/base/utils/test_switch_workspace.py
@@ -0,0 +1,141 @@
+import importlib
+import os
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, call, patch
+
+from pype.base.utils.workspace import (
+    _create_temporary_workspace,
+    _reset_workspace,
+    switch_workspace,
+)
+
+
+def test_reset_workspace():
+    old_workspace = "old_workspace"
+    old_sys_path = ["a", "b"]
+    main_lib = MagicMock()
+
+    v3 = MagicMock()
+    main_lib.__dict__ = {"var1": MagicMock(), "var2": MagicMock(), "var3": v3}
+    main_to_remove = {"var1", "var2"}
+
+    with patch("pype.base.utils.workspace.os") as mock_os, patch("pype.base.utils.workspace.sys") as mock_sys:
+        _reset_workspace(old_workspace, old_sys_path, main_lib, main_to_remove)
+
+    mock_os.chdir.assert_called_once_with(old_workspace)
+
+    assert mock_sys.path == old_sys_path
+    assert main_lib.__dict__ == {"var3": v3}
+
+
+def make_moduled_variable(module: str):
+    var = MagicMock()
+    var.__module__ = module
+    return var
+
+
+def test_create_temporary_workspace():
+    target_workspace = Path("directory")
+    extra_files = ["a.py", "example/b.py", "dummy_file.txt"]  # the last file will be ignored.
+    main_lib = MagicMock()
+
+    v1 = MagicMock()
+    v2 = MagicMock()
+    v3 = MagicMock()
+    main_lib.__dict__ = {"var1": v1, "var2": v2, "var3": v3}
+
+    a_lib = MagicMock()
+    a_lib.__dict__ = {
+        "var1": make_moduled_variable("directory.a"),  # will be ignored because of name
+        "var_new_1": make_moduled_variable("directory.other_module"),  # will be ignored because of module
+        "var_new_2": make_moduled_variable("directory.a"),  # will be added
+    }
+    example_b_lib = MagicMock()
+    example_b_lib.__dict__ = {
+        "var_new_2": make_moduled_variable("directory.a"),  # will be ignored since a_lib has it
+        "var_new_3": make_moduled_variable("directory.example.b"),  # will be added
+    }
+    libs = [a_lib, example_b_lib]
+
+    with patch("pype.base.utils.workspace.os.chdir") as mock_chdir, patch(
+        "pype.base.utils.workspace.sys"
+    ) as mock_sys, patch("pype.base.utils.workspace.importlib.import_module", side_effect=libs) as mock_import:
+        result = _create_temporary_workspace(target_workspace, extra_files, main_lib)
+
+        mock_chdir.assert_called_once_with(target_workspace)
+        mock_sys.path.append.assert_has_calls(
+            [
+                call(str(target_workspace)),
+                call(os.path.join(".", "a.py")),
+                call(os.path.join(".", "example/b.py")),
+            ],
+            any_order=True,
+        )
+
+        mock_import.assert_has_calls([call("directory.a"), call("directory.example.b")])
+
+        assert main_lib.__dict__ == {
+            "var1": v1,
+            "var2": v2,
+            "var3": v3,
+            "var_new_2": a_lib.__dict__["var_new_2"],
+            "var_new_3": example_b_lib.__dict__["var_new_3"],
+        }
+
+        assert result == {"var_new_2", "var_new_3"}
+
+
+def test_switch_workspace():
+    path = Path(".")
+    extra_files = ["a.py"]
+
+    with patch("pype.base.utils.workspace._create_temporary_workspace") as mock_create, patch(
+        "pype.base.utils.workspace._reset_workspace"
+    ) as mock_reset, patch("pype.base.utils.workspace.importlib.import_module") as mock_importlib:
+        cwd = os.getcwd()
+        c_sys_path = sys.path.copy()
+
+        with switch_workspace(path, extra_files):
+            pass
+
+        mock_create.assert_called_once_with(path, extra_files, mock_importlib.return_value)
+        mock_reset.assert_called_once_with(cwd, c_sys_path, mock_importlib.return_value, mock_create.return_value)
+        mock_importlib.assert_called_once_with("__main__")
+
+
+def _is_dir_in_sys_path(path: Path) -> bool:
+    return str(path.absolute()) in [str(Path(p).absolute()) for p in sys.path]
+
+
+def test_switch_workspace_integration():
+    path = Path(__file__).parent / "dummy_folder"
+    extra_files = ["a.py"]
+
+    cwd = os.getcwd()
+    c_sys_path = sys.path.copy()
+    main_lib = importlib.import_module("__main__")
+
+    try:
+        assert not _is_dir_in_sys_path(path)
+        assert "probably_not_known_form_of_hello" not in main_lib.__dict__
+        with switch_workspace(path, extra_files):
+            assert os.getcwd() != cwd
+            assert Path(os.getcwd()).absolute() == path.absolute()
+
+            # make sure the new directory is added to path, temporarily
+            assert _is_dir_in_sys_path(path)
+
+            # make sure the 'hello' function is imported correctly
+            # see dummy_folder/a.py for the function
+            assert "probably_not_known_form_of_hello" in main_lib.__dict__
+            result = main_lib.__dict__["probably_not_known_form_of_hello"]()
+            assert result == "Hello!"
+
+        assert os.getcwd() == cwd
+        assert sys.path == c_sys_path
+        assert not _is_dir_in_sys_path(path)
+        assert "probably_not_known_form_of_hello" not in main_lib.__dict__
+    finally:
+        sys.path = c_sys_path
+        os.chdir(cwd)
diff --git a/tests/fastapi/test_app.py b/tests/fastapi/test_app.py
new file mode 100644
index 0000000..ce13cdd
--- /dev/null
+++ b/tests/fastapi/test_app.py
@@ -0,0 +1,230 @@
+import json
+from datetime import datetime
+from pathlib import Path
+from unittest.mock import MagicMock, call, patch
+
+from data.data_sink import DataSink
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from pytest import fixture
+
+from pype.base.data.dataset import DataSet
+from pype.base.deploy.inference import Inferencer
+from pype.base.experiment import Experiment
+from pype.base.pipeline.type_checker import DataSetModel
+from pype.fastapi.deploy.app import PypeApp, write_in_background
+from tests.shared_fixtures import dummy_experiment
+from tests.utils import (
+    AnyArg,
+    DummyDataModel,
+    DummyDataSet,
+    DummyDataSink,
+    get_dummy_data,
+)
+
+dummy_experiment
+
+
+@fixture(scope="module")
+def experiment_path(dummy_experiment: Experiment):
+    dummy_experiment.run()
+
+    yield dummy_experiment.output_folder
+
+
+@fixture
+def app(experiment_path: Path):
+    return PypeApp("dummy-app", experiment_path)
+
+
+@fixture
+def tracking():
+    return {
+        "x": DummyDataSink(),
+        "y": DummyDataSink(),
+    }
+
+
+@fixture
+def app_with_tracking(experiment_path: Path, tracking: dict[str, DataSink]):
+    return PypeApp("dummy-app", experiment_path, tracking_servers=tracking)
+
+
+@fixture()
+def test_client(app: PypeApp) -> TestClient:
+    return TestClient(app.create_app())
+
+
+@fixture()
+def test_client_with_tracking(app_with_tracking: PypeApp) -> TestClient:
+    return TestClient(app_with_tracking.create_app())
+
+
+class Test_create_app:
+    def test_loading(self, app: PypeApp):
+        with patch.object(app, "_load_model") as mock_load, patch.object(
+            app, "_verify_tracking_servers"
+        ) as mock_verify, patch("pype.fastapi.deploy.app.getLogger") as mock_get_logger:
+            mock_inferencer = mock_load.return_value
+            mock_inferencer.input_type_checker.get_pydantic_types.return_value = DummyDataSet
+            mock_inferencer.output_type_checker.get_pydantic_types.return_value = DummyDataSet
+
+            created_app = app.create_app()
+            assert isinstance(created_app, FastAPI)
+
+            mock_load.assert_called_once_with()
+            mock_inferencer.input_type_checker.get_pydantic_types.assert_called_once_with()
+            mock_inferencer.output_type_checker.get_pydantic_types.assert_called_once_with()
+
+            mock_get_logger.assert_called_once_with(app.name)
+
+            # input and output
+            mock_verify.assert_called_once_with(DummyDataSet, DummyDataSet, mock_get_logger.return_value)
+
+
+class Test_verify_tracking_servers:
+    def test_verify_tracking_servers_without_tracking(self, app: PypeApp):
+        # This should just run, since nothing is checked.
+        app._verify_tracking_servers(MagicMock(), MagicMock(), MagicMock())
+
+    def test_verify_tracking_servers_with_tracking(self, app_with_tracking: PypeApp):
+        # This should just run, since nothing is checked.
+        logger = MagicMock()
+        app_with_tracking._verify_tracking_servers(DummyDataSet, DummyDataSet, logger)
+
+        logger.warning.assert_not_called()
+
+    def test_verify_tracking_servers_with_tracking_raise_warning(self, app_with_tracking: PypeApp):
+        # This should just run, since nothing is checked.
+        class DummyDataSet(DataSetModel):
+            x: DummyDataModel
+            z: DummyDataModel
+
+        logger = MagicMock()
+        app_with_tracking._verify_tracking_servers(DummyDataSet, DummyDataSet, logger)
+
+        logger.warning.assert_called_once_with(
+            f"No dataset named `y` found in the fields of input or output DataSetModels"
+        )
+
+
+class Test_handle_tracking:
+    def test_handle_tracking_without_tracking(self, app: PypeApp):
+        data = get_dummy_data(10, 2, 1)
+
+        # should just run.
+        app._handle_tracking(DataSet(x=data["x"]), DataSet(y=data["y"]), MagicMock(), MagicMock())
+
+    def test_handle_tracking_with_tracking(self, app_with_tracking: PypeApp, tracking: dict[str, MagicMock]):
+        data = get_dummy_data(10, 2, 1).read()
+
+        background_tasks = MagicMock()
+        # should just run.
+        app_with_tracking._handle_tracking(DataSet(x=data["x"]), DataSet(y=data["y"]), MagicMock(), background_tasks)
+
+        background_tasks.add_task.assert_has_calls(
+            [
+                call(write_in_background, sink=tracking["x"], data=data["x"]),
+                call(write_in_background, sink=tracking["y"], data=data["y"]),
+            ],
+            any_order=True,
+        )
+
+    def test_handle_tracking_handles_crashes(self, app_with_tracking: PypeApp, tracking: dict[str, MagicMock]):
+        data = get_dummy_data(10, 2, 1).read()
+
+        logger = MagicMock()
+        background_tasks = MagicMock()
+        # This will call add_task to throw an error, causing only 1 background task to be executed.
+        background_tasks.add_task.side_effect = [ConnectionError("Dummy connection error!"), None]
+
+        app_with_tracking._handle_tracking(DataSet(x=data["x"]), DataSet(y=data["y"]), logger, background_tasks)
+
+        background_tasks.add_task.assert_has_calls(
+            [
+                call(write_in_background, sink=tracking["x"], data=data["x"]),
+                call(write_in_background, sink=tracking["y"], data=data["y"]),
+            ],
+            any_order=True,
+        )
+
+        logger.error.assert_called_once_with(f"Encountered error while sending data to x: Dummy connection error!")
+
+
+def test_write_in_background():
+    sink = MagicMock()
+    data = MagicMock()
+    write_in_background(sink, data)
+
+    sink.write.assert_called_once_with(data)
+
+
+class Test_app:
+    def test_home_page(self, test_client: TestClient):
+        response = test_client.get("/")
+        assert response.status_code == 200
+
+        content = json.loads(response.content.decode())
+        assert content == "Welcome to the Pype FastAPI app for dummy-app"
+
+    def test_predict(self, test_client: TestClient, experiment_path: Path):
+        x = [1, 2, 3, 4]
+        response = test_client.post("/predict", json={"x": {"data": x}})
+        assert response.status_code == 200
+
+        content = json.loads(response.content.decode())
+        assert "y" in content
+        assert "data" in content["y"]
+        prediction = content["y"]["data"]
+
+        assert isinstance(prediction, list)
+        for p in prediction:
+            assert isinstance(p, (float, int))
+
+        # compare to local reading
+        inferencer = Inferencer.from_folder(experiment_path)
+        y_true = inferencer.predict(DataSet(x=x))["y"]
+
+        assert y_true == prediction
+
+    def test_predict_with_tracking(self, app_with_tracking: PypeApp):
+        x = [1.0, 2.0, 3.0, 4.0]
+
+        with patch.object(app_with_tracking, "_handle_tracking") as mock_handle, patch(
+            "pype.fastapi.deploy.app.getLogger"
+        ) as mock_get_logger:
+            app = app_with_tracking.create_app()
+            test_client = TestClient(app)
+
+            response = test_client.post("/predict", json={"x": {"data": x}})
+        assert response.status_code == 200
+
+        content = json.loads(response.content.decode())
+        prediction = content["y"]["data"]
+
+        # Cannot check `background` argument
+        mock_handle.assert_called_once_with(DataSet(x=x), DataSet(y=prediction), mock_get_logger.return_value, AnyArg())
+
+    def test_predict_with_tracking_integration(
+        self, test_client_with_tracking: TestClient, tracking: dict[str, DummyDataSink]
+    ):
+        x = [1.0, 2.0, 3.0, 4.0]
+
+        response = test_client_with_tracking.post("/predict", json={"x": {"data": x}})
+        assert response.status_code == 200
+
+        content = json.loads(response.content.decode())
+        prediction = content["y"]["data"]
+
+        # Wait for max 2 seconds
+        current_time = datetime.now()
+        succeeded = False
+        while not succeeded and ((datetime.now() - current_time).seconds < 2):
+            try:
+                assert tracking["x"].data == x
+                assert tracking["y"].data == prediction
+                succeeded = True
+            except AssertionError as e:
+                # Potentially expected, just continue
+                pass
+        assert succeeded, "Did not manage to receive background task result within 2 seconds!"
diff --git a/tests/hyperopt/__init__.py b/tests/hyperopt/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/hyperopt/test_optimise.py b/tests/hyperopt/test_optimise.py
new file mode 100644
index 0000000..2e82f73
--- /dev/null
+++ b/tests/hyperopt/test_optimise.py
@@ -0,0 +1,272 @@
+from unittest.mock import MagicMock, patch
+
+from hyperopt import Trials, hp
+from pytest import mark
+
+from pype.base.experiment import Experiment
+from pype.hyperopt.optimise import create_optimisation_function, optimise_experiment
+from tests.shared_fixtures import dummy_experiment
+from tests.utils import AnyArg, pytest_assert
+
+dummy_experiment
+
+
+class Test_create_optimisation_function:
+    def test_create_optimisation_function(self):
+        exp = MagicMock()
+        exp_copy = MagicMock()
+        exp.copy.return_value = exp_copy
+
+        metrics = {"test": {"diff": 9, "other": 2}, "train": {"diff": 3}}
+        exp_copy.run.return_value = metrics
+
+        opt_func = create_optimisation_function(
+            exp,
+            ("test", "diff"),
+            minimise_target=True,
+            seed=1,
+        )
+
+        parameters_opt = {"model__a": 0}
+        performance_opt = opt_func(parameters_opt)
+
+        exp.copy.assert_called_once_with(parameters_opt, seed=1)
+        exp_copy.run.assert_called_once_with()
+
+        assert performance_opt == 9
+
+    def test_create_optimisation_function_integration(self, dummy_experiment: Experiment):
+        exp = dummy_experiment
+        exp.model.a = 0
+        baseline_performance = exp.run()
+
+        opt_func = create_optimisation_function(
+            exp,
+            ("test", "diff"),
+            minimise_target=True,
+            seed=1,
+        )
+
+        parameters_opt = {"model__a": 0}
+        performance_opt = opt_func(parameters_opt)
+
+        assert performance_opt == baseline_performance["test"]["diff"]
+
+
+class Test_optimise_experiment:
+    def test_optimise_experiment(self):
+        exp = MagicMock()
+        exp_copy = MagicMock()
+        exp.copy.return_value = exp_copy
+
+        search_space = {
+            "model__a": hp.randint("a", -10, 10),
+        }
+        metric = ("test", "diff")
+
+        with patch("pype.hyperopt.optimise.create_optimisation_function") as mock_create, patch(
+            "pype.hyperopt.optimise.fmin"
+        ) as mock_fmin, patch("pype.hyperopt.optimise.space_eval") as mock_space, patch(
+            "pype.hyperopt.optimise.Trials.best_trial"
+        ) as mock_best_trial:
+            best_trial_id, performance, best_arguments, trials = optimise_experiment(
+                experiment_template=exp,
+                search_space=search_space,
+                target_metric=metric,
+                max_evals=10,
+                minimise_target=True,
+                trial_type="normal",
+                training_seed=1,
+                fmin_arg=2390,
+            )
+
+            mock_create.assert_called_once_with(exp, metric, True, seed=1)
+            mock_fmin.assert_called_once_with(
+                mock_create.return_value, search_space, max_evals=10, trials=AnyArg(), fmin_arg=2390
+            )
+            mock_space.assert_called_once_with(search_space, mock_fmin.return_value)
+
+            assert best_trial_id == mock_best_trial["tid"]
+            assert performance == mock_best_trial["result"]["loss"]
+
+            assert best_arguments == mock_space.return_value
+            assert isinstance(trials, Trials)
+
+    def test_trials_mongo(self):
+        exp = MagicMock()
+        exp_copy = MagicMock()
+        exp.copy.return_value = exp_copy
+
+        search_space = {
+            "model__a": hp.randint("a", -10, 10),
+        }
+        metric = ("test", "diff")
+        mongo_url = "mongo://db/jobs"
+        exp_key = "experiment_key"
+
+        with patch("pype.hyperopt.optimise.create_optimisation_function"), patch("pype.hyperopt.optimise.fmin"), patch(
+            "pype.hyperopt.optimise.space_eval"
+        ), patch("pype.hyperopt.optimise.MongoTrials") as mock_trial:
+            optimise_experiment(
+                experiment_template=exp,
+                search_space=search_space,
+                target_metric=metric,
+                max_evals=10,
+                minimise_target=True,
+                trial_type="mongo",
+                mongo_url="mongo://db/jobs",
+                mongo_exp_key=exp_key,
+                training_seed=1,
+                fmin_arg=2390,
+            )
+            mock_trial.assert_called_once_with(mongo_url, exp_key=exp_key)
+
+    @mark.parametrize(
+        ["error", "url"],
+        [
+            ["mongo_url can not be None if you want to run hyperopt with Mongo.", None],
+            ["hyperopt currently required the DB name to end in /jobs.", "something"],
+        ],
+    )
+    def test_trials_mongo_asserts(self, url: str, error: str):
+        exp = MagicMock()
+        exp_copy = MagicMock()
+        exp.copy.return_value = exp_copy
+
+        search_space = {
+            "model__a": hp.randint("a", -10, 10),
+        }
+        metric = ("test", "diff")
+        mongo_url = "mongo://db/jobs"
+        exp_key = "experiment_key"
+
+        with patch("pype.hyperopt.optimise.create_optimisation_function"), patch("pype.hyperopt.optimise.fmin"), patch(
+            "pype.hyperopt.optimise.space_eval"
+        ), patch("pype.hyperopt.optimise.MongoTrials") as mock_trial, pytest_assert(AssertionError, error):
+            optimise_experiment(
+                experiment_template=exp,
+                search_space=search_space,
+                target_metric=metric,
+                max_evals=10,
+                minimise_target=True,
+                trial_type="mongo",
+                mongo_url=url,
+                mongo_exp_key=exp_key,
+                training_seed=1,
+                fmin_arg=2390,
+            )
+            mock_trial.assert_called_once_with(mongo_url, exp_key=exp_key)
+
+    def test_trials_normal(self):
+        exp = MagicMock()
+        exp_copy = MagicMock()
+        exp.copy.return_value = exp_copy
+
+        search_space = {
+            "model__a": hp.randint("a", -10, 10),
+        }
+        metric = ("test", "diff")
+
+        with patch("pype.hyperopt.optimise.create_optimisation_function"), patch("pype.hyperopt.optimise.fmin"), patch(
+            "pype.hyperopt.optimise.space_eval"
+        ), patch("pype.hyperopt.optimise.Trials") as mock_trial:
+            optimise_experiment(
+                experiment_template=exp,
+                search_space=search_space,
+                target_metric=metric,
+                max_evals=10,
+                minimise_target=True,
+                trial_type="normal",
+                training_seed=1,
+                fmin_arg=2390,
+            )
+            mock_trial.assert_called_once_with()
+
+    def test_trials_spark(self):
+        exp = MagicMock()
+        exp_copy = MagicMock()
+        exp.copy.return_value = exp_copy
+
+        search_space = {
+            "model__a": hp.randint("a", -10, 10),
+        }
+        metric = ("test", "diff")
+
+        with patch("pype.hyperopt.optimise.create_optimisation_function"), patch("pype.hyperopt.optimise.fmin"), patch(
+            "pype.hyperopt.optimise.space_eval"
+        ), pytest_assert(
+            ValueError,
+            "Spark trials have not been implemented yet with auto-setup."
+            + "Provide a SparkTrials object to the trials argument instead.",
+        ):
+            optimise_experiment(
+                experiment_template=exp,
+                search_space=search_space,
+                target_metric=metric,
+                max_evals=10,
+                minimise_target=True,
+                trial_type="spark",
+                training_seed=1,
+                fmin_arg=2390,
+            )
+
+    def test_trials_provided(self):
+        exp = MagicMock()
+        exp_copy = MagicMock()
+        exp.copy.return_value = exp_copy
+
+        search_space = {
+            "model__a": hp.randint("a", -10, 10),
+        }
+        metric = ("test", "diff")
+        trials = MagicMock()
+
+        with patch("pype.hyperopt.optimise.create_optimisation_function") as mock_create, patch(
+            "pype.hyperopt.optimise.fmin"
+        ) as mock_fmin, patch("pype.hyperopt.optimise.space_eval"):
+            optimise_experiment(
+                experiment_template=exp,
+                search_space=search_space,
+                target_metric=metric,
+                max_evals=10,
+                minimise_target=True,
+                trial_type="normal",
+                training_seed=1,
+                trials=trials,
+                fmin_arg=2390,
+            )
+            mock_fmin.assert_called_once_with(
+                mock_create.return_value, search_space, max_evals=10, trials=trials, fmin_arg=2390
+            )
+
+    def test_maxisime_target_result(self):
+        pass
+
+    def test_optimise_experiment_integration(self, dummy_experiment: Experiment):
+        exp = dummy_experiment
+        # prevent logging
+        exp.experiment_logger = MagicMock()
+        search_space = {
+            "model__a": hp.randint("a", -10, 10),
+        }
+        metric = ("test", "diff")
+        n_evals = 6
+
+        best_trial_id, performance, best_arguments, trials = optimise_experiment(
+            experiment_template=exp,
+            search_space=search_space,
+            target_metric=metric,
+            max_evals=n_evals,
+            minimise_target=True,
+            trial_type="normal",
+            training_seed=1,
+        )
+
+        assert best_trial_id is not None
+        assert n_evals >= best_trial_id and best_trial_id >= 0
+        assert trials is not None
+        assert performance == trials.best_trial["result"]["loss"]
+
+        # Higher a leads to better performance. Highest a should have been selected.
+        max_a = max([t["misc"]["vals"]["a"] for t in trials.trials])
+        assert max_a == best_arguments["model__a"]
diff --git a/tests/mlflow/__init__.py b/tests/mlflow/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/mlflow/deploy/__init__.py b/tests/mlflow/deploy/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/mlflow/deploy/test_load_experiment.py b/tests/mlflow/deploy/test_load_experiment.py
new file mode 100644
index 0000000..d451d95
--- /dev/null
+++ b/tests/mlflow/deploy/test_load_experiment.py
@@ -0,0 +1,121 @@
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from pytest import mark
+
+from pype.mlflow.deploy.load_experiment import (
+    _download_and_load,
+    load_experiment_from_mlflow,
+)
+from tests.utils import pytest_assert
+
+
+def dummy_run_info(run_id: str) -> MagicMock:
+    mock = MagicMock()
+    mock.run_id = run_id
+    return mock
+
+
+def dummy_experiment(run_id: str) -> MagicMock:
+    mock = MagicMock()
+    mock.experiment_id = run_id
+    return mock
+
+
+class Test_load_experiment_from_mlflow:
+    def test_assertion_valid_experiment(self):
+        url = "some_url"
+        name = "some name"
+        run_id = "jkldsf9032sd"
+        with patch("pype.mlflow.deploy.load_experiment.get_experiment_by_name") as mock_get_experiment_by_name, patch(
+            "pype.mlflow.deploy.load_experiment._download_and_load"
+        ):
+            # experiment not found
+            mock_get_experiment_by_name.return_value = None
+
+            with pytest_assert(AssertionError, f"Experiment {name} does not exist in {url}."):
+                load_experiment_from_mlflow(url, name, run_id)
+
+    def test_assertion_valid_run_id(self):
+        url = "some_url"
+        name = "some name"
+        run_id = "jkldsf9032sd"
+        experiment_id = "238jlkdfala932la@@"
+        with patch("pype.mlflow.deploy.load_experiment.get_experiment_by_name") as mock_get_experiment_by_name, patch(
+            "pype.mlflow.deploy.load_experiment.list_run_infos"
+        ) as mock_list_run_infos, patch("pype.mlflow.deploy.load_experiment._download_and_load"):
+            # experiment not found
+            mock_get_experiment_by_name.return_value = dummy_experiment(experiment_id)
+            mock_list_run_infos.return_value = [dummy_run_info("others")]
+
+            with pytest_assert(AssertionError, f"Run ID {run_id} is not present in the given experiment."):
+                load_experiment_from_mlflow(url, name, run_id)
+
+    @mark.parametrize(["directory"], [["str-directory"], [Path("path-directory")]])
+    def test_with_folder(self, directory: str | Path):
+        url = "some_url"
+        name = "some name"
+        run_id = "jkldsf9032sd"
+        experiment_id = "238jlkdfala932la@@"
+        with patch("pype.mlflow.deploy.load_experiment.get_experiment_by_name") as mock_get_experiment_by_name, patch(
+            "pype.mlflow.deploy.load_experiment.list_run_infos"
+        ) as mock_list_run_infos, patch(
+            "pype.mlflow.deploy.load_experiment.set_tracking_uri"
+        ) as mock_set_tracking_uri, patch(
+            "pype.mlflow.deploy.load_experiment._download_and_load"
+        ) as mock_download:
+            mock_get_experiment_by_name.return_value = dummy_experiment(experiment_id)
+            mock_list_run_infos.return_value = [dummy_run_info(run_id), dummy_run_info("others")]
+
+            result = load_experiment_from_mlflow(url, name, run_id, directory)
+
+            mock_set_tracking_uri.assert_called_once_with(url)
+            mock_get_experiment_by_name.assert_called_once_with(name)
+            mock_list_run_infos.assert_called_once_with(experiment_id)
+
+            mock_download.assert_called_once_with(run_id, Path(directory))
+            assert result == mock_download.return_value
+
+    def test_with_tmp_folder(self):
+        url = "some_url"
+        name = "some name"
+        run_id = "jkldsf9032sd"
+        experiment_id = "238jlkdfala932la@@"
+        with patch("pype.mlflow.deploy.load_experiment.get_experiment_by_name") as mock_get_experiment_by_name, patch(
+            "pype.mlflow.deploy.load_experiment.list_run_infos"
+        ) as mock_list_run_infos, patch(
+            "pype.mlflow.deploy.load_experiment.set_tracking_uri"
+        ) as mock_set_tracking_uri, patch(
+            "pype.mlflow.deploy.load_experiment._download_and_load"
+        ) as mock_download, patch(
+            "pype.mlflow.deploy.load_experiment.TemporaryDirectory"
+        ) as mock_tmp_dir:
+            tmp_path = "tmp_path"
+            mock_tmp_dir.return_value.__enter__.return_value = tmp_path
+
+            mock_get_experiment_by_name.return_value = dummy_experiment(experiment_id)
+            mock_list_run_infos.return_value = [dummy_run_info(run_id), dummy_run_info("others")]
+
+            result = load_experiment_from_mlflow(url, name, run_id)
+
+            mock_set_tracking_uri.assert_called_once_with(url)
+            mock_get_experiment_by_name.assert_called_once_with(name)
+            mock_list_run_infos.assert_called_once_with(experiment_id)
+
+            mock_tmp_dir.assert_called_once_with(prefix="mlflow_model")
+
+            mock_download.assert_called_once_with(run_id, Path(tmp_path))
+            assert result == mock_download.return_value
+
+
+def test_download_and_load():
+    run_id = "run_id"
+    directory = Path("folder_with_model")
+    with patch("pype.mlflow.deploy.load_experiment.download_artifacts") as mock_download, patch(
+        "pype.mlflow.deploy.load_experiment.Inferencer.from_folder"
+    ) as mock_load:
+        result = _download_and_load(run_id=run_id, directory=directory)
+
+        mock_download.assert_called_once_with(f"runs:/{run_id}/", dst_path="folder_with_model")
+        mock_load.assert_called_once_with(directory)
+        assert result == mock_load.return_value
diff --git a/tests/mlflow/logger/__init__.py b/tests/mlflow/logger/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/mlflow/logger/test_mlflow_logger.py b/tests/mlflow/logger/test_mlflow_logger.py
new file mode 100644
index 0000000..3d9554d
--- /dev/null
+++ b/tests/mlflow/logger/test_mlflow_logger.py
@@ -0,0 +1,174 @@
+from pathlib import Path
+from unittest.mock import MagicMock, call, patch
+
+from git import InvalidGitRepositoryError
+from pytest import fixture
+
+from pype.mlflow.logger.mlflow_logger import MlflowLogger
+from tests.utils import pytest_assert
+
+
+class Test_mlflow_logger:
+    @fixture
+    def logger(self) -> MlflowLogger:
+        return MlflowLogger("example", "http://localhost:5000")
+
+    def test_enter(self, logger: MlflowLogger):
+        assert logger.run is None
+        with patch("pype.mlflow.logger.mlflow_logger.set_tracking_uri") as mock_set_tracking_uri, patch(
+            "pype.mlflow.logger.mlflow_logger.get_experiment_by_name"
+        ) as mock_get_experiment_by_name, patch(
+            "pype.mlflow.logger.mlflow_logger.create_experiment"
+        ) as mock_create_experiment, patch(
+            "pype.mlflow.logger.mlflow_logger.set_experiment"
+        ) as mock_set_experiment, patch(
+            "pype.mlflow.logger.mlflow_logger.start_run"
+        ) as mock_start_run, patch.object(
+            logger, "log_branch"
+        ) as mock_log_branch:
+            mock_get_experiment_by_name.return_value = []  # just not None
+            logger.__enter__()
+
+            mock_set_tracking_uri.assert_called_once_with(logger.uri)
+            mock_get_experiment_by_name.assert_called_once_with(logger.name)
+            mock_create_experiment.assert_not_called()
+            mock_set_experiment.assert_called_once_with(experiment_name=logger.name)
+
+            mock_start_run.assert_called_once_with()
+            mock_log_branch.assert_called_once_with()
+            run = mock_start_run.return_value
+            run.__enter__.assert_called_once_with()
+            active_run = run.__enter__.return_value
+
+            assert logger.run == active_run
+
+    def test_enter_new_experiment(self, logger: MlflowLogger):
+        assert logger.run is None
+        with patch("pype.mlflow.logger.mlflow_logger.set_tracking_uri") as mock_set_tracking_uri, patch(
+            "pype.mlflow.logger.mlflow_logger.get_experiment_by_name"
+        ) as mock_get_experiment_by_name, patch(
+            "pype.mlflow.logger.mlflow_logger.create_experiment"
+        ) as mock_create_experiment, patch(
+            "pype.mlflow.logger.mlflow_logger.set_experiment"
+        ) as mock_set_experiment, patch(
+            "pype.mlflow.logger.mlflow_logger.start_run"
+        ) as mock_start_run, patch.object(
+            logger, "log_branch"
+        ) as mock_log_branch:
+            mock_get_experiment_by_name.return_value = None
+            logger.__enter__()
+
+            mock_set_tracking_uri.assert_called_once_with(logger.uri)
+            mock_get_experiment_by_name.assert_called_once_with(logger.name)
+            mock_create_experiment.assert_called_once_with(logger.name, None)
+            mock_set_experiment.assert_called_once_with(experiment_name=logger.name)
+
+            mock_start_run.assert_called_once_with()
+            mock_log_branch.assert_called_once_with()
+            run = mock_start_run.return_value
+            run.__enter__.assert_called_once_with()
+            active_run = run.__enter__.return_value
+
+            assert logger.run == active_run
+
+    def test_log_branch(self, logger: MlflowLogger):
+        branch_name = "branch/name"
+        repo = MagicMock()
+        repo.active_branch.name = branch_name
+        with patch.object(logger, "_find_encapsulating_repo", return_value=repo) as mock_find, patch(
+            "pype.mlflow.logger.mlflow_logger.set_tag"
+        ) as mock_set_tag:
+            logger.run = MagicMock()
+            logger.log_branch()
+
+            mock_find.assert_called_once_with(Path("."))
+            mock_set_tag.assert_called_once_with("git_branch", branch_name)
+
+    def test_log_branch_check_run(self, logger: MlflowLogger):
+        with patch.object(logger, "_find_encapsulating_repo"), patch(
+            "pype.mlflow.logger.mlflow_logger.set_tag"
+        ), pytest_assert(AssertionError, "Please start the experiment first"):
+            logger.log_branch()
+
+    def test_log_branch_no_repo_found(self, logger: MlflowLogger):
+        repo = None
+        with patch.object(logger, "_find_encapsulating_repo", return_value=repo) as mock_find, patch(
+            "pype.mlflow.logger.mlflow_logger.set_tag"
+        ) as mock_set_tag:
+            logger.run = MagicMock()
+            logger.log_branch()
+
+            mock_find.assert_called_once_with(Path("."))
+            mock_set_tag.assert_called_once_with("git_branch", "unknown_branch")
+
+    def test_find_encapsulating_repo(self, logger: MlflowLogger):
+        repo = MagicMock()
+        directory = Path("a_path") / "sub_dir"
+        with patch("pype.mlflow.logger.mlflow_logger.Repo", return_value=repo) as mock_repo_class:
+            result = logger._find_encapsulating_repo(directory)
+            mock_repo_class.assert_called_once_with(directory.absolute())
+
+            assert result == repo
+
+    def test_find_encapsulating_repo_find_repo_up(self, logger: MlflowLogger):
+        repo = MagicMock()
+        directory = Path("a_path") / "sub_dir"
+        with patch(
+            "pype.mlflow.logger.mlflow_logger.Repo", side_effect=[InvalidGitRepositoryError, repo]
+        ) as mock_repo_class:
+            result = logger._find_encapsulating_repo(directory)
+            mock_repo_class.assert_has_calls([call(directory.absolute()), call(directory.absolute().parent)])
+
+            assert result == repo
+
+    def test_find_encapsulating_repo_find_no_repo(self, logger: MlflowLogger):
+        directory = Path("/") / "1_dir" / "2_dir"
+        with patch(
+            "pype.mlflow.logger.mlflow_logger.Repo", side_effect=[InvalidGitRepositoryError] * 2
+        ) as mock_repo_class:
+            result = logger._find_encapsulating_repo(directory)
+            mock_repo_class.assert_has_calls(
+                [
+                    call(Path("/") / "1_dir" / "2_dir"),
+                    call(Path("/") / "1_dir"),
+                ]
+            )
+
+            assert result == None
+
+    def test_exit(self, logger: MlflowLogger):
+        mock_run = MagicMock()
+        logger = mock_run
+        a = MagicMock()
+        b = MagicMock()
+        c = MagicMock()
+        logger.__exit__(a, b, c)
+        mock_run.__exit__.assert_called_once_with(a, b, c)
+
+    def test_exit_assert_run_started(self, logger: MlflowLogger):
+        a = MagicMock()
+        b = MagicMock()
+        c = MagicMock()
+        with pytest_assert(AssertionError, "Run not started, cannot exit!"):
+            logger.__exit__(a, b, c)
+
+    def test_log_metrics(self, logger: MlflowLogger):
+        with patch("pype.mlflow.logger.mlflow_logger.log_metrics") as mock_log_metrics:
+            metrics = {"a": 1, "b": 3}
+            logger._log_metrics(metrics)
+
+            mock_log_metrics.assert_called_once_with(metrics)
+
+    def test_log_parameters(self, logger: MlflowLogger):
+        with patch("pype.mlflow.logger.mlflow_logger.log_params") as mock_log_params:
+            parameters = {"a": 1, "b": 3}
+            logger.log_parameters(parameters)
+
+            mock_log_params.assert_called_once_with(parameters)
+
+    def test_log_file(self, logger: MlflowLogger):
+        with patch("pype.mlflow.logger.mlflow_logger.log_artifact") as mock_log_artifact:
+            file = "some file"
+            logger.log_file(file)
+
+            mock_log_artifact.assert_called_once_with(file)
diff --git a/tests/shared_fixtures.py b/tests/shared_fixtures.py
new file mode 100644
index 0000000..7ce3d45
--- /dev/null
+++ b/tests/shared_fixtures.py
@@ -0,0 +1,42 @@
+import shutil
+from pathlib import Path
+from typing import Iterable
+
+from pytest import fixture
+
+from pype.base.experiment.experiment import Experiment
+from pype.base.logger.local_logger import LocalLogger
+from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
+from tests.utils import (
+    DummyModel,
+    get_dummy_data,
+    get_dummy_evaluator,
+    get_dummy_pipeline,
+    get_dummy_type_checkers,
+)
+
+
+@fixture(scope="module")
+def dummy_experiment() -> Iterable[Experiment]:
+    train = get_dummy_data(20, 1, 0)
+    test = get_dummy_data(5, 2, -1)
+
+    output_folder = Path(__file__).parent / "tmp_folder"
+
+    input_checker, output_checker = get_dummy_type_checkers()
+
+    yield Experiment(
+        data_sources={"train": train, "test": test},
+        model=DummyModel(inputs=["x"], outputs=["y"]),
+        pipeline=get_dummy_pipeline(),
+        logger=LocalLogger(),
+        evaluator=get_dummy_evaluator(),
+        serialiser=JoblibSerialiser(),
+        output_folder=output_folder,
+        input_type_checker=input_checker,
+        output_type_checker=output_checker,
+        parameters={"param_1": 2},
+    )
+
+    # make sure we purge the output folder again after a test is done.
+    shutil.rmtree(output_folder)
diff --git a/tests/sklearn/__init__.py b/tests/sklearn/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/sklearn/data/__init__.py b/tests/sklearn/data/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/sklearn/data/test_data_frame_source.py b/tests/sklearn/data/test_data_frame_source.py
new file mode 100644
index 0000000..c99e388
--- /dev/null
+++ b/tests/sklearn/data/test_data_frame_source.py
@@ -0,0 +1,12 @@
+import pandas as pd
+from pandas.testing import assert_frame_equal
+
+from pype.sklearn.data.data_frame_source import DataFrameSource
+
+
+def test_DataFrameSource():
+    df = pd.DataFrame({"x": [1, 2, 3]})
+
+    source = DataFrameSource(df.copy())
+
+    assert_frame_equal(df, source.read())
diff --git a/tests/sklearn/model/__init__.py b/tests/sklearn/model/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/sklearn/model/test_linear_regression_model.py b/tests/sklearn/model/test_linear_regression_model.py
new file mode 100644
index 0000000..f070faf
--- /dev/null
+++ b/tests/sklearn/model/test_linear_regression_model.py
@@ -0,0 +1,18 @@
+from unittest.mock import MagicMock
+
+from sklearn.linear_model import LinearRegression
+
+from pype.sklearn.model.linear_regression_model import LinearRegressionModel
+
+
+class Test_LinearRegressionModel:
+    def test_init_new_model(self):
+        model = LinearRegressionModel(["x"], ["y"])
+
+        assert isinstance(model.model, LinearRegression)
+
+    def test_init_existing_model(self):
+        sklearn_model = MagicMock()
+        model = LinearRegressionModel(["x"], ["y"], model=sklearn_model)
+
+        assert model.model == sklearn_model
diff --git a/tests/sklearn/model/test_logistic_regression_model.py b/tests/sklearn/model/test_logistic_regression_model.py
new file mode 100644
index 0000000..9798193
--- /dev/null
+++ b/tests/sklearn/model/test_logistic_regression_model.py
@@ -0,0 +1,18 @@
+from unittest.mock import MagicMock
+
+from sklearn.linear_model import LogisticRegression
+
+from pype.sklearn.model.logistic_regression_model import LogisticRegressionModel
+
+
+class Test_LogisticRegressionModel:
+    def test_init_new_model(self):
+        model = LogisticRegressionModel(["x"], ["y"])
+
+        assert isinstance(model.model, LogisticRegression)
+
+    def test_init_existing_model(self):
+        sklearn_model = MagicMock()
+        model = LogisticRegressionModel(["x"], ["y"], model=sklearn_model)
+
+        assert model.model == sklearn_model
diff --git a/tests/sklearn/model/test_sklearn_model.py b/tests/sklearn/model/test_sklearn_model.py
new file mode 100644
index 0000000..af670c5
--- /dev/null
+++ b/tests/sklearn/model/test_sklearn_model.py
@@ -0,0 +1,103 @@
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pandas as pd
+from pytest import fixture
+from sklearn.linear_model import LinearRegression
+
+from pype.base.data.dataset import DataSet
+from pype.sklearn.model.linear_regression_model import LinearRegressionModel
+from pype.sklearn.model.sklearn_model import SklearnModel
+
+
+class Test_SklearnModel:
+    @fixture
+    def sklearn_model(self):
+        return MagicMock()
+
+    @fixture
+    def model(self, sklearn_model: MagicMock):
+        return LinearRegressionModel(inputs=["x"], outputs=["y"], model=sklearn_model, seed=43)
+
+    def test_set_seed(self, model: SklearnModel):
+        with patch("pype.sklearn.model.sklearn_model.np.random.seed") as mock_seed:
+            model.set_seed()
+        mock_seed.assert_called_once_with(43)
+
+    def test_fit(self, model: SklearnModel, sklearn_model: MagicMock):
+        ds1 = MagicMock()
+        ds2 = MagicMock()
+        model._fit(ds1, ds2)
+
+        sklearn_model.fit.assert_called_once_with(ds1, ds2)
+
+    def test_transform(self, model: SklearnModel, sklearn_model: MagicMock):
+        ds1 = MagicMock()
+        ds2 = MagicMock()
+        model._transform(ds1, ds2)
+
+        sklearn_model.predict.assert_called_once_with(ds1, ds2)
+
+    def test_save(self, model: SklearnModel, sklearn_model: MagicMock):
+        with patch("pype.sklearn.model.sklearn_model.JoblibSerialiser.serialise") as mock_serialise:
+            folder = Path("folder")
+            model._save(folder)
+
+            mock_serialise.assert_called_once_with(sklearn_model, folder / model.SKLEARN_MODEL_FILE)
+
+    def test_load(self, model: SklearnModel, sklearn_model: MagicMock):
+        with patch(
+            "pype.sklearn.model.sklearn_model.JoblibSerialiser.deserialise", return_value=sklearn_model
+        ) as mock_deserialise:
+            folder = Path("folder")
+            inputs = ["x"]
+            outputs = ["y"]
+            result: SklearnModel = model._load(folder, inputs, outputs)
+
+            mock_deserialise.assert_called_once_with(folder / model.SKLEARN_MODEL_FILE)
+            assert result.model == sklearn_model
+            assert result.inputs == inputs
+            assert result.outputs == outputs
+
+    def test_get_parameters_from_object(self):
+        class DummyModel(SklearnModel[LinearRegression]):
+            def _init_model(self, args: dict[str, Any]) -> LinearRegression:
+                return LinearRegression(**args)
+
+        parser = MagicMock()
+        model = DummyModel(inputs=["x"], outputs=["y"])
+
+        with patch("pype.sklearn.model.sklearn_model.add_args_to_parser_for_class") as mock_add_args:
+            model.get_parameters(parser)
+            mock_add_args.assert_called_once_with(
+                parser, LinearRegression, "model", [], excluded_args=["seed", "inputs", "outputs", "model"]
+            )
+
+    def test_get_parameters_from_class(self):
+        class DummyModel(SklearnModel[LinearRegression]):
+            def _init_model(self, args: dict[str, Any]) -> LinearRegression:
+                return LinearRegression(**args)
+
+        parser = MagicMock()
+
+        with patch("pype.sklearn.model.sklearn_model.add_args_to_parser_for_class") as mock_add_args:
+            DummyModel.get_parameters(parser)
+            mock_add_args.assert_called_once_with(
+                parser, LinearRegression, "model", [], excluded_args=["seed", "inputs", "outputs", "model"]
+            )
+
+    def test_integration(self):
+        model = LinearRegressionModel(["inputs"], ["outputs"])
+
+        x = pd.DataFrame({"x": [1, 2, 3, 4], "y": [3, 5, 6, 7]})
+        y = pd.DataFrame({"z": x["x"] * 1.5 - 2 * x["y"] + 2.4})
+        data = DataSet(inputs=x.copy(), outputs=y.copy())
+
+        model.fit(data)
+
+        predictions = model.transform(data)
+
+        assert isinstance(predictions["outputs"], np.ndarray)
+        np.testing.assert_array_almost_equal(y, predictions["outputs"])
diff --git a/tests/sklearn/pipeline/__init__.py b/tests/sklearn/pipeline/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/sklearn/pipeline/test_numpy_type_checker.py b/tests/sklearn/pipeline/test_numpy_type_checker.py
new file mode 100644
index 0000000..c77c0d5
--- /dev/null
+++ b/tests/sklearn/pipeline/test_numpy_type_checker.py
@@ -0,0 +1,115 @@
+import numpy as np
+from pydantic import ValidationError
+from pytest import mark
+
+from pype.base.pipeline.type_checker import (
+    DataModel,
+    DataSetModel,
+    DataSetTypeChecker,
+    TypeCheckerPipe,
+)
+from pype.sklearn.pipeline.numpy_type_checker import NumpyData, NumpyTypeChecker
+from tests.utils import pytest_assert
+
+
+class Test_NumpyData:
+    def test_convert(self):
+        model = NumpyData(data=[[1, 2, 3], [4, 5, 6], [6, 7, 8]])
+        expected = np.array([[1, 2, 3], [4, 5, 6], [6, 7, 8]])
+
+        result = model.convert()
+
+        np.testing.assert_array_equal(expected, result)
+
+    def test_to_model(self):
+        array = np.array([[1, 2, 3], [2, 3, 4]])
+        result = NumpyData.to_model(array)
+
+        expected = NumpyData(
+            data=[[1, 2, 3], [2, 3, 4]],
+        )
+
+        assert expected == result
+
+
+class Test_NumpyTypeChecker:
+    def test_fit(self):
+        type_checker = NumpyTypeChecker()
+        array = np.array([[1, 2, 3], [3, 4, 5]])
+        assert array.shape[0] != array.shape[1]
+
+        type_checker.fit(array)
+
+        assert type_checker.dims == (3,)
+        assert type_checker.dtype == int
+
+    def test_transform(self):
+        type_checker = NumpyTypeChecker()
+
+        array = np.array([[1, 2, 3], [3, 4, 5]])
+        type_checker.fit(array)
+
+        # works
+        result = type_checker.transform(array)
+
+        # errors #
+        # other type of object
+        with pytest_assert(AssertionError, "Please provide a numpy array!"):
+            type_checker.transform([])
+
+        # other shape
+        other_shape = np.array([[1, 2], [3, 4]])
+        with pytest_assert(
+            AssertionError, f"Dimensions of numpy arrays do not add up: {other_shape.shape[1:]} vs {array.shape[1:]}"
+        ):
+            type_checker.transform(other_shape)
+
+        # other type
+        other_type = np.array([[1, 2, 3], [3, 4, 5.0]])
+        with pytest_assert(AssertionError, f"Dtype of data does not add up: {float} vs {int}"):
+            type_checker.transform(other_type)
+
+    @mark.parametrize(
+        ["input_array", "expected"],
+        [
+            [np.array([1, 2, 3]), int],
+            [np.array([1.0, 2, 3]), float],
+            [np.array([False, True]), bool],
+            [np.array(["1", "2", "3"]), str],
+        ],
+    )
+    def test_convert_dtype(self, input_array: np.ndarray, expected: type):
+        type_checker = NumpyTypeChecker()
+
+        result = type_checker._convert_dtype(input_array.dtype)
+        assert result == expected
+
+    def test_get_pydantic_types(self):
+        type_checker = NumpyTypeChecker()
+        array = np.array([[1, 2, 3], [7, 6, 5]])
+        type_checker.fit(array)
+
+        NumpySpecificType = type_checker.get_pydantic_type()
+
+        # succeed
+        NumpySpecificType(data=[[1, 2, 3], [9, 0, 3]])
+
+        # also succeed
+        data = [[1, 2, 3], [9, 0, 3.0]]
+        result = NumpySpecificType(data=data).convert()
+
+        np.testing.assert_array_equal(np.array(data), result)
+
+    def test_get_pydantic_types_strict_on_bool(self):
+        type_checker = NumpyTypeChecker()
+        array = np.array([[True, False]])
+        type_checker.fit(array)
+
+        NumpySpecificType = type_checker.get_pydantic_type()
+
+        # succeed
+        NumpySpecificType(data=[[False, False]])
+
+        # fails
+        with pytest_assert(ValidationError):
+            NumpySpecificType(data=[[1, 2, 3]])
diff --git a/tests/sklearn/pipeline/test_pandas_type_checker copy.py b/tests/sklearn/pipeline/test_pandas_type_checker copy.py
new file mode 100644
index 0000000..f2a5294
--- /dev/null
+++ b/tests/sklearn/pipeline/test_pandas_type_checker copy.py
@@ -0,0 +1,181 @@
+from datetime import datetime
+from typing import Type
+
+import pandas as pd
+from pandas.api import types as pandas_types
+from pandas.api.types import (
+    is_bool_dtype,
+    is_datetime64_any_dtype,
+    is_float_dtype,
+    is_integer_dtype,
+    is_object_dtype,
+    is_string_dtype,
+)
+from pandas.testing import assert_frame_equal
+from pydantic import create_model
+from pytest import fixture
+
+from pype.sklearn.pipeline.pandas_type_checker import PandasData, PandasTypeChecker
+from tests.utils import pytest_assert
+
+
+class Test_PandasData:
+    @fixture
+    def model_class(self) -> Type[PandasData]:
+        return create_model(
+            "PandasData",
+            x=(list[int] | dict[str | int, int], ...),
+            y=(list[str] | dict[str | int, str], ...),
+            __base__=PandasData,
+        )
+
+    def test_convert(self, model_class: Type[PandasData]):
+        model = model_class(x=[1, 2, 3, 4], y=["a", "b", "c", "e"])
+
+        expected = pd.DataFrame(
+            {
+                "x": [1, 2, 3, 4],
+                "y": ["a", "b", "c", "e"],
+            }
+        )
+
+        result = model.convert()
+
+        assert_frame_equal(expected, result)
+
+    def test_to_model(self, model_class: Type[PandasData]):
+        df = pd.DataFrame(
+            {
+                "x": [1, 2, 3, 4],
+                "y": ["a", "b", "c", "e"],
+            }
+        )
+        result = model_class.to_model(df)
+
+        expected = model_class(
+            x=[1, 2, 3, 4],
+            y=["a", "b", "c", "e"],
+        )
+
+        assert expected == result
+
+
+class Test_PandasTypeChecker:
+    def test_fit(self):
+        type_checker = PandasTypeChecker()
+        df = pd.DataFrame(
+            {
+                "x": [1, 2, 3, 4],
+                "y": ["a", "b", "c", "e"],
+            }
+        )
+
+        type_checker.fit(df)
+        expected_dict = {
+            "x": (int, is_integer_dtype),
+            "y": (str, is_object_dtype),
+        }
+        assert type_checker.raw_types == expected_dict
+
+    def test_transform(self):
+        type_checker = PandasTypeChecker()
+        df = pd.DataFrame(
+            {
+                "x": [1, 2, 3, 4],
+                "y": ["a", "b", "c", "e"],
+            }
+        )
+
+        type_checker.fit(df)
+        # works
+        result = type_checker.transform(df)
+
+        # errors #
+        # other type of object
+        with pytest_assert(AssertionError, "Please provide a pandas DataFrame!"):
+            type_checker.transform([])
+
+        # missing columns
+        missing_df = pd.DataFrame({"x": [1, 2, 3]})  # missing
+        with pytest_assert(AssertionError, "Not all columns are present."):
+            type_checker.transform(missing_df)
+
+        other_type = df.copy()
+        other_type["x"] = other_type["x"].astype(str)
+
+        with pytest_assert(AssertionError, "Dtypes did not match up for col x."):
+            type_checker.transform(other_type)
+
+    def test_transform_reduces_columns(self):
+        type_checker = PandasTypeChecker()
+        df = pd.DataFrame(
+            {
+                "x": [1, 2, 3, 4],
+                "y": ["a", "b", "c", "e"],
+            }
+        )
+
+        type_checker.fit(df)
+        # works
+        extra_col_df = df.copy()
+        extra_col_df["z"] = 1
+        result = type_checker.transform(extra_col_df)
+
+        assert "z" not in df.columns
+
+    def test_convert_raw_types(self):
+        type_checker = PandasTypeChecker()
+        df = pd.DataFrame(
+            {
+                "x": [1, 2, 3],
+                "y": [1, 2, 3.0],
+                "z": ["1", "2", "3"],
+                "a": [True, False, True],
+                "b": [1, "1", 1],
+                "c": [datetime(2018, 1, 1, 1), datetime(2018, 1, 1, 0), datetime(2018, 1, 1, 2)],
+            }
+        )
+
+        result = type_checker._convert_raw_types(dict(df.dtypes))
+        expected = {
+            "x": (int, pandas_types.is_integer_dtype),
+            "y": (float, pandas_types.is_float_dtype),
+            "z": (str, pandas_types.is_object_dtype),
+            "a": (bool, pandas_types.is_bool_dtype),
+            "b": (str, pandas_types.is_object_dtype),
+            "c": (datetime, pandas_types.is_datetime64_any_dtype),
+        }
+        assert result == expected
+
+    def test_get_pydantic_types(self):
+        type_checker = PandasTypeChecker()
+        df = pd.DataFrame(
+            {
+                "x": [1, 2, 3, 4],
+                "y": ["a", "b", "c", "e"],
+            }
+        )
+        type_checker.fit(df.copy())
+
+        PandasSpecificType = type_checker.get_pydantic_type()
+
+        # succeed
+        result = PandasSpecificType(x=[1, 2, 3, 4], y=["a", "b", "c", "e"]).convert()
+
+        assert_frame_equal(df, result)
+
+    # def test_get_pydantic_types_strict_on_bool(self):
+    #     type_checker = PandasTypeChecker()
+    #     array = np.array([
+    #         [True, False]
+    #     ])
+    #     type_checker.fit(array)
+
+    #     PandasSpecificType = type_checker.get_pydantic_type()
+
+    #     # succeed
+    #     PandasSpecificType(data=[[False, False]])
+
+    #     # fails
+    #     with pytest_assert(ValidationError):
+    #         PandasSpecificType(data=[[1, 2 ,3]])
diff --git a/tests/spark/__init__.py b/tests/spark/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/spark/data/__init__.py b/tests/spark/data/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/spark/data/test_spark_dataframe_source.py b/tests/spark/data/test_spark_dataframe_source.py
new file mode 100644
index 0000000..f44e133
--- /dev/null
+++ b/tests/spark/data/test_spark_dataframe_source.py
@@ -0,0 +1,13 @@
+from unittest.mock import MagicMock
+
+from pype.spark.data import SparkDataFrameSource
+
+
+class Test_SparkDataFrameSource:
+    def test(self):
+        df = MagicMock()
+
+        source = SparkDataFrameSource(df)
+        result = source.read()
+
+        assert result == df
diff --git a/tests/spark/data/test_spark_read_source.py b/tests/spark/data/test_spark_read_source.py
new file mode 100644
index 0000000..ce37515
--- /dev/null
+++ b/tests/spark/data/test_spark_read_source.py
@@ -0,0 +1,21 @@
+from unittest.mock import MagicMock
+
+from pype.spark.data import SparkReadSource
+
+
+class Test_SparkReadSource:
+    def test(self):
+        mock_spark = MagicMock()
+
+        file = "/a/path"
+        formats = "a format"
+        options = {"a": 3, "b": 4}
+
+        source = SparkReadSource(mock_spark, file, formats, options)
+        result = source.read()
+
+        mock_format = mock_spark.read.format
+        mock_load = mock_format.return_value.load
+        mock_format.assert_called_once_with(formats)
+        mock_load.assert_called_once_with(file, a=3, b=4)
+        assert result == mock_load.return_value
diff --git a/tests/spark/data/test_spark_sql_source.py b/tests/spark/data/test_spark_sql_source.py
new file mode 100644
index 0000000..00f514b
--- /dev/null
+++ b/tests/spark/data/test_spark_sql_source.py
@@ -0,0 +1,15 @@
+from unittest.mock import MagicMock
+
+from pype.spark.data import SparkSqlSource
+
+
+class Test_SparkSqlSource:
+    def test(self):
+        session = MagicMock()
+        query = "select * from table"
+
+        source = SparkSqlSource(session, query)
+        result = source.read()
+
+        session.sql.assert_called_once_with(query)
+        assert result == session.sql.return_value
diff --git a/tests/spark/model/__init__.py b/tests/spark/model/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/spark/model/test_spark_model.py b/tests/spark/model/test_spark_model.py
new file mode 100644
index 0000000..e847134
--- /dev/null
+++ b/tests/spark/model/test_spark_model.py
@@ -0,0 +1,258 @@
+import json
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import List
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pandas as pd
+from pandas.testing import assert_frame_equal
+from pyspark.ml.feature import VectorAssembler
+from pyspark.ml.regression import LinearRegression
+from pyspark.sql import SparkSession
+from pytest import mark
+
+from pype.base.data import DataSet
+from pype.spark.model import LinearSparkModel
+from tests.spark.utils import spark_session
+from tests.utils import pytest_assert
+
+spark_session
+
+
+class Test_SparkModel:
+    # Tested using LinearSparkModel for initialisation purposes
+
+    @mark.parametrize(
+        ["inputs", "outputs"],
+        [
+            [[], ["x"]],
+            [["x"], []],
+            [["1", "2"], ["x"]],
+            [["x"], ["1", "2"]],
+            [["x"], ["y"]],
+        ],
+    )
+    def test_init_assert(self, inputs: List[str], outputs: List[str]):
+        with pytest_assert(
+            AssertionError, "SparkML only requires a single DataFrame as input and output, the same one."
+        ):
+            LinearSparkModel(
+                inputs=inputs,
+                outputs=outputs,
+            )
+
+    def test_init(self):
+        with patch.object(LinearSparkModel, "_init_model") as mock_init_model:
+            LinearSparkModel(inputs=["x"], outputs=["x"], arg1=1, arg2=4)
+
+            mock_init_model.assert_called_once_with({"arg1": 1, "arg2": 4})
+
+    def test_init_with_predictor(self):
+        with patch.object(LinearSparkModel, "_init_model") as mock_init_model:
+            LinearSparkModel(inputs=["x"], outputs=["x"], predictor=MagicMock())
+
+            mock_init_model.assert_not_called()
+
+    def test_get_annotated_class(self):
+        model = LinearSparkModel(inputs=["x"], outputs=["x"], predictor=MagicMock())
+
+        result = model._get_annotated_class()
+
+        assert result == LinearRegression
+
+    def test_fit(self):
+        model = LinearSparkModel(inputs=["x"], outputs=["x"], predictor=MagicMock())
+        x = MagicMock()
+        data = DataSet(x=x, y=MagicMock())
+
+        with patch.object(model, "_fit") as mock_fit:
+            result = model.fit(data)
+
+        assert result == model
+        mock_fit.assert_called_once_with(x)
+
+    def test__fit_assertions(self):
+        model = LinearSparkModel(inputs=["x"], outputs=["x"], predictor=MagicMock())
+        x = MagicMock()
+        y = MagicMock()
+
+        with pytest_assert(AssertionError, "SparkML needs a single DataFrame as input, got 2"):
+            model._fit(x, y)
+
+    def test__fit(self):
+        predictor = MagicMock()
+        model = LinearSparkModel(inputs=["x"], outputs=["x"], predictor=predictor)
+        x = MagicMock()
+
+        model._fit(x)
+
+        predictor.fit.assert_called_once_with(x)
+        assert model.model == predictor.fit.return_value
+
+    def test_transform_assertions_excess_data(self):
+        model = LinearSparkModel(inputs=["x"], outputs=["x"], predictor=MagicMock())
+        x = MagicMock()
+        y = MagicMock()
+
+        with pytest_assert(AssertionError, "SparkML needs a single DataFrame as input, got 2"):
+            model._transform(x, y)
+
+    def test_transform_assertions_unfitted(self):
+        model = LinearSparkModel(inputs=["x"], outputs=["x"], predictor=MagicMock())
+        x = MagicMock()
+
+        with pytest_assert(AssertionError, "Please fit this model before transforming data."):
+            model._transform(x)
+
+    @mark.parametrize(
+        ["name", "prediction_col"],
+        [
+            ["return full set", None],
+            ["return only column", "pred_col"],
+        ],
+    )
+    def test_transform(self, name: str, prediction_col: str | None):
+        predictor = MagicMock()
+        model = LinearSparkModel(
+            inputs=["x"],
+            outputs=["x"],
+            predictor=predictor,
+            output_col=prediction_col,
+        )
+        x = MagicMock()
+        data = DataSet(x=x, y=MagicMock())
+
+        model.fit(data)
+        result = model.transform(data)
+
+        model = predictor.fit.return_value
+        model.transform.assert_called_once_with(x)
+
+        predictions = model.transform.return_value
+
+        if prediction_col is None:
+            assert result["x"] == predictions
+        else:
+            predictions.select.assert_called_once_with(prediction_col)
+            assert result["x"] == predictions.select.return_value
+
+    def test_save_assertions(self):
+        model = LinearSparkModel(
+            inputs=["x"],
+            outputs=["x"],
+            predictor=MagicMock(),
+        )
+
+        with pytest_assert(
+            AssertionError, "Please fit this model before transforming data."
+        ), TemporaryDirectory() as tmp_dir:
+            model.save(tmp_dir)
+
+    def test_save(self):
+        predictor = MagicMock()
+        model = MagicMock()
+        pype_model = LinearSparkModel(
+            inputs=["x"],
+            outputs=["x"],
+            predictor=predictor,
+            model=model,
+            output_col=None,
+        )
+
+        with TemporaryDirectory() as tmp_dir, patch(
+            "pype.spark.model.spark_model.JoblibSerialiser.serialise"
+        ) as mock_serialise:
+            tmp_dir = Path(tmp_dir)
+            pype_model._save(tmp_dir)
+
+            with open(tmp_dir / pype_model.PYPE_MODEL_CONFIG, "r") as f:
+                config = json.load(f)
+                assert config == {"output_col": None}
+
+            model.save.assert_called_once_with(str(tmp_dir / pype_model.SPARK_MODEL_PATH))
+            predictor.save.assert_called_once_with(str(tmp_dir / pype_model.SPARK_PREDICTOR_PATH))
+            mock_serialise.assert_called_once_with(type(model), str(tmp_dir / pype_model.SPARK_MODEL_CLASS_PATH))
+
+    def test_load(self):
+        with TemporaryDirectory() as tmp_dir, patch.object(
+            LinearSparkModel, "_get_annotated_class"
+        ) as mock_annotate, patch("pype.spark.model.spark_model.JoblibSerialiser.deserialise") as mock_deserialise:
+            tmp_dir = Path(tmp_dir)
+
+            with open(tmp_dir / LinearSparkModel.PYPE_MODEL_CONFIG, "w") as f:
+                json.dump({"output_col": None}, f)
+
+            result = LinearSparkModel._load(tmp_dir, ["x"], ["x"])
+
+            mock_annotate.assert_called_once_with()
+            mock_pred_class = mock_annotate.return_value
+            mock_pred_class.load.assert_called_once_with(str(tmp_dir / LinearSparkModel.SPARK_PREDICTOR_PATH))
+
+            mock_deserialise.assert_called_once_with(str(tmp_dir / LinearSparkModel.SPARK_MODEL_CLASS_PATH))
+            mock_model_class = mock_deserialise.return_value
+            mock_model_class.load.assert_called_once_with(str(tmp_dir / LinearSparkModel.SPARK_MODEL_PATH))
+
+            assert result.inputs == ["x"]
+            assert result.outputs == ["x"]
+            assert result.output_col is None
+            assert result.predictor == mock_pred_class.load.return_value
+            assert result.model == mock_model_class.load.return_value
+
+    def test_get_parameters(self):
+        parser = MagicMock()
+
+        with patch("pype.spark.model.spark_model.add_args_to_parser_for_class") as mock_add:
+            LinearSparkModel.get_parameters(parser)
+
+        mock_add.assert_called_once_with(
+            parser, LinearRegression, "model", [], excluded_args=["seed", "inputs", "outputs", "model"]
+        )
+
+    def test_integration(self, spark_session: SparkSession):
+        df = pd.DataFrame(
+            {
+                "x": [
+                    1,
+                    2,
+                    3,
+                    4,
+                ],
+                "y": [9, 4, 6, 1],
+                "response": [3, 5, 6, 7],
+            }
+        )
+        spark_df = spark_session.createDataFrame(df)
+        spark_df = VectorAssembler(inputCols=["x", "y"], outputCol="vector").transform(spark_df)
+
+        dataset = DataSet(df=spark_df)
+
+        model = LinearSparkModel(
+            inputs=["df"],
+            outputs=["df"],
+            output_col="prediction_col",
+            predictionCol="prediction_col",
+            featuresCol="vector",
+            labelCol="response",
+            fitIntercept=False,
+        )
+
+        model.fit(dataset)
+        transformed = model.transform(dataset)
+
+        assert "df" in transformed
+        assert ["prediction_col"] == transformed["df"].columns
+
+        transformed_df = transformed["df"].toPandas()
+        assert transformed_df.shape == (df.shape[0], 1)
+        assert np.all(~transformed_df.isna())
+
+        # saving/loading
+        with TemporaryDirectory() as tmp_dir:
+            tmp_dir = Path(tmp_dir)
+
+            model.save(tmp_dir)
+            loaded_model = LinearSparkModel.load(tmp_dir)
+            loaded_preds = loaded_model.transform(dataset)["df"].toPandas()
+
+            assert_frame_equal(transformed_df, loaded_preds)
diff --git a/tests/spark/pipeline/__init__.py b/tests/spark/pipeline/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/spark/pipeline/test_spark_pipe.py b/tests/spark/pipeline/test_spark_pipe.py
new file mode 100644
index 0000000..b388550
--- /dev/null
+++ b/tests/spark/pipeline/test_spark_pipe.py
@@ -0,0 +1,120 @@
+import warnings
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pandas as pd
+from pandas.testing import assert_frame_equal
+from pyspark.ml import Estimator, Transformer
+from pyspark.ml.feature import MinMaxScaler, VectorAssembler
+from pyspark.sql import SparkSession
+from pytest import fixture
+
+from pype.base.data.dataset import DataSet
+from pype.spark.pipeline.spark_pipe import SparkPipe
+from tests.spark.utils import spark_session
+from tests.utils import pytest_assert
+
+spark_session
+
+
+class DummyEstimator(Estimator):
+    def _fit(self, data: Any):
+        pass
+
+    def _transform(self):
+        pass
+
+
+class DummyTransformer(Transformer):
+    def _transform(self):
+        pass
+
+
+class Test_SparkPipe:
+    @fixture
+    def data(self) -> DataSet:
+        return DataSet(x=MagicMock(), z=MagicMock())
+
+    def test_fit_assert(self, data: DataSet):
+        class Dummy:
+            pass
+
+        pipe = SparkPipe("test", Dummy, ["x"], ["y"])
+
+        with pytest_assert(
+            ValueError, f"In a SparkPipe, the operator must be a Transformer or Estimator. Got: {Dummy}"
+        ):
+            pipe.fit(data)
+
+    def test_fit_transformer(self, data: DataSet):
+
+        pipe = SparkPipe("test", DummyTransformer, ["x"], ["y"])
+
+        pipe.fit(data)
+        assert pipe.operator == pipe.fitted
+
+    def test_fit_estimator(self, data: DataSet):
+
+        pipe = SparkPipe("test", DummyEstimator, ["x"], ["y"])
+
+        with patch.object(DummyEstimator, "fit") as mock_fit:
+            pipe.fit(data)
+            mock_fit.assert_called_once_with(data["x"])
+
+    def test_transform(self, data: DataSet):
+        pipe = SparkPipe("test", DummyEstimator, ["x"], ["y"])
+
+        pipe.fitted = MagicMock()
+
+        assert "y" not in data
+        result = pipe.transform(data)
+        assert "y" in result
+
+        pipe.fitted.transform.assert_called_once_with(data["x"])
+        assert result["y"] == pipe.fitted.transform.return_value
+
+    def test_set_state(self):
+        warnings.warn("SparkPipe's __set_state__ is not tested atm.")
+
+    def test_get_state(self):
+        pipe = SparkPipe("test", DummyEstimator, ["x"], ["y"])
+
+        state = pipe.__getstate__()
+
+        assert "fitted" not in state
+        assert "operator" not in state
+
+    def test_get_state_from_fitted(self, data: DataSet):
+        pipe = SparkPipe("test", DummyEstimator, ["x"], ["y"])
+        pipe.fit(data)
+
+        state = pipe.__getstate__()
+
+        assert "fitted" not in state
+        assert "operator" not in state
+
+    def test_str(self):
+        # just test if it runs.
+        pipe = SparkPipe("test", DummyEstimator, ["x"], ["y"])
+        str(pipe)
+
+    def test_integration(self, spark_session: SparkSession):
+        x = pd.DataFrame({"col1": [1, 2, 3, 4]})
+        spark_x = VectorAssembler(inputCols=["col1"], outputCol="transformed").transform(
+            spark_session.createDataFrame(x)
+        )
+        data = DataSet(x=spark_x, z=MagicMock())
+
+        pipe = SparkPipe("test", MinMaxScaler, ["x"], ["y"], kw_args={"inputCol": "transformed", "outputCol": "scaled"})
+        pipe.fit(data)
+
+        assert "y" not in data
+        transformed = pipe.transform(data)
+        assert "y" in transformed
+        transformed_df = transformed["y"].toPandas()
+
+        expected = pd.DataFrame(
+            {"col1": [1, 2, 3, 4], "transformed": [(1,), (2,), (3,), (4,)], "scaled": [(0,), (1 / 3,), (2 / 3,), (1,)]}
+        )
+
+        assert_frame_equal(expected, transformed_df)
diff --git a/tests/spark/pipeline/test_spark_type_checker.py b/tests/spark/pipeline/test_spark_type_checker.py
new file mode 100644
index 0000000..06a527d
--- /dev/null
+++ b/tests/spark/pipeline/test_spark_type_checker.py
@@ -0,0 +1,150 @@
+from unittest.mock import MagicMock, call, patch
+
+import pandas as pd
+from pandas.testing import assert_frame_equal
+from pydantic import create_model
+from pyspark.sql import SparkSession
+from pytest import mark
+
+from pype.spark.pipeline.spark_type_checker import SparkData, SparkTypeChecker
+from tests.spark.utils import spark_session
+from tests.utils import pytest_assert
+
+spark_session
+
+
+class Test_SparkData:
+    def test_convert(self, spark_session: SparkSession):
+        data_type = {
+            "a": (list[int] | dict[str | int, int], ...),
+            "b": (list[str] | dict[str | int, str], ...),
+        }
+
+        CustomSparkData = create_model("CustomSparkData", **data_type, __base__=SparkData)
+
+        data_obj = CustomSparkData(a=[1, 2, 3], b=["a", "2", "f"])
+
+        converted = data_obj.convert()
+
+        assert converted.columns == ["a", "b"]
+        converted_df: pd.DataFrame = converted.toPandas()
+
+        expected = pd.DataFrame(
+            {
+                "a": [1, 2, 3],
+                "b": ["a", "2", "f"],
+            }
+        )
+
+        assert_frame_equal(expected, converted_df)
+
+    def test_to_model(self, spark_session: SparkSession):
+        data_type = {
+            "a": (list[int] | dict[str | int, int], ...),
+            "b": (list[str] | dict[str | int, str], ...),
+        }
+        CustomSparkData = create_model("CustomSparkData", **data_type, __base__=SparkData)
+
+        df = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "2", "f"]})
+        spark_df = spark_session.createDataFrame(df)
+
+        result = CustomSparkData.to_model(spark_df)
+
+        assert result.a == [1, 2, 3]
+        assert result.b == ["a", "2", "f"]
+
+
+class Test_SparkTypeChecker:
+    def test_fit(self):
+        data = MagicMock()
+        dtypes = {"a": int, "b": float}
+        data.dtypes = dtypes
+        checker = SparkTypeChecker()
+
+        with patch.object(checker, "_convert_dtypes") as mock_convert:
+            result = checker.fit(data)
+
+        mock_convert.assert_called_once_with(dtypes)
+        assert result == checker
+
+    def test_convert_dtypes(self):
+        checker = SparkTypeChecker()
+        data = {"a": "2", "b": "9"}
+        returns = [2, 3]
+
+        with patch.object(checker, "_convert_dtype", side_effect=returns) as mock_convert:
+            result = checker._convert_dtypes(data)
+
+        mock_convert.assert_has_calls([call("2"), call("9")])
+
+        assert result == {"a": 2, "b": 3}
+
+    @mark.parametrize(
+        ["string_type", "expected"],
+        [
+            ["str", str],
+            ["string", str],
+            ["int", int],
+            ["bigint", int],
+            ["float", float],
+            ["double", float],
+            ["bool", bool],
+        ],
+    )
+    def test_convert_dtype(self, string_type: str, expected: type):
+        checker = SparkTypeChecker()
+        result = checker._convert_dtype(string_type)
+
+        assert result == (string_type, expected)
+
+    def test_convert_dtype_assertion(self):
+        checker = SparkTypeChecker()
+        with pytest_assert(ValueError, "- not supported"):
+            checker._convert_dtype("-")
+
+    def test_transform_assert_missing_columns(self, spark_session: SparkSession):
+        df = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "2", "f"]})
+        spark_df = spark_session.createDataFrame(df)
+
+        missing_df = pd.DataFrame({"a": [1, 2, 3]})
+        spark_missing_df = spark_session.createDataFrame(missing_df)
+
+        checker = SparkTypeChecker()
+        checker.fit(spark_df)
+
+        with pytest_assert(AssertionError, "`b` is missing from the dataset"):
+            checker.transform(spark_missing_df)
+
+    def test_transform_check_dtype(self, spark_session: SparkSession):
+        df = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "2", "f"]})
+        spark_df = spark_session.createDataFrame(df)
+
+        missing_df = pd.DataFrame({"a": [1, 2, 3], "b": [2.2, 3.4, 5.6]})
+        spark_missing_df = spark_session.createDataFrame(missing_df)
+
+        checker = SparkTypeChecker()
+        checker.fit(spark_df)
+
+        with pytest_assert(AssertionError, "Dtypes did not match up for col b: Expected string, got double"):
+            checker.transform(spark_missing_df)
+
+    def test_transform_pass(self, spark_session: SparkSession):
+        df = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "2", "f"]})
+        spark_df = spark_session.createDataFrame(df)
+
+        checker = SparkTypeChecker()
+        checker.fit(spark_df)
+        checker.transform(spark_df)
+
+    def test_get_pydantic_type(self, spark_session: SparkSession):
+        df = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "2", "f"]})
+        spark_df = spark_session.createDataFrame(df)
+
+        checker = SparkTypeChecker()
+        checker.fit(spark_df)
+
+        Model = checker.get_pydantic_type()
+
+        converted = Model.to_model(spark_df)
+        assert converted.a == df["a"].to_list()
+        assert converted.b == df["b"].to_list()
diff --git a/tests/spark/serialisation/__init__.py b/tests/spark/serialisation/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/spark/serialisation/test_spark_serialiser.py b/tests/spark/serialisation/test_spark_serialiser.py
new file mode 100644
index 0000000..34641b5
--- /dev/null
+++ b/tests/spark/serialisation/test_spark_serialiser.py
@@ -0,0 +1,66 @@
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pandas as pd
+from pandas.testing import assert_frame_equal
+from pyspark.ml.feature import StandardScaler, VectorAssembler
+from pyspark.sql import SparkSession
+
+from pype.base.data import DataSet
+from pype.base.pipeline import Pipeline
+from pype.spark.pipeline import SparkPipe
+from pype.spark.serialisation import SparkSerialiser
+from tests.spark.utils import spark_session
+
+spark_session
+
+
+class Test_SparkSerialiser:
+    def test_serialise_pipeline(self, spark_session: SparkSession):
+        data = DataSet(
+            x=spark_session.createDataFrame(pd.DataFrame({"a": [1, 2, 3.0], "target": [2, 3, 4.0]})),
+            helper=spark_session.createDataFrame(pd.DataFrame({"a": [2, 3, 4]})),
+        )
+
+        pipeline = Pipeline(
+            [
+                SparkPipe(
+                    "assemble", VectorAssembler, ["x"], ["x"], kw_args={"inputCols": ["a"], "outputCol": "a_gathered"}
+                ),
+                SparkPipe(
+                    "scale", StandardScaler, ["x"], ["x"], kw_args={"inputCol": "a_gathered", "outputCol": "a_scaled"}
+                ),
+            ]
+        )
+
+        pipeline.fit(data)
+
+        serialiser = SparkSerialiser()
+
+        with TemporaryDirectory() as tmp_dir:
+            tmp_dir = Path(tmp_dir)
+            file = tmp_dir / "tmp"
+            serialiser.serialise(pipeline, file)
+
+            loaded = serialiser.deserialise(file)
+
+            old_data = pipeline.transform(data)
+            new_data = loaded.transform(data)
+
+            assert len(old_data) == len(new_data)
+            for name, old_value in old_data.items():
+                new_value = new_data[name]
+                assert_frame_equal(old_value.toPandas(), new_value.toPandas())
+
+    def test_serialise_joblib(self):
+        serialiser = SparkSerialiser()
+
+        with TemporaryDirectory() as tmp_dir:
+            tmp_dir = Path(tmp_dir)
+            file = tmp_dir / "tmp"
+            a = Path("a") / "b"
+
+            serialiser.serialise(a, file)
+            loaded = serialiser.deserialise(file)
+
+            assert a == loaded
diff --git a/tests/spark/utils.py b/tests/spark/utils.py
new file mode 100644
index 0000000..983f5b0
--- /dev/null
+++ b/tests/spark/utils.py
@@ -0,0 +1,46 @@
+import numpy as np
+import pandas as pd
+from pandas.testing import assert_frame_equal
+from pyspark import SparkConf
+from pyspark.sql import SparkSession
+from pytest import fixture, mark, param
+
+
+@fixture(
+    scope="session",
+    params=[param("spark", marks=[mark.spark, mark.filterwarnings("ignore:distutils Version classes are deprecated")])],
+)
+def spark_session() -> SparkSession:
+    """Creates a SparkSession fixture for the entire test session.
+
+    This is setup such that that SparkSession should only be initialised once.
+    Since it can take a long time for Spark to start, this should save us time.
+
+    If you do not want to test Spark-related functionality, you can skip these
+    tests by running:
+        `python -m pytest -m "not spark"`
+
+    This SHOULD disable all tests using this fixture through the mark in the
+    fixture's parameters.
+
+    Returns:
+        SparkSession: A Spark session to use, shared between tests to save
+            boot time.
+    """
+    print("Launching Spark session: please wait...")
+    config = (
+        SparkConf()
+        .set("spark.sql.execution.arrow.enabled", "False")
+        .set("spark.sql.execution.arrow.pyspark.enabled", "False")
+    )
+    return SparkSession.builder.config(conf=config).getOrCreate()
+
+
+def test_equal_dataframes(
+    expected: pd.DataFrame,
+    result: pd.DataFrame,
+):
+    assert expected.shape[1] == result.shape[1]
+    assert np.all(np.isin(expected.columns, result.columns))
+
+    assert_frame_equal(expected, result[expected.columns])
diff --git a/tests/tensorflow/__init__.py b/tests/tensorflow/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/tensorflow/data/__init__.py b/tests/tensorflow/data/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/tensorflow/data/test_tensor_source.py b/tests/tensorflow/data/test_tensor_source.py
new file mode 100644
index 0000000..90fdcbf
--- /dev/null
+++ b/tests/tensorflow/data/test_tensor_source.py
@@ -0,0 +1,12 @@
+import tensorflow as tf
+from numpy.testing import assert_array_equal
+
+from pype.tensorflow.data.tensor_source import TensorSource
+
+
+def test_TensorSource():
+    tensor = tf.convert_to_tensor([1, 2, 3, 4])
+
+    source = TensorSource(tensor)
+
+    assert_array_equal(tensor.numpy(), source.read().numpy())
diff --git a/tests/tensorflow/model/__init__.py b/tests/tensorflow/model/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/tensorflow/model/test_keras_pype_model.py b/tests/tensorflow/model/test_keras_pype_model.py
new file mode 100644
index 0000000..1799d10
--- /dev/null
+++ b/tests/tensorflow/model/test_keras_pype_model.py
@@ -0,0 +1,126 @@
+import shutil
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock, call, patch
+
+import numpy as np
+import tensorflow as tf
+from keras import Model
+from keras.losses import MeanAbsoluteError
+from keras.optimizers import Adam
+from pytest import mark
+
+from pype.base.data import DataSet
+from pype.base.model import Model as PypeModel
+from pype.tensorflow.model import KerasPypeModel, MLPKeras
+
+
+class DummyKerasModel(KerasPypeModel[MLPKeras]):
+    def _init_model(self, args: dict[str, Any]) -> MLPKeras:
+        return MLPKeras(10, 2, 2)
+
+
+class Test_keras_pype_model:
+    # tested through MLPPypeModel since that can be instantiated.
+    @mark.parametrize(
+        ["model"],
+        [
+            [None],
+            [MLPKeras(10, 2, 2)],
+        ],
+    )
+    def test_init(self, model):
+        arguments = {"a": None}
+        with patch.object(DummyKerasModel, "_init_model") as mock_init:
+            pype_model = DummyKerasModel(["x"], ["y"], MagicMock(), MagicMock(), model=model, **arguments)
+
+            if model is None:
+                mock_init.assert_called_once_with(arguments)
+            else:
+                mock_init.assert_not_called()
+
+    def test_get_annotated_class(self):
+        result = DummyKerasModel._get_annotated_class()
+
+        assert result == MLPKeras
+
+    def test_set_seed(self):
+        seed = 3
+        pype_model = DummyKerasModel(["x"], ["y"], MagicMock(), MagicMock(), seed=seed)
+
+        with patch("pype.tensorflow.model.keras_pype_model.set_seed") as mock_seed:
+            pype_model.set_seed()
+            mock_seed.assert_called_once_with(seed)
+
+    def test_fit(self):
+        model = MagicMock()
+        loss = MagicMock()
+        optimizer_class = MagicMock()
+        metrics = [MagicMock(), MagicMock()]
+        pype_model = DummyKerasModel(
+            ["x"], ["y"], loss=loss, optimizer_class=optimizer_class, model=model, metrics=metrics
+        )
+
+        x = MagicMock()
+        y = MagicMock()
+        pype_model._fit(x, y)
+
+        model.compile.assert_called_once_with(optimizer=optimizer_class.return_value, loss=loss, metrics=metrics)
+
+        model.fit.assert_called_once()
+
+    def test_transform(self):
+        model = MagicMock()
+        loss = MagicMock()
+        optimizer_class = MagicMock()
+        pype_model = DummyKerasModel(["x"], ["y"], loss=loss, optimizer_class=optimizer_class, model=model)
+
+        x = MagicMock()
+        y = MagicMock()
+        pype_model._transform(x, y)
+
+        model.call.assert_called_once_with(x, y)
+
+    def test_integration_with_save(self):
+        folder = Path(__file__).parent / "tmp"
+        folder.mkdir()
+        ds = DataSet(
+            x=tf.zeros((25, 5)),
+            y=tf.zeros(25),
+        )
+        loss = MeanAbsoluteError()
+
+        try:
+            model = MLPKeras(layer_size=20, n_layers=3, output_size=1)
+            pype_model = DummyKerasModel(["x"], ["y"], loss, Adam, model=model)
+
+            pype_model.fit(ds)
+            initial_result = pype_model.transform(ds)
+
+            pype_model.save(folder)
+
+            loaded = DummyKerasModel.load(folder)
+            loaded_result = loaded.transform(ds)
+
+            np.testing.assert_array_almost_equal(initial_result["y"].numpy(), loaded_result["y"].numpy())
+        finally:
+            shutil.rmtree(folder)
+
+    def test_get_parameters(self):
+        parser = MagicMock()
+
+        with patch("pype.tensorflow.model.keras_pype_model.add_args_to_parser_for_class") as mock_args:
+            DummyKerasModel.get_parameters(parser)
+
+            mock_args.assert_has_calls(
+                [
+                    call(parser, MLPKeras, "model", [Model], excluded_args=[]),
+                    call(
+                        parser,
+                        DummyKerasModel,
+                        "model",
+                        [PypeModel],
+                        excluded_args=["seed", "inputs", "outputs", "model", "loss", "optimizer_class"],
+                    ),
+                ]
+            )
diff --git a/tests/tensorflow/model/test_mlp_keras.py b/tests/tensorflow/model/test_mlp_keras.py
new file mode 100644
index 0000000..6603b26
--- /dev/null
+++ b/tests/tensorflow/model/test_mlp_keras.py
@@ -0,0 +1,17 @@
+import tensorflow as tf
+
+from pype.tensorflow.model.mlp_keras import MLPKeras
+
+
+def test_mlp_keras():
+    n_hidden = 2
+    output_size = 3
+    size_hidden = 20
+    hidden_act = "relu"
+
+    model = MLPKeras(output_size, n_hidden, size_hidden, hidden_act)
+    assert len(model._layers) == n_hidden + 1
+
+    data = tf.random.normal((5, 17))
+    result = model.call(data)
+    assert result.shape == (5, output_size)
diff --git a/tests/tensorflow/model/test_mlp_pype_keras_model.py b/tests/tensorflow/model/test_mlp_pype_keras_model.py
new file mode 100644
index 0000000..a5a9e66
--- /dev/null
+++ b/tests/tensorflow/model/test_mlp_pype_keras_model.py
@@ -0,0 +1,36 @@
+import tensorflow as tf
+from keras.losses import MeanAbsoluteError
+from keras.optimizers import Nadam
+
+from pype.base.data.dataset import DataSet
+from pype.tensorflow.model.mlp_pype_model import MLPPypeModel
+
+
+def test_mlp_pype_keras_model():
+    n = 1000
+    x = tf.random.normal((n, 4))
+    y = x[:, 0] + x[:, 1] * -0.4 + x[:, 2] * 2 + x[:, 3] * -0.4
+
+    model = MLPPypeModel(
+        ["x"],
+        ["y"],
+        epochs=10,
+        batch_size=25,
+        optimizer_class=Nadam,
+        learning_rate=0.01,
+        loss=MeanAbsoluteError(),
+        output_size=1,
+        n_layers=1,
+        layer_size=5,
+    )
+    dataset = DataSet(x=x, y=y)
+
+    old_predictions = model.transform(dataset)["y"]
+    model.fit(dataset)
+
+    fit_predictions = model.transform(dataset)["y"]
+
+    old_error = MeanAbsoluteError()(y, old_predictions)
+    fit_error = MeanAbsoluteError()(y, fit_predictions)
+
+    assert fit_error < old_error
diff --git a/tests/tensorflow/pipeline/__init__.py b/tests/tensorflow/pipeline/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/tensorflow/pipeline/test_tensor_type_checker.py b/tests/tensorflow/pipeline/test_tensor_type_checker.py
new file mode 100644
index 0000000..d7a4f0d
--- /dev/null
+++ b/tests/tensorflow/pipeline/test_tensor_type_checker.py
@@ -0,0 +1,116 @@
+import numpy as np
+import tensorflow as tf
+from pydantic import ValidationError
+from pytest import mark
+
+from pype.tensorflow.pipeline.tensor_checker import (
+    TensorflowData,
+    TensorflowTypeChecker,
+)
+from tests.utils import pytest_assert
+
+
+class Test_TensorflowData:
+    def test_convert(self):
+        model = TensorflowData(data=[[1, 2, 3], [4, 5, 6], [6, 7, 8]])
+        expected = tf.convert_to_tensor([[1, 2, 3], [4, 5, 6], [6, 7, 8]])
+
+        result = model.convert()
+
+        assert isinstance(result, tf.Tensor)
+        np.testing.assert_array_equal(expected.numpy(), result.numpy())
+
+    def test_to_model(self):
+        array = tf.convert_to_tensor([[1, 2, 3], [2, 3, 4]])
+        result = TensorflowData.to_model(array)
+
+        expected = TensorflowData(
+            data=[[1, 2, 3], [2, 3, 4]],
+        )
+
+        assert expected == result
+
+
+class Test_TensorflowTypeChecker:
+    def test_fit(self):
+        type_checker = TensorflowTypeChecker()
+        array = tf.convert_to_tensor([[1, 2, 3], [3, 4, 5]])
+        assert array.shape[0] != array.shape[1]
+
+        type_checker.fit(array)
+
+        assert type_checker.dims == (3,)
+        assert type_checker.dtype == int
+
+    def test_transform(self):
+        type_checker = TensorflowTypeChecker()
+
+        array = tf.convert_to_tensor([[1, 2, 3], [3, 4, 5]])
+        type_checker.fit(array)
+
+        # works
+        result = type_checker.transform(array)
+
+        # errors #
+        # other type of object
+        with pytest_assert(AssertionError, "Please provide a Tensorflow tensor!"):
+            type_checker.transform([])
+
+        # other shape
+        other_shape = tf.convert_to_tensor([[1, 2], [3, 4]])
+        with pytest_assert(
+            AssertionError,
+            f"Dimensions of Tensorflow tensors do not add up: {other_shape.shape[1:]} vs {array.shape[1:]}",
+        ):
+            type_checker.transform(other_shape)
+
+        # other type
+        other_type = tf.convert_to_tensor([[1, 2, 3], [3, 4, 5.0]])
+        with pytest_assert(AssertionError, f"Dtype of data does not add up: {float} vs {int}"):
+            type_checker.transform(other_type)
+
+    @mark.parametrize(
+        ["input_array", "expected"],
+        [
+            [tf.convert_to_tensor([1, 2, 3]), int],
+            [tf.convert_to_tensor([1.0, 2, 3]), float],
+            [tf.convert_to_tensor([False, True]), bool],
+            [tf.convert_to_tensor(["1", "2", "3"]), str],
+        ],
+    )
+    def test_convert_dtype(self, input_array: tf.Tensor, expected: type):
+        type_checker = TensorflowTypeChecker()
+
+        result = type_checker._convert_dtype(input_array.dtype)
+        assert result == expected
+
+    def test_get_pydantic_types(self):
+        type_checker = TensorflowTypeChecker()
+        array = tf.convert_to_tensor([[1, 2, 3], [7, 6, 5]])
+        type_checker.fit(array)
+
+        TensorflowSpecificType = type_checker.get_pydantic_type()
+
+        # succeed
+        TensorflowSpecificType(data=[[1, 2, 3], [9, 0, 3]])
+
+        # also succeed
+        data = [[1, 2, 3], [9, 0, 3.0]]
+        result = TensorflowSpecificType(data=data).convert()
+
+        assert isinstance(result, tf.Tensor)
+        np.testing.assert_array_equal(tf.convert_to_tensor(data).numpy(), result.numpy())
+
+    def test_get_pydantic_types_strict_on_bool(self):
+        type_checker = TensorflowTypeChecker()
+        array = tf.convert_to_tensor([[True, False]])
+        type_checker.fit(array)
+
+        TensorflowSpecificType = type_checker.get_pydantic_type()
+
+        # succeed
+        TensorflowSpecificType(data=[[False, False]])
+
+        # fails
+        with pytest_assert(ValidationError):
+            TensorflowSpecificType(data=[[1, 2, 3]])
diff --git a/tests/test_utils.py b/tests/utils.py
similarity index 69%
rename from tests/test_utils.py
rename to tests/utils.py
index 28cf90f..eced4f6 100644
--- a/tests/test_utils.py
+++ b/tests/utils.py
@@ -1,36 +1,55 @@
 from contextlib import contextmanager
 from pathlib import Path
 
+from data.data_sink import DataSink
+
 from pype.base.data.data_source import DataSource
 from pype.base.data.dataset_source import DataSetSource
 from pype.base.evaluate.evaluator import Evaluator
-from pype.base.experiment.experiment import Experiment
-from pype.base.logger.local_logger import LocalLogger
 from pype.base.model import Model
 from pype.base.pipeline.operator import Operator
 from pype.base.pipeline.pipe import Pipe
 from pype.base.pipeline.pipeline import Pipeline
-from pype.base.pipeline.type_checker import DataModel, TypeChecker, TypeCheckerPipe
-from pype.base.serialiser.joblib_serialiser import JoblibSerialiser
+from pype.base.pipeline.type_checker import (
+    DataModel,
+    DataSetModel,
+    TypeChecker,
+    TypeCheckerPipe,
+)
+from tests.utils_training_support import reverse
+
+
+class AnyArg:
+    def __eq__(self, __o: object) -> bool:
+        return True
 
 
 @contextmanager
-def pytest_assert(error, message: str):
+def pytest_assert(error_class, message: str | None = None, exact: bool = True):
     try:
         yield
         raise ValueError("No error was raised!")
-    except error as e:
-        assert e.args[0] == message
+    except error_class as e:
+        if message is not None:
+            error_message: str = e.args[0]
+            if exact:
+                assert error_message == message
+            else:
+                assert message in error_message
 
 
 class DummyModel(Model[list[int | float]]):
     mean_file = "mean.txt"
 
+    def __init__(self, inputs: list[str], outputs: list[str], seed: int = 1, a: int = 3) -> None:
+        super().__init__(inputs, outputs, seed)
+        self.a = a
+
     def set_seed(self) -> None:
         pass
 
     def _fit(self, x: list[int | float], y: list[int | float]) -> None:
-        self.prediction = sum(y) / len(y)
+        self.prediction = sum(y) / len(y) + self.a
 
     def _transform(self, x: list[int | float]) -> list[int | float]:
         return [self.prediction for _ in x]
@@ -55,6 +74,9 @@ def __init__(self, l) -> None:
     def read(self) -> list[float]:
         return self.l
 
+    def __eq__(self, __o: object) -> bool:
+        return self.l == __o.l
+
 
 def get_dummy_data(n: int, x_offset: int, y_offset: int) -> DataSetSource:
     return DataSetSource(
@@ -63,12 +85,20 @@ def get_dummy_data(n: int, x_offset: int, y_offset: int) -> DataSetSource:
     )
 
 
+class DummyDataSink(DataSink[list[float]]):
+    def __init__(self) -> None:
+        self.data = None
+
+    def write(self, data: list[float]) -> None:
+        self.data = data
+
+
 class DummyOperator(Operator[list[float]]):
     def fit(self, x: list[float]) -> "Operator":
         return self
 
     def transform(self, x: list[float]) -> list[float]:
-        return [i - 1 for i in x]
+        return reverse([i - 1 for i in x])
 
 
 def get_dummy_pipeline() -> Pipeline:
@@ -90,13 +120,18 @@ def to_model(cls, data: list[float]) -> "DataModel":
         return cls(data=data)
 
 
+class DummyDataSet(DataSetModel):
+    x: DummyDataModel
+    y: DummyDataModel
+
+
 class DummyTypeChecker(TypeChecker):
     def fit(self, data: list[float]) -> "Operator":
         return super().fit(data)
 
     def transform(self, data: list[float]) -> list[float]:
-        assert isinstance(data, list)
-        assert isinstance(data[0], float) or isinstance(data[0], int)
+        assert isinstance(data, list), "Provide a list!"
+        assert isinstance(data[0], float) or isinstance(data[0], int), "Provide a list with ints/floats!"
         return data
 
     def get_pydantic_type(self) -> type[DataModel]:
@@ -107,25 +142,3 @@ def get_dummy_type_checkers() -> tuple[TypeCheckerPipe, TypeCheckerPipe]:
     return TypeCheckerPipe("input", ["x"], type_checker_classes=[(list, DummyTypeChecker)]), TypeCheckerPipe(
         "output", ["y"], type_checker_classes=[(list, DummyTypeChecker)]
     )
-
-
-def get_dummy_experiment() -> Experiment:
-    train = get_dummy_data(20, 1, 0)
-    test = get_dummy_data(5, 2, -1)
-
-    output_folder = Path(__file__).parent / "tmp_folder"
-
-    input_checker, output_checker = get_dummy_type_checkers()
-
-    return Experiment(
-        data_sources={"train": train, "test": test},
-        model=DummyModel(inputs=["x"], outputs=["y"]),
-        pipeline=get_dummy_pipeline(),
-        logger=LocalLogger(),
-        evaluator=get_dummy_evaluator(),
-        serialiser=JoblibSerialiser(),
-        output_folder=output_folder,
-        input_type_checker=input_checker,
-        output_type_checker=output_checker,
-        parameters={"param_1": 2},
-    )
diff --git a/tests/utils_training_support.py b/tests/utils_training_support.py
new file mode 100644
index 0000000..8a45155
--- /dev/null
+++ b/tests/utils_training_support.py
@@ -0,0 +1,5 @@
+"""Custom 'function' to test importing / exporting data files."""
+
+
+def reverse(x: list[float]) -> list[float]:
+    return x[::-1]
diff --git a/tests/xgboost/__init__.py b/tests/xgboost/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/xgboost/test_classifier.py b/tests/xgboost/test_classifier.py
new file mode 100644
index 0000000..9e2ae2e
--- /dev/null
+++ b/tests/xgboost/test_classifier.py
@@ -0,0 +1,74 @@
+import shutil
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from data.dataset import DataSet
+from pytest import fixture
+
+from pype.xgboost.model import XGBClassifierModel
+
+
+@fixture
+def dataset() -> DataSet:
+    n = 3
+    x = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5] * n,
+            "b": [20, 3, 2, 3, 20] * n,
+        }
+    )
+
+    y = pd.DataFrame(
+        {
+            "target": [1, 0, 0, 0, 1] * n,
+        }
+    )
+    return DataSet(x=x, y=y)
+
+
+@fixture
+def model() -> XGBClassifierModel:
+    return XGBClassifierModel(
+        inputs=["x"],
+        outputs=["y"],
+        n_estimators=10,
+        max_depth=3,
+        learning_rate=0.01,
+        min_child_weight=1,  # Please do not change this. Weird bugs may happen
+    )
+
+
+def error(target, prediction):
+    return float((target == prediction).mean())
+
+
+def test_xgb_regressor_training(dataset: DataSet, model: XGBClassifierModel):
+    # training works if trained model produces better performance than mean
+
+    baseline = error(dataset["y"]["target"], 0)
+
+    model.fit(dataset)
+
+    predictions = model.transform(dataset)["y"]
+    performance = error(dataset["y"]["target"], predictions)
+
+    assert performance > baseline
+
+
+def test_xgb_regressor_saving_and_loading(dataset: DataSet, model: XGBClassifierModel):
+    # test if we can save and load the model
+    target_path = Path("tmp_xgboost")
+
+    try:
+        model.fit(dataset)
+        predictions_old = model.transform(dataset)["y"]
+
+        model.save(target_path)
+        new_model = XGBClassifierModel.load(target_path)
+        predictions_new = new_model.transform(dataset)["y"]
+
+        np.testing.assert_array_equal(predictions_old, predictions_new)
+    finally:
+        shutil.rmtree(target_path)
+    assert not target_path.exists()
diff --git a/tests/xgboost/test_regressor.py b/tests/xgboost/test_regressor.py
new file mode 100644
index 0000000..09195eb
--- /dev/null
+++ b/tests/xgboost/test_regressor.py
@@ -0,0 +1,72 @@
+import shutil
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from data.dataset import DataSet
+from pytest import fixture
+
+from pype.xgboost.model import XGBRegressorModel
+
+
+@fixture
+def dataset() -> DataSet:
+    x = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5],
+            "b": [9, 3, 2, 3, 8],
+        }
+    )
+
+    y = pd.DataFrame(
+        {
+            "target": [11, 2, 3, 4, 15],
+        }
+    )
+    return DataSet(x=x, y=y)
+
+
+@fixture
+def model() -> XGBRegressorModel:
+    return XGBRegressorModel(
+        inputs=["x"],
+        outputs=["y"],
+        n_estimators=20,
+        objective="reg:squarederror",
+        max_depth=3,
+    )
+
+
+def error(target, prediction):
+    return float((target - prediction).abs().sum())
+
+
+def test_xgb_regressor_training(dataset: DataSet, model: XGBRegressorModel):
+    # training works if trained model produces better performance than mean
+
+    baseline = error(dataset["y"]["target"], dataset["y"]["target"].mean())
+
+    model.fit(dataset)
+
+    predictions = model.transform(dataset)["y"]
+    performance = error(dataset["y"]["target"], predictions)
+
+    assert performance < baseline
+
+
+def test_xgb_regressor_saving_and_loading(dataset: DataSet, model: XGBRegressorModel):
+    # test if we can save and load the model
+    target_path = Path("tmp_xgboost")
+
+    try:
+        model.fit(dataset)
+        predictions_old = model.transform(dataset)["y"]
+
+        model.save(target_path)
+        new_model = XGBRegressorModel.load(target_path)
+        predictions_new = new_model.transform(dataset)["y"]
+
+        np.testing.assert_array_equal(predictions_old, predictions_new)
+    finally:
+        shutil.rmtree(target_path)
+    assert not target_path.exists()