Skip to content

Commit 4820118

Browse files
gjbexgjbex
committed
Add Cython + dask example
1 parent 301c1ea commit 4820118

File tree

11 files changed

+293
-0
lines changed

11 files changed

+293
-0
lines changed

‎source-code/dask/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,4 @@ CSV or HDF5 files.
3030
* `dask_sum_aarays.py`: somewhat artificial example of a Dask computation
3131
on `numpy` arrays.
3232
* `dask_sum_aarays.pbs`: PBS script to execute `dask_sum_aarays.py`.
33+
* `julia_set`: example of combining Cython and dask.

‎source-code/dask/julia_set/Makefile‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
PYX_FILES = julia_cython.pyx julia_cython_omp.pyx
2+
C_FILES = $(patsubst %.pyx,%.c,$(PYX_FILES))
3+
JULIA_LIB = $(patsubst %.c,%.so,$(C_FILES))
4+
5+
all: $(JULIA_LIB)
6+
7+
$(JULIA_LIB): $(PYX_FILES)
8+
python setup.py build_ext --inplace
9+
10+
clean:
11+
python setup.py clean
12+
rm -f $(C_FILES) $(JULIA_LIB)

‎source-code/dask/julia_set/README.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Julia set
2+
3+
Implementatoin of the computation of the Julia set using Cython and Dask.
4+
5+
## What is it?
6+
7+
1. `julia_cython.pyx`: Cython implementation of computing a range of values.
8+
1. `julia_cython_omp.pyx`: Cython OpenMP implementation of computing a range of
9+
values.
10+
1. `setup.py`: Python build file for the Cython modules.
11+
1. `Makefile`: make file to build the Cython modules.
12+
1. `julia_set_dask.py`: Python driver program.
13+
1. `julia_set_dask.pbs`: PBS script to run the computation.
14+
1. `launch_scheduler.sh`: Bash script to launch the scheduler.
15+
1. `launch_worker.sh`: Bash script to launch a worker.

‎source-code/dask/julia_set/conda.sh‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# >>> conda initialize >>>
2+
# !! Contents within this block are managed by 'conda init' !!
3+
__conda_setup="$('/data/leuven/301/vsc30140/miniconda3/bin/conda' 'shell.bash' 'hook' 2> /dev/null)"
4+
if [ $? -eq 0 ]; then
5+
eval "$__conda_setup"
6+
else
7+
if [ -f "/data/leuven/301/vsc30140/miniconda3/etc/profile.d/conda.sh" ]; then
8+
. "/data/leuven/301/vsc30140/miniconda3/etc/profile.d/conda.sh"
9+
else
10+
export PATH="/data/leuven/301/vsc30140/miniconda3/bin:$PATH"
11+
fi
12+
fi
13+
unset __conda_setup
14+
# <<< conda initialize <<<

‎source-code/dask/julia_set/julia_cython.pyx‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
'''Module containing function for computing Julia sets'''
2+
3+
4+
def julia_set(domain, iterations):
5+
_julia_set(memoryview(domain), memoryview(iterations), 2.0, 255)
6+
return iterations
7+
8+
9+
cdef _julia_set(double complex[:] domain, int[:] iterations,
10+
double max_norm, int max_iters):
11+
'''Compute the Julia set on a complex domain.
12+
13+
Positional arguments:
14+
domain -- complex domain to compute a 1-D array
15+
iterations -- number of iterations as a 1-D array
16+
max_norm -- maximum complex norm to iterate to.
17+
max_iters -- maximum number of iterations
18+
'''
19+
cdef int i
20+
cdef complex z
21+
for i in range(len(domain)):
22+
z = domain[i]
23+
while (iterations[i] <= max_iters and
24+
z.real*z.real + z.imag*z.imag <= max_norm*max_norm):
25+
z = z**2 - 0.622772 + 0.42193j
26+
iterations[i] += 1

