pvw-spark.py

from __future__ import print_function

import os
import sys
from pyspark import SparkContext

sc = SparkContext()

# Define the address of the PMI server and the number of MPI workers

hostname = os.uname()[1]
hydra_proxy_port = os.getenv("HYDRA_PROXY_PORT")
pmi_port = hostname + ":" + hydra_proxy_port

partitions = 4

# Prepare a list of environmental variables

env = []
for id in range(0, partitions):
    kvs = {
        'PMI_PORT' : pmi_port,
        'PMI_ID' : id,
        'Center': [id, 0, 0],
        'Resolution': (id + 1) * 6
    }
    env.append(kvs)

# Create the rdd collection associated with the MPI workers

rdd = sc.parallelize(env, partitions)

def allreduce(kvs):
    # -------------------------------------------------------------------------
    # Setup environment for MPI and pvbatch
    # -------------------------------------------------------------------------

    import os
    os.environ["PMI_PORT"] = kvs["PMI_PORT"]
    os.environ["PMI_ID"] = str(kvs["PMI_ID"])
    os.environ["PV_ALLOW_BATCH_INTERACTION"] = "1"
    os.environ["DISPLAY"] = ":0"

    # -------------------------------------------------------------------------
    # Load data on each node from Spark
    # -------------------------------------------------------------------------

    from vtk.vtkPVVTKExtensionsCore import vtkDistributedTrivialProducer
    from vtk.vtkFiltersSources import vtkSphereSource

    source = vtkSphereSource()
    source.SetCenter(kvs['Center'])
    source.SetPhiResolution(kvs['Resolution'])
    source.SetThetaResolution(kvs['Resolution'])
    source.Update()
    dataset = source.GetOutput()

    vtkDistributedTrivialProducer.SetGlobalOutput('Spark', dataset)

    # -------------------------------------------------------------------------
    # Start Visualizer in parallel
    # -------------------------------------------------------------------------

    import paraview
    paraview.options.batch = True

    from vtk.vtkPVClientServerCoreCore import vtkProcessModule
    from paraview     import simple
    from paraview.web import wamp      as pv_wamp
    from paraview.web import protocols as pv_protocols
    from vtk.web      import server

    pm = vtkProcessModule.GetProcessModule()

    class Options(object):
        debug = False
        nosignalhandlers = True
        host = 'localhost'
        port = 9753
        timeout = 300
        content = '/home/chaudhary/tools/sparkmpi/visualizer/dist'
        forceFlush = False
        sslKey = ''
        sslCert = ''
        ws = 'ws'
        lp = 'lp'
        hp = 'hp'
        nows = False
        nobws = False
        nolp = False
        fsEndpoints = ''
        uploadPath = None
        testScriptPath = None
        baselineImgDir = ''
        useBrowser = 'nobrowser'
        tmpDirectory = '.'
        testImgFile = None

    class _VisualizerServer(pv_wamp.PVServerProtocol):
        dataDir = '/data'
        groupRegex = "[0-9]+\\.[0-9]+\\.|[0-9]+\\."
        excludeRegex = "^\\.|~$|^\\$"
        allReaders = True
        viewportScale=1.0
        viewportMaxWidth=2560
        viewportMaxHeight=1440

        def initialize(self):
            # Bring used components
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebFileListing(_VisualizerServer.dataDir, "Home", _VisualizerServer.excludeRegex, _VisualizerServer.groupRegex))
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebProxyManager(baseDir=_VisualizerServer.dataDir, allowUnconfiguredReaders=_VisualizerServer.allReaders))
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebColorManager())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebMouseHandler())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebViewPort(_VisualizerServer.viewportScale, _VisualizerServer.viewportMaxWidth, _VisualizerServer.viewportMaxHeight))
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebViewPortImageDelivery())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebViewPortGeometryDelivery())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebTimeHandler())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebSelectionHandler())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebWidgetManager())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebKeyValuePairStore())
            self.registerVtkWebProtocol(pv_protocols.ParaViewWebSaveData(baseSavePath=_VisualizerServer.dataDir))

            # Disable interactor-based render calls
            simple.GetRenderView().EnableRenderOnInteraction = 0
            simple.GetRenderView().Background = [0,0,0]

            # Update interaction mode
            pxm = simple.servermanager.ProxyManager()
            interactionProxy = pxm.GetProxy('settings', 'RenderViewInteractionSettings')
            interactionProxy.Camera3DManipulators = ['Rotate', 'Pan', 'Zoom', 'Pan', 'Roll', 'Pan', 'Zoom', 'Rotate', 'Zoom']

    args = Options()

    # -------------------------------------------------------------------------
    # Only MPI Process 0 is going to start the web server
    # -------------------------------------------------------------------------

    if pm.GetPartitionId() == 0:
        producer = simple.DistributedTrivialProducer()
        producer.UpdateDataset = ''
        producer.UpdateDataset = 'Spark'
        server.start_webserver(options=args, protocol=_VisualizerServer)
        pm.GetGlobalController().TriggerBreakRMIs()

    # -------------------------------------------------------------------------
    # Send back results to Spark
    # -------------------------------------------------------------------------

    out = {
       'rank' : pm.GetPartitionId(),
       'size': pm.GetNumberOfLocalPartitions()
    }
    return out

# Run MPI application on Spark workers and collect the results

results = rdd.map(allreduce).collect()
for out in results:
    print ("rank: ", out['rank'], 'size:', out['size'])