<jupyter_start><jupyter_text>Customize your own GPU Kernels in gQuantThe gQuant is designed to accelerate quantitive finance workflows on the GPU. The acceleration on GPU is facilitated by using cuDF dataframes organized into a computation graph. The cuDF project is a continously evolving library that provides a pandas-like API. Sometimes the data scientists are facing a few challenges that cannot be easily solved:    1. The quantitative work needs customized logic to manipulate the data, and there are no direct methods within cuDF to support this logic.    2. Each cuDF dataframe method call launches the GPU kernel once. For performance crtical task, it is sometimes required to wrap lots of computation steps together in a single GPU kernel to reduce the kernel launch overheads.The solution is to build customized GPU kernels to implement them. The code and examples below illustrate a variety of approaches to implement customized GPU kernels in Python.<jupyter_code>import sys; sys.path.insert(0, '..')
# Load necessary Python modules
import sys
from gquant.dataframe_flow import TaskSpecSchema, TaskGraph, MetaData
from gquant.dataframe_flow import Node, NodePorts, PortsSpecSchema
from gquant.dataframe_flow import ConfSchema
import cudf
import numpy as np
from numba import cuda
import cupy
import math
import dask
import dask_cudf<jupyter_output><empty_output><jupyter_text>Define a utility function to verify the results:<jupyter_code>def verify(ground_truth, computed):
    max_difference = (ground_truth - computed).abs().max()
    # print('Max Difference: {}'.format(max_difference))
    assert(max_difference < 1e-8)
    return max_difference
task_graph = TaskGraph()<jupyter_output><empty_output><jupyter_text>Example Problem: Calculating the distance of points to the originThe sample problem is to take a list of points in 2-D space and compute their distance to the origin.We start by creating a source `Node` in the graph that generates a cuDF dataframe containing some configurable number of random points. A custom node is defined by inheriting from the `Node` class and overriding methods `init`, `meta_setup`, `ports_setup`, `conf_schema`, `process`.The `ports_setup` must return an instance of `NodePorts` which encapsulates the ports specs. Ports specs are dictionaries with port attributes/options per `PortsSpecSchema`.In the case of the `PointNode` below the input port is an empty dictionary, since no inputs are required, and the output has two ports `points_df_out` and `points_ddf_out`. It can output two types of dataframe frames depends who connects it.The `process` method receives a input dictionary where keys are input ports and values are input data. It return a dictionary where the keys correspond to the output ports. The `meta_setup` is used to compute the output meta information. It returns a dictionary where keys correspond to the output ports.The `conf_schema` is used to define the Node configuration [JSON schema](https://json-schema.org/). gQuantlab UI uses [RJSF](https://github.com/rjsf-team/react-jsonschema-form) project to generate HTML form elements based on the JSON schema. [RJSF playground](https://rjsf-team.github.io/react-jsonschema-form/)  is a good place to learn how to write JSON schema and visualize it. `conf_schema` returns `ConfSchema` which encapsulate the JSON schema and UI schema together.The `column` and `port_types` information sometimes are determined dynamically. gQuant provides a few utility functions to help get dynamical graph information. `self.get_connected_inports()` will return a dictionay where keys are connected inport names and values are inport types. `self.get_input_meta()` will return a dictionary where keys are connected inport names and values are column name/type paris from the parent node. `self.outport_connected(port_name)` method returns a boolean if the output port `port_name` is connected. The `PointNode` uses it to determine what kind of computation it needs to do depending on the connection.<jupyter_code>class PointNode(Node):

    def ports_setup(self):
        input_ports = {}
        output_ports = {
            'points_df_out': {
                PortsSpecSchema.port_type: cudf.DataFrame
            },
            'points_ddf_out': {
                PortsSpecSchema.port_type: dask_cudf.DataFrame
            },
        }
        return NodePorts(inports=input_ports, outports=output_ports)

    def conf_schema(self):
        json = {
            "title": "PointNode configure",
            "type": "object",
            "properties": {
                "npts":  {
                    "type": "number",
                    "description": "number of data points",
                    "minimum": 10
                },
                "npartitions":  {
                    "type": "number",
                    "description": "num of partitions in the Dask dataframe",
                    "minimum": 1
                }

            },
            "required": ["npts", "npartitions"],
        }

        ui = {
            "npts": {"ui:widget": "updown"},
            "npartitions": {"ui:widget": "updown"}
        }
        return ConfSchema(json=json, ui=ui)

    def init(self):
        pass
        
    def meta_setup(self):
        columns_out = {
            'points_df_out': {
                'x': 'float64',
                'y': 'float64'
            },
            'points_ddf_out': {
                'x': 'float64',
                'y': 'float64'
            }
        }
        return MetaData(inports={}, outports=columns_out)

    def process(self, inputs):
        npts = self.conf['npts']
        df = cudf.DataFrame()
        df['x'] = np.random.rand(npts)
        df['y'] = np.random.rand(npts)
        output = {}
        if self.outport_connected('points_df_out'):
            output.update({'points_df_out': df})
        if self.outport_connected('points_ddf_out'):
            npartitions = self.conf['npartitions']
            ddf = dask_cudf.from_cudf(df, npartitions=npartitions)
            output.update({'points_ddf_out': ddf})
        return output<jupyter_output><empty_output><jupyter_text>The distance can be computed via cuDF methods. We define the `DistanceNode` to calculate the euclidean distance and add a `distance_cudf` column to the output dataframe. We will use that as the ground truth to compare and verify results later. Additionally, the distance node calculates absolute distance (Manhattan distance) in another output port. The compuation is done depending which output is connected.<jupyter_code>class DistanceNode(Node):

    def ports_setup(self):
        port_type = PortsSpecSchema.port_type
        input_ports = {
            'points_df_in': {
                port_type: [cudf.DataFrame, dask_cudf.DataFrame]
            }
        }

        output_ports = {
            'distance_df': {
                port_type: [cudf.DataFrame, dask_cudf.DataFrame]
            },
            'distance_abs_df': {
                PortsSpecSchema.port_type:  [cudf.DataFrame, dask_cudf.DataFrame]
            }
        }
        input_connections = self.get_connected_inports()
        if 'points_df_in' in input_connections:
            types = input_connections['points_df_in']
            # connected, use the types passed in from parent
            return NodePorts(inports={'points_df_in': {port_type: types}},
                             outports={'distance_df': {port_type: types},
                                       'distance_abs_df': {port_type: types},
                                       })
        else:
            return NodePorts(inports=input_ports, outports=output_ports)

    def conf_schema(self):
        return ConfSchema()

    def init(self):
        self.delayed_process = True


    def meta_setup(self):
        req_cols = {
            'x': 'float64',
            'y': 'float64'
        }
        required = {
            'points_df_in': req_cols,
        }
        input_meta = self.get_input_meta()
        output_cols = ({
                'distance_df': {
                    'distance_cudf': 'float64',
                    'x': 'float64',
                    'y': 'float64'
                },
                'distance_abs_df': {
                    'distance_abs_cudf': 'float64',
                    'x': 'float64',
                    'y': 'float64'
                }
            })
        if 'points_df_in' in input_meta:
            col_from_inport = input_meta['points_df_in']
            # additional ports
            output_cols['distance_df'].update(col_from_inport)
            output_cols['distance_abs_df'].update(col_from_inport)
        return MetaData(inports=required, outports=output_cols)

    def process(self, inputs):
        df = inputs['points_df_in']
        output = {}
        if self.outport_connected('distance_df'):
            copy_df = df.copy()
            copy_df['distance_cudf'] = (df['x'] ** 2 + df['y'] ** 2).sqrt()
            output.update({'distance_df': copy_df})
        if self.outport_connected('distance_abs_df'):
            copy_df = df.copy()
            copy_df['distance_abs_cudf'] = df['x'].abs() + df['y'].abs()
            output.update({'distance_abs_df': copy_df})
        return output<jupyter_output><empty_output><jupyter_text>Having these two nodes, we can construct a simple task graph to compute the distance.<jupyter_code># Task specifications.
module_name = 'custom_nodes'

points_tspec = {
    TaskSpecSchema.task_id: 'points_task',
    TaskSpecSchema.node_type: PointNode,
    TaskSpecSchema.conf: {'npts': 1000},
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {},
}

cudf_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_cudf',
    TaskSpecSchema.node_type: DistanceNode,
    TaskSpecSchema.conf: {},
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'points_df_in': 'points_task.points_df_out'
    }
}

