Tutorial: Distributed Compute with Arrays

arrays
tutorials
python
scalable compute
Task graphs are a powerful feature of TileDB, allowing you to distribute work across numerous machines in a parallel, serverless manner.
How to run this tutorial

You can run this tutorial only on TileDB Cloud. However, TileDB Cloud has a free tier. We strongly recommend that you sign up and run everything there, as that requires no installations or deployment.

This tutorial shows distributed compute on TileDB Cloud using task graphs. It assumes you have already completed the Get Started section. You will learn how to ingest data and perform array computations, all at scale in a distributed (i.e., parallel across multiple machines) fashion. For more information on distributed compute, visit the Key Concepts: Distributed Compute section.

Distributed array ingestion

In this section, you will learn how to write to an array in parallel with a task graph containing four write nodes, each writing to a different partition of the array.

First, import the necessary libraries, set the array URI for TileDB Cloud, and delete any previously created arrays with the same name.

# Import necessary libraries
import os.path

import numpy as np
import tiledb
import tiledb.cloud
from tiledb.cloud.compute import Delayed, DelayedArrayUDF

# You should set the appropriate environment variables with your keys.
# Get the keys from the environment variables.
tiledb_token = os.environ["TILEDB_REST_TOKEN"]
tiledb_account = os.environ["TILEDB_ACCOUNT"]

# Get the bucket and region from environment variables
s3_bucket = os.environ["S3_BUCKET"]

# Loging with your TileDB Cloud token
tiledb.cloud.login(token=tiledb_token)

# Set array URI
array_name = "distributed_computing"
array_uri = "tiledb://" + tiledb_account + "/" + array_name

# The following context will carry the TileDB Cloud credentials
cloud_ctx = tiledb.cloud.Ctx()

# Delete array, if it already exists
with tiledb.scope_ctx(cloud_ctx):
    if tiledb.array_exists(array_uri):
        tiledb.Array.delete_array(array_uri)

Next, create a dense array by specifying its schema (the case of sparse arrays is similar). The only difference between TileDB Cloud and TileDB Open-source when creating and registering arrays is that the TileDB Cloud URI should be of the form: tiledb://<account>/s3://<bucket>/<array_name>. TileDB Cloud understands that you are trying to create an array in S3 URI s3://<bucket>/<array_name> and register it under <account>. After you create and register the array, you can access the array as tiledb://<account>/<array_name>.

# Create the two dimensions
d1 = tiledb.Dim(name="d1", domain=(1, 4), tile=2, dtype=np.int32)
d2 = tiledb.Dim(name="d2", domain=(1, 4), tile=2, dtype=np.int32)

# Create a domain using the two dimensions
dom = tiledb.Domain(d1, d2)

# Create an attribute
a = tiledb.Attr(name="a", dtype=np.int32)

# Create the array schema, setting `sparse=False` to indicate a dense array.
sch = tiledb.ArraySchema(domain=dom, sparse=False, attrs=[a])

# Create and register the array on TileDB Cloud
array_uri_reg = "tiledb://" + tiledb_account + "/" + s3_bucket + "/" + array_name
tiledb.Array.create(array_uri_reg, sch, ctx=cloud_ctx)

Subsequently, you need a function that takes as input an array URI, a slice, and some data, and writes the data to the corresponding array slice. This will be called by each of the task graph nodes with different slice and data arguments to perform the entire array ingestion in parallel.

# Create a generic UDF that writes the input data to
# the input slice of the array
def write_slice(array_uri, slice, data):
    with tiledb.open(array_uri, mode="w", ctx=tiledb.cloud.Ctx()) as A:
        A[slice[0], slice[1]] = {"a": data}

You are now ready to create and run the task graph that will perform the parallel ingestion. There will be four nodes that perform the writes in parallel, all connected to a sink node that will run an empty function and complete the graph execution.

# Write the following data to the array.
#    [1, 2, 3, 4],
#    [5, 6, 7, 8],
#    [9, 10, 11, 12],
#    [13, 14, 15, 16]
#
# The following code will create a task graph
# with 4 nodes/workers, each writing to the
# array a part of the array, namely:

# 1. Write [[1, 2], [5,6]] in subarray [1,2], [1,2]
write1 = Delayed(write_slice, name="write_1")(
    array_uri, [slice(1, 3), slice(1, 3)], np.array([[1, 2], [5, 6]], dtype=np.int32)
)

# 2. Write [[3, 4], [7,8]] in subarray [1,2], [3,4]
write2 = Delayed(write_slice, name="write_2")(
    array_uri, [slice(1, 3), slice(3, 5)], np.array([[3, 4], [7, 8]], dtype=np.int32)
)

# 3. Write [[9, 10], [13,14]] in subarray [3,4], [1,2]
write3 = Delayed(write_slice, name="write_3")(
    array_uri, [slice(3, 5), slice(1, 3)], np.array([[9, 10], [13, 14]], dtype=np.int32)
)

# 4. Write [[11, 12], [15,16]] in subarray [3,4], [3,4]
write4 = Delayed(write_slice, name="write_4")(
    array_uri,
    [slice(3, 5), slice(3, 5)],
    np.array([[11, 12], [15, 16]], dtype=np.int32),
)


# Create a sink node with an empty function to connect all other nodes
def empty_func(write1, write2, write3, write4):
    pass


sink = Delayed(func_exec=empty_func, name="sink")(write1, write2, write3, write4)

# Visualize the graph and see it working in action
sink.visualize()

# Compute the whole task graph
sink.compute()

Read the entire array to verify that the parallel ingestion was successful.

# Read the entire array to make sure everything was
# written correctly
with tiledb.open(array_uri, ctx=tiledb.cloud.Ctx()) as A:
    print(A[:])
OrderedDict([('a', array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12],
       [13, 14, 15, 16]], dtype=int32))])

Distributed array processing

This section continues the same code example as above to perform parallel computations over the array. Specifically, you will compute the sum over all the elements of the array you ingested. You will again use a task graph with 4 nodes, each reading a different array slice and computing a partial sum. These four nodes will be connected to a sink node that will receive the four partial results, and add them together to produce the final sum over the entire array.

First, you need to create a function that will act as an array user-defined function (UDF); visit the Use User-Defined Functions with Array Data tutorial for dedicated examples on array UDFs.

# Create an array UDF that sums the values in a given slice of the array
def sum_slice(data):
    return data["a"].sum()

Next, create the task graph. Notice that this time the task graph consists of array UDFs, as opposed to parallel ingestion that used generic UDFs in each graph node.

# The following code will compute the sum of all values
# in the array using a task graph with 4 nodes/workers.
# It is similar to ingestion, in that each node will
# focus on a subarray partition of the array.

# 1. Return sum of subarray [1,2], [1,2]
sum1 = DelayedArrayUDF(array_uri, sum_slice, name="sum_1")(
    [(1, 2), (1, 2)], attrs=["a"]
)

# 2. Write [[3, 4], [7,8]] in subarray [1,2], [3,4]
sum2 = DelayedArrayUDF(array_uri, sum_slice, name="sum_2")(
    [(1, 2), (3, 4)], attrs=["a"]
)

# 3. Write [[9, 10], [13,14]] in subarray [3,4], [1,2]
sum3 = DelayedArrayUDF(array_uri, sum_slice, name="sum_3")(
    [(3, 4), (1, 2)], attrs=["a"]
)

# 4. Write [[11, 12], [15,16]] in subarray [3,4], [3,4]
sum4 = DelayedArrayUDF(array_uri, sum_slice, name="sum_4")(
    [(3, 4), (3, 4)], attrs=["a"]
)


# Create a sink node with sum function to connect all other nodes
def sum_all(sum1, sum2, sum3, sum4):
    return sum1 + sum2 + sum3 + sum4


sink = Delayed(func_exec=sum_all, name="sink")(sum1, sum2, sum3, sum4)

# Visualize the graph and see it working in action
sink.visualize()

# Compute the whole task graph
sink.compute()

Clean up by removing the created array.

# Delete array, if it already exists
with tiledb.scope_ctx(cloud_ctx):
    if tiledb.array_exists(array_uri):
        tiledb.Array.delete_array(array_uri)