How-To: Distribute Python Functions Across DFM Sites#

This guide shows you how to take existing Python functions and distribute them across multiple DFM sites, then execute a pipeline that uses these distributed functions.

In this first tutorial, you will use the local DFM target to run these functions in different processes as a model of a fully distributed federation.

Scenario#

Assume you have a Python script called plot_gradient.py that contains three functions:

  • create_gradient(shape, min, max) -> numpy.ndarray - Creates a NumPy array of given shape

  • subset(array, index) -> numpy.ndarray - Extracts a slice from an array given array indices

  • plot2D(array) -> PIL.Image - Converts a 2D array to a grayscale image

It should look something like this:

import numpy as np
from PIL import Image


def create_gradient(shape: tuple[int, ...], min: float = 0.0, max: float = 1.0) -> np.ndarray:
    ndims = len(shape)
    grads = [np.arange(s) for s in shape]
    grad = sum(g[tuple(slice(None) if i == j else np.newaxis for i in range(ndims))] / (shape[j] - 1) for j, g in enumerate(grads)) / ndims
    return np.asarray((max - min) * grad + min)


def subset(array: np.ndarray, index: list[int, ...]) -> np.ndarray:
    return array[tuple(index)]


def plot2D(array: np.ndarray) -> Image.Image:
    if len(np.shape(array)) != 2:
        raise ValueError("Array must be 2D to be plotted")
    normalized = ((array - array.min()) / (array.max() - array.min()) * 255).astype(np.uint8)
    return Image.fromarray(normalized, mode='L')


if __name__ == '__main__':
    gradient3D = create_gradient((200, 300, 400), min=-100, max=100)
    gradient2D = subset(gradient3D, [100])
    image = plot2D(gradient2D)
    image.save('gradient.jpg', quality=90)

Try this script on your system. It should produce a file called gradient.jpg that looks like this:

gradient

Goal: Distribute these functions across multiple DFM sites and run them as a pipeline from a Jupyter notebook.

This How-To is divided into the following parts:

  1. Setting up a basic Data Federation and testing it locally

  2. Setting your federation to use NVIDIA Flare and testing with POC mode