Repository: globus-labs/FLoX-prototype (language: python)

Issue:
added support for running average



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 595230a..7db1d11 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,9 @@ env
 *.npy
 build/
 pyflox.egg-info/
+build/
+dist/
+data/
 
+*.pyc
+.vscode/settings.json
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 49768af..21978a2 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -28,8 +28,8 @@ repos:
   hooks:
     - id: pyupgrade
       args: ["--py36-plus"]
-- repo: https://github.com/PyCQA/flake8
-  rev: 5.0.4
-  hooks:
-    - id: flake8
-      args: [flox/]
\ No newline at end of file
+# - repo: https://github.com/PyCQA/flake8
+#   rev: 5.0.4
+#   hooks:
+#     - id: flake8
+#       args: [flox/]
\ No newline at end of file
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
new file mode 100644
index 0000000..5f6a6c5
--- /dev/null
+++ b/.readthedocs.yaml
@@ -0,0 +1,25 @@
+
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/conf.py
+
+# Build documentation with MkDocs
+#mkdocs:
+#  configuration: mkdocs.yml
+
+# Optionally build your docs in additional formats such as PDF
+#formats:
+#  - pdf
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  version: 3.7
+  install:
+    - requirements: docs/requirements.txt
\ No newline at end of file
diff --git a/Makefile b/Makefile
index ef3617d..439b4b7 100644
--- a/Makefile
+++ b/Makefile
@@ -4,7 +4,11 @@ typehint:
 
 .PHONY: test
 test:
-	pytest tests/
+	pytest flox/tests/ --disable-warnings
+
+.PHONY: test_unit
+test_unit:
+	pytest flox/tests/unit --disable-warnings
 
 .PHONY: lint
 lint:
diff --git a/README.rst b/README.rst
index e3ee85e..1f33d4a 100644
--- a/README.rst
+++ b/README.rst
@@ -43,18 +43,16 @@ on the *Controller*:
    (.venv) $ pip install pyflox
 
 To be able to include your *endpoints* into the Federated Learning process,
-you need to install `funcx-endpoint <https://funcx.readthedocs.io/en/latest/endpoints.html>`_:
+you need to install `funcx-endpoint <https://funcx.readthedocs.io/en/latest/endpoints.html>`_ on the *endpoints*:
 
 .. code-block:: console
 
    (.endpoint_venv) $ python3 -m pipx install funcx_endpoint
 
-Finally, install compatible versions of Tensorflow on both the Controller and endpoints.
-As of now, FLoX only supports Tensorflow, although support for PyTorch will be added soon.
-
-.. code-block:: console
-
-   (.venv) $ pip install tensorflow==1.14.0
+Finally, install compatible versions of `PyTorch <https://qengineering.eu/install-pytorch-on-raspberry-pi-4.html) & [Torchvision](https://qengineering.eu/install-pytorch-on-raspberry-pi-4.html#:~:text=rm%20%2Drf%20~/pytorch-,TorchVision.,-Install%20torchvision%20on>`_
+or `Tensorflow <https://qengineering.eu/install-tensorflow-2.1.0-on-raspberry-pi-4.html>`_ on both the Controller and endpoints.
+As of now, FLoX supports Tensorflow and PyTorch, although you can add support for your frameworks by creating a new
+ModelTrainer. See ``flox.logic`` and examples of ModelTrainers in ``flox.model_trainers``.
 
 *Note*: ``funcX-endpoint`` is only supported on Linux.
 FLoX Controller functionality is supported on MacOS, Linux and Windows.
