How does index assignment to nodes work in Spark with mapPartitionsWithIndex()?
<p>I am trying to coordinate GPU execution on a Spark cluster. In order to achieve this I need each task/partition to only use a specific GPU slot per system. Each system has 4 GPUs, and the easiest way I have found to achieve this is by doing a mapPartitionsWithIndex() on the rdd with the data, and then using the index as the gpu slot.</p>

<p>My question is if I can depend on the indexes to always be assigned in order to the worker nodes? Is there any documentation that refers to this?</p>

<p>Quick example:</p>

<pre><code>from pyspark import SparkContext
from pyspark.sql import SQLContext

sc = SparkContext(appName="sample")
sqlContext = SQLContext(sc)

def print_partition_info(idx, part):
    print "Index: {} - Partition data: {}".format(idx, part)
    for p in part:
        yield p

data = [1,1,1,2,2,2, 20, 30, 40, 100]
partitions = 4
rdd = sc.parallelize(data, partitions)
# rdd = rdd.coalesce(len(partition_keys))
rdd2 = rdd.mapPartitionsWithIndex(print_partition_info)
</code></pre>

<p>Lets pretend the printed output of running an action on rdd2 is  </p>

<pre><code>Index: 0 - Partition data: [1,1]  
Index: 1 - Partition data: [1,2]  
Index: 2 - Partition data: [2,2]  
Index: 3 - Partition data: [20,30,40,100] 
</code></pre>

<p>When these partitions get sent out to workers (say 2 workers), will they always be in order as in...<br>
Worker 1 partitions = 1, 2<br>
Worker 2 partitions = 3, 4  </p>

<p>Expanding on that will the partitions always be ordered for say a cluster of 10 or 50 nodes?</p>

<p>Finally here is a slide that seems to support this technique:  </p>

<p><a href="http://i.stack.imgur.com/VDyM1.png" rel="nofollow"><img src="http://i.stack.imgur.com/VDyM1.png" alt="enter image description here"></a>
<a href="http://www.slideshare.net/continuumio/gpu-computing-with-apache-spark-and-python" rel="nofollow" title="Slides">slides - check #52</a></p>

<p>Short answer: No. In a cluster, the RessourceManager (YARN most of the time) will use a worker if it is available, and it is not always the case when your system is multi-user, or if you have already started a job in a subset of your cluster. So you can't bind a worker with an index. </p>

<p>Thus, I am pretty sure that index 0 will bind to the first partition in your data, and so on. An expert will correct me if I'm wrong.</p>

python apache-spark pyspark partitioning rdd