out_spec = {
    TaskSpecSchema.task_id: '',
    TaskSpecSchema.node_type: "Output_Collector",
    TaskSpecSchema.conf: {},
    TaskSpecSchema.inputs: {
        'in0': 'points_task.points_df_out',
        'in1': 'distance_by_cudf.distance_df',
        'in2': 'distance_by_cudf.distance_abs_df'
    }
}

task_list = [points_tspec, cudf_distance_tspec, out_spec]
task_graph = TaskGraph(task_list)<jupyter_output><empty_output><jupyter_text>We can visualize the task graph with and without ports.<jupyter_code>print('WITHOUT PORTS')
task_graph.draw(show='ipynb')
print('WITH PORTS')
task_graph.draw(show='ipynb', show_ports=True)<jupyter_output>WITH PORTS<jupyter_text>We can draw the graph in an interactive widget. First, let's register the dynamically defined gQuant nodes so the client knows about them. Note, this step is only needed if we would like to interact with gQuant by Jupyterlab UI.<jupyter_code>TaskGraph.register_lab_node(module_name, PointNode)
TaskGraph.register_lab_node(module_name, DistanceNode)<jupyter_output><empty_output><jupyter_text>Draw the widget:<jupyter_code>task_graph.draw()<jupyter_output><empty_output><jupyter_text>The next step is to run the task graph to obtain the distances. The output is identified by the `id` of the distance node:<jupyter_code>r = task_graph.run()
points_df = r['points_task.points_df_out']
task_graph.run(formated=True)<jupyter_output><jupyter_text>Customized Kernel with Numba libraryNumba is an excellent python library used for accelerating numerical computations. Numba supports CUDA GPU programming by directly compiling a restricted subset of Python code into CUDA kernels and device functions. The Numba GPU kernel is written in Python and translated (JIT just-in-time compiled) into GPU code at runtime. This is achieved by decorating a Python function with `@cuda.jit`.  Just like a C/C++ CUDA GPU kernel, the `distance_kernel` function is called by thousands of threads in the GPU. The thread id is computed by `threadIdx.x`, `blockId.x` and `blockDim.x` built-in variables. Please check the [CUDA programming guild](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.htmlthread-hierarchy) for details. A cuDF series can be converted to GPU arrays compatible with the Numba library via.to_gpu_array` API. The next step is to define a Node that calls this Numba kernel to compute the distance and save the result into `distance_numba` column in the output dataframe.<jupyter_code>import rmm
@cuda.jit
def distance_kernel(x, y, distance, array_len):
    # ii - overall thread index
    ii = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    if ii < array_len:
        distance[ii] = math.sqrt(x[ii] ** 2 + y[ii] ** 2)


class NumbaDistanceNode(Node):

    def ports_setup(self):
        port_type = PortsSpecSchema.port_type
        input_ports = {
            'points_df_in': {
                port_type: [cudf.DataFrame,
                            dask_cudf.DataFrame]
            }
        }

        output_ports = {
            'distance_df': {
                port_type: [cudf.DataFrame,
                            dask_cudf.DataFrame]
            }
        }

        input_connections = self.get_connected_inports()
        if 'points_df_in' in input_connections:
            types = input_connections['points_df_in']
            # connected
            return NodePorts(inports={'points_df_in': {port_type: types}},
                             outports={'distance_df': {port_type: types}})
        else:
            return NodePorts(inports=input_ports, outports=output_ports)
  
    def init(self):
        self.delayed_process = True


    def meta_setup(self,):
        required_cols = {'x': 'float64',
                    'y': 'float64'}
        required = {
            'points_df_in': required_cols,
            'distance_df': required_cols
        }
        input_meta = self.get_input_meta()
        output_cols = ({
                'distance_df': {
                    'distance_numba': 'float64',
                    'x': 'float64',
                    'y': 'float64'
                }
            })
        if 'points_df_in' in input_meta:
            col_from_inport = input_meta['points_df_in']
            # additional ports
            output_cols['distance_df'].update(col_from_inport)
        return MetaData(inports=required, outports=output_cols)

    def conf_schema(self):
        return ConfSchema()

    def process(self, inputs):
        df = inputs['points_df_in']
        number_of_threads = 16
        number_of_blocks = ((len(df) - 1) // number_of_threads) + 1
        # Inits device array by setting 0 for each index.
        # df['distance_numba'] = 0.0
        darr = rmm.device_array(len(df))
        distance_kernel[(number_of_blocks,), (number_of_threads,)](
            df['x'],
            df['y'],
            darr,
            len(df))
        df['distance_numba'] = darr
        return {'distance_df': df}<jupyter_output><empty_output><jupyter_text>The `self.delayed_process = True` flag in the `meta_setup` is necesary to enable the logic in the `Node` class for handling `dask_cudf` dataframes in order to use Dask (for distributed computation i.e. multi-gpu in examples later on). The `dask_cudf` dataframe does not support GPU customized kernels directly. The `to_delayed` and `from_delayed` low level interfaces of `dask_cudf` enable this support. The gQuant framework handles `dask_cudf` dataframes automatically under the hood when we set this flag.  Customized Kernel by CuPy libraryCuPy is an alternative to Numba. Numba JIT compiles Python code into GPU device code at runtime. There are some limitations in how Numba can be used as well as JIT compilation latency overhead. When a Python process calls a Numba GPU kernel for the first time Numba has to compile the Python code, and each time a new Python process is started the GPU kernel has to be recompiled. If advanced features of CUDA are needed and latency is important, CuPy is an alternative library that can be used to compile C/C++ CUDA code. CuPy caches the GPU device code on disk (default location `$(HOME)/.cupy/kernel_cache` which can be changed via `CUPY_CACHE_DIR` environment variable) thus eliminating compilation latency for subsequent Python processes.`CuPy` GPU kernel is esentially a C/C++ GPU kernel. Below we define the `compute_distance` kernel using `CuPy`: Using gQuant we can now define a Node that calls this CuPy kernel to compute the distance and save the results into `distance_cupy` column of a `cudf` dataframe.<jupyter_code>kernel_string = r'''
    extern "C" __global__
    void compute_distance(const double* x, const double* y,
            double* distance, int arr_len) {
        int tid = blockDim.x * blockIdx.x + threadIdx.x;
        if (tid < arr_len){
        distance[tid] = sqrt(x[tid]*x[tid] + y[tid]*y[tid]);
        }
    }
'''


class CupyDistanceNode(Node):

    def ports_setup(self):
        port_type = PortsSpecSchema.port_type
        input_ports = {
            'points_df_in': {
                port_type: [cudf.DataFrame,
                            dask_cudf.DataFrame]
            }
        }

        output_ports = {
            'distance_df': {
                port_type: [cudf.DataFrame,
                            dask_cudf.DataFrame]
            }
        }

        input_connections = self.get_connected_inports()
        if 'points_df_in' in input_connections:
            types = input_connections['points_df_in']
            # connected
            return NodePorts(inports={'points_df_in': {port_type: types}},
                             outports={'distance_df': {port_type: types}})
        else:
            return NodePorts(inports=input_ports, outports=output_ports)

    def init(self):
        self.delayed_process = True


    def meta_setup(self,):
        cols_required = {'x': 'float64',
                         'y': 'float64'}
        required = {
            'points_df_in': cols_required,
            'distance_df': cols_required
        }
        input_meta = self.get_input_meta()
        output_cols = ({
                'distance_df': {
                    'distance_cupy': 'float64',
                    'x': 'float64',
                    'y': 'float64'
                }
            })
        if 'points_df_in' in input_meta:
            col_from_inport = input_meta['points_df_in']
            # additional ports
            output_cols['distance_df'].update(col_from_inport)
        return MetaData(inports=required, outports=output_cols)

    def conf_schema(self):
        return ConfSchema()

    def get_kernel(self):
        raw_kernel = cupy.RawKernel(kernel_string, 'compute_distance')
        return raw_kernel

    def process(self, inputs):
        df = inputs['points_df_in']
        cupy_x = cupy.asarray(df['x'])
        cupy_y = cupy.asarray(df['y'])
        number_of_threads = 16
        number_of_blocks = (len(df) - 1) // number_of_threads + 1
        dis = cupy.ndarray(len(df), dtype=cupy.float64)
        self.get_kernel()((number_of_blocks,), (number_of_threads,),
                          (cupy_x, cupy_y, dis, len(df)))
        df['distance_cupy'] = dis

        return {'distance_df': df}<jupyter_output><empty_output><jupyter_text>The `self.delayed_process = True` flag is added for the same reason as with `DistanceNumbaNode` i.e. to support `dask_cudf` data frames.Let's register these two added new nodes<jupyter_code>TaskGraph.register_lab_node(module_name, NumbaDistanceNode)
TaskGraph.register_lab_node(module_name, CupyDistanceNode)<jupyter_output><empty_output><jupyter_text>Computing using the Nodes with customized GPU kernelsFirst we construct the computation graph for gQuant.<jupyter_code># For comparison to above re-use points dataframe instead
# of rand generating each time when running the task-graph.

numba_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_numba',
    TaskSpecSchema.node_type: NumbaDistanceNode,
    TaskSpecSchema.conf: {}, 
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'points_df_in': 'points_task.points_df_out'
    },
}

cupy_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_cupy',
    TaskSpecSchema.node_type: CupyDistanceNode,
    TaskSpecSchema.conf: {},
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'points_df_in': 'points_task.points_df_out'
    },
}

out_spec = {
    TaskSpecSchema.task_id: '',
    TaskSpecSchema.node_type: "Output_Collector",
    TaskSpecSchema.conf: {},
    TaskSpecSchema.inputs: {
        'in0': 'distance_by_cudf.distance_df',
        'in1': 'distance_by_numba.distance_df',
        'in2': 'distance_by_cupy.distance_df'
    }
}

task_list = [
    points_tspec,
    cudf_distance_tspec,
    numba_distance_tspec,
    cupy_distance_tspec,
    out_spec
]
task_graph = TaskGraph(task_list)

task_graph.draw(show='ipynb', show_ports=True)<jupyter_output><empty_output><jupyter_text>Then we run the task graph interatively by using the gQuantlab widget<jupyter_code>task_graph.draw()<jupyter_output><empty_output><jupyter_text>Or, run it programmatically<jupyter_code>out_list = [
    'distance_by_cudf.distance_df',
    'distance_by_numba.distance_df',
    'distance_by_cupy.distance_df'
]
cache_load = {"points_task": {"load": {'points_df_out': points_df}}}
(df_w_cudf, df_w_numba, df_w_cupy) = task_graph.run(out_list, replace=cache_load)
task_graph.run(out_list, replace=cache_load, formated=True)<jupyter_output><jupyter_text>Use `verify` function defined above to verify the results:<jupyter_code>mdiff = verify(df_w_cudf['distance_cudf'], df_w_numba['distance_numba'])
print('Max Difference cudf to numba: {}'.format(mdiff))
mdiff = verify(df_w_cudf['distance_cudf'], df_w_cupy['distance_cupy'])
print('Max Difference cudf to cupy: {}'.format(mdiff))<jupyter_output>Max Difference cudf to numba: 2.220446049250313e-16
Max Difference cudf to cupy: 2.220446049250313e-16<jupyter_text>To illustrate multi-input nodes let's create a verify node.<jupyter_code>class VerifyNode(Node):

    def ports_setup(self):
        input_ports = {
            'df1': {
                PortsSpecSchema.port_type: [cudf.DataFrame,
                                            dask_cudf.DataFrame]
            },
            'df2': {
                PortsSpecSchema.port_type: [cudf.DataFrame,
                                            dask_cudf.DataFrame]
            }
        }
        output_ports = {
            'max_diff': {
                PortsSpecSchema.port_type: float
            }
        }

        connections = self.get_connected_inports()   
        for key in input_ports:
            if key in connections:
                # connected
                types = connections[key]
                input_ports[key].update({PortsSpecSchema.port_type: types})
        return NodePorts(inports=input_ports, outports=output_ports)

    def meta_setup(self):
        required ={
            "df1": {},
            "df2": {}
        }
        return MetaData(inports=required, outports={'max_diff': {}})

    def conf_schema(self):
        json = {
            "title": "VerifyNode configure",
            "type": "object",
            "properties": {
                "df1_col":  {
                    "type": "string",
                    "description": "dataframe1 column name"
                },
                "df2_col":  {
                    "type": "string",
                    "description": "dataframe2 column name"
                }
            },
            "required": ["df1_col", "df2_col"],
        }

        ui = {
            "df1_col": {"ui:widget": "text"},
            "df2_col": {"ui:widget": "text"}
        }
        return ConfSchema(json=json, ui=ui)

    def process(self, inputs):
        df1 = inputs['df1']
        df2 = inputs['df2']
        col_df1 = self.conf['df1_col']
        col_df2 = self.conf['df2_col']

        df1_col = df1[col_df1]
        if isinstance(df1, dask_cudf.DataFrame):
            # df1_col = df1_col.compute()
            pass

        df2_col = df2[col_df2]
        if isinstance(df2, dask_cudf.DataFrame):
            # df2_col = df2_col.compute()
            pass

        max_difference = (df1_col - df2_col).abs().max()

        if isinstance(max_difference, dask.dataframe.core.Scalar):
            max_difference = float(max_difference.compute())
        max_difference = float(max_difference)
        return {'max_diff': max_difference}<jupyter_output><empty_output><jupyter_text>Register the `VerifyNode`:<jupyter_code>TaskGraph.register_lab_node(module_name, VerifyNode)<jupyter_output><empty_output><jupyter_text>Define the full Taskgraph:<jupyter_code>verify_tspec = {
    TaskSpecSchema.task_id: 'verify_cudf_to_numba',
    TaskSpecSchema.node_type: VerifyNode,
    TaskSpecSchema.conf: {
        'df1_col': 'distance_cudf',
        'df2_col': 'distance_numba'
    },    
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'df1': 'distance_by_cudf.distance_df',
        'df2': 'distance_by_numba.distance_df'
    }
}

verify_tspec2 = {
    TaskSpecSchema.task_id: 'verify_cudf_to_cupy',
    TaskSpecSchema.node_type: VerifyNode,
    TaskSpecSchema.conf: {
        'df1_col': 'distance_cudf',
        'df2_col': 'distance_cupy'
    },
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'df1': 'distance_by_cudf.distance_df',
        'df2': 'distance_by_cupy.distance_df'
    }
}
out_spec = {
    TaskSpecSchema.task_id: '',
    TaskSpecSchema.node_type: "Output_Collector",
    TaskSpecSchema.conf: {},
    TaskSpecSchema.inputs: {
        'in0': 'verify_cudf_to_numba.max_diff',
        'in1': 'verify_cudf_to_cupy.max_diff'
    }
}

task_list = [
    points_tspec,
    cudf_distance_tspec,
    numba_distance_tspec,
    cupy_distance_tspec,
    out_spec,
    verify_tspec, 
    verify_tspec2
]
task_graph = TaskGraph(task_list)
task_graph.draw(show='ipynb', show_ports=True)
task_graph.draw()
(max_cudf_to_numba_diff, max_cudf_to_cupy_diff) = task_graph.run([
    'verify_cudf_to_numba.max_diff',
    'verify_cudf_to_cupy.max_diff'
])
print('Max Difference cudf to numba: {}'.format(max_cudf_to_numba_diff))
print('Max Difference cudf to cupy: {}'.format(max_cudf_to_cupy_diff))<jupyter_output>Max Difference cudf to numba: 2.220446049250313e-16
Max Difference cudf to cupy: 2.220446049250313e-16<jupyter_text>Dask distributed computationUsing Dask and `dask-cudf` we can run the Nodes with customized GPU kernels on distributed dataframes. Under the hood of the `Node` class the Dask delayed processing API is handled for cudf dataframes when the `self.delayed_process = True` flag is set.We first start a distributed Dask environment. When a dask client is instantiated it registers itself as the default Dask scheduler (). Therefore all subsequent Dask distibuted dataframe operations will run in distributed fashion.<jupyter_code>from dask_cuda import LocalCUDACluster
from dask.distributed import Client

cluster = LocalCUDACluster()
client = Client(cluster)
client<jupyter_output><empty_output><jupyter_text>The Dask status page can be displayed in a web browser at `:8787`. The ip-address corresponds to the machine where the dask cluster (scheduler) was launched. Most likely same ip-address as where this jupyter notebook is running. Using the Dask status page is convenient for monitoring dask distributed processing.  The next step is to partition the `cudf` dataframe into a `dask_cudf` dataframe. Here we make the number of partitions corresponding to the number of workers:<jupyter_code>class DistributedNode(Node):

    def ports_setup(self):
        input_ports = {
            'points_df_in': {
                PortsSpecSchema.port_type: cudf.DataFrame
            }
        }

        output_ports = {
            'points_ddf_out': {
                PortsSpecSchema.port_type: dask_cudf.DataFrame
            }
        }

        return NodePorts(inports=input_ports, outports=output_ports)

    def init(self):
        pass

    def meta_setup(self,):
        cols_required = {
            'x': 'float64',
            'y': 'float64'
        }
        required = {
            'points_df_in': cols_required,
            'points_ddf_out': cols_required
        }
        input_meta = self.get_input_meta()
        output_cols = ({
                'points_ddf_out': {
                    'x': 'float64',
                    'y': 'float64'
                }
            })
        if 'points_df_in' in input_meta:
            col_from_inport = input_meta['points_df_in']
            # additional ports
            output_cols['points_ddf_out'].update(col_from_inport)
        return MetaData(inports=required, outports=output_cols)

    def conf_schema(self):
        json = {
            "title": "DistributedNode configure",
            "type": "object",
            "properties": {
                "npartitions":  {
                    "type": "number",
                    "description": "num of partitions in the Dask dataframe",
                    "minimum": 1
                }
            },
            "required": ["npartitions"],
        }

        ui = {
            "npartitions": {"ui:widget": "updown"}
        }
        return ConfSchema(json=json, ui=ui)

    def process(self, inputs):
        npartitions = self.conf['npartitions']
        df = inputs['points_df_in']
        ddf = dask_cudf.from_cudf(df, npartitions=npartitions)
        return {'points_ddf_out': ddf}<jupyter_output><empty_output><jupyter_text>Register it:<jupyter_code>TaskGraph.register_lab_node(module_name, DistributedNode)<jupyter_output><empty_output><jupyter_text>We add this distribution node to the computation graph to convert `cudf` dataframes into `dask-cudf` dataframes. The `dask-cudf` dataframes are handled automatically in gQuant when `self.delayed_process=True` within a `Node` implementation (setup in `meta_setup`). When using nodes with ports with `self.delayed_process=True` setting, it is required that all input and output ports be of type `cudf.DataFrame`. Otherwise don't set  `self.delayed_process` and one can write custom logic to handle distributed dataframes (refer to `VerifyNode` abover for an example where `dask_cudf` dataframes are handled directly within the process method).<jupyter_code>npartitions = len(client.scheduler_info()['workers'])


distribute_tspec = {
    TaskSpecSchema.task_id: 'distributed_points',
    TaskSpecSchema.node_type: DistributedNode,
    TaskSpecSchema.conf: {'npartitions': npartitions},
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'points_df_in': 'points_task.points_df_out'
    }
}

dask_cudf_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_cudf',
    TaskSpecSchema.node_type: DistanceNode,
    TaskSpecSchema.conf: {},
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'points_df_in': 'distributed_points.points_ddf_out'
    }
}

dask_numba_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_numba',
    TaskSpecSchema.node_type: NumbaDistanceNode,
    TaskSpecSchema.conf: {},
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'points_df_in': 'distributed_points.points_ddf_out'
    }
}

dask_cupy_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_cupy',
    TaskSpecSchema.node_type: CupyDistanceNode,
    TaskSpecSchema.conf: {},
    TaskSpecSchema.module: module_name,
    TaskSpecSchema.inputs: {
        'points_df_in': 'distributed_points.points_ddf_out'
    }
}

out_spec = {
    TaskSpecSchema.task_id: '',
    TaskSpecSchema.node_type: "Output_Collector",
    TaskSpecSchema.conf: {},
    TaskSpecSchema.inputs: {
        'in0': 'distributed_points.points_ddf_out',
        'in1': 'distance_by_cudf.distance_df',
        'in2': 'distance_by_numba.distance_df',
        'in3': 'distance_by_cupy.distance_df'
    }
}


task_list = [
    points_tspec,
    distribute_tspec,
    dask_cudf_distance_tspec,
    dask_numba_distance_tspec,
    dask_cupy_distance_tspec,
    out_spec
]

task_graph = TaskGraph(task_list)
task_graph.draw(show='ipynb', show_ports=True)
task_graph.draw()<jupyter_output><empty_output><jupyter_text>Run the distributed computation programmatically:<jupyter_code>out_list = [
    'distributed_points.points_ddf_out',
    'distance_by_cudf.distance_df',
    'distance_by_numba.distance_df',
    'distance_by_cupy.distance_df'
]

(points_ddf, ddf_w_cudf, ddf_w_numba, ddf_w_cupy) = task_graph.run(out_list)
df_w_cudf = ddf_w_cudf.compute()
df_w_numba = ddf_w_numba.compute()
df_w_cupy = ddf_w_cupy.compute()<jupyter_output><empty_output><jupyter_text>Verify the results:<jupyter_code>verify_cudf_numba_tspec = verify_tspec.copy()
verify_cudf_cupy_tspec = verify_tspec2.copy()

task_graph.extend(
    [verify_cudf_numba_tspec,
     verify_cudf_cupy_tspec],
    replace=True)
task_graph.draw(show='ipynb', show_ports=True)

# Use results above and avoid re-running dask
replace_spec = {
    'distance_by_cudf': {
        TaskSpecSchema.load: {
            'distance_df': ddf_w_cudf
        }
    },
    'distance_by_numba': {
        TaskSpecSchema.load: {
            'distance_df': ddf_w_numba
        }
    },
    'distance_by_cupy': {
        TaskSpecSchema.load: {
            'distance_df': ddf_w_cupy
        }
    }
}

(max_cudf_to_numba_diff, max_cudf_to_cupy_diff) = task_graph.run(
    ['verify_cudf_to_numba.max_diff',
     'verify_cudf_to_cupy.max_diff'],
    replace=replace_spec
)

print('HEAD points_ddf:\n{}\n'.format(points_ddf.head()))
print('HEAD df_w_cudf:\n{}\n'.format(ddf_w_cudf.head()))
print('HEAD df_w_numba:\n{}\n'.format(ddf_w_numba.head()))
print('HEAD df_w_cupy:\n{}\n'.format(ddf_w_cupy.head()))
print('Max Difference cudf to numba: {}'.format(max_cudf_to_numba_diff))
print('Max Difference cudf to cupy: {}'.format(max_cudf_to_cupy_diff))<jupyter_output>HEAD points_ddf:
          x         y
0  0.017881  0.460670
1  0.184132  0.864031
2  0.241513  0.956813
3  0.336786  0.975579
4  0.767782  0.330257

HEAD df_w_cudf:
          x         y  distance_cudf
0  0.017881  0.460670       0.461017
1  0.184132  0.864031       0.883433
2  0.241513  0.956813       0.986823
3  0.336786  0.975579       1.032075
4  0.767782  0.330257       0.835799

HEAD df_w_numba:
          x         y  distance_numba
0  0.017881  0.460670        0.461017
1  0.184132  0.864031        0.883433
2  0.241513  0.956813        0.986823
3  0.336786  0.975579        1.032075
4  0.767782  0.330257        0.835799

HEAD df_w_cupy:
          x         y  distance_cupy
0  0.017881  0.460670       0.461017
1  0.184132  0.864031       0.883433
2  0.241513  0.956813       0.986823
3  0.336786  0.975579       1.032075
4  0.767782  0.330257       0.835799

Max Difference cudf to numba: 2.220446049250313e-16
Max Difference cudf to cupy: 2.220446049250313e-16<jupyter_text>One limitation to be aware of when using customized kernels within Nodes in the Dask environment, is that each GPU kernel works on one partition of the dataframe. Therefore if the computation depends on other partitions of the dataframe the approach above does not work.  Saving Custom Nodes and KernelsThe gQuant examples already implement a number of `Nodes`. These can be found in `gquant.plugin_nodes` submodules.The customized kernels and nodes can be saved to your own python modules for future re-use instead of having to re-define them at runtime. The nodes we defined above were to a written to a python module "custom_port_nodes.py" (the `DistanceNode` was simplified to ommit the absolute distance calculation). We will re-run our workflow importing the Nodes from the custom module we wrote out.When defining the tasks we specify `filepath` for the path to the python module that has the Node definition. Notice, that the `node_type` is specified as a string instead of class. The string is the class name of the node that will be imported for running a task.<jupyter_code>npartitions = len(client.scheduler_info()['workers'])

points_tspec = {
    TaskSpecSchema.task_id: 'points_task',
    TaskSpecSchema.node_type: 'PointNode',
    TaskSpecSchema.filepath: '/home/quant/gQuant/notebooks/custom_port_nodes.py',
    TaskSpecSchema.conf: {'npts': 1000},
    TaskSpecSchema.inputs: {},
}

distribute_tspec = {
    TaskSpecSchema.task_id: 'distributed_points',
    TaskSpecSchema.node_type: 'DistributedNode',
    TaskSpecSchema.filepath: '/home/quant/gQuant/notebooks/custom_port_nodes.py',
    TaskSpecSchema.conf: {'npartitions': npartitions},
    TaskSpecSchema.inputs: {
        'points_df_in': 'points_task.points_df_out'
    }
}

dask_cudf_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_cudf',
    TaskSpecSchema.node_type: 'DistanceNode',
    TaskSpecSchema.filepath: '/home/quant/gQuant/notebooks/custom_port_nodes.py',
    TaskSpecSchema.conf: {},
    TaskSpecSchema.inputs: {
        'points_df_in': 'distributed_points.points_ddf_out'
    }
}

dask_numba_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_numba',
    TaskSpecSchema.node_type: 'NumbaDistanceNode',
    TaskSpecSchema.filepath: '/home/quant/gQuant/notebooks/custom_port_nodes.py',
    TaskSpecSchema.conf: {},
    TaskSpecSchema.inputs: {
        'points_df_in': 'distributed_points.points_ddf_out'
    }
}

dask_cupy_distance_tspec = {
    TaskSpecSchema.task_id: 'distance_by_cupy',
    TaskSpecSchema.node_type: 'CupyDistanceNode',
    TaskSpecSchema.filepath: '/home/quant/gQuant/notebooks/custom_port_nodes.py',
    TaskSpecSchema.conf: {},
    TaskSpecSchema.inputs: {
        'points_df_in': 'distributed_points.points_ddf_out'
    }
}

verify_cudf_to_numba_tspec = {
    TaskSpecSchema.task_id: 'verify_cudf_to_numba',
    TaskSpecSchema.node_type: 'VerifyNode',
    TaskSpecSchema.filepath: '/home/quant/gQuant/notebooks/custom_port_nodes.py',
    TaskSpecSchema.conf: {
        'df1_col': 'distance_cudf',
        'df2_col': 'distance_numba'
    },    
    TaskSpecSchema.inputs: {
        'df1': 'distance_by_cudf.distance_df',
        'df2': 'distance_by_numba.distance_df'
    }
}

verify_cudf_to_cupy_tspec = {
    TaskSpecSchema.task_id: 'verify_cudf_to_cupy',
    TaskSpecSchema.node_type: 'VerifyNode',
    TaskSpecSchema.filepath: '/home/quant/gQuant/notebooks/custom_port_nodes.py',
    TaskSpecSchema.conf: {
        'df1_col': 'distance_cudf',
        'df2_col': 'distance_cupy'
    },    
    TaskSpecSchema.inputs: {
        'df1': 'distance_by_cudf.distance_df',
        'df2': 'distance_by_cupy.distance_df'
    }
}

task_list = [
    points_tspec,
    distribute_tspec,
    dask_cudf_distance_tspec,
    dask_numba_distance_tspec,
    dask_cupy_distance_tspec,
    verify_cudf_to_numba_tspec,
    verify_cudf_to_cupy_tspec
]

task_graph = TaskGraph(task_list)
task_graph.draw(show='ipynb', show_ports=True)
task_graph.draw()
out_list = [
    'distance_by_cudf.distance_df',
    'distance_by_numba.distance_df',
    'distance_by_cupy.distance_df',
    'verify_cudf_to_numba.max_diff',
    'verify_cudf_to_cupy.max_diff'
]

(ddf_w_cudf, ddf_w_numba, ddf_w_cupy,
 mdiff_cudf_to_numba, mdiff_cudf_to_cupy) = task_graph.run(out_list)

print('HEAD df_w_cudf:\n{}\n'.format(ddf_w_cudf.head()))
print('HEAD df_w_numba:\n{}\n'.format(ddf_w_numba.head()))
print('HEAD df_w_cupy:\n{}\n'.format(ddf_w_cupy.head()))
print('Max Difference cudf to numba: {}'.format(mdiff_cudf_to_numba))
print('Max Difference cudf to cupy: {}'.format(mdiff_cudf_to_cupy))<jupyter_output>HEAD df_w_cudf:
          x         y  distance_cudf
0  0.438593  0.505563       0.669296
1  0.714780  0.531261       0.890589
2  0.174467  0.734742       0.755171
3  0.964969  0.986593       1.380048
4  0.631274  0.819636       1.034558

HEAD df_w_numba:
          x         y  distance_numba
0  0.438593  0.505563        0.669296
1  0.714780  0.531261        0.890589
2  0.174467  0.734742        0.755171
3  0.964969  0.986593        1.380048
4  0.631274  0.819636        1.034558

HEAD df_w_cupy:
          x         y  distance_cupy
0  0.438593  0.505563       0.669296
1  0.714780  0.531261       0.890589
2  0.174467  0.734742       0.755171
3  0.964969  0.986593       1.380048
4  0.631274  0.819636       1.034558

Max Difference cudf to numba: 2.220446049250313e-16
Max Difference cudf to cupy: 2.220446049250313e-16<jupyter_text>The final illustration is how to save and load a task graph to a file for re-use.<jupyter_code>task_graph.save_taskgraph('custom_wflow.gq.yaml')<jupyter_output><empty_output><jupyter_text>The gQuant TaskGraph file is created and saved. You can double click on it to open it up in the JupyterLab to edit it.Or you can display it by gQuant widget and play with it interactively<jupyter_code>task_graph = TaskGraph.load_taskgraph('custom_wflow.gq.yaml')
task_graph.draw()<jupyter_output><empty_output><jupyter_text>Of course you can run it by callign `run` method.<jupyter_code># update npartitions in case the scheduler is running with
# different number of workers than what was saved.
npartitions = len(client.scheduler_info()['workers'])
replace_spec = {
    'distributed_points': {
        TaskSpecSchema.conf: {'npartitions': npartitions},
    }
}

out_list = [
    'distance_by_cudf.distance_df',
    'distance_by_numba.distance_df',
    'distance_by_cupy.distance_df',
    'verify_cudf_to_numba.max_diff',
    'verify_cudf_to_cupy.max_diff'
]

task_graph.run(out_list, replace=replace_spec, formated=True)


print('HEAD df_w_cudf:\n{}\n'.format(ddf_w_cudf.head()))
print('HEAD df_w_numba:\n{}\n'.format(ddf_w_numba.head()))
print('HEAD df_w_cupy:\n{}\n'.format(ddf_w_cupy.head()))<jupyter_output>HEAD df_w_cudf:
          x         y  distance_cudf
0  0.438593  0.505563       0.669296
1  0.714780  0.531261       0.890589
2  0.174467  0.734742       0.755171
3  0.964969  0.986593       1.380048
4  0.631274  0.819636       1.034558

HEAD df_w_numba:
          x         y  distance_numba
0  0.438593  0.505563        0.669296
1  0.714780  0.531261        0.890589
2  0.174467  0.734742        0.755171
3  0.964969  0.986593        1.380048
4  0.631274  0.819636        1.034558

HEAD df_w_cupy:
          x         y  distance_cupy
0  0.438593  0.505563       0.669296
1  0.714780  0.531261       0.890589
2  0.174467  0.734742       0.755171
3  0.964969  0.986593       1.380048
4  0.631274  0.819636       1.034558<jupyter_text>ConclusionUsing customized GPU kernels allows data scientists to implement and incorporate advanced algorithms. We demonstrated implementations using Numba and CuPy.The Numba approach enables data scientists to write GPU kernels directly in the Python language. Numba is easy to use for implementing and accelerating computations. However there is some overhead incurred for compiling the kernels whenever the Numba GPU kernels are used for the first time in a Python process. Currently Numba library only supports primitive data types. Some advanced CUDA programming features, such as function pointers and function recursions are not supported. The Cupy method is very flexible, because data scientists are writing C/C++ GPU kernels with CUDA directly. All the CUDA programming features are supported. CuPy compiles the kernel and caches the device code to the filesystem. The launch overhead is low. Also, the GPU kernel is built statically resulting in runtime efficiency. However it might be harder for data scientists to use, because C/C++ programming is more complicated. Below is a brief summary comparison table:| Methods | Development Difficulty | Flexibility | Efficiency | Latency ||---|---|---|---|---|| Numba method | medium | medium | low | high || CuPy method | hard | high  | high | low |We recommend that the data scientists select the approach appropriate for their task taking into consideration the efficiency, latency, difficulty and flexibility of their workflow. In this blog, we showed how to wrap the customized GPU kernels in gQuant nodes. Also, by taking advantage of having the gQuant handle the low-level Dask interfaces for the developer, we demonstrated how to use the gQuant workflow with Dask distributed computations.<jupyter_code># Clean up

# Shutdown the Dask cluster
client.close()
cluster.close()
import IPython
app = IPython.Application.instance()
app.kernel.do_shutdown(True)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09386583491546877, 'fraction_numerical': 0.03085164090046108, 'mean_word_length': 3.02813182811362, 'pattern_counts': {'":': 44, '<': 116, '<?xml version=': 0, '>': 113, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '30', 'dir': 'jupyter-structured-clean-dedup', 'id': '39970', 'license': 'Apache-2.0', 'n_tokens_mistral': 15062, 'n_tokens_neox': 14221, 'n_words': 3438, 'path': 'notebooks/05_customize_nodes_with_ports.ipynb', 'repo_name': 'goosen78/gQuant', 'size': '44244'}