diff --git a/build/lib/flox/flox.py b/build/lib/flox/flox.py
deleted file mode 100644
index d585eee..0000000
--- a/build/lib/flox/flox.py
+++ /dev/null
@@ -1,409 +0,0 @@
-import numpy as np
-import time
-
-from funcx.sdk.client import FuncXClient
-from funcx.sdk.executor import FuncXExecutor
-
-
-def get_edge_weights(sample_counts):
-    """
-    Takes an array of numbers and returns their fractions of the total number of samples
-    Can be used to find weights for the weighted_average 
-
-    Parameters
-    ----------
-    sample_counts: numpy array of integers
-
-    Returns
-    -------
-    fractions: numpy array
-    
-    """
-    total = sum(sample_counts)
-    fractions = sample_counts/total
-    return fractions
-
-def eval_model(m, x, y, silent=False):
-    """
-    Evaluate the model on a datset
-
-    Parameters
-    ----------
-    m: Tensorflow model
-
-    x: numpy array
-        dataset entries, e.g. x_test
-    
-    y: numpy array
-        labels for the entries, e.g., y_test
-    
-    """
-    score = m.evaluate(x, y, verbose=0)
-    if not silent:
-        print("Test loss:", score[0])
-        print("Test accuracy:", score[1])
-    return score[0], score[1]
-
-def training_function(json_model_config, 
-                      global_model_weights, 
-                      num_samples=None,
-                      epochs=10,     
-                      data_source="keras",
-                      keras_dataset="mnist",
-                      preprocess=True,
-                      path_dir='/home/pi/datasets', 
-                      x_train_name="mnist_x_train.npy", 
-                      y_train_name="mnist_y_train.npy",
-                      input_shape=(32, 28, 28, 1),
-                      loss="categorical_crossentropy",
-                      optimizer="adam",
-                      metrics=["accuracy"]                                               
-):
-    """
-    This function gets deployed to the given endpoints with corresponding parameters. 
-    Returns the updated model weights and number of samples it was trained on.
-
-    Parameters
-    ----------
-    json_model_config: str
-        configuration of the TF model retrieved using model.to_json()
-
-    global_model_weights: numpy array
-        a numpy array with weights of the TF model
-
-    num_samples: int
-        if data_source="keras", randomly samples n data points from (x_train, y_train)
-
-    epochs: int
-        the number of epochs to train the model for
-
-    data_source: str
-        the function supports three data sources: "local", "keras"
-        for "local" and "keras", see get_local_data and get_keras_data functions
-
-    keras_dataset: str
-        specifies one of the default keras datasets to use if using "keras" data source
-        ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-
-    preprocess: boolean
-            if True, will attempt to preprocess your data from "keras" data sources
-
-    path_dir: str
-        path to the folder with x_train and y_train; needed when data_sourse="local"
-
-    x_train_name: str
-        filename for x_train; needed when data_sourse="local"
-
-    y_train_name: str
-        file name for y_train; needed when data_sourse="local"
-
-    input_shape: tupple
-        input shape for the provided model
-    
-    loss: str
-        loss for TF's model.fit() function
-
-    optimizer: str
-        optimizer for for TF's model.fit() function
-
-    metrics: str/list 
-        metrics for TF's model.fit() function. E.g, metrics=["accuracy"],
-
-    Returns
-    -------
-    model_weights: numpy array
-        updated TF model weights after training on the local data
-
-    samples_counts: int
-        the number of samples the model was trained on. 
-        Can be used to find the weighted average by # of samples seen
-
-    """
-    # import all the dependencies required for funcX functions)
-    from tensorflow import keras
-    import numpy as np
-    import os
-
-    # retrieve (and optionally process) the data
-    if data_source == 'keras':
-        available_datasets = ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-        dataset_mapping= {
-            'mnist': keras.datasets.mnist,
-            'fashion_mnist': keras.datasets.fashion_mnist,
-            'cifar10': keras.datasets.cifar10,
-            'cifar100': keras.datasets.cifar100,
-            'imdb': keras.datasets.imdb,
-            'reuters': keras.datasets.reuters,
-            'boston_housing': keras.datasets.boston_housing
-        }
-        image_datasets = ['mnist', 'fashion_mnist', 'cifar10', 'cifar100']
-
-        # check if the Keras dataset exists
-        if keras_dataset not in available_datasets:
-            raise Exception(f"Please select one of the built-in Keras datasets: {available_datasets}")
-
-        else:
-            # load the data
-            (x_train, y_train), _ = dataset_mapping[keras_dataset].load_data()
-
-            # take a random set of images
-            if num_samples:
-                idx = np.random.choice(np.arange(len(x_train)), num_samples, replace=True)
-                x_train = x_train[idx]
-                y_train = y_train[idx]
-
-            if preprocess:
-                # do default image processing for built-in Keras images    
-                if keras_dataset in image_datasets:
-                    # Scale images to the [0, 1] range
-                    x_train = x_train.astype("float32") / 255
-
-                    # Make sure images have shape (num_samples, x, y, 1) if working with MNIST images
-                    if x_train.shape[-1] not in [1, 3]:
-                        x_train = np.expand_dims(x_train, -1)
-
-                    # convert class vectors to binary class matrices
-                    if keras_dataset == 'cifar100':
-                        num_classes=100
-                    else:
-                        num_classes=10
-                        
-                    y_train = keras.utils.to_categorical(y_train, num_classes)
-
-    elif data_source == 'local':        
-        # construct the path
-        x_train_path_file = os.sep.join([path_dir, x_train_name])
-        y_train_path_file = os.sep.join([path_dir, y_train_name])
-
-        # load the files
-        with open(x_train_path_file, 'rb') as f:
-            x_train = np.load(f)
-            
-        with open(y_train_path_file, 'rb') as f:
-            y_train = np.load(f)
-
-        # if preprocess is True & the function is valid, preprocess the data
-        if preprocess:
-            # check if a valid function was given
-            depth = input_shape[3]
-            image_size_y = input_shape[2]
-            image_size_x = input_shape[1]
-
-            # take a limited number of samples, if indicated
-            if num_samples:
-                idx = np.random.choice(np.arange(len(x_train)), num_samples, replace=True)
-                x_train = x_train[idx]
-                y_train = y_train[idx]
-            
-            # reshape and scale to pixel values to 0-1
-            x_train = x_train.reshape(len(x_train), image_size_x, image_size_y, depth)
-            x_train = x_train / 255.0
-
-    else:
-        raise Exception("Please choose one of data sources: ['local', 'keras']")
-
-    # create the model
-    model = keras.models.model_from_json(json_model_config)
-
-    # compile the model and set weights to the global model
-    model.compile(loss=loss, optimizer=optimizer, metrics=metrics)
-
-    try:
-        model.set_weights(global_model_weights)
-    except:
-        # some older TF versions require the model to be 'built' first
-        model.build(input_shape=input_shape)
-        model.set_weights(global_model_weights)
-
-    # train the model on the local data and extract the weights
-    model.fit(x_train, y_train, epochs=epochs)
-    model_weights = model.get_weights()
-
-    # transform to a numpy array
-    np_model_weights = np.asarray(model_weights, dtype=object)
-
-    # return the updated weights and number of samples the model was trained on
-    return {"model_weights":np_model_weights, "samples_count": x_train.shape[0]}
-
-def federated_learning(global_model, 
-                      endpoint_ids, 
-                      num_samples=100,
-                      epochs=5,
-                      loops=1,
-                      time_interval=0,
-                      aggregation_mode="weighted_average",
-                      data_source: str = "keras",
-                      keras_dataset = "mnist",  
-                      preprocess=False,
-                      path_dir='/home/pi/datasets', 
-                      x_train_name="mnist_x_train.npy", 
-                      y_train_name="mnist_y_train.npy",
-                      input_shape=(32, 28, 28, 1),
-                      loss="categorical_crossentropy",
-                      optimizer="adam", 
-                      metrics=["accuracy"],
-                      evaluation_function=eval_model,
-                      x_test=None,
-                      y_test=None
-                      ):
-    """
-    Facilitates Federated Learning for *loops* rounds. 
-
-    Parameters
-    ----------
-    global_model: TF model object
-        compiled TF model that will be deployed for training on the endpoints
-
-    endpoint_ids: list of str
-        a list with endpoint_ids to include in the FL process. 
-        The ids can be retrieved by running 'funcx-endpoint list' on participating devices
-
-    num_samples: int or list
-        indicates how many samples to get for training on endpoints
-        if int, applies the same num_samples to all endpoints. 
-        if list, it will use the corresponding number of samples for each device
-        the list should have the same number of entries as the number of endpoints
-
-
-    epochs: int or list
-        indicates how many epochs to use for training on endpoints
-        if int, applies the same number of epochs to all endpoints. 
-        if list, it will use the corresponding number of epochs for each device
-        the list should have the same number of entries as the number of endpoints
-
-    loops: int
-        defines how many FL rounds to run. Each round consists of deploying the mode, training
-        aggregating the updates, and reassigning new weights to the model. 
-
-    time_interval: int
-        defines the pause between FL rounds in seconds (default=0). Can be useful if you want to run 
-        a round every minute/hour/day.
-
-    aggregation_mode: str
-        defines the aggregation mode from a choice of a simple 'average' or a 'weighted_average'
-        the 'weighted_average' aggregates the updates based on how many samples each updates has 
-        been trained on. This gives more weight to updates from devices that have seen more data.
-
-    data_source: str
-        the function supports three data sources: "local", "keras"
-        for "local" and "keras", see get_local_data and get_keras_data functions
-
-    keras_dataset: str
-        specifies one of the default keras datasets to use if using "keras" data source
-        ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-
-    preprocess: boolean
-            if True, will attempt to preprocess your data from "keras" data sources
-
-    path_dir: str
-        path to the folder with x_train and y_train; needed when data_sourse="local"
-
-    x_train_name: str
-        filename for x_train; needed when data_sourse="local"
-
-    y_train_name: str
-        file name for y_train; needed when data_sourse="local"
-
-    input_shape: tupple
-        input shape for the provided model
-    
-    loss: str
-        loss for TF's model.fit() function
-
-    optimizer: str
-        optimizer for for TF's model.fit() function
-
-    metrics: str/list 
-        metrics for TF's model.fit() function. E.g, metrics=["accuracy"],
-
-    evaluation_function: function
-        if supplied, evaluates the model on x_test and y_test
-
-    x_test: list/numpy array/tensors
-        x_test data for testing
-
-    y_test: list
-        y_test labels for x_test
-
-
-    Returns
-    -------
-    global_model: TF model
-        the original model but with updated weights after all the FL rounds
-
-    """
-    # instantiate the FuncXExecutor
-    fx = FuncXExecutor(FuncXClient())
-
-    # if num_samples or epochs is an int, convert to list so the same number can be applied to all endpoints
-    if type(num_samples) == int:
-        num_samples = [num_samples]*len(endpoint_ids)
-
-    if type(epochs) == int:
-        epochs = [epochs]*len(endpoint_ids)
-    
-    # start running FL loops
-    for i in range(loops):
-
-        # get the model's architecture and weights
-        json_config = global_model.to_json()
-        gm_weights = global_model.get_weights()
-        gm_weights_np = np.asarray(gm_weights, dtype=object)
-
-        #fx = FuncXExecutor(FuncXClient())
-        tasks = []
-
-        # submit the corresponding parameters to each endpoint for a round of FL 
-        for e, num_s, num_epoch, path_d in zip(endpoint_ids, num_samples, epochs, path_dir): 
-            tasks.append(fx.submit(training_function, 
-                                   json_model_config=json_config, 
-                                    global_model_weights=gm_weights_np, 
-                                    num_samples=num_s,
-                                    epochs=num_epoch,
-                                    data_source=data_source,
-                                    keras_dataset=keras_dataset,
-                                    preprocess=preprocess,
-                                    path_dir=path_d,
-                                    x_train_name=x_train_name,
-                                    y_train_name=y_train_name,
-                                    input_shape=input_shape,
-                                    loss=loss,
-                                    optimizer=optimizer,
-                                    metrics=metrics,
-                                    endpoint_id=e))
-        
-        # extract model updates from each endpoints once they are available
-        model_weights = [t.result()["model_weights"] for t in tasks]
-        
-        # aggregate the updates using simple 'average' or 'weighted_average'
-        if aggregation_mode == "average":
-            average_weights = np.mean(model_weights, axis=0)
-
-        elif aggregation_mode == "weighted_average":
-            # get the weights of model updates based on how much data they have been trained on
-            sample_counts = np.array([t.result()["samples_count"] for t in tasks])
-            edge_weights = get_edge_weights(sample_counts)
-            
-            # find weighted average
-            average_weights = np.average(model_weights, weights=edge_weights, axis=0)
-
-        else:
-            raise Exception(f"Federated mode {aggregation_mode} is not recognized. \
-                 Please select one of the available modes: ['average', 'weighted_average']")
-            
-        # assign the updated weights to the global_model
-        global_model.set_weights(average_weights)
-
-        print(f'Epoch {i}, Trained Federated Model')
-
-        # if the all paramters are supplied, evaluate the model
-        if x_test is not None and y_test is not None and evaluation_function and callable(evaluation_function):
-            loss_eval, accuracy = evaluation_function(global_model, x_test, y_test)
-
-        # if time_interval is supplied, wait for *time_interval* seconds
-        if time_interval > 0:
-            time.sleep(time_interval)
-
-    return global_model
-
-
diff --git a/build/lib/flox/utils.py b/build/lib/flox/utils.py
deleted file mode 100644
index d6aa3f9..0000000
--- a/build/lib/flox/utils.py
+++ /dev/null
@@ -1,78 +0,0 @@
-from tensorflow import keras
-import numpy as np
-
-def get_test_data(keras_dataset='mnist', num_samples=None, preprocess=True, preprocessing_function=None, **kwargs):
-    """
-    Returns (x_test, y_test) of a chosen built-in Keras dataset. 
-    Also preprocesses the image datasets (mnist, fashion_mnist, cifar10, cifar100) by default.
-
-    Parameters
-    ----------
-    keras_dataset: str
-        one of the available Keras datasets: 
-        ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-
-    num_samples: int 
-        randomly samples n data points from (x_test, y_test). Set to None by default.
-
-    preprocess: boolean
-        if True, preprocesses (x_test, y_test) 
-
-    preprocessing_function: function
-        a custom user-provided function that processes (x_test, y_test) and outputs 
-        a tuple (x_test, y_test)
-
-    Returns
-    -------
-    (x_test, y_test): tuple of numpy arrays
-        testing data
-
-    """
-    available_datasets = ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-    dataset_mapping= {
-        'mnist': keras.datasets.mnist,
-        'fashion_mnist': keras.datasets.fashion_mnist,
-        'cifar10': keras.datasets.cifar10,
-        'cifar100': keras.datasets.cifar100,
-        'imdb': keras.datasets.imdb,
-        'reuters': keras.datasets.reuters,
-        'boston_housing': keras.datasets.boston_housing
-    }
-    image_datasets = ['mnist', 'fashion_mnist', 'cifar10', 'cifar100']
-
-    # check if the dataset exists
-    if keras_dataset not in available_datasets:
-        raise Exception(f"Please select one of the built-in Keras datasets: {available_datasets}")
-
-    else:
-        _, (x_test, y_test) = dataset_mapping[keras_dataset].load_data()
-
-        # take a random set of images
-        if num_samples:
-            idx = np.random.choice(np.arange(len(x_test)), num_samples, replace=True)
-            x_test = x_test[idx]
-            y_test = y_test[idx]
-
-        if preprocess:
-            if preprocessing_function and callable(preprocessing_function):
-                (x_test, y_test) = preprocessing_function(x_test, y_test)
-
-            else:
-                # do default image processing for built-in Keras images    
-                if keras_dataset in image_datasets:
-                    # Scale images to the [0, 1] range
-                    x_test = x_test.astype("float32") / 255
-
-                    # Make sure images have shape (num_samples, x, y, 1) if working with MNIST images
-                    if x_test.shape[-1] not in [1, 3]:
-                        x_test = np.expand_dims(x_test, -1)
-
-                    # convert class vectors to binary class matrices
-                    if keras_dataset == 'cifar100':
-                        num_classes=100
-                    else:
-                        num_classes=10
-                        
-                    y_test = keras.utils.to_categorical(y_test, num_classes)
-
-        return (x_test, y_test)
\ No newline at end of file
diff --git a/build/lib/flox_module/flox.py b/build/lib/flox_module/flox.py
deleted file mode 100644
index 7879b1d..0000000
--- a/build/lib/flox_module/flox.py
+++ /dev/null
@@ -1,411 +0,0 @@
-import numpy as np
-from tensorflow import keras
-from tensorflow.keras import layers
-import time
-
-from funcx.sdk.client import FuncXClient
-from funcx.sdk.executor import FuncXExecutor
-
-
-def get_edge_weights(sample_counts):
-    """
-    Takes an array of numbers and returns their fractions of the total number of samples
-    Can be used to find weights for the weighted_average 
-
-    Parameters
-    ----------
-    sample_counts: numpy array of integers
-
-    Returns
-    -------
-    fractions: numpy array
-    
-    """
-    total = sum(sample_counts)
-    fractions = sample_counts/total
-    return fractions
-
-def eval_model(m, x, y, silent=False):
-    """
-    Evaluate the model on a datset
-
-    Parameters
-    ----------
-    m: Tensorflow model
-
-    x: numpy array
-        dataset entries, e.g. x_test
-    
-    y: numpy array
-        labels for the entries, e.g., y_test
-    
-    """
-    score = m.evaluate(x, y, verbose=0)
-    if not silent:
-        print("Test loss:", score[0])
-        print("Test accuracy:", score[1])
-    return score[0], score[1]
-
-def training_function(json_model_config, 
-                      global_model_weights, 
-                      num_samples=None,
-                      epochs=10,     
-                      data_source="keras",
-                      keras_dataset="mnist",
-                      preprocess=True,
-                      path_dir='/home/pi/datasets', 
-                      x_train_name="mnist_x_train.npy", 
-                      y_train_name="mnist_y_train.npy",
-                      input_shape=(32, 28, 28, 1),
-                      loss="categorical_crossentropy",
-                      optimizer="adam",
-                      metrics=["accuracy"]                                               
-):
-    """
-    This function gets deployed to the given endpoints with corresponding parameters. 
-    Returns the updated model weights and number of samples it was trained on.
-
-    Parameters
-    ----------
-    json_model_config: str
-        configuration of the TF model retrieved using model.to_json()
-
-    global_model_weights: numpy array
-        a numpy array with weights of the TF model
-
-    num_samples: int
-        if data_source="keras", randomly samples n data points from (x_train, y_train)
-
-    epochs: int
-        the number of epochs to train the model for
-
-    data_source: str
-        the function supports three data sources: "local", "keras"
-        for "local" and "keras", see get_local_data and get_keras_data functions
-
-    keras_dataset: str
-        specifies one of the default keras datasets to use if using "keras" data source
-        ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-
-    preprocess: boolean
-            if True, will attempt to preprocess your data from "keras" data sources
-
-    path_dir: str
-        path to the folder with x_train and y_train; needed when data_sourse="local"
-
-    x_train_name: str
-        filename for x_train; needed when data_sourse="local"
-
-    y_train_name: str
-        file name for y_train; needed when data_sourse="local"
-
-    input_shape: tupple
-        input shape for the provided model
-    
-    loss: str
-        loss for TF's model.fit() function
-
-    optimizer: str
-        optimizer for for TF's model.fit() function
-
-    metrics: str/list 
-        metrics for TF's model.fit() function. E.g, metrics=["accuracy"],
-
-    Returns
-    -------
-    model_weights: numpy array
-        updated TF model weights after training on the local data
-
-    samples_counts: int
-        the number of samples the model was trained on. 
-        Can be used to find the weighted average by # of samples seen
-
-    """
-    # import all the dependencies required for funcX functions)
-    from tensorflow import keras
-    import numpy as np
-    import os
-
-    # retrieve (and optionally process) the data
-    if data_source == 'keras':
-        available_datasets = ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-        dataset_mapping= {
-            'mnist': keras.datasets.mnist,
-            'fashion_mnist': keras.datasets.fashion_mnist,
-            'cifar10': keras.datasets.cifar10,
-            'cifar100': keras.datasets.cifar100,
-            'imdb': keras.datasets.imdb,
-            'reuters': keras.datasets.reuters,
-            'boston_housing': keras.datasets.boston_housing
-        }
-        image_datasets = ['mnist', 'fashion_mnist', 'cifar10', 'cifar100']
-
-        # check if the Keras dataset exists
-        if keras_dataset not in available_datasets:
-            raise Exception(f"Please select one of the built-in Keras datasets: {available_datasets}")
-
-        else:
-            # load the data
-            (x_train, y_train), _ = dataset_mapping[keras_dataset].load_data()
-
-            # take a random set of images
-            if num_samples:
-                idx = np.random.choice(np.arange(len(x_train)), num_samples, replace=True)
-                x_train = x_train[idx]
-                y_train = y_train[idx]
-
-            if preprocess:
-                # do default image processing for built-in Keras images    
-                if keras_dataset in image_datasets:
-                    # Scale images to the [0, 1] range
-                    x_train = x_train.astype("float32") / 255
-
-                    # Make sure images have shape (num_samples, x, y, 1) if working with MNIST images
-                    if x_train.shape[-1] not in [1, 3]:
-                        x_train = np.expand_dims(x_train, -1)
-
-                    # convert class vectors to binary class matrices
-                    if keras_dataset == 'cifar100':
-                        num_classes=100
-                    else:
-                        num_classes=10
-                        
-                    y_train = keras.utils.to_categorical(y_train, num_classes)
-
-    elif data_source == 'local':        
-        # construct the path
-        x_train_path_file = os.sep.join([path_dir, x_train_name])
-        y_train_path_file = os.sep.join([path_dir, y_train_name])
-
-        # load the files
-        with open(x_train_path_file, 'rb') as f:
-            x_train = np.load(f)
-            
-        with open(y_train_path_file, 'rb') as f:
-            y_train = np.load(f)
-
-        # if preprocess is True & the function is valid, preprocess the data
-        if preprocess:
-            # check if a valid function was given
-            depth = input_shape[3]
-            image_size_y = input_shape[2]
-            image_size_x = input_shape[1]
-
-            # take a limited number of samples, if indicated
-            if num_samples:
-                idx = np.random.choice(np.arange(len(x_train)), num_samples, replace=True)
-                x_train = x_train[idx]
-                y_train = y_train[idx]
-            
-            # reshape and scale to pixel values to 0-1
-            x_train = x_train.reshape(len(x_train), image_size_x, image_size_y, depth)
-            x_train = x_train / 255.0
-
-    else:
-        raise Exception("Please choose one of data sources: ['local', 'keras']")
-
-    # create the model
-    model = keras.models.model_from_json(json_model_config)
-
-    # compile the model and set weights to the global model
-    model.compile(loss=loss, optimizer=optimizer, metrics=metrics)
-
-    try:
-        model.set_weights(global_model_weights)
-    except:
-        # some older TF versions require the model to be 'built' first
-        model.build(input_shape=input_shape)
-        model.set_weights(global_model_weights)
-
-    # train the model on the local data and extract the weights
-    model.fit(x_train, y_train, epochs=epochs)
-    model_weights = model.get_weights()
-
-    # transform to a numpy array
-    np_model_weights = np.asarray(model_weights, dtype=object)
-
-    # return the updated weights and number of samples the model was trained on
-    return {"model_weights":np_model_weights, "samples_count": x_train.shape[0]}
-
-def federated_learning(global_model, 
-                      endpoint_ids, 
-                      num_samples=100,
-                      epochs=5,
-                      loops=1,
-                      time_interval=0,
-                      aggregation_mode="weighted_average",
-                      data_source: str = "keras",
-                      keras_dataset = "mnist",  
-                      preprocess=False,
-                      path_dir='/home/pi/datasets', 
-                      x_train_name="mnist_x_train.npy", 
-                      y_train_name="mnist_y_train.npy",
-                      input_shape=(32, 28, 28, 1),
-                      loss="categorical_crossentropy",
-                      optimizer="adam", 
-                      metrics=["accuracy"],
-                      evaluation_function=eval_model,
-                      x_test=None,
-                      y_test=None
-                      ):
-    """
-    Facilitates Federated Learning for *loops* rounds. 
-
-    Parameters
-    ----------
-    global_model: TF model object
-        compiled TF model that will be deployed for training on the endpoints
-
-    endpoint_ids: list of str
-        a list with endpoint_ids to include in the FL process. 
-        The ids can be retrieved by running 'funcx-endpoint list' on participating devices
-
-    num_samples: int or list
-        indicates how many samples to get for training on endpoints
-        if int, applies the same num_samples to all endpoints. 
-        if list, it will use the corresponding number of samples for each device
-        the list should have the same number of entries as the number of endpoints
-
-
-    epochs: int or list
-        indicates how many epochs to use for training on endpoints
-        if int, applies the same number of epochs to all endpoints. 
-        if list, it will use the corresponding number of epochs for each device
-        the list should have the same number of entries as the number of endpoints
-
-    loops: int
-        defines how many FL rounds to run. Each round consists of deploying the mode, training
-        aggregating the updates, and reassigning new weights to the model. 
-
-    time_interval: int
-        defines the pause between FL rounds in seconds (default=0). Can be useful if you want to run 
-        a round every minute/hour/day.
-
-    aggregation_mode: str
-        defines the aggregation mode from a choice of a simple 'average' or a 'weighted_average'
-        the 'weighted_average' aggregates the updates based on how many samples each updates has 
-        been trained on. This gives more weight to updates from devices that have seen more data.
-
-    data_source: str
-        the function supports three data sources: "local", "keras"
-        for "local" and "keras", see get_local_data and get_keras_data functions
-
-    keras_dataset: str
-        specifies one of the default keras datasets to use if using "keras" data source
-        ['mnist', 'fashion_mnist', 'cifar10', 'cifar100', 'imdb', 'reuters', 'boston_housing']
-
-    preprocess: boolean
-            if True, will attempt to preprocess your data from "keras" data sources
-
-    path_dir: str
-        path to the folder with x_train and y_train; needed when data_sourse="local"
-
-    x_train_name: str
-        filename for x_train; needed when data_sourse="local"
-
-    y_train_name: str
-        file name for y_train; needed when data_sourse="local"
-
-    input_shape: tupple
-        input shape for the provided model
-    
-    loss: str
-        loss for TF's model.fit() function
-
-    optimizer: str
-        optimizer for for TF's model.fit() function
-
-    metrics: str/list 
-        metrics for TF's model.fit() function. E.g, metrics=["accuracy"],
-
-    evaluation_function: function
-        if supplied, evaluates the model on x_test and y_test
-
-    x_test: list/numpy array/tensors
-        x_test data for testing
-
-    y_test: list
-        y_test labels for x_test
-
-
-    Returns
-    -------
-    global_model: TF model
-        the original model but with updated weights after all the FL rounds
-
-    """
-    # instantiate the FuncXExecutor
-    fx = FuncXExecutor(FuncXClient())
-
-    # if num_samples or epochs is an int, convert to list so the same number can be applied to all endpoints
-    if type(num_samples) == int:
-        num_samples = [num_samples]*len(endpoint_ids)
-
-    if type(epochs) == int:
-        epochs = [epochs]*len(endpoint_ids)
-    
-    # start running FL loops
-    for i in range(loops):
-
-        # get the model's architecture and weights
-        json_config = global_model.to_json()
-        gm_weights = global_model.get_weights()
-        gm_weights_np = np.asarray(gm_weights, dtype=object)
-
-        #fx = FuncXExecutor(FuncXClient())
-        tasks = []
-
-        # submit the corresponding parameters to each endpoint for a round of FL 
-        for e, num_s, num_epoch, path_d in zip(endpoint_ids, num_samples, epochs, path_dir): 
-            tasks.append(fx.submit(training_function, 
-                                   json_model_config=json_config, 
-                                    global_model_weights=gm_weights_np, 
-                                    num_samples=num_s,
-                                    epochs=num_epoch,
-                                    data_source=data_source,
-                                    keras_dataset=keras_dataset,
-                                    preprocess=preprocess,
-                                    path_dir=path_d,
-                                    x_train_name=x_train_name,
-                                    y_train_name=y_train_name,
-                                    input_shape=input_shape,
-                                    loss=loss,
-                                    optimizer=optimizer,
-                                    metrics=metrics,
-                                    endpoint_id=e))
-        
-        # extract model updates from each endpoints once they are available
-        model_weights = [t.result()["model_weights"] for t in tasks]
-        
-        # aggregate the updates using simple 'average' or 'weighted_average'
-        if aggregation_mode == "average":
-            average_weights = np.mean(model_weights, axis=0)
-
-        elif aggregation_mode == "weighted_average":
-            # get the weights of model updates based on how much data they have been trained on
-            sample_counts = np.array([t.result()["samples_count"] for t in tasks])
-            edge_weights = get_edge_weights(sample_counts)
-            
-            # find weighted average
-            average_weights = np.average(model_weights, weights=edge_weights, axis=0)
-
-        else:
-            raise Exception(f"Federated mode {aggregation_mode} is not recognized. \
-                 Please select one of the available modes: ['average', 'weighted_average']")
-            
-        # assign the updated weights to the global_model
-        global_model.set_weights(average_weights)
-
-        print(f'Epoch {i}, Trained Federated Model')
-
-        # if the all paramters are supplied, evaluate the model
-        if x_test is not None and y_test is not None and evaluation_function and callable(evaluation_function):
-            loss_eval, accuracy = evaluation_function(global_model, x_test, y_test)
-
-        # if time_interval is supplied, wait for *time_interval* seconds
-        if time_interval > 0:
-            time.sleep(time_interval)
-
-    return global_model
-
-
diff --git a/dist/pyflox-0.1.5.tar.gz b/dist/pyflox-0.1.5.tar.gz
deleted file mode 100644
index f07d061..0000000
Binary files a/dist/pyflox-0.1.5.tar.gz and /dev/null differ
diff --git a/docs/_static/pydata-custom.css b/docs/_static/pydata-custom.css
new file mode 100644
index 0000000..b7f903a
--- /dev/null
+++ b/docs/_static/pydata-custom.css
@@ -0,0 +1,10 @@
+/*Tweaks to the Pydata default CSS */
+
+/*No yellow background highlight when targeted by summary tables */
+/*dt:target { background-color: #f8f8f8; border: 1px solid black, }*/
+dt:target { background: transparent;}
+/*More space between H1s and signatures in API reference*/
+h1 { margin-bottom: 40px; }
+
+/*No line underneath summary table headings (clashes with line above first member)*/
+p.rubric { border-bottom: 0px; }
\ No newline at end of file
diff --git a/docs/_static/readthedocs-custom.css b/docs/_static/readthedocs-custom.css
new file mode 100644
index 0000000..414380d
--- /dev/null
+++ b/docs/_static/readthedocs-custom.css
@@ -0,0 +1,29 @@
+/* Override nav bar color */
+/*.wy-side-nav-search {
+    background-color: #fbfbb6;
+}
+.wy-side-nav-search > a {
+    color: #b2355c
+}*/
+
+/* Override text bar color */
+/*.caption-text {
+    color: #b2355c;
+}*/
+
+/* Override code signature colour */
+/*.rst-content dl:not(.docutils) dt {
+    background: #fbfbb6;
+    color: #b2355c;
+    border-top: solid 3px #b2355c;
+}*/
+
+/* Override hyperlink colour */
+/* a {
+    color: #b2355c;
+}*/
+
+/* Make content width wider*/
+.wy-nav-content {
+    max-width: 60% !important;
+}
\ No newline at end of file
diff --git a/docs/_templates/custom-class-template.rst b/docs/_templates/custom-class-template.rst
new file mode 100644
index 0000000..d64b80d
--- /dev/null
+++ b/docs/_templates/custom-class-template.rst
@@ -0,0 +1,34 @@
+{{ fullname | escape | underline}}
+
+.. currentmodule:: {{ module }}
+
+.. autoclass:: {{ objname }}
+   :members:
+   :show-inheritance:
+   :inherited-members:
+   :special-members: __call__, __add__, __mul__
+
+   {% block methods %}
+   {% if methods %}
+   .. rubric:: {{ _('Methods') }}
+
+   .. autosummary::
+      :nosignatures:
+   {% for item in methods %}
+      {%- if not item.startswith('_') %}
+      ~{{ name }}.{{ item }}
+      {%- endif -%}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
+   {% block attributes %}
+   {% if attributes %}
+   .. rubric:: {{ _('Attributes') }}
+
+   .. autosummary::
+   {% for item in attributes %}
+      ~{{ name }}.{{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
\ No newline at end of file
diff --git a/docs/_templates/custom-module-template.rst b/docs/_templates/custom-module-template.rst
new file mode 100644
index 0000000..dd90e32
--- /dev/null
+++ b/docs/_templates/custom-module-template.rst
@@ -0,0 +1,66 @@
+{{ fullname | escape | underline}}
+
+.. automodule:: {{ fullname }}
+
+   {% block attributes %}
+   {% if attributes %}
+   .. rubric:: Module attributes
+
+   .. autosummary::
+      :toctree:
+   {% for item in attributes %}
+      {{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
+   {% block functions %}
+   {% if functions %}
+   .. rubric:: {{ _('Functions') }}
+
+   .. autosummary::
+      :toctree:
+      :nosignatures:
+   {% for item in functions %}
+      {{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
+   {% block classes %}
+   {% if classes %}
+   .. rubric:: {{ _('Classes') }}
+
+   .. autosummary::
+      :toctree:
+      :template: custom-class-template.rst
+      :nosignatures:
+   {% for item in classes %}
+      {{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
+   {% block exceptions %}
+   {% if exceptions %}
+   .. rubric:: {{ _('Exceptions') }}
+
+   .. autosummary::
+      :toctree:
+   {% for item in exceptions %}
+      {{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
+{% block modules %}
+{% if modules %}
+.. autosummary::
+   :toctree:
+   :template: custom-module-template.rst
+   :recursive:
+{% for item in modules %}
+   {{ item }}
+{%- endfor %}
+{% endif %}
+{% endblock %}
\ No newline at end of file
diff --git a/docs/api.rst b/docs/api.rst
new file mode 100644
index 0000000..c5b9975
--- /dev/null
+++ b/docs/api.rst
@@ -0,0 +1,6 @@
+.. autosummary::
+   :toctree: _autosummary
+   :template: custom-module-template.rst
+   :recursive:
+
+   flox
\ No newline at end of file
diff --git a/docs/conf.py b/docs/conf.py
new file mode 100644
index 0000000..29d5c55
--- /dev/null
+++ b/docs/conf.py
@@ -0,0 +1,100 @@
+# Configuration file for the Sphinx documentation builder.
+
+# -- Project information
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(".."))  # Source code dir relative to this file
+
+project = "FLoX"
+copyright = "2022, Globus Labs"
+author = "Globus Labs"
+
+release = "0.1"
+version = "0.1.0"
+
+# -- General configuration ---------------------------------------------------
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    "sphinx.ext.autodoc",  # Core Sphinx library for auto html doc generation from docstrings
+    "sphinx.ext.autosummary",  # Create neat summary tables for modules/classes/methods etc
+    "sphinx.ext.intersphinx",  # Link to other project's documentation (see mapping below)
+    "sphinx.ext.viewcode",  # Add a link to the Python source code for classes, functions etc.
+    "sphinx_autodoc_typehints",  # Automatically document param types (less noise in class signature)
+    "nbsphinx",  # Integrate Jupyter Notebooks and Sphinx
+    "IPython.sphinxext.ipython_console_highlighting",
+]
+
+# Mappings for sphinx.ext.intersphinx. Projects have to have Sphinx-generated doc! (.inv file)
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3/", None),
+}
+
+autosummary_generate = True  # Turn on sphinx.ext.autosummary
+autoclass_content = "both"  # Add __init__ doc (ie. params) to class summaries
+html_show_sourcelink = (
+    False  # Remove 'view source code' from top of page (for html, not python)
+)
+autodoc_inherit_docstrings = True  # If no docstring, inherit from base class
+set_type_checking_flag = True  # Enable 'expensive' imports for sphinx_autodoc_typehints
+nbsphinx_allow_errors = True  # Continue through Jupyter errors
+# autodoc_typehints = "description" # Sphinx-native method. Not as good as sphinx_autodoc_typehints
+add_module_names = False  # Remove namespaces from class/method signatures
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# Exclusions
+# To exclude a module, use autodoc_mock_imports. Note this may increase build time, a lot.
+# (Also, when installing on readthedocs.org, we omit installing Tensorflow and
+# Tensorflow Probability so mock them here instead.)
+# autodoc_mock_imports = [
+# 'tensorflow',
+# 'tensorflow_probability',
+# ]
+# To exclude a class, function, method or attribute, use autodoc-skip-member. (Note this can also
+# be used in reverse, ie. to re-include a particular member that has been excluded.)
+# 'Private' and 'special' members (_ and __) are excluded using the Jinja2 templates; from the main
+# doc by the absence of specific autoclass directives (ie. :private-members:), and from summary
+# tables by explicit 'if-not' statements. Re-inclusion is effective for the main doc though not for
+# the summary tables.
+# def autodoc_skip_member_callback(app, what, name, obj, skip, options):
+#     # This would exclude the Matern12 class and to_default_float function:
+#     exclusions = ('Matern12', 'to_default_float')
+#     # This would re-include __call__ methods in main doc, previously excluded by templates:
+#     inclusions = ('__call__')
+#     if name in exclusions:
+#         return True
+#     elif name in inclusions:
+#         return False
+#     else:
+#         return skip
+# def setup(app):
+#     # Entry point to autodoc-skip-member
+#     app.connect("autodoc-skip-member", autodoc_skip_member_callback)
+
+# -- Options for HTML output -------------------------------------------------
+
+# Readthedocs theme
+# on_rtd is whether on readthedocs.org, this line of code grabbed from docs.readthedocs.org...
+on_rtd = os.environ.get("READTHEDOCS", None) == "True"
+if not on_rtd:  # only import and set the theme if we're building docs locally
+    import sphinx_rtd_theme
+
+    html_theme = "sphinx_rtd_theme"
+    html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
+html_css_files = ["readthedocs-custom.css"]  # Override some CSS settings
+
+# Pydata theme
+# html_theme = "pydata_sphinx_theme"
+# html_logo = "_static/logo-company.png"
+# html_theme_options = { "show_prev_next": False}
+# html_css_files = ['pydata-custom.css']
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ["_static"]
diff --git a/docs/source/index.rst b/docs/index.rst
similarity index 90%
rename from docs/source/index.rst
rename to docs/index.rst
index 349ac52..3346d8e 100644
--- a/docs/source/index.rst
+++ b/docs/index.rst
@@ -12,6 +12,7 @@ Contents
 
 .. toctree::
 
-   install
+   Homepage <self>
    usage
-   API reference <_autosummary/flox>
\ No newline at end of file
+   install
+   logic
\ No newline at end of file
diff --git a/docs/install.rst b/docs/install.rst
new file mode 100644
index 0000000..8c26db6
--- /dev/null
+++ b/docs/install.rst
@@ -0,0 +1,28 @@
+.. _installation:
+
+Installation
+------------
+*Controller* is a device from which Federated Learning is facilitated, such as your laptop.
+*Endpoints* are devices that participate in Federated Learning, such as Raspberry Pis.
+
+Use the package manager `pip <https://pip.pypa.io/en/stable/>`_ to install flox
+on the *Controller*:
+
+.. code-block:: console
+
+   (.venv) $ pip install pyflox
+
+To be able to include your *endpoints* into the Federated Learning process,
+you need to install `funcx-endpoint <https://funcx.readthedocs.io/en/latest/endpoints.html>`_ on the *endpoints*:
+
+.. code-block:: console
+
+   (.endpoint_venv) $ python3 -m pipx install funcx_endpoint
+
+Finally, install compatible versions of `PyTorch <https://qengineering.eu/install-pytorch-on-raspberry-pi-4.html) & [Torchvision](https://qengineering.eu/install-pytorch-on-raspberry-pi-4.html#:~:text=rm%20%2Drf%20~/pytorch-,TorchVision.,-Install%20torchvision%20on>`_
+or `Tensorflow <https://qengineering.eu/install-tensorflow-2.1.0-on-raspberry-pi-4.html>`_ on both the Controller and endpoints.
+As of now, FLoX supports Tensorflow and PyTorch, although you can add support for your frameworks by creating a new
+ModelTrainer. See ``flox.logic`` and examples of ModelTrainers in ``flox.model_trainers``.
+
+*Note*: ``funcX-endpoint`` is only supported on Linux.
+FLoX Controller functionality is supported on MacOS, Linux and Windows.
\ No newline at end of file
diff --git a/docs/logic.rst b/docs/logic.rst
new file mode 100644
index 0000000..7cd444b
--- /dev/null
+++ b/docs/logic.rst
@@ -0,0 +1,145 @@
+.. _logic:
+
+FLoX Client-Controller Logic
+============================
+
+FLoX is designed to be a highly modular and customizable framework for
+serverless, FL processes. It is built on top of a general 10-step
+abstraction for FL processes grouped into controller-side and
+client-side logical steps. FLoX considers two abstract classes,
+``FloxControllerLogic`` and ``FloxClientLogic``, for
+wrapping/implementing the logic needed for the controller-side and
+client-side, respectively. Each step is run on either the Controller (S)
+or the clients (C). All 7 steps are listed below with the respective
+function that corresponds with them:
+
+1.  **(S)** Model Initialization: ``on_model_init``
+2.  **(S)** Model Sharing: ``on_model_broadcast``
+3.  **(C)** Receiving Model on Client: ``on_model_receive``
+4.  **(C)** Data Fetching: ``on_data_retrieve``
+5.  **(C)** Local Model Training: ``on_model_fit``
+6.  **(C)** Model Parameter Submission: ``on_model_send``
+7.  **(S)** Receiving Model on Controller: ``on_model_receive``
+8.  **(S)** Model Aggregation: ``on_model_aggregate``
+9.  **(S)** Model Updating: ``on_model_update``
+10. **(S)** Model Evaluation: ``on_model_evaluate``
+
+Below is a (``mermaid.js``) figure showing the sequence of the logic
+steps. Then, we further describe each of the logic steps.
+
+.. code:: mermaid
+
+   flowchart
+       id1[[on_model_init]]
+       id2[/on_model_broadcast\]
+
+       id3[\on_model_receive/]
+       id4[(on_data_retrieve)]
+       id5{on_model_fit}
+       id6[on_model_send]
+
+       id7[on_model_receive]
+       id8[on_model_aggregate]
+       id9[on_model_update]
+       id10[on_model_evaluate]
+
+
+
+       subgraph controller
+           id1-->id2
+           id7-->id8
+           id8-->id9
+           id9-->id10
+           id10-->id2
+       end
+
+       subgraph client
+           direction TB
+           id3-->id4
+           id4-->id5
+           id5-->id6
+       end
+
+       id2-->|model params|id3
+       id6-->|model params|id7
+
+1. Model Initialization
+-----------------------
+``Controller.on_model_init()`` is where one would provide initial
+setup scripts that need to be run only once rather than needing to run
+every FL round. In ``flox.controllers.TensorflowController`` we use
+``.on_model_init()`` to set up variables that were not provided by the
+user but will be reused in all of the FL rounds going forward.
+
+2. Model Sharing
+----------------
+Model sharing happend in ``Controller.on_model_broadcast`` where one
+would assemble all variables into a config file, encrypt data is
+necessary, and then send the tasks to the *clients* using an Executor,
+such as FuncXExecutor. The method should return a list of futures/tasks
+that can later be parsed out by ``Controller.on_model_receive()`` once
+*clients* return the results.
+
+3. Receiving Model on Client
+----------------------------
+Once *clients* receive the model and the config with necessary
+parameters, ``Client.on_model_receive()`` is responsible for the initial
+actions such as decrypting the data if it wsa encrypted.
+
+4. Data Retrieval
+-----------------
+``Client.on_data_retrieve()`` is where *clients* retrieve and prepare
+their data for training.
+
+5. Local Model Training
+-----------------------
+``Client.on_data_fit()`` is where the training process is defined and
+executed.
+
+6. Model Parameter Submission
+-----------------------------
+When *clients* retrieved local data and made updates to the global
+model, the new model weights are returned to the *Controller* and
+``Client.on_data_send()`` is for things like encryption of data before
+it is sent back.
+
+7. Receiving Model on Controller
+--------------------------------
+Once the *Controller* receives the results back from the *clients*,
+``Controller.on_model_receive()`` parses the results and decrypts if
+necessary.
+
+8. Model Aggregation
+--------------------
+``Controller.on_model_aggregate()`` takes the parsed results from
+``Controller.on_model_receive()`` and aggregates weights from the
+endpoints.
+
+9. Model Updating
+-----------------
+``Controller.on_model_update()`` simply takes the new weights from
+``Controller.on_model_aggregate()`` and assigns them to the global
+model.
+
+10. Model Evaluation
+--------------------
+Finally, ``Controller.on_model_evaluate()`` evaluates the model using
+a user-provided testing dataset, reports the results, and then the
+entire loop from Step 2 to Step 10 is repeated for as many rounds as was
+specified by the user.
+
+We implemented the abstract base classes in
+``flox.logic.base_client.py`` and ``flox.logic.base_controller.py``. We
+also implemented a base class for Machine Learning Model Trainers,
+located in ``flox.logic.base_model_trainer.py``. We are providing
+practical examples on top of these classes to illustrate how all of
+these steps come together:
+
+- ``flox.examples.quickstart_pytorch`` makes
+use of ``PyTorchController``, ``PyTorchClient``, and ``PyTorchTrainer``
+to run a Federated Learning workflow on PyTorch.
+
+- ``flox.examples.quickstart_tensorflow`` makes use of
+``TensorflowController``, ``TensorflowClient``, and
+``TensorflowTrainer`` to run a Federated Learning workflow on
+Tensorflow.
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 0000000..173cee6
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,10 @@
+sphinx
+sphinx-rtd-theme
+pydata-sphinx-theme
+sphinx-autodoc-typehints
+nbsphinx
+ipython
+jupytext
+jupyter
+matplotlib
+numpy
\ No newline at end of file
diff --git a/docs/source/conf.py b/docs/source/conf.py
deleted file mode 100644
index 9f16558..0000000
--- a/docs/source/conf.py
+++ /dev/null
@@ -1,36 +0,0 @@
-# Configuration file for the Sphinx documentation builder.
-
-# -- Project information
-
-project = "FLoX"
-copyright = "2022, Globus Labs"
-author = "Globus Labs"
-
-release = "0.1"
-version = "0.1.0"
-
-# -- General configuration
-
-extensions = [
-    "sphinx.ext.duration",
-    "sphinx.ext.doctest",
-    "sphinx.ext.autodoc",
-    "sphinx.ext.autosummary",
-    "sphinx.ext.intersphinx",
-    "sphinx.ext.autosectionlabel",
-]
-
-intersphinx_mapping = {
-    "python": ("https://docs.python.org/3/", None),
-    "sphinx": ("https://www.sphinx-doc.org/en/master/", None),
-}
-intersphinx_disabled_domains = ["std"]
-
-templates_path = ["_templates"]
-
-# -- Options for HTML output
-
-html_theme = "sphinx_rtd_theme"
-
-# -- Options for EPUB output
-epub_show_urls = "footnote"
diff --git a/docs/source/install.rst b/docs/source/install.rst
deleted file mode 100644
index 80ebf8f..0000000
--- a/docs/source/install.rst
+++ /dev/null
@@ -1,10 +0,0 @@
-.. _installation:
-
-Installation
-------------
-
-To use Flox, first install it using pip:
-
-.. code-block:: console
-
-   (.venv) $ pip install pyflox
\ No newline at end of file
diff --git a/docs/source/usage.rst b/docs/usage.rst
similarity index 100%
rename from docs/source/usage.rst
rename to docs/usage.rst
diff --git a/flox/clients/PyTorchClient.py b/flox/clients/PyTorchClient.py
new file mode 100644
index 0000000..72b7dae
--- /dev/null
+++ b/flox/clients/PyTorchClient.py
@@ -0,0 +1,71 @@
+from flox.logic import FloxClientLogic
+
+
+class PyTorchClient(FloxClientLogic):
+    def on_model_receive():
+        pass
+
+    def on_data_retrieve(self, config):
+        import torch
+        import torchvision
+        import torchvision.transforms as transforms
+
+        transform = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+            ]
+        )
+
+        batch_size = 32 if "batch_size" not in config.keys() else config["batch_size"]
+        dataset_name = config["dataset_name"]
+        num_workers = 8 if "num_workers" not in config.keys() else config["num_workers"]
+        root = "./data" if "data_root" not in config.keys() else config["data_root"]
+
+        # create train DataLoader
+        trainset = dataset_name(
+            root=root, train=True, download=True, transform=transform
+        )
+
+        train_split_len = (
+            len(trainset)
+            if "num_samples" not in config.keys()
+            else config["num_samples"]
+        )
+
+        train_subpart = torch.utils.data.random_split(
+            trainset, [train_split_len, len(trainset) - train_split_len]
+        )[0]
+        trainloader = torch.utils.data.DataLoader(
+            train_subpart, batch_size=batch_size, shuffle=True, num_workers=num_workers
+        )
+
+        # create train DataLoader
+        testset = dataset_name(
+            root=root, train=False, download=True, transform=transform
+        )
+        test_split_len = (
+            len(trainset)
+            if "num_samples" not in config.keys()
+            else config["num_samples"]
+        )
+
+        test_subpart = torch.utils.data.random_split(
+            testset, [test_split_len, len(testset) - test_split_len]
+        )[0]
+        testloader = torch.utils.data.DataLoader(
+            test_subpart, batch_size=batch_size, shuffle=False, num_workers=num_workers
+        )
+
+        return trainloader, testloader
+
+    def on_model_fit(self, ModelTrainer, trainloader, config):
+        model_weights = ModelTrainer.fit(trainloader, config)
+
+        return model_weights
+
+    def run_round(self, config, ModelTrainer):
+        trainloader, testloader = self.on_data_retrieve(config)
+        fit_results = self.on_model_fit(ModelTrainer, trainloader, config)
+
+        return {"model_weights": fit_results, "samples_count": config["num_samples"]}
diff --git a/flox/clients/TensorflowClient.py b/flox/clients/TensorflowClient.py
new file mode 100644
index 0000000..fffb0f7
--- /dev/null
+++ b/flox/clients/TensorflowClient.py
@@ -0,0 +1,138 @@
+from flox.logic import FloxClientLogic
+
+
+class TensorflowClient(FloxClientLogic):
+    def on_model_receive(self):
+        pass
+
+    def on_data_retrieve(self, config):
+        import os
+
+        import numpy as np
+        from tensorflow import keras
+
+        if config["data_source"] == "keras":
+            available_datasets = [
+                "mnist",
+                "fashion_mnist",
+                "cifar10",
+                "cifar100",
+                "imdb",
+                "reuters",
+                "boston_housing",
+            ]
+            dataset_mapping = {
+                "mnist": keras.datasets.mnist,
+                "fashion_mnist": keras.datasets.fashion_mnist,
+                "cifar10": keras.datasets.cifar10,
+                "cifar100": keras.datasets.cifar100,
+                "imdb": keras.datasets.imdb,
+                "reuters": keras.datasets.reuters,
+                "boston_housing": keras.datasets.boston_housing,
+            }
+            image_datasets = ["mnist", "fashion_mnist", "cifar10", "cifar100"]
+
+            # check if the Keras dataset exists
+            if config["dataset_name"] not in available_datasets:
+                raise Exception(
+                    f"Please select one of the built-in Keras datasets: {available_datasets}"
+                )
+
+            else:
+                # load the data
+                (x_train, y_train), _ = dataset_mapping[
+                    config["dataset_name"]
+                ].load_data()
+
+                # take a random set of images
+                if config["num_samples"]:
+                    idx = np.random.choice(
+                        np.arange(len(x_train)), config["num_samples"], replace=True
+                    )
+                    x_train = x_train[idx]
+                    y_train = y_train[idx]
+
+                if config["preprocess"]:
+                    # do default image processing for built-in Keras images
+                    if config["dataset_name"] in image_datasets:
+                        # Scale images to the [0, 1] range
+                        x_train = x_train.astype("float32") / 255
+
+                        # Make sure images have shape (num_samples, x, y, 1) if working with MNIST images
+                        if x_train.shape[-1] not in [1, 3]:
+                            x_train = np.expand_dims(x_train, -1)
+
+                        # convert class vectors to binary class matrices
+                        if config["dataset_name"] == "cifar100":
+                            num_classes = 100
+                        else:
+                            num_classes = 10
+
+                        y_train = keras.utils.to_categorical(y_train, num_classes)
+
+        elif config["data_source"] == "local":
+            # construct the path
+            x_train_path_file = os.sep.join(
+                [config["path_dir"], config["x_train_filename"]]
+            )
+            y_train_path_file = os.sep.join(
+                [config["path_dir"], config["y_train_filename"]]
+            )
+
+            # load the files
+            with open(x_train_path_file, "rb") as f:
+                x_train = np.load(f)
+
+            with open(y_train_path_file, "rb") as f:
+                y_train = np.load(f)
+
+            # if preprocess is True & the function is valid, preprocess the data
+            if config["preprocess"]:
+                # check if a valid function was given
+                depth = config["input_shape"][3]
+                image_size_y = config["input_shape"][2]
+                image_size_x = config["input_shape"][1]
+
+                # take a limited number of samples, if indicated
+                if config["num_samples"]:
+                    idx = np.random.choice(
+                        np.arange(len(x_train)), config["num_samples"], replace=True
+                    )
+                    x_train = x_train[idx]
+                    y_train = y_train[idx]
+
+                # reshape and scale to pixel values to 0-1
+                x_train = x_train.reshape(
+                    len(x_train), image_size_x, image_size_y, depth
+                )
+                x_train = x_train / 255.0
+
+        else:
+            raise Exception("Please choose one of data sources: ['local', 'keras']")
+
+        return x_train, y_train
+
+    def on_model_fit(self, ModelTrainer, config, x_train, y_train):
+        import numpy as np
+
+        model = ModelTrainer.create_model(config["architecture"])
+        ModelTrainer.compile_model(model)
+        ModelTrainer.set_weights(model, config["weights"])
+
+        ModelTrainer.fit(model, x_train, y_train, epochs=config["epochs"])
+        model_weights = ModelTrainer.get_weights(model)
+
+        # transform to a numpy array
+        np_model_weights = np.asarray(model_weights, dtype=object)
+
+        return np_model_weights
+
+    def on_model_send(self):
+        pass
+
+    def run_round(self, config, ModelTrainer):
+        x_train, y_train = self.on_data_retrieve(config)
+
+        fit_results = self.on_model_fit(ModelTrainer, config, x_train, y_train)
+
+        return {"model_weights": fit_results, "samples_count": x_train.shape[0]}
diff --git a/build/lib/flox/__init__.py b/flox/clients/__init__.py
similarity index 100%
rename from build/lib/flox/__init__.py
rename to flox/clients/__init__.py
diff --git a/flox/common/__init__.py b/flox/common/__init__.py
new file mode 100644
index 0000000..2e662fa
--- /dev/null
+++ b/flox/common/__init__.py
@@ -0,0 +1,11 @@
+from .typing import (
+    ConfigFile,
+    EvaluateRes,
+    FitIns,
+    FuturesList,
+    NDArray,
+    NDArrays,
+    ResultsList,
+    UpdateIns,
+    XYData,
+)
diff --git a/flox/common/logging_config.py b/flox/common/logging_config.py
new file mode 100644
index 0000000..e4fcc4f
--- /dev/null
+++ b/flox/common/logging_config.py
@@ -0,0 +1,233 @@
+"""
+This module contains logging configuration for flox.
+"""
+from __future__ import annotations
+
+import logging
+import logging.config
+import logging.handlers
+import os
+import pathlib
+import re
+import sys
+
+log = logging.getLogger(__name__)
+
+DEFAULT_FORMAT = (
+    "%(created)f %(asctime)s %(levelname)s %(processName)s-%(process)d "
+    "%(threadName)s-%(thread)d %(name)s:%(lineno)d %(funcName)s "
+    "%(message)s"
+)
+
+_ital = "\033[3m"
+_redb = "\033[41m"
+_teal = "\033[32m"
+_yel = "\033[33m"
+_byel = "\033[93m"
+_yelb = "\033[43m"
+_purp = "\033[35m"
+_cyan = "\033[36m"
+_gray = "\033[37m"
+_grayonb = "\033[37;40m"
+_r = "\033[m"
+_C_BASE = (
+    f"{_teal}%(created)f{_r} {_yel}%(asctime)s{_r} {_ital}%(levelname)s{_r}"
+    " %(processName)s-%(process)d %(threadName)s-%(thread)d"
+    f" %(name)s:{_cyan}%(lineno)d{_r} {_purp}%(funcName)s{_r}"
+)
+COLOR_ERROR = _redb
+COLOR_WARNING = _yelb
+COLOR_INFO = _gray
+COLOR_DEBUG = _grayonb
+C_ERROR_FMT = _C_BASE + f" {COLOR_ERROR}%(message)s{_r}"
+C_WARNING_FMT = _C_BASE + f" {COLOR_WARNING}%(message)s{_r}"
+C_INFO_FMT = _C_BASE + f" {COLOR_INFO}%(message)s{_r}"
+C_DEBUG_FMT = _C_BASE + f" {COLOR_DEBUG}%(message)s{_r}"
+
+
+class FloxConsoleFormatter(logging.Formatter):
+    """
+    For internal use only.
+    This formatter handles output to standard streams in the following way:
+    if 'debug' is False (default):
+        info messages and below are treated as "user output" and are minimally decorated
+        warning messages and up are given a full format, so are debug messages
+    if 'debug' is True:
+        all messages are given the full format
+    """
+
+    _u = "[0-9A-Fa-f]"  # convenience
+    _uuid_re = f"{_u}{{8}}-{_u}{{4}}-{_u}{{4}}-{_u}{{4}}-{_u}{{12}}"
+    # match uuids for colorization that have not otherwise already been colorized
+    uuid_re = re.compile(rf"(?<!\dm)({_uuid_re})")
+
+    def __init__(
+        self,
+        debug: bool = False,
+        no_color: bool = False,
+        fmt: str = "",
+        datefmt: str = "%Y-%m-%d %H:%M:%S",
+    ) -> None:
+        super().__init__()
+
+        self.use_color = debug and not no_color and sys.stderr.isatty()
+
+        if fmt:
+            d_fmt, i_fmt, w_fmt, e_fmt = fmt, fmt, fmt, fmt
+        else:
+            d_fmt, i_fmt, w_fmt, e_fmt = (
+                C_DEBUG_FMT,
+                C_INFO_FMT,
+                C_WARNING_FMT,
+                C_ERROR_FMT,
+            )
+
+        if not self.use_color:
+            ansi_re = re.compile("\033.*?m")
+            d_fmt = ansi_re.sub("", d_fmt)
+            i_fmt = ansi_re.sub("", i_fmt)
+            w_fmt = ansi_re.sub("", w_fmt)
+            e_fmt = ansi_re.sub("", e_fmt)
+
+        if debug:
+            self._error_formatter = logging.Formatter(fmt=e_fmt, datefmt=datefmt)
+            self._warning_formatter = logging.Formatter(fmt=w_fmt, datefmt=datefmt)
+            self._debug_formatter = logging.Formatter(fmt=d_fmt, datefmt=datefmt)
+            self._info_formatter = logging.Formatter(fmt=i_fmt, datefmt=datefmt)
+        else:
+            self._info_formatter = logging.Formatter(fmt="> %(message)s")
+            self._warning_formatter = logging.Formatter(fmt=w_fmt, datefmt=datefmt)
+            self._error_formatter = self._warning_formatter
+            self._debug_formatter = self._warning_formatter
+
+    def format(self, record: logging.LogRecord):
+        if self.use_color:
+            # Highlight all UUIDs
+            if record.levelno > logging.WARNING:
+                end_coloring = COLOR_ERROR
+            elif record.levelno > logging.INFO:
+                end_coloring = COLOR_WARNING
+            elif record.levelno > logging.DEBUG:
+                end_coloring = COLOR_INFO
+            else:
+                end_coloring = COLOR_DEBUG
+
+            repl = f"{_byel}\\1{_r}{end_coloring}"
+            try:
+                record.msg = self.uuid_re.sub(repl, record.msg)
+                if isinstance(record.args, dict):
+                    for k, v in record.args.items():
+                        record.args[k] = self.uuid_re.sub(repl, str(v))
+                elif record.args:
+                    args = tuple(self.uuid_re.sub(repl, str(a)) for a in record.args)
+                    record.args = args
+            except Exception as exc:
+                # Basically, inform, but ignore
+                print(f"Unable to colorize log message: {exc}")
+
+        if record.levelno > logging.WARNING:
+            return self._error_formatter.format(record)
+        elif record.levelno > logging.INFO:
+            return self._warning_formatter.format(record)
+        elif record.levelno > logging.DEBUG:
+            return self._info_formatter.format(record)
+        return self._debug_formatter.format(record)
+
+
+def _get_file_dict_config(
+    logfile: str, console_enabled: bool, debug: bool, no_color: bool
+) -> dict:
+    # ensure that the logdir exists
+    logdir = os.path.dirname(logfile)
+    os.makedirs(logdir, exist_ok=True)
+    log_handlers = ["logfile"]
+    if console_enabled:
+        log_handlers.append("console")
+
+    return {
+        "version": 1,
+        "formatters": {
+            "streamfmt": {
+                "()": "flox.common.logging_config.FloxConsoleFormatter",
+                "debug": debug,
+                "no_color": no_color,
+            },
+            "filefmt": {
+                "format": DEFAULT_FORMAT,
+                "datefmt": "%Y-%m-%d %H:%M:%S",
+            },
+        },
+        "handlers": {
+            "console": {
+                "class": "logging.StreamHandler",
+                "level": "DEBUG",
+                "formatter": "streamfmt",
+            },
+            "logfile": {
+                "class": "logging.handlers.RotatingFileHandler",
+                "level": "DEBUG",
+                "filename": logfile,
+                "formatter": "filefmt",
+                "maxBytes": 100 * 1024 * 1024,
+                "backupCount": 1,
+            },
+        },
+        "loggers": {
+            "": {
+                "level": "DEBUG" if debug else "INFO",
+                "handlers": log_handlers,
+            },
+        },
+    }
+
+
+def _get_stream_dict_config(debug: bool, no_color: bool) -> dict:
+    return {
+        "version": 1,
+        "formatters": {
+            "streamfmt": {
+                "()": "flox.common.logging_config.FloxConsoleFormatter",
+                "debug": debug,
+                "no_color": no_color,
+            },
+        },
+        "handlers": {
+            "console": {
+                "class": "logging.StreamHandler",
+                "level": "DEBUG" if debug else "INFO",
+                "formatter": "streamfmt",
+            }
+        },
+        "loggers": {
+            "": {
+                "level": "DEBUG",
+                "handlers": ["console"],
+            },
+        },
+    }
+
+
+class FloxLogger(logging.Logger):
+    TRACE = logging.DEBUG - 5
+
+    def trace(self, msg, *args, **kwargs):
+        self.log(FloxLogger.TRACE, msg, args, **kwargs)
+
+
+logging.setLoggerClass(FloxLogger)
+logger = logging.getLogger(__name__)
+
+
+def setup_logging(
+    *,
+    logfile: pathlib.Path | str | None = None,
+    console_enabled: bool = True,
+    debug: bool = False,
+    no_color: bool = False,
+) -> None:
+    if logfile is not None:
+        config = _get_file_dict_config(str(logfile), console_enabled, debug, no_color)
+    else:
+        config = _get_stream_dict_config(debug, no_color)
+
+    logging.config.dictConfig(config)
diff --git a/flox/common/typing.py b/flox/common/typing.py
new file mode 100644
index 0000000..39c04d6
--- /dev/null
+++ b/flox/common/typing.py
@@ -0,0 +1,74 @@
+"""FLoX Type Definitions"""
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy.typing as npt
+
+NDArray = npt.NDArray[Any]
+NDArrays = List[NDArray]
+
+Scalar = Union[bool, bytes, float, int, str]
+
+Metrics = Dict[str, Scalar]
+
+
+@dataclass
+class EvaluateRes:
+    """Evaluate response from a client."""
+
+    loss: float
+    metrics: Dict[str, Scalar]
+
+
+@dataclass
+class FuturesList:
+    """List of futures. If using FuncXExecutor, this would most likely be a list of futures
+    funcX returns after you submit functions to endpoints"""
+
+    tasks: List[Any]
+
+
+@dataclass
+class ResultsList:
+    """FL results extracted from FuturesList, formatted as a dictionary. For example:
+    ResultsList = {
+        "model_weights": model_weights,
+        "samples_count": samples_count,
+        "bias_weights": fractions,
+    }
+    """
+
+    results: Dict[str, Scalar]
+
+
+@dataclass
+class UpdateIns:
+    """List of inputs for on_model_update"""
+
+    weights: NDArrays
+
+
+@dataclass
+class ConfigFile:
+    """Dictionary with the necessary inputs for functions like on_data_retrieve"""
+
+    config: Dict[str, Scalar]
+
+
+@dataclass
+class XYData:
+    """x_train/test and y_train/test data for model training or testing"""
+
+    x_data: NDArrays
+    y_data: Union[NDArray, NDArrays]
+
+
+@dataclass
+class FitIns:
+    """Parameters for fitting the model"""
+
+    model: Any
+    ModelTrainer: Any
+    config: ConfigFile
+    x_train: NDArrays
+    y_train: Union[NDArray, NDArrays]
diff --git a/flox/controllers/MainController.py b/flox/controllers/MainController.py
new file mode 100644
index 0000000..dccdd73
--- /dev/null
+++ b/flox/controllers/MainController.py
@@ -0,0 +1,482 @@
+import concurrent.futures
+import logging
+import time
+from collections import deque
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+from funcx import FuncXClient, FuncXExecutor
+
+from flox.common.logging_config import setup_logging
+from flox.common.typing import NDArrays
+from flox.logic import BaseModelTrainer, FloxClientLogic, FloxControllerLogic
+
+setup_logging(debug=True)
+logger = logging.getLogger(__name__)
+
+
+class MainController(FloxControllerLogic):
+    """MainController implements the main controller functionality while allowing for
+    extensibility by inhereting functionality from it
+
+    endpoint_ids: List[str]
+        a list with endpoint_ids to include in the FL process.
+        The ids can be retrieved by running 'funcx-endpoint list' on participating devices.
+        If using a non-funcx executor, pass *unique* arbitrary string values, such as ["fake_ep1", "fake_ep2"]
+
+    client_logic: FloxClientLogic instance
+        a class that implements abstract class methods from FloxClientLogic
+
+    model_trainer: BaseModelTrainer instance
+        a class that implements abstract class methods from BaseModelTrainer
+
+    # TODO: should I move this parameter under those client_logic classes that actually use them?
+    global_model: ML model object
+       ML model that will be deployed for training on the endpoints.
+       However, some ML frameworks might instead benefit from storing the model as
+       an attribute of model_trainer, so global_model defaults to None and is not required.
+
+    executor: concurrent.futures.Executor instance
+        a class similar to concurrent.futures.Executor class that implements methods
+        .submit() and returns a future-like class. MainController provides two executors
+        to choose from (FuncXExecutor, ThreadPoolExecutor) although users may provide custom executors
+
+    executor_type: str
+        you can provide "funcx" or "local" if you wish to use one of the offered executors
+        (FuncXExecutor, ThreadPoolExecutor). Otherwise, it will default to "local"
+
+    num_samples: int or list
+        indicates how many samples to get for training on endpoints
+        if int, applies the same num_samples to all endpoints.
+        if list, it will use the corresponding number of samples for each device
+        the list should have the same number of entries as the number of endpoints
+
+    epochs: int or list
+        indicates how many epochs to use for training on endpoints
+        if int, applies the same number of epochs to all endpoints.
+        if list, it will use the corresponding number of epochs for each device
+        the list should have the same number of entries as the number of endpoints
+
+    rounds: int
+        defines how many FL rounds to run. Each round consists of deploying the model, training,
+        aggregating the updates, and reassigning new weights to the model.
+
+    timeout: int
+        defines the timeout after which the MainController will stop trying to retrieve tasks,
+        assuming that the device has dropped out from training midway.
+
+    data_source: str
+        type of the data source supported by the provided client_logic instance
+
+    dataset_name: str
+        the dataset name supported by the provided client_logic instance
+
+    # TODO: should I move this parameter under those client_logic classes that actually use them?
+    preprocess: boolean
+        used to indicate whether to preprocess data on the client side, if supported by client_logic
+
+    path_dir: str
+        path to the folder with x_train_filename and y_train_filename; needed when data_sourse="local"
+
+    x_train_filename: str
+        filename for x_train; needed when data_sourse="local"
+
+    y_train_filename: str
+        file name for y_train; needed when data_sourse="local"
+
+    # TODO: should I move this parameter under those client_logic classes that actually use them?
+    input_shape: tupple
+        input shape for the provided model
+
+    x_test: list/numpy array/tensors
+        x_test data for testing
+
+    y_test: list
+        y_test labels for x_test
+
+    running_average: bool
+        if True, will use `tasks_to_running_average` to aggregate the model weights
+        as a simple average on the fly as they are returned from the Executor.
+        This speeds up the total processing time when you have many endpoints and aggregating
+        all of them in a single operation takes too long
+
+    """
+
+    AVAILABLE_EXECUTORS = {"local": ThreadPoolExecutor, "funcx": FuncXExecutor}
+
+    def __init__(
+        self,
+        endpoint_ids: List[str],
+        client_logic: FloxClientLogic,
+        model_trainer: BaseModelTrainer,
+        global_model: NDArrays = None,
+        executor: concurrent.futures.Executor = None,
+        executor_type: str = "local",
+        num_samples: Union[int, List[int]] = 100,
+        epochs: Union[int, List[int]] = 5,
+        rounds: int = 3,
+        path_dir: Union[str, List[int]] = ["."],
+        x_test=None,
+        y_test=None,
+        data_source: str = None,
+        dataset_name=None,
+        preprocess: bool = True,
+        x_train_filename: str = None,
+        y_train_filename: str = None,
+        input_shape: Tuple[int] = None,
+        timeout: int = 60,
+        running_average: bool = False,
+    ):
+        self.endpoint_ids = endpoint_ids
+        self.num_samples = num_samples
+        self.epochs = epochs
+        self.rounds = rounds
+        self.client_logic = client_logic
+        self.global_model = global_model
+        self.model_trainer = model_trainer
+        self.executor = executor
+        self.executor_type = executor_type
+        self.path_dir = path_dir
+        self.x_test = x_test
+        self.y_test = y_test
+        self.data_source = data_source
+        self.dataset_name = dataset_name
+        self.preprocess = preprocess
+        self.x_train_filename = x_train_filename
+        self.y_train_filename = y_train_filename
+        self.input_shape = input_shape
+        self.timeout = timeout
+        self.running_average = running_average
+
+    def on_model_init(self) -> None:
+        """Does initial Controller setup before running the main Federated Learning loop"""
+        # if num_samples or epochs is an int, convert to list so the same number can be applied to all endpoints
+        if type(self.num_samples) == int:
+            self.num_samples = [self.num_samples] * len(self.endpoint_ids)
+
+        if type(self.epochs) == int:
+            self.epochs = [self.epochs] * len(self.endpoint_ids)
+
+        if type(self.path_dir) == str:
+            self.path_dir = [self.path_dir] * len(self.endpoint_ids)
+
+        if not self.executor:
+            logger.debug(
+                f"No executor was provided, trying to retrieve the provided executor type {self.executor_type} from the list of available executors: {self.AVAILABLE_EXECUTORS}"
+            )
+            try:
+                self.executor = self.AVAILABLE_EXECUTORS[self.executor_type]
+                logger.debug(f"The selected executor is {self.executor}")
+            except KeyError as exp:
+                logger.debug(
+                    f"Could not find {self.executor_type} in the list of available executors: {self.AVAILABLE_EXECUTORS}. Please provide your own executor or select from {self.AVAILABLE_EXECUTORS.keys()}"
+                )
+                self.executor = self.AVAILABLE_EXECUTORS["local"]
+                logger.debug(
+                    f"Defaulting to {self.AVAILABLE_EXECUTORS['local']} executor"
+                )
+
+        self.funcx_client = FuncXClient(http_timeout=60)
+
+        self.endpoints_statuses = {}
+
+    def create_config(self, *args, **kwargs) -> Dict:
+        """Creates a config dictionary that will be broadcasted in self.on_model_broadcast()
+        to the endpoints. The methods requires the user to extend it since different ML frameworks
+        require a different set of arguments.
+
+        Returns
+        -------
+        config: Dict
+            a dictionary of parameters required for use by client_logic.run_round()
+
+        """
+        raise NotImplementedError("Method not implemented")
+
+    def on_model_broadcast(self) -> deque:
+        """Sends the model and config to endpoints for FL training.
+
+        Returns
+        -------
+        tasks: deque[future]
+            return a list of futures. The futures should support methods .done() and .result()
+
+        """
+        # define list storage for results
+        tasks = deque()
+
+        logger.debug(f"Launching the {self.executor} executor")
+        with self.executor() as executor:
+            # submit the corresponding parameters to each endpoint for a round of FL
+            for ep, num_s, num_epoch, path_d in zip(
+                self.endpoint_ids, self.num_samples, self.epochs, self.path_dir
+            ):
+                logger.info(f"Starting to broadcast a task to endpoint {ep}")
+                try:
+                    ep_status = self.funcx_client.get_endpoint_status(ep)["status"]
+                except Exception as exp:
+                    logger.warning(
+                        f"Could not check the status of the endpoint {ep}, the error is: {exp}"
+                    )
+                    ep_status = "error"
+
+                if ep_status != "online" and self.executor == FuncXExecutor:
+                    logger.warning(f"Endpoint {ep} is not online, it's {ep_status}!")
+                else:
+                    config = self.create_config(num_s, num_epoch, path_d)
+                    executor.endpoint_id = ep
+
+                    if self.executor_type == "local":
+                        task = executor.submit(
+                            self.client_logic.run_round, config, self.model_trainer
+                        )
+
+                    elif self.executor_type == "funcx":
+                        task = executor.submit(
+                            self.client_logic.run_round,
+                            self.client_logic,  # funcxExecutor requires the class submitted as well, while the ThreadPoolExecutor does not
+                            config,
+                            self.model_trainer,
+                        )
+                    else:
+                        raise ValueError(
+                            f"{self.executor_type} is invalid executor type, choose one of [local, funcx]"
+                        )
+
+                    logger.info(f"Deployed the task to endpoint {ep}")
+                    tasks.append(task)
+
+                self.endpoints_statuses[ep] = ep_status
+
+            self.task_start_time = (
+                time.time()
+            )  # how would this work for multiple endpoints?
+
+            if len(tasks) == 0:
+                logger.error(
+                    f"The tasks queue is empty, here are the endpoints' statuses: {self.endpoints_statuses}"
+                )
+                raise ValueError(
+                    f"The tasks queue is empty, no tasks were submitted for training!"
+                )
+
+        return tasks
+
+    def on_model_receive(self, tasks: deque) -> Dict:
+        """Processes returned tasks from on_model_broadcast.
+
+        Parameters
+        ----------
+        tasks: List[futures]
+            A list of tasks/futures with results of the FL training returned from the endpoints.
+            If using FuncXExecutor/ThreadPoolExecutor, this would be a list of futures
+            funcX/ThreadPoolExecutor returns after you submit functions to endpoints.
+
+        Returns
+        -------
+        results: Dict
+            FL results extracted from tasks, formatted as a dictionary. For example:
+            results = {
+                "model_weights": model_weights,
+                "samples_count": samples_count,
+                "bias_weights": fractions,
+            }
+        """
+        model_weights = []
+        samples_count = []
+        endpoint_result_order = []
+
+        logger.info("Starting to retrieve results from endpoints")
+        while tasks and (time.time() - self.task_start_time) < self.timeout:
+            t = tasks.popleft()
+            if t.done():
+                res = t.result()
+                model_weights.append(res["model_weights"])
+                samples_count.append(res["samples_count"])
+                endpoint_result_order.append(t)
+            else:
+                tasks.append(t)
+                logger.info(f"Retrieved results from endpoints {t}")
+
+        samples_count = np.array(samples_count)
+        total = sum(samples_count)
+        fractions = samples_count / total
+        logger.info("Finished retrieving all results from the endpoints")
+        return {
+            "model_weights": model_weights,
+            "samples_count": samples_count,
+            "bias_weights": fractions,
+            "endpoint_result_order": endpoint_result_order,
+        }
+
+    def on_model_aggregate(self, results: Dict) -> NDArrays:
+        """Aggregates weights.
+
+        Parameters
+        ----------
+        results: Dict
+            FL results returned by on_model_receive, formatted as a dictionary. For example:
+            results = {
+                "model_weights": model_weights,
+                "samples_count": samples_count,
+                "bias_weights": fractions,
+            }
+
+        Returns
+        -------
+        average_weights: NDArrays
+            ML model weights for a single model in the form of Numpy Arrays.
+
+        """
+        average_weights = np.average(
+            results["model_weights"], weights=results["bias_weights"], axis=0
+        )
+        logger.info("Finished aggregating weights")
+        return average_weights
+
+    def on_model_update(self, updated_weights: NDArrays) -> None:
+        """Updates the model's weights with new weights.
+        The method might need to be overriden depending on the Machine Learning framework as use.
+
+        Parameters
+        ----------
+        updated_weights: NDArrays
+            ML model weights for a single model in the form of Numpy Arrays.
+
+        """
+        self.model_trainer.set_weights(self.global_model, updated_weights)
+        logger.info("Updated the global model's weights")
+
+    def on_model_evaluate(self):
+        """Evaluates the given model using provided test data.
+        The method might need to be overriden depending on the Machine Learning framework as use.
+
+        Returns
+        -------
+        EvaluateRes
+            A class consisting of two attributes showing the loss and metrics of the evaluated
+            model on the dataset:
+                loss: float
+                metrics: Dict[str, Scalar]
+        """
+        if self.x_test is not None and self.y_test is not None:
+            logger.info("Starting evaluation")
+            results = self.model_trainer.evaluate(
+                self.global_model, self.x_test, self.y_test
+            )
+            logger.info(f"Evaluation results: {results}")
+            return results
+        else:
+            logger.warning("Skipping evaluation, no x_test and/or y_test provided")
+            return False
+
+    def run_federated_learning(self):
+        """The main Federated Learning loop that ties all the functions together.
+        Runs <self.rounds> rounds of federated learning"""
+        self.on_model_init()
+        # start running FL loops
+        for i in range(self.rounds):
+            # broadcast the model
+            tasks = self.on_model_broadcast()
+
+            if self.running_average:
+                updated_weights = self.tasks_to_running_average(tasks)[
+                    "running_average_weights"
+                ]
+            else:
+                # process & decrypt the results
+                results = self.on_model_receive(tasks)
+
+                # aggregate the weights
+                updated_weights = self.on_model_aggregate(results)
+
+            # update the model's weights
+            self.on_model_update(updated_weights)
+
+            # evaluate the model
+            logger.info(f"Round {i} evaluation results: ")
+            self.on_model_evaluate()
+
+    def tasks_to_running_average(self, tasks: deque):
+        """Serves as a modified combination of on_model_receive and on_model_aggregate.
+        Aggregates the model weights
+        as a simple average on the fly as they are returned from the Executor.
+        This speeds up the total processing time when you have many endpoints and aggregating
+        all of them in a single operation takes too long. Note that the aggregated model
+        will be slightly different from what you would get if you used on_model_aggregate
+        on all of the weights in a single operation. This is because of the reiterative
+        floating point averaging.
+
+        Parameters
+        ----------
+        tasks: List[futures]
+            A list of tasks/futures with results of the FL training returned from the endpoints.
+            If using FuncXExecutor/ThreadPoolExecutor, this would be a list of futures
+            funcX/ThreadPoolExecutor returns after you submit functions to endpoints.
+
+        Returns
+        -------
+        results: Dict
+            FL results extracted from tasks, formatted as a dictionary. For example:
+            results = {
+                "model_weights": model_weights,
+                "samples_count": samples_count,
+                "bias_weights": fractions,
+                "running_average_weights": running_average
+            }
+        """
+        model_weights = []
+        samples_count = []
+        endpoint_result_order = []
+
+        n_aggregated = 0
+        n_total = 0
+        running_average = None
+
+        logger.info("Starting to retrieve results from endpoints")
+        while tasks and (time.time() - self.task_start_time) < self.timeout:
+            t = tasks.popleft()
+            if t.done():
+                res = t.result()
+                new_weights = res["model_weights"]
+                model_weights.append(new_weights)
+                samples_count.append(res["samples_count"])
+                endpoint_result_order.append(t)
+                # update running average
+                if running_average is None:
+                    logger.debug(
+                        "the running average is NONE, instantiating it for the first time"
+                    )
+                    running_average = new_weights
+                    n_aggregated += 1
+                    n_total += 1
+                else:
+                    n_total += 1
+                    running_weights = [
+                        (n_aggregated / n_total),
+                        (n_total - n_aggregated) / n_total,
+                    ]
+                    logger.info(
+                        f"The weights are {running_weights} and sum up to {sum(running_weights)}"
+                    )
+                    running_average = np.average(
+                        [running_average, new_weights], weights=running_weights, axis=0
+                    )
+                    n_aggregated += 1
+
+            else:
+                tasks.append(t)
+                logger.info(f"Retrieved results from endpoints {t}")
+
+        samples_count = np.array(samples_count)
+        total = sum(samples_count)
+        fractions = samples_count / total
+        logger.info("Finished retrieving all results from the endpoints")
+        return {
+            "model_weights": model_weights,
+            "samples_count": samples_count,
+            "bias_weights": fractions,
+            "endpoint_result_order": endpoint_result_order,
+            "running_average_weights": running_average,
+        }
diff --git a/flox/controllers/PyTorchController.py b/flox/controllers/PyTorchController.py
new file mode 100644
index 0000000..3c42c3a
--- /dev/null
+++ b/flox/controllers/PyTorchController.py
@@ -0,0 +1,30 @@
+import logging
+
+from flox.controllers.MainController import MainController
+
+logger = logging.getLogger(__name__)
+
+
+class PyTorchController(MainController):
+    def __init__(self, *args, testloader=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.testloader = testloader
+
+    def create_config(self, num_s, num_epoch, path_d):
+        config = {
+            "num_samples": num_s,
+            "epochs": num_epoch,
+            "path_dir": path_d,
+            "dataset_name": self.dataset_name,
+            "preprocess": self.preprocess,
+        }
+
+        return config
+
+    def on_model_update(self, updated_weights) -> None:
+        self.model_trainer.set_weights(updated_weights)
+
+    def on_model_evaluate(self):
+        results = self.model_trainer.evaluate(self.testloader)
+        print(results)
+        return results
diff --git a/flox/controllers/TensorflowController.py b/flox/controllers/TensorflowController.py
new file mode 100644
index 0000000..5216a59
--- /dev/null
+++ b/flox/controllers/TensorflowController.py
@@ -0,0 +1,37 @@
+from flox.controllers.MainController import MainController
+
+
+class TensorflowController(MainController):
+    def create_config(self, num_s, num_epoch, path_d):
+        model_architecture = self.model_trainer.get_architecture(self.global_model)
+        model_weights = self.model_trainer.get_weights(self.global_model)
+
+        config = {
+            "num_samples": num_s,
+            "epochs": num_epoch,
+            "path_dir": path_d,
+            "data_source": self.data_source,
+            "dataset_name": self.dataset_name,
+            "preprocess": self.preprocess,
+            "architecture": model_architecture,
+            "weights": model_weights,
+            "x_train_filename": self.x_train_filename,
+            "y_train_filename": self.y_train_filename,
+            "input_shape": self.input_shape,
+        }
+
+        return config
+
+    def on_model_update(self, updated_weights) -> None:
+        self.model_trainer.set_weights(self.global_model, updated_weights)
+
+    def on_model_evaluate(self):
+        if self.x_test is not None and self.y_test is not None:
+            results = self.model_trainer.evaluate(
+                self.global_model, self.x_test, self.y_test
+            )
+            print(results)
+            return results
+        else:
+            print("Skipping evaluation, no x_test and/or y_test provided")
+            return False
diff --git a/build/lib/flox_module/__init__.py b/flox/controllers/__init__.py
similarity index 100%
rename from build/lib/flox_module/__init__.py
rename to flox/controllers/__init__.py
diff --git a/flox/examples/quickstart_pytorch/README.md b/flox/examples/quickstart_pytorch/README.md
new file mode 100644
index 0000000..b855481
--- /dev/null
+++ b/flox/examples/quickstart_pytorch/README.md
@@ -0,0 +1,38 @@
+# Quickstart PyTorch Tutorial
+
+
+This example is in development and includes only basic setup and execution instructions.
+
+### Controller instructions
+
+On your *Controller*, follow these instructions:
+
+1. Create a conda or virtualenv environment
+
+2. Install FLoX with ``pip install -e .`` since the latest version of the library is not out yet.
+
+3. Install [PyTorch](https://pytorch.org/): ``pip install torch==1.12.0`` and ``pip install torchvision==0.13.0``
+
+4. Create a dist file with ``python setup.py bdist_wheel`` and transfer the created dist file to the Clients
+This step will be eliminated once the library is published on PyPI.
+
+5. Configure your *Clients* as per instructions below
+
+6. Run ``python flox/examples/quickstart_pytorch/pytorch_funcx.py``
+
+### Clients instructions
+
+On your *Clients*:
+1. Make sure you are using a 64-bit OS, this will make PyTorch installation easier
+
+2. Install the FLoX dependencies from the dist file you created on the Controller: ``pip install pyflox-0.1.5-py3-none-any.whl``
+
+3. Install [PyTorch](https://qengineering.eu/install-pytorch-on-raspberry-pi-4.html) & [Torchvision](https://qengineering.eu/install-pytorch-on-raspberry-pi-4.html#:~:text=rm%20%2Drf%20~/pytorch-,TorchVision.,-Install%20torchvision%20on)
+
+4. Install ``funcx-endpoint`` with ``pip install funcx-endpoint``
+
+5. Configure (``funcx-endpoint configure ep1``) and start (``funcx-endpoint start ep1``) an endpoint.
+
+6. Paste the endpoint's id (``funcx-endpoint list``) and paste it into ``pytorch_funcx.py`` on the *Controller* at the endpoint definition line
+
+
diff --git a/flox/examples/quickstart_pytorch/pytorch_funcx.py b/flox/examples/quickstart_pytorch/pytorch_funcx.py
new file mode 100644
index 0000000..d04d45f
--- /dev/null
+++ b/flox/examples/quickstart_pytorch/pytorch_funcx.py
@@ -0,0 +1,126 @@
+import logging
+from concurrent.futures import ThreadPoolExecutor
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+from funcx import FuncXExecutor
+from torch import Tensor
+
+from flox.clients.PyTorchClient import PyTorchClient
+from flox.controllers.PyTorchController import PyTorchController
+from flox.model_trainers.PyTorchTrainer import PyTorchTrainer
+
+logger = logging.getLogger(__name__)
+
+
+def get_test_data(config):
+    import torch
+    import torchvision
+    import torchvision.transforms as transforms
+
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ]
+    )
+
+    batch_size = config.get("batch_size", 32)
+    dataset_name = config["dataset_name"]
+    num_workers = config.get("num_workers", 8)
+    root = config.get("data_root", "./data")
+
+    # create train DataLoader
+    trainset = dataset_name(root=root, train=True, download=True, transform=transform)
+
+    train_split_len = (
+        len(trainset) if "num_samples" not in config.keys() else config["num_samples"]
+    )
+
+    train_subpart = torch.utils.data.random_split(
+        trainset, [train_split_len, len(trainset) - train_split_len]
+    )[0]
+    trainloader = torch.utils.data.DataLoader(
+        train_subpart, batch_size=batch_size, shuffle=True, num_workers=num_workers
+    )
+
+    # create test DataLoader
+    testset = dataset_name(root=root, train=False, download=True, transform=transform)
+    test_split_len = (
+        len(trainset) if "num_samples" not in config.keys() else config["num_samples"]
+    )
+
+    test_subpart = torch.utils.data.random_split(
+        testset, [test_split_len, len(testset) - test_split_len]
+    )[0]
+    testloader = torch.utils.data.DataLoader(
+        test_subpart, batch_size=batch_size, shuffle=False, num_workers=num_workers
+    )
+
+    return trainloader, testloader
+
+
+def main():
+    class Net(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv1 = nn.Conv2d(3, 6, 5)
+            self.pool = nn.MaxPool2d(2, 2)
+            self.conv2 = nn.Conv2d(6, 16, 5)
+            self.fc1 = nn.Linear(16 * 5 * 5, 120)
+            self.fc2 = nn.Linear(120, 84)
+            self.fc3 = nn.Linear(84, 10)
+
+        def forward(self, x):
+            import torch
+
+            x = self.pool(torch.nn.functional.relu(self.conv1(x)))
+            x = self.pool(torch.nn.functional.relu(self.conv2(x)))
+            x = torch.flatten(x, 1)  # flatten all dimensions except batch
+            x = torch.nn.functional.relu(self.fc1(x))
+            x = torch.nn.functional.relu(self.fc2(x))
+            x = self.fc3(x)
+            return x
+
+    net = Net()
+    torch_trainer = PyTorchTrainer(net)
+    torch_client = PyTorchClient()
+
+    data_config = {
+        "num_samples": 1000,
+        "batch_size": 32,
+        "dataset_name": torchvision.datasets.CIFAR10,
+        "num_workers": 4,
+    }
+
+    _, testloader = get_test_data(data_config)
+
+    ep1 = "fb93a1c2-a8d7-49f3-ad59-375f4e298784"
+
+    eps = [ep1]
+    logger.info(f"Endpoints: {eps}")
+
+    flox_controller = PyTorchController(
+        endpoint_ids=eps,
+        num_samples=100,
+        epochs=1,
+        rounds=1,
+        client_logic=torch_client,
+        model_trainer=torch_trainer,
+        executor_type="funcx",  # choose "funcx" for FuncXExecutor, "local" for ThreadPoolExecutor
+        path_dir=".",
+        testloader=testloader,
+        dataset_name=torchvision.datasets.CIFAR10,
+        preprocess=True,
+    )
+
+    logger.info("STARTING FL TORCH FLOW...")
+    flox_controller.run_federated_learning()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/flox/examples/quickstart_tensorflow/README.md b/flox/examples/quickstart_tensorflow/README.md
new file mode 100644
index 0000000..b16c927
--- /dev/null
+++ b/flox/examples/quickstart_tensorflow/README.md
@@ -0,0 +1,37 @@
+
+# Quickstart Tensorflow Tutorial
+
+This example is in development and includes only basic setup and execution instructions.
+
+### Controller instructions
+
+On your *Controller*, follow these instructions:
+
+1. Create a conda or virtualenv environment with ``python=3.7``: ``conda create --name tf python=3.7``
+
+2. Install FLoX with ``pip install -e .`` since the latest version of the library is not out yet.
+
+3. Install [Tensorflow 2.1.0](https://www.tensorflow.org/install/pip>): ``pip install tensorflow==2.1.0``.
+Tensorflow Verstion between *Controller* and *Clients* need match.
+We use 2.1.0 because it's the version that was easily installable on our Rasperry Pi 4.
+
+4. Create a dist file with ``python setup.py bdist_wheel`` and transfer the created dist file to the Clients
+This step will be eliminated once the library is published on PyPI.
+
+5. Configure your *Clients* as per instructions below
+
+6. Run ``python flox/examples/quickstart_tensorflow/tf_funcx.py``
+
+### Clients instructions
+
+On your *Clients*:
+1. Install the FLoX dependencies from the dist file you created on the Controller: ``pip install pyflox-0.1.5-py3-none-any.whl``
+
+2. Install [Tensorflow 2.1.0](https://qengineering.eu/install-tensorflow-2.1.0-on-raspberry-pi-4.html)
+
+3. Install ``funcx-endpoint`` with ``pip install funcx-endpoint``
+
+4. Configure (``funcx-endpoint configure ep1``) and start (``funcx-endpoint start ep1``) an endpoint.
+
+5. Paste the endpoint's id (``funcx-endpoint list``) and paste it into ``pytorch_funcx.py`` on the *Controller* at the endpoint definition line
+
diff --git a/flox/examples/quickstart_tensorflow/tf_funcx.py b/flox/examples/quickstart_tensorflow/tf_funcx.py
new file mode 100644
index 0000000..24cc975
--- /dev/null
+++ b/flox/examples/quickstart_tensorflow/tf_funcx.py
@@ -0,0 +1,74 @@
+import logging
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+import tensorflow as tf
+from funcx import FuncXExecutor
+from tensorflow import keras
+
+from flox.clients.TensorflowClient import TensorflowClient
+from flox.controllers.TensorflowController import TensorflowController
+from flox.model_trainers.TensorflowTrainer import TensorflowTrainer
+from flox.utils import get_test_data
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+
+    ep1 = "c7487b2b-b129-47e2-989b-5a9ac361befc"
+
+    eps = [ep1]
+    logger.info(f"Endpoints: {eps}")
+
+    # `fashion_mnist` images are grayscale, 28 x 28 pixels in size
+    input_shape = (28, 28, 1)
+    # there are 10 classes in the dataset
+    num_classes = 10
+
+    # define the model architecture
+    global_model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=input_shape),
+            tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(num_classes, activation="softmax"),
+        ]
+    )
+
+    # compile the model
+    global_model.compile(
+        loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
+    )
+    x_test, y_test = get_test_data(keras_dataset="fashion_mnist", num_samples=2000)
+
+    tf_trainer = TensorflowTrainer()
+    tf_client = TensorflowClient()
+
+    flox_controller = TensorflowController(
+        endpoint_ids=eps,
+        num_samples=200,
+        epochs=2,
+        rounds=1,
+        client_logic=tf_client,
+        global_model=global_model,
+        executor_type="funcx",  # choose "funcx" for FuncXExecutor, "local" for ThreadPoolExecutor
+        model_trainer=tf_trainer,
+        path_dir=".",
+        x_test=x_test,
+        y_test=y_test,
+        data_source="keras",
+        dataset_name="fashion_mnist",
+        preprocess=True,
+    )
+
+    logger.info("STARTING FL FLOW...")
+    flox_controller.run_federated_learning()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/flox/examples/quickstart_tensorflow/tf_funcx_local_dataset.py b/flox/examples/quickstart_tensorflow/tf_funcx_local_dataset.py
new file mode 100644
index 0000000..e38621d
--- /dev/null
+++ b/flox/examples/quickstart_tensorflow/tf_funcx_local_dataset.py
@@ -0,0 +1,98 @@
+import os
+import pickle
+
+import numpy as np
+import tensorflow as tf
+from tensorflow import keras
+
+from flox.clients.TensorflowClient import TensorflowClient
+from flox.controllers.TensorflowController import TensorflowController
+from flox.model_trainers.TensorflowTrainer import TensorflowTrainer
+from flox.utils import get_test_data
+
+
+def process_data(train_image, train_label, num_samples=None):
+    depth = 3
+    image_size_y = 32
+    image_size_x = 32
+
+    if num_samples:
+        idx = np.random.choice(np.arange(len(train_image)), num_samples, replace=True)
+        train_image = train_image[idx]
+        train_label = train_label[idx]
+
+    train_image = train_image.reshape(
+        len(train_image), image_size_x, image_size_y, depth
+    )
+    train_image = train_image / 255.0
+
+    return (train_image, train_label)
+
+
+def main():
+
+    # ep1 = "fe49ba41-9654-4d1b-8266-fd2f8197b242sgsdg"
+    ep1 = "d917e715-bafa-466c-b12b-b9ce8e251257"
+    eps = [ep1]
+    print(f"Endpoints: {eps}")
+
+    # `fashion_mnist` images are grayscale, 28 x 28 pixels in size
+    input_shape = (32, 32, 3)
+    # there are 10 classes in the dataset
+    num_classes = 10
+
+    # define the model architecture
+    global_model = keras.Sequential(
+        [
+            keras.Input(shape=input_shape),
+            tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(num_classes, activation="softmax"),
+        ]
+    )
+
+    # compile the model
+    global_model.compile(
+        loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
+    )
+    with open("data/test_data_animal10_32.pkl", "rb") as file:
+        x_test, y_test = pickle.load(file)
+
+    x_test, y_test = process_data(x_test, y_test)
+
+    tf_trainer = TensorflowTrainer(loss="sparse_categorical_crossentropy")
+    tf_client = TensorflowClient()
+
+    flox_controller = TensorflowController(
+        endpoint_ids=eps,
+        num_samples=200,
+        epochs=1,
+        rounds=1,
+        client_logic=tf_client,
+        global_model=global_model,
+        model_trainer=tf_trainer,
+        executor_type="funcx",
+        data_source="local",
+        path_dir="/home/pi/datasets",
+        x_train_filename="x_animal10_32.npy",
+        y_train_filename="y_animal10_32.npy",
+        input_shape=(32, 32, 32, 3),
+        preprocess=True,
+        x_test=x_test,
+        y_test=y_test,
+        timeout=5,
+    )
+
+    print("STARTING FL FLOW...")
+    flox_controller.run_federated_learning()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/flox/examples/quickstart_tensorflow/tf_running_avg.py b/flox/examples/quickstart_tensorflow/tf_running_avg.py
new file mode 100644
index 0000000..8552fd2
--- /dev/null
+++ b/flox/examples/quickstart_tensorflow/tf_running_avg.py
@@ -0,0 +1,78 @@
+import logging
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+import tensorflow as tf
+from funcx import FuncXExecutor
+from tensorflow import keras
+
+from flox.clients.TensorflowClient import TensorflowClient
+from flox.controllers.TensorflowController import TensorflowController
+from flox.model_trainers.TensorflowTrainer import TensorflowTrainer
+from flox.utils import get_test_data
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+
+    ep1 = "c7487b2b-b129-47e2-989b-5a9ac361befc"
+    ep2 = "d7487b2b-b129-47e2-989b-5a9ac361befd"
+    ep3 = "f7487b2b-b129-47e2-989b-5a9ac361beff"
+    ep4 = "e7487b2b-b129-47e2-989b-5a9ac361befe"
+
+    eps = [ep1, ep2, ep3, ep4]
+    logger.info(f"Endpoints: {eps}")
+
+    # `fashion_mnist` images are grayscale, 28 x 28 pixels in size
+    input_shape = (28, 28, 1)
+    # there are 10 classes in the dataset
+    num_classes = 10
+
+    # define the model architecture
+    global_model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=input_shape),
+            tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(num_classes, activation="softmax"),
+        ]
+    )
+
+    # compile the model
+    global_model.compile(
+        loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
+    )
+    x_test, y_test = get_test_data(keras_dataset="fashion_mnist", num_samples=2000)
+
+    tf_trainer = TensorflowTrainer()
+    tf_client = TensorflowClient()
+
+    flox_controller = TensorflowController(
+        endpoint_ids=eps,
+        num_samples=100,
+        epochs=1,
+        rounds=3,
+        client_logic=tf_client,
+        global_model=global_model,
+        executor_type="local",  # choose "funcx" for FuncXExecutor, "local" for ThreadPoolExecutor
+        model_trainer=tf_trainer,
+        path_dir=".",
+        x_test=x_test,
+        y_test=y_test,
+        data_source="keras",
+        dataset_name="fashion_mnist",
+        preprocess=True,
+        running_average=True,
+    )
+
+    logger.info("STARTING FL FLOW...")
+    flox_controller.run_federated_learning()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/flox/logic/README.rst b/flox/logic/README.rst
new file mode 100644
index 0000000..d4cb30f
--- /dev/null
+++ b/flox/logic/README.rst
@@ -0,0 +1,143 @@
+FLoX Client-Controller Logic
+============================
+
+FLoX is designed to be a highly modular and customizable framework for
+serverless, FL processes. It is built on top of a general 10-step
+abstraction for FL processes grouped into controller-side and
+client-side logical steps. FLoX considers two abstract classes,
+``FloxControllerLogic`` and ``FloxClientLogic``, for
+wrapping/implementing the logic needed for the controller-side and
+client-side, respectively. Each step is run on either the Controller (S)
+or the clients (C). All 7 steps are listed below with the respective
+function that corresponds with them:
+
+1.  **(S)** Model Initialization: ``on_model_init``
+2.  **(S)** Model Sharing: ``on_model_broadcast``
+3.  **(C)** Receiving Model on Client: ``on_model_receive``
+4.  **(C)** Data Fetching: ``on_data_retrieve``
+5.  **(C)** Local Model Training: ``on_model_fit``
+6.  **(C)** Model Parameter Submission: ``on_model_send``
+7.  **(S)** Receiving Model on Controller: ``on_model_receive``
+8.  **(S)** Model Aggregation: ``on_model_aggregate``
+9.  **(S)** Model Updating: ``on_model_update``
+10. **(S)** Model Evaluation: ``on_model_evaluate``
+
+Below is a (``mermaid.js``) figure showing the sequence of the logic
+steps. Then, we further describe each of the logic steps.
+
+.. code:: mermaid
+
+   flowchart
+       id1[[on_model_init]]
+       id2[/on_model_broadcast\]
+
+       id3[\on_model_receive/]
+       id4[(on_data_retrieve)]
+       id5{on_model_fit}
+       id6[on_model_send]
+
+       id7[on_model_receive]
+       id8[on_model_aggregate]
+       id9[on_model_update]
+       id10[on_model_evaluate]
+
+
+
+       subgraph controller
+           id1-->id2
+           id7-->id8
+           id8-->id9
+           id9-->id10
+           id10-->id2
+       end
+
+       subgraph client
+           direction TB
+           id3-->id4
+           id4-->id5
+           id5-->id6
+       end
+
+       id2-->|model params|id3
+       id6-->|model params|id7
+
+1. Model Initialization
+-----------------------
+``Controller.on_model_init()`` is where one would provide initial
+setup scripts that need to be run only once rather than needing to run
+every FL round. In ``flox.controllers.TensorflowController`` we use
+``.on_model_init()`` to set up variables that were not provided by the
+user but will be reused in all of the FL rounds going forward.
+
+2. Model Sharing
+----------------
+Model sharing happend in ``Controller.on_model_broadcast`` where one
+would assemble all variables into a config file, encrypt data is
+necessary, and then send the tasks to the *clients* using an Executor,
+such as FuncXExecutor. The method should return a list of futures/tasks
+that can later be parsed out by ``Controller.on_model_receive()`` once
+*clients* return the results.
+
+3. Receiving Model on Client
+----------------------------
+Once *clients* receive the model and the config with necessary
+parameters, ``Client.on_model_receive()`` is responsible for the initial
+actions such as decrypting the data if it wsa encrypted.
+
+4. Data Retrieval
+-----------------
+``Client.on_data_retrieve()`` is where *clients* retrieve and prepare
+their data for training.
+
+5. Local Model Training
+-----------------------
+``Client.on_data_fit()`` is where the training process is defined and
+executed.
+
+6. Model Parameter Submission
+-----------------------------
+When *clients* retrieved local data and made updates to the global
+model, the new model weights are returned to the *Controller* and
+``Client.on_data_send()`` is for things like encryption of data before
+it is sent back.
+
+7. Receiving Model on Controller
+--------------------------------
+Once the *Controller* receives the results back from the *clients*,
+``Controller.on_model_receive()`` parses the results and decrypts if
+necessary.
+
+8. Model Aggregation
+--------------------
+``Controller.on_model_aggregate()`` takes the parsed results from
+``Controller.on_model_receive()`` and aggregates weights from the
+endpoints.
+
+9. Model Updating
+-----------------
+``Controller.on_model_update()`` simply takes the new weights from
+``Controller.on_model_aggregate()`` and assigns them to the global
+model.
+
+10. Model Evaluation
+--------------------
+Finally, ``Controller.on_model_evaluate()`` evaluates the model using
+a user-provided testing dataset, reports the results, and then the
+entire loop from Step 2 to Step 10 is repeated for as many rounds as was
+specified by the user.
+
+We implemented the abstract base classes in
+``flox.logic.base_client.py`` and ``flox.logic.base_controller.py``. We
+also implemented a base class for Machine Learning Model Trainers,
+located in ``flox.logic.base_model_trainer.py``. We are providing
+practical examples on top of these classes to illustrate how all of
+these steps come together:
+
+- ``flox.examples.quickstart_pytorch`` makes
+use of ``PyTorchController``, ``PyTorchClient``, and ``PyTorchTrainer``
+to run a Federated Learning workflow on PyTorch.
+
+- ``flox.examples.quickstart_tensorflow`` makes use of
+``TensorflowController``, ``TensorflowClient``, and
+``TensorflowTrainer`` to run a Federated Learning workflow on
+Tensorflow.
diff --git a/flox/logic/__init__.py b/flox/logic/__init__.py
new file mode 100644
index 0000000..f2d2a70
--- /dev/null
+++ b/flox/logic/__init__.py
@@ -0,0 +1,5 @@
+from .base_client import FloxClientLogic
+from .base_controller import FloxControllerLogic
+from .base_model_trainer import BaseModelTrainer
+
+__all__ = ["BaseModelTrainer", "FloxClientLogic", "FloxControllerLogic"]
diff --git a/flox/logic/base_client.py b/flox/logic/base_client.py
new file mode 100644
index 0000000..6219b0e
--- /dev/null
+++ b/flox/logic/base_client.py
@@ -0,0 +1,94 @@
+"""Abstract Base Class for FLoX Clients"""
+from flox.common import ConfigFile, FitIns, NDArrays, ResultsList, XYData
+
+
+class FloxClientLogic:
+    """Abstract base class for FLoX Clients"""
+
+    def on_model_receive(self) -> None:
+        """Parses & decrypts the received data from a controller"""
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_data_retrieve(self, config: ConfigFile) -> XYData:
+        """Retrieves data for training the model.
+
+        Parameters
+        ----------
+        config: ConfigFile
+            a dictionary with values necessary to retrieve the data, such as
+            the data source and preprocessing parameters
+
+        Returns
+        -------
+        XYData
+            data for evaluating the model. This can take different forms depending on the
+            ML framework you use. For Tensorflow, it would look as x_test and y_test, while
+            for PyTorch, it would look as a single DataLoader instance.
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_model_fit(self, ins: FitIns) -> NDArrays:
+        """Fit the provided global model using a local dataset
+
+        Parameters
+        ----------
+        FitIns
+            Parameters for fitting the model. This might take different forms depending on the
+            ML framework. For Tensorflow, it might look like this:
+                model: keras.Sequential
+                ModelTrainer: Class
+                config: ConfigFile
+                x_train: NDArrays
+                y_train: Union[NDArray, NDArrays]
+
+        Returns
+        -------
+        NDArrays
+            new model weights as Numpy arrays
+
+        """
+
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_model_send(self) -> ResultsList:
+        """Final data processing before sending the results back, such as encryption.
+
+        Returns
+        -------
+        ResultsList
+            FL results formatted as a dictionary. For example:
+            ResultsList = {
+                "model_weights": model_weights,
+                "samples_count": samples_count,
+                "bias_weights": fractions,
+            }
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def run_round(self, ModelTrainer, config: ConfigFile) -> ResultsList:
+        """Combines the rest of the functions to run a single Federated Learning round
+
+        Parameters
+        ----------
+        ModelTrainer
+            instance of a class based on BaseModelTrainer that implements
+            .fit(), .evaluate(), .get_weights(), and .set_weights()
+
+        config: ConfigFile
+            a dictionary with all the necessary parameters for retrieving data and
+            training the model
+
+        Returns
+        -------
+        ResultsList
+            FL results formatted as a dictionary. For example:
+            ResultsList = {
+                "model_weights": model_weights,
+                "samples_count": samples_count,
+                "bias_weights": fractions,
+            }
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
diff --git a/flox/logic/base_controller.py b/flox/logic/base_controller.py
new file mode 100644
index 0000000..f8c9d21
--- /dev/null
+++ b/flox/logic/base_controller.py
@@ -0,0 +1,103 @@
+"""Abstract Base Class for FLoX Controllers"""
+from flox.common import EvaluateRes, FuturesList, NDArrays, ResultsList
+
+
+class FloxControllerLogic:
+    """Abstract base class for FLoX Controller logic"""
+
+    def on_model_init(self) -> None:
+        """Does initial Controller setup before running the main Federated Learning loop"""
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_model_broadcast(self) -> FuturesList:
+        """Sends the model and config to endpoints for FL training.
+
+        Returns
+        -------
+        FuturesList
+            A list of tasks/futures with results of the FL training returned from the endpoints.
+            If using FuncXExecutor, this would most likely be a list of futures
+            funcX returns after you submit functions to endpoints.
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_model_receive(self, ins: FuturesList) -> ResultsList:
+        """Processes returned tasks from on_model_broadcast.
+
+        Parameters
+        ----------
+        FuturesList
+            A list of tasks/futures with results of the FL training returned from the endpoints.
+            If using FuncXExecutor, this would most likely be a list of futures
+            funcX returns after you submit functions to endpoints.
+
+        Returns
+        -------
+        ResultsList
+            FL results extracted from FuturesList, formatted as a dictionary. For example:
+            ResultsRes = {
+                "model_weights": model_weights,
+                "samples_count": samples_count,
+                "bias_weights": fractions,
+            }
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_model_aggregate(self, ins: ResultsList) -> NDArrays:
+        """Aggregates weights.
+
+        Parameters
+        ----------
+        ResultsList
+            FL results extracted from FuturesList, formatted as a dictionary. For example:
+            ResultsRes = {
+                "model_weights": model_weights,
+                "samples_count": samples_count,
+                "bias_weights": fractions,
+            }
+
+        Returns
+        -------
+        NDArrays
+            ML model weights for a single model in the form of Numpy Arrays.
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_model_update(self, weights: NDArrays) -> None:
+        """Updates the model's weights with new weights
+
+        Parameters
+        ----------
+        weights: NDArrays
+            ML model weights for a single model in the form of Numpy Arrays.
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def on_model_evaluate(self, test_data, model=None) -> EvaluateRes:
+        """Evaluates the given model using test_data
+
+        Parameters
+        ----------
+        test_data
+            data for evaluating the model. This can take different forms depending on the
+            ML framework you use. For Tensorflow, it would look as x_test and y_test, while
+            for PyTorch, it would look as a single DataLoader instance.
+
+        model
+            The machine learning model for evaluation. Whether you pass it as a parameter
+            would depend on the ML framework. With Tensorflow, we cannot save and transfer
+            the model as a class attribute, while for PyTorch it was possible so we do not need
+            to pass it as a paramete.
+
+        Returns
+        -------
+        EvaluateRes
+            A class consisting of two attributes showing the loss and metrics of the evaluated
+            model on the dataset:
+                loss: float
+                metrics: Dict[str, Scalar]
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
diff --git a/flox/logic/base_model_trainer.py b/flox/logic/base_model_trainer.py
new file mode 100644
index 0000000..cd1dd42
--- /dev/null
+++ b/flox/logic/base_model_trainer.py
@@ -0,0 +1,44 @@
+"""Abstract Base Class for FLoX Model Trainers"""
+from flox.common import NDArrays
+
+
+class BaseModelTrainer:
+    """Abstract base class for FLoX ML model trainers"""
+
+    def fit(self) -> NDArrays:
+        """Fits the model using training data
+
+        Returns
+        ----------
+        NDArrays
+            the new model weights in the Numpy array form
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def evaluate(self):
+        """Evaluates the model using testing data"""
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def get_weights(self) -> NDArrays:
+        """Returns the weights of the model as a Numpy array
+
+
+        Returns
+        -------
+        NDArrays
+            the model's weights as Numpy arrays
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
+
+    def set_weights(self, new_weights: NDArrays) -> None:
+        """Sets the model's weights to the new weights
+
+        Parameters
+        ----------
+        new_weights: NDArrays
+            new model weights in the Numpy array form
+
+        """
+        raise NotImplementedError("Abstract class method. Cannot be called directly.")
diff --git a/flox/model_trainers/PyTorchTrainer.py b/flox/model_trainers/PyTorchTrainer.py
new file mode 100644
index 0000000..dcadef4
--- /dev/null
+++ b/flox/model_trainers/PyTorchTrainer.py
@@ -0,0 +1,91 @@
+"""PyTorch ML ModelTrainer Class"""
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from flox.common import EvaluateRes, NDArrays
+from flox.logic import BaseModelTrainer
+
+
+class PyTorchTrainer(BaseModelTrainer):
+    """PyTorch ML ModelTrainer Class"""
+
+    def __init__(
+        self, model, device=None, criterion=nn.CrossEntropyLoss(), optimizer=None
+    ) -> None:
+        self.model = model
+        if not device:
+            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        else:
+            self.device = device
+
+        self.criterion = criterion
+        if optimizer:
+            self.optimizer = optimizer
+        else:
+            self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
+
+    def fit(self, trainloader, config):
+        epochs = config["epochs"]
+        for epoch in range(epochs):  # loop over the dataset multiple times
+            running_loss = 0.0
+            for i, data in enumerate(trainloader, 0):
+                images, labels = data[0].to(self.device), data[1].to(self.device)
+
+                # zero the parameter gradients
+                self.optimizer.zero_grad()
+
+                # forward + backward + optimize
+                outputs = self.model(images)
+                loss = self.criterion(outputs, labels)
+                loss.backward()
+                self.optimizer.step()
+
+                # print statistics
+                running_loss += loss.item()
+                if i % 2000 == 1999:  # print every 2000 mini-batches
+                    print(
+                        "[%d, %5d] loss: %.3f" % (epoch + 1, i + 1, running_loss / 2000)
+                    )
+                    running_loss = 0.0
+
+        return self.get_weights()
+
+    def evaluate(self, testloader):
+        import torch
+
+        correct = 0
+        total = 0
+        loss = 0.0
+        with torch.no_grad():
+            for data in testloader:
+                images, labels = data[0].to(self.device), data[1].to(self.device)
+                outputs = self.model(images)
+                loss += self.criterion(outputs, labels).item()
+                _, predicted = torch.max(outputs.data, 1)  # pylint: disable=no-member
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+        accuracy = correct / total
+        return EvaluateRes(loss=float(loss), metrics={"accuracy": float(accuracy)})
+
+    def get_weights(self) -> NDArrays:
+        """Get model weights as a list of NumPy ndarrays."""
+        return [val.cpu().numpy() for _, val in self.model.state_dict().items()]
+
+    def set_weights(self, weights: NDArrays) -> None:
+        import torch
+
+        state_dict = OrderedDict(
+            {
+                k: torch.tensor(v)
+                for k, v in zip(self.model.state_dict().keys(), weights)
+            }
+        )
+        self.model.load_state_dict(state_dict, strict=True)
+
+    # def create_dataloader(self, x_data, y_data, batch_size=4, shuffle=True) -> torch.utils.data.DataLoader:
+    #     combined_data = torch.hstack((x_data, y_data))
+    #     data_loader = torch.utils.data.DataLoader(combined_data, batch_size=batch_size, shuffle=shuffle)
+    #     return data_loader
diff --git a/flox/model_trainers/TensorflowTrainer.py b/flox/model_trainers/TensorflowTrainer.py
new file mode 100644
index 0000000..1cfc3be
--- /dev/null
+++ b/flox/model_trainers/TensorflowTrainer.py
@@ -0,0 +1,61 @@
+"""TensorFlow ML ModelTrainer Class"""
+from tensorflow.keras.models import Sequential
+
+from flox.common import EvaluateRes, NDArrays
+from flox.logic import BaseModelTrainer
+
+
+class TensorflowTrainer(BaseModelTrainer):
+    """TensorFlow ML ModelTrainer Class"""
+
+    def __init__(
+        self,
+        loss="categorical_crossentropy",
+        optimizer="adam",
+        metrics=["accuracy"],
+    ) -> None:
+        self.loss = loss
+        self.optimizer = optimizer
+        self.metrics = metrics
+
+    def get_weights(self, model: Sequential) -> NDArrays:
+        import numpy as np
+
+        model_weights = model.get_weights()
+        model_weights_numpy = np.asarray(model_weights, dtype=object)
+        return model_weights_numpy
+
+    def set_weights(self, model: Sequential, weights: NDArrays) -> None:
+        model.set_weights(weights)
+
+    def fit(
+        self, model: Sequential, x_train: NDArrays, y_train: NDArrays, epochs: int = 10
+    ) -> None:
+        model.fit(x_train, y_train, epochs=epochs)
+
+    def evaluate(
+        self, model: Sequential, x_test: NDArrays, y_test: NDArrays
+    ) -> EvaluateRes:
+        scores = model.evaluate(x_test, y_test, verbose=0)
+        loss, accuracy = scores[0], scores[1]
+        return EvaluateRes(loss=float(loss), metrics={"accuracy": float(accuracy)})
+
+    def get_architecture(self, model: Sequential) -> object:
+        """Returns Tensorflow's model architecture as a JSON file"""
+        return model.to_json()
+
+    def create_model(self, architecture) -> Sequential:
+        """Creates a keras.Sequential model from a JSON architecture file"""
+        from tensorflow import keras
+
+        model = keras.models.model_from_json(architecture)
+        return model
+
+    def build_model(self, model: Sequential, input_shape) -> None:
+        """Builds a keras.Sequential model, requires input_shape. Some older Tensorflow versions
+        require the model to be built before we can set new weights"""
+        model.build(input_shape=input_shape)
+
+    def compile_model(self, model: Sequential) -> None:
+        "Compiles keras.Sequential model"
+        model.compile(loss=self.loss, optimizer=self.optimizer, metrics=self.metrics)
diff --git a/flox/model_trainers/__init__.py b/flox/model_trainers/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pyflox.egg-info/PKG-INFO b/pyflox.egg-info/PKG-INFO
deleted file mode 100644
index 7506c0d..0000000
--- a/pyflox.egg-info/PKG-INFO
+++ /dev/null
@@ -1,65 +0,0 @@
-Metadata-Version: 2.1
-Name: pyflox
-Version: 0.1.5
-Summary: Library for serverless Federated Learning experiments.
-Home-page: https://github.com/nikita-kotsehub/FLoX
-Author: Nikita Kotsehub
-Author-email: mykyta.kotsehub@gmail.com
-License: MIT
-Download-URL: https://github.com/nikita-kotsehub/FLoX/archive/refs/tags/v0.1.3-test.tar.gz
-Keywords: federated_learning,serverless,edge_devices
-Platform: UNKNOWN
-Classifier: Development Status :: 3 - Alpha
-Classifier: Intended Audience :: Science/Research
-Classifier: Topic :: Software Development :: Build Tools
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3.0
-Classifier: Programming Language :: Python :: 3.4
-Classifier: Programming Language :: Python :: 3.5
-Classifier: Programming Language :: Python :: 3.6
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Description-Content-Type: text/markdown
-License-File: LICENSE.TXT
-
-# FLoX
-
-FLoX (**F**ederated **L**earning on func**X**) is a Python library for serverless Federated Learning experiments.
-
-This is initial documentation that will be soon expanded. 
-
-## Installation
-
-Use the package manager [pip](https://pip.pypa.io/en/stable/) to install flox.
-
-```bash
-# seems like flox is taken, will update the name soon
-pip install flox
-```
-
-## Usage
-For a full example, see this [Google Colab tutorial](https://colab.research.google.com/drive/19X1N8E5adUrmeE10Srs1hSQqCCecv23m?usp=sharing).
-
-```python
-from flox.flox import federated_learning
-
-# performs 5 rounds of Federated Learning train global_model on given endpoint_ids
-# uses 10 epochs and 100 samples from fashion_mnist dataset for training
-federated_learning(global_model=global_model, 
-                  endpoint_ids=endpoint_ids,
-                  loops=5,
-                  epochs=10,
-                  keras_dataset="fashion_mnist", 
-                  num_samples=100, 
-                  input_shape=(32, 28, 28, 1))
-```
-
-## Contributing
-Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
-
-## License
-[MIT](https://choosealicense.com/licenses/mit/)
-
diff --git a/pyflox.egg-info/SOURCES.txt b/pyflox.egg-info/SOURCES.txt
deleted file mode 100644
index 17b58a2..0000000
--- a/pyflox.egg-info/SOURCES.txt
+++ /dev/null
@@ -1,12 +0,0 @@
-LICENSE.TXT
-README.md
-setup.cfg
-setup.py
-flox/__init__.py
-flox/flox.py
-flox/utils.py
-pyflox.egg-info/PKG-INFO
-pyflox.egg-info/SOURCES.txt
-pyflox.egg-info/dependency_links.txt
-pyflox.egg-info/requires.txt
-pyflox.egg-info/top_level.txt
\ No newline at end of file
diff --git a/pyflox.egg-info/dependency_links.txt b/pyflox.egg-info/dependency_links.txt
deleted file mode 100644
index 8b13789..0000000
--- a/pyflox.egg-info/dependency_links.txt
+++ /dev/null
@@ -1,1 +0,0 @@
-
diff --git a/pyflox.egg-info/requires.txt b/pyflox.egg-info/requires.txt
deleted file mode 100644
index 6558555..0000000
--- a/pyflox.egg-info/requires.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-numpy
-funcx
-parsl
diff --git a/pyflox.egg-info/top_level.txt b/pyflox.egg-info/top_level.txt
deleted file mode 100644
index e8a6ce2..0000000
--- a/pyflox.egg-info/top_level.txt
+++ /dev/null
@@ -1,1 +0,0 @@
-flox
diff --git a/setup.py b/setup.py
index 455d9aa..24005e1 100644
--- a/setup.py
+++ b/setup.py
@@ -4,14 +4,14 @@
 from setuptools import find_packages, setup
 
 this_directory = Path(__file__).parent
-long_description = (this_directory / "README.md").read_text()
+long_description = (this_directory / "README.rst").read_text()
 
 setup(
     name="pyflox",
-    packages=find_packages(include=["flox"]),
+    packages=find_packages(include=["flox", "flox.*"]),
     version="0.1.5",
     description="Library for serverless Federated Learning experiments.",
-    readme="README.md",
+    readme="README.rst",
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/nikita-kotsehub/FLoX",


Test patch:
diff --git a/flox/tests/conftest.py b/flox/tests/conftest.py
new file mode 100644
index 0000000..aa3283d
--- /dev/null
+++ b/flox/tests/conftest.py
@@ -0,0 +1,63 @@
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+import tensorflow as tf
+from funcx import FuncXExecutor
+
+from flox.clients.TensorflowClient import TensorflowClient
+from flox.controllers.TensorflowController import TensorflowController
+from flox.model_trainers.TensorflowTrainer import TensorflowTrainer
+from flox.utils import get_test_data
+
+
+@pytest.fixture
+def tf_controller():
+    # `fashion_mnist` images are grayscale, 28 x 28 pixels in size
+    input_shape = (28, 28, 1)
+    # there are 10 classes in the dataset
+    num_classes = 10
+
+    # define the model architecture
+    global_model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=input_shape),
+            tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(num_classes, activation="softmax"),
+        ]
+    )
+
+    # compile the model
+    global_model.compile(
+        loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
+    )
+    x_test, y_test = get_test_data(keras_dataset="fashion_mnist", num_samples=2000)
+
+    tf_trainer = TensorflowTrainer()
+    tf_client_logic = TensorflowClient()
+
+    eps = ["fake_endpoint_1", "fake_endpoint_2", "fake_endpoint_3"]
+
+    tf_controller = TensorflowController(
+        endpoint_ids=eps,
+        num_samples=200,
+        epochs=2,
+        rounds=2,
+        client_logic=tf_client_logic,
+        global_model=global_model,
+        executor_type="local",  # choose "funcx" for FuncXExecutor, "local" for ThreadPoolExecutor
+        model_trainer=tf_trainer,
+        path_dir=".",
+        x_test=x_test,
+        y_test=y_test,
+        data_source="keras",
+        dataset_name="fashion_mnist",
+        preprocess=True,
+    )
+
+    return tf_controller
diff --git a/flox/tests/integration/test_tf_local_flow.py b/flox/tests/integration/test_tf_local_flow.py
new file mode 100644
index 0000000..6027e13
--- /dev/null
+++ b/flox/tests/integration/test_tf_local_flow.py
@@ -0,0 +1,11 @@
+def test_broadcast_local(tf_controller):
+    """The simple federated learning process runs without throwing errors"""
+    tf_controller.run_federated_learning()
+    assert True
+
+
+def test_broadcast_local_running_average(tf_controller):
+    """The running_average version of federated learning process runs without throwing errors"""
+    tf_controller.running_average = True
+    tf_controller.run_federated_learning()
+    assert True
diff --git a/flox/tests/unit/test_main_controller.py b/flox/tests/unit/test_main_controller.py
new file mode 100644
index 0000000..8481e15
--- /dev/null
+++ b/flox/tests/unit/test_main_controller.py
@@ -0,0 +1,94 @@
+from concurrent.futures import ThreadPoolExecutor
+
+from funcx import FuncXExecutor
+
+
+def test_init_num_samples_int(tf_controller):
+    """MainController.on_model_init() correctly initiates num_samples
+    when only a single integer is provided"""
+    samples_int = 1000
+    tf_controller.num_samples = samples_int
+    tf_controller.on_model_init()
+    assert tf_controller.num_samples == [1000, 1000, 1000]
+
+
+def test_init_num_samples_list(tf_controller):
+    """MainController.on_model_init() correctly initiates num_samples
+    when a list of integers is provided"""
+
+    samples_list = [200, 500, 1000]
+    tf_controller.num_samples = samples_list
+    tf_controller.on_model_init()
+    assert tf_controller.num_samples == samples_list
+
+
+def test_init_epochs_int(tf_controller):
+    """MainController.on_model_init() correctly initiates number of epochs
+    when only a single integer is provided"""
+    epoch_int = 5
+    tf_controller.epochs = epoch_int
+    tf_controller.on_model_init()
+    assert tf_controller.epochs == [5, 5, 5]
+
+
+def test_init_epochs_list(tf_controller):
+    """MainController.on_model_init() correctly initiates number of epochs
+    when a list of integers is provided"""
+
+    epochs_list = [5, 10, 15]
+    tf_controller.epochs = epochs_list
+    tf_controller.on_model_init()
+    assert tf_controller.epochs == epochs_list
+
+
+def test_init_path_str(tf_controller):
+    """MainController.on_model_init() correctly initiates the data directory paths
+    when only a single string is provided"""
+    path_dir = "."
+    tf_controller.path_dir = path_dir
+    tf_controller.on_model_init()
+    assert tf_controller.path_dir == [".", ".", "."]
+
+
+def test_init_path_list(tf_controller):
+    """MainController.on_model_init() correctly initiates the data directory paths
+    when a list of strings is provided"""
+
+    path_dir_list = [".", "./hello/", "./world/"]
+    tf_controller.path_dir = path_dir_list
+    tf_controller.on_model_init()
+    assert tf_controller.path_dir == path_dir_list
+
+
+def test_init_executor_default(tf_controller):
+    """MainController.on_model_init() correctly initiates the executor when the executor
+    and executor_type are not provided"""
+    tf_controller.executor = None
+
+    tf_controller.on_model_init()
+
+    assert tf_controller.executor_type == "local"
+    assert tf_controller.executor == ThreadPoolExecutor
+
+
+def test_init_executor_provided_type(tf_controller):
+    """MainController.on_model_init() correctly initiates the executor when executor_type is provided
+    but no executor is provided."""
+    tf_controller.executor = None
+    tf_controller.executor_type = "funcx"
+
+    tf_controller.on_model_init()
+
+    assert tf_controller.executor_type == "funcx"
+    assert tf_controller.executor == FuncXExecutor
+
+
+def test_init_executor_provided_executor(tf_controller):
+    """MainController.on_model_init() correctly initiates the executor when a custom executor is provided."""
+    tf_controller.executor = "fake_executor"
+    tf_controller.executor_type = "local"
+
+    tf_controller.on_model_init()
+
+    assert tf_controller.executor_type == "local"
+    assert tf_controller.executor == "fake_executor"