‎source-code/dask/julia_set/julia_cython_omp.pyx‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
'''Module containing function for computing Julia sets'''
2+
3+
cimport cython
4+
from cython.parallel import prange
5+
6+
def julia_set(domain, iterations, max_norm, max_iters):
7+
_julia_set(memoryview(domain), memoryview(iterations),
8+
2.0, 255)
9+
return iterations
10+
11+
12+
@cython.boundscheck(False)
13+
cdef _julia_set(double complex[:] domain, int[:] iterations,
14+
double max_norm, int max_iters):
15+
'''Compute the Julia set on a complex domain.
16+
17+
Positional arguments:
18+
domain -- complex domain to compute a 1-D array
19+
iterations -- number of iterations as a 1-D array
20+
max_norm -- maximum complex norm to iterate to.
21+
max_iters -- maximum number of iterations
22+
'''
23+
cdef int i, length = len(domain)
24+
cdef complex z
25+
with nogil:
26+
for i in prange(length, schedule='guided'):
27+
z = domain[i]
28+
while (iterations[i] <= max_iters and
29+
z.real*z.real + z.imag*z.imag <= max_norm*max_norm):
30+
z = z**2 - 0.622772 + 0.42193j
31+
iterations[i] += 1

‎source-code/dask/julia_set/julia_set_dask.pbs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#!/usr/bin/env bash
2+
#PBS -l nodes=2:ppn=36
3+
#PBS -l walltime=00:30:00
4+
#PBS -m ae
5+
#PBS -M geertjan.bex@uhasselt.be
6+
7+
cd $PBS_O_WORKDIR
8+
9+
source "conda.sh"
10+
conda activate dask
11+
12+
if [ $? -ne 0 ]
13+
then
14+
(>&2 echo '### error: conda environment not sourced successfully' )
15+
fi
16+
17+
scheduler_node=$(hostname)
18+
scheduler_port=8786
19+
scheduler="${scheduler_node}:${scheduler_port}"
20+
worker_nodes=$(uniq $PBS_NODEFILE)
21+
worker_launcher="$(pwd)/launch_worker.sh"
22+
23+
echo "launching scheduler: ${scheduler}"
24+
./launch_scheduler.sh
25+
26+
echo 'waitling till scheduler launched...'
27+
sleep 15
28+
29+
echo 'starting workers...'
30+
for worker in $worker_nodes;
31+
do
32+
echo "launching worker on ${worker}"
33+
ssh $worker $worker_launcher "$(pwd)" "${scheduler}" "${PBS_JOBID}" &
34+
done
35+
36+
echo 'waiting for workers to start and connect'
37+
sleep 15
38+
39+
40+
echo 'starting computation'
41+
42+
time python julia_set_dask.py --host ${scheduler_node} \
43+
--port ${scheduler_port} \
44+
--implementation cython

‎source-code/dask/julia_set/julia_set_dask.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
#!/usr/bin/env python
2+
3+
from argparse import ArgumentParser
4+
from distributed import Client, Future
5+
import numpy as np
6+
import os
7+
import sys
8+
import time
9+
10+
11+
def init_julia(re, im, n):
12+
'''Initialize the complex domain.
13+
14+
Positional arguments:
15+
re -- minimum and maximum real value as 2-tuple
16+
im -- minimum and maximum imaginary value as 2-tuple
17+
n -- number of real and imaginary points as 2-tuple
18+
'''
19+
re_vals, im_vals = np.meshgrid(
20+
np.linspace(re[0], re[1], n[0]),
21+
np.linspace(im[0], im[1], n[1])
22+
)
23+
domain = re_vals + im_vals*1j
24+
return domain.flatten()
25+
26+
27+
def init_pyx(dask_worker):
28+
import pyximport
29+
pyximport.install()
30+
sys.path.insert(0, os.getcwd())
31+
# sys.path.insert(0, '/scratch/leuven/301/vsc30140/julia_set/')
32+
from julia_cython import julia_set
33+
34+
35+
def init_omp_pyx(dask_worker):
36+
import pyximport
37+
pyximport.install()
38+
sys.path.insert(0, os.getcwd())
39+
# sys.path.insert(0, '/scratch/leuven/301/vsc30140/julia_set/')
40+
from julia_cython_omp import julia_set
41+
42+
43+
44+
if __name__ == '__main__':
45+
arg_parser = ArgumentParser(description='Compute julia set')
46+
arg_parser.add_argument('--re_min', type=float, default=-1.8,
47+
help='minimum real value')
48+
arg_parser.add_argument('--re_max', type=float, default=1.8,
49+
help='maximum real value')
50+
arg_parser.add_argument('--im_min', type=float, default=-1.8,
51+
help='minimum imaginary value')
52+
arg_parser.add_argument('--im_max', type=float, default=1.8,
53+
help='maximum imaginary value')
54+
arg_parser.add_argument('--max_norm', type=float, default=2.0,
55+
help='maximum complex norm for z')
56+
arg_parser.add_argument('--n_re', type=int, default=100,
57+
help='number of points on the real axis')
58+
arg_parser.add_argument('--n_im', type=int, default=100,
59+
help='number of points on the imaginary axis')
60+
arg_parser.add_argument('--max_iters', type=int, default=300,
61+
help='maximum number of iterations')
62+
arg_parser.add_argument('--implementation', default='python',
63+
choices=['python', 'cython', 'cython_omp'],
64+
help='implementation to use')
65+
arg_parser.add_argument('--partitions', type=int, default=100,
66+
help='number of partitions for dask workers')
67+
arg_parser.add_argument('--host', required=True,
68+
help='hostname of the dask scheduler')
69+
arg_parser.add_argument('--port', type=int, required=True,
70+
help='port of the dask scheduler')
71+
options = arg_parser.parse_args()
72+
client = Client(f'{options.host}:{options.port:d}')
73+
if options.implementation == 'python':
74+
from julia_python import julia_set
75+
elif options.implementation == 'cython':
76+
from julia_cython import julia_set
77+
client.register_worker_callbacks(init_pyx)
78+
elif options.implementation == 'cython_omp':
79+
from julia_cython_omp import julia_set
80+
client.register_worker_callbacks(init_omp_pyx)
81+
else:
82+
msg = '{0} version not implemented\n'
83+
sys.stderr.write(msg.format(options.implementation))
84+
sys.exit(1)
85+
86+
domain = init_julia(
87+
(options.re_min, options.re_max),
88+
(options.im_min, options.im_max),
89+
(options.n_re, options.n_im)
90+
)
91+
domains = np.array_split(domain, options.partitions)
92+
iterations = np.array_split(np.zeros(options.n_re*options.n_im,
93+
dtype=np.int32), options.partitions)
94+
start_time = time.time()
95+
futures = client.map(julia_set, domains, iterations)
96+
results = client.gather(futures)
97+
end_time = time.time()
98+
print('compute time = {0:.6f} s'.format(end_time - start_time))
99+
np.savetxt('julia.txt', np.concatenate(results).reshape(options.n_re,
100+
options.n_im))

‎source-code/dask/julia_set/launch_scheduler.sh‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
#!/usr/bin/env bash
2+
3+
source "conda.sh"
4+
conda activate dask 2> /dev/null
5+
if [ $? -ne 0 ]
6+
then
7+
(>&2 echo '### error: conda environment not sourced correctly' )
8+
fi
9+
10+
nohup dask-scheduler &> "scheduler-${PBS_JOBID}.log" &

‎source-code/dask/julia_set/launch_worker.sh‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
#!/usr/bin/env bash
2+
3+
if [ $# -ne 3 ]
4+
then
5+
(>&2 echo '### error: expecting WORK_DIR, SCHEDULER and JOBID as arguments' )
6+
exit 1
7+
fi
8+
9+
WORK_DIR=$1
10+
if [ ! -d "${WORK_DIR}" ]
11+
then
12+
(>&2 echo "### error: WORK_DIR '${WORK_DIR}' does not exist" )
13+
exit 2
14+
fi
15+
SCHEDULER=$2
16+
JOBID=$3
17+
18+
cd "${WORK_DIR}"
19+
20+
source "conda.sh"
21+
conda activate dask 2> /dev/null
22+
if [ $? -ne 0 ]
23+
then
24+
(>&2 echo '### error: conda environment not sourced correctly' )
25+
fi
26+
27+
nohup dask-worker "${SCHEDULER}" &> "worker-$(hostname)-${JOBID}.log" &

0 commit comments

Comments
 (0)