Skip to content

Commit 4064c2e

Browse files
author
Feng Wang
committed
2 parents 6b1876c + a5e2632 commit 4064c2e

11 files changed

Lines changed: 179 additions & 119 deletions

File tree

File renamed without changes.
File renamed without changes.
File renamed without changes.

‎cluster_tests/ritoy/ritoy.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
""" Radio interferometry toy
2+
3+
This mimics the overall structure and workload of our processing.
4+
5+
Tim Cornwell 9 Sept 2017
6+
realtimcornwell@gmail.com
7+
"""
8+
import numpy
9+
from dask import delayed
10+
from distributed import Client
11+
12+
# Make some randomly located points on 2D plane
13+
def init_sparse(n, margin=0.1):
14+
numpy.random.seed(8753193)
15+
return numpy.array([numpy.random.uniform(margin, 1.0 - margin, n),
16+
numpy.random.uniform(margin, 1.0 - margin, n)]).reshape([n, 2])
17+
18+
19+
# Put the points onto a grid and FFT
20+
def grid_and_invert_data(sparse_data, shape):
21+
grid = numpy.zeros(shape, dtype='complex')
22+
loc = numpy.round(shape * sparse_data).astype('int')
23+
for i in range(0, sparse_data.shape[0]):
24+
grid[loc[i,:]] = 1.0
25+
return numpy.fft.fft(grid).real
26+
27+
if __name__ == '__main__':
28+
import sys
29+
import time
30+
31+
start=time.time()
32+
33+
# Process nchunks each of length len_chunk 2d points, making a psf of size shape
34+
len_chunk = 16384*8
35+
nchunks = 256*4
36+
nreduce = 16*4
37+
shape=[1024, 1024]
38+
skip = 1
39+
40+
# We pass in the scheduler from the invoking script
41+
if len(sys.argv) > 1:
42+
scheduler = sys.argv[1]
43+
client = Client(scheduler)
44+
else:
45+
client = Client()
46+
47+
sparse_graph_list = [delayed(init_sparse)(len_chunk) for i in range(nchunks)]
48+
psf_graph_list = [delayed(grid_and_invert_data)(s, shape) for s in sparse_graph_list]
49+
sum_psf_graph_rank1 = [delayed(numpy.sum)(psf_graph_list[i:i+nreduce]) for i in range(0, nchunks, nreduce)]
50+
sum_psf_graph = delayed(numpy.sum)(sum_psf_graph_rank1)
51+
52+
future = client.compute(sum_psf_graph)
53+
psf = future.result()
54+
print(numpy.max(psf))
55+
56+
client.close()
57+
print("*** Successfully reached end in %.1f seconds ***" % (time.time() - start))
58+
59+
exit(0)
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
#!/bin/bash
2+
#!
3+
#! Dask job script for P3
4+
#! Tim Cornwell
5+
#!
6+
7+
#!#############################################################
8+
#!#### Modify the options in this section as appropriate ######
9+
#!#############################################################
10+
11+
#! sbatch directives begin here ###############################
12+
#! Name of the job:
13+
#SBATCH -J SDP_ARL
14+
#! Which project should be charged:
15+
#SBATCH -A SKA-SDP
16+
#! How many whole nodes should be allocated?
17+
#SBATCH --nodes=8
18+
#! How many (MPI) tasks will there be in total? (<= nodes*16)
19+
#SBATCH --ntasks=16
20+
#! Memory limit
21+
##SBATCH --mem 50000
22+
#! How much wallclock time will be required?
23+
#SBATCH --time=00:20:00
24+
#! What types of email messages do you wish to receive?
25+
#SBATCH --mail-type=FAIL
26+
#! Uncomment this to prevent the job from being requeued (e.g. if
27+
#! interrupted by node failure or system downtime):
28+
##SBATCH --no-requeue
29+
#! Do not change:
30+
#SBATCH -p compute
31+
#! What types of email messages do you wish to receive?
32+
#SBATCH --mail-type=FAIL
33+
#! Uncomment this to prevent the job from being requeued (e.g. if
34+
#! interrupted by node failure or system downtime):
35+
##SBATCH --no-requeue
36+
37+
#! Modify the settings below to specify the application's environment, location
38+
#! and launch method:
39+
40+
#! Optionally modify the environment seen by the application
41+
#! (note that SLURM reproduces the environment at submission irrespective of ~/.bashrc):
42+
module purge # Removes all modules still loaded
43+
44+
#! Set up python
45+
export PYTHONPATH=$PYTHONPATH:$ARL
46+
echo "PYTHONPATH is ${PYTHONPATH}"
47+
48+
echo -e "Running python: `which python`"
49+
echo -e "Running dask-scheduler: `which dask-scheduler`"
50+
51+
cd $SLURM_SUBMIT_DIR
52+
echo -e "Changed directory to `pwd`.\n"
53+
54+
module load slurm
55+
56+
JOBID=${SLURM_JOB_ID}
57+
echo ${SLURM_JOB_NODELIST}
58+
59+
#! Create a hostfile:
60+
scontrol show hostnames $SLURM_JOB_NODELIST | uniq > hostfile.$JOBID
61+
62+
scheduler=$(head -1 hostfile.$JOBID)
63+
hostIndex=0
64+
for host in `cat hostfile.$JOBID`; do
65+
echo "Working on $host ...."
66+
if [ "$hostIndex" = "0" ]; then
67+
echo "run dask-scheduler"
68+
ssh $host dask-scheduler --port=8786 &
69+
sleep 5
70+
fi
71+
echo "run dask-worker"
72+
ssh $host dask-worker --host ${host} --nprocs 2 --nthreads 1 \
73+
--memory-limit 100GB $scheduler:8786 &
74+
sleep 1
75+
hostIndex="1"
76+
done
77+
echo "Scheduler and workers now running"
78+
79+
CMD="python ./ritoy.py ${scheduler}:8786 > ritoy_${JOBID}.out"
80+
echo "About to execute $CMD"
81+
82+
eval $CMD
File renamed without changes.
File renamed without changes.
File renamed without changes.

‎processing_library/util/array_functions.py‎

Lines changed: 27 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,8 @@
11
"""Useful array functions.
22
33
"""
4-
from numba import jit, types as nbtypes, c16, f8, i8,float64, int64, int32,int8
4+
# import numba
55
import numpy
6-
import warnings
7-
from numba import NumbaPerformanceWarning
8-
warnings.simplefilter('ignore', NumbaPerformanceWarning)
96

107

118
# @numba.jit([numba.types.Tuple((numba.float64[:], numba.float64[:]))
@@ -49,40 +46,6 @@ def average_chunks_jit(arr, wts, chunksize):
4946
return chunks, weights
5047

5148

52-
@jit([
53-
nbtypes.Tuple((c16[:], f8[:]))(c16[:], f8[:], i8),
54-
], nopython=True)
55-
def average_chunks_complex(arr, wts, chunksize):
56-
""" Average the array arr with weights by chunks
57-
58-
Array len does not have to be multiple of chunksize
59-
The function is used for processing arr with complex data type.
60-
61-
:param arr: 1D array of values
62-
:param wts: 1D array of weights
63-
:param chunksize: averaging size
64-
:return: 1D array of averaged data_models, 1d array of weights
65-
"""
66-
if chunksize <= 1:
67-
return arr, wts
68-
69-
mask = numpy.zeros(((len(arr) - 1) // chunksize + 1, arr.shape[0]), dtype=arr.dtype)
70-
maskwts = numpy.zeros_like(mask, dtype=wts.dtype)
71-
for enumerate_id, i in enumerate(range(0, len(arr), chunksize)):
72-
mask[enumerate_id, i:i + chunksize] = 1+1j
73-
maskwts[enumerate_id, i:i + chunksize] = 1
74-
tmp = wts * arr
75-
chunks = mask.dot(tmp)
76-
weights = maskwts.dot(wts)
77-
chunks[weights > 0.0] = chunks[weights > 0.0] / weights[weights > 0.0]
78-
79-
return chunks, weights
80-
81-
@jit([nbtypes.Tuple((float64[:], float64[:]))(float64[:], float64[:], int64),
82-
nbtypes.Tuple((f8[:],float64[:]))(f8[:], float64[:], int64),
83-
nbtypes.Tuple((f8[:],float64[:]))(f8[:], float64[:], int32),
84-
nbtypes.Tuple((f8[:],float64[:]))(f8[:], float64[:], int8)
85-
], nopython=True)
8649
def average_chunks(arr, wts, chunksize):
8750
""" Average the array arr with weights by chunks
8851
@@ -108,24 +71,13 @@ def average_chunks(arr, wts, chunksize):
10871

10972
# Codes optimized
11073

111-
# mask = numpy.zeros(((len(arr)-1)//chunksize + 1, arr.shape[0]), dtype=bool)
112-
# for enumerate_id, i in enumerate(range(0, len(arr), chunksize)):
113-
# mask[enumerate_id,i:i+chunksize]= 1
114-
# chunks = mask.dot(wts*arr)
115-
# weights = mask.dot(wts)
116-
# # chunks[weights > 0.0] = chunks[weights > 0.0] / weights[weights > 0.0]
117-
# numpy.putmask(chunks, weights>0.0, chunks/weights)
118-
119-
# Numba
120-
mask = numpy.zeros(((len(arr)-1)//chunksize + 1, arr.shape[0]), dtype=arr.dtype)
121-
maskwts = numpy.zeros_like(mask,dtype=wts.dtype)
122-
for enumerate_id, i in enumerate(range(0, len(arr), chunksize)):
123-
mask[enumerate_id,i:i+chunksize]= 1
124-
maskwts[enumerate_id,i:i+chunksize]= 1
125-
tmp = wts*arr
126-
chunks = mask.dot(tmp)
127-
weights = maskwts.dot(wts)
128-
chunks[weights > 0.0] = chunks[weights > 0.0] / weights[weights > 0.0]
74+
mask = numpy.zeros(((len(arr)-1)//chunksize + 1, arr.shape[0]), dtype=bool)
75+
for enumerate_id,i in enumerate(range(0, len(arr), chunksize)):
76+
mask[enumerate_id,i:i+chunksize]=1
77+
chunks = mask.dot(wts*arr)
78+
weights = mask.dot(wts)
79+
# chunks[weights > 0.0] = chunks[weights > 0.0] / weights[weights > 0.0]
80+
numpy.putmask(chunks, weights>0.0, chunks/weights)
12981

13082
return chunks, weights
13183

@@ -144,66 +96,25 @@ def average_chunks2(arr, wts, chunksize):
14496
# assert arr.shape == wts.shape, "Shapes of arrays must be the same"
14597
# It is possible that there is a dangling null axis on wts
14698
wts = wts.reshape(arr.shape)
147-
#
148-
# For numba
149-
if arr.dtype=='c16':
150-
l0 = len(average_chunks_complex(arr[:, 0].flatten(), wts[:, 0].flatten(), chunksize[0])[0])
151-
l1 = len(average_chunks_complex(arr[0, :].flatten(), wts[0, :].flatten(), chunksize[1])[0])
152-
153-
tempchunks = numpy.zeros([arr.shape[0], l1], dtype=arr.dtype)
154-
tempwt = numpy.zeros([arr.shape[0], l1])
155-
156-
tempchunks *= tempwt
157-
for i in range(arr.shape[0]):
158-
result = average_chunks_complex(arr[i, :], wts[i, :], chunksize[1])
159-
tempchunks[i, :], tempwt[i, :] = result[0].flatten(), result[1].flatten()
160-
161-
chunks = numpy.zeros([l0, l1], dtype=arr.dtype)
162-
weights = numpy.zeros([l0, l1])
163-
164-
for i in range(l1):
165-
result = average_chunks_complex(tempchunks[:, i], tempwt[:, i], chunksize[0])
166-
chunks[:, i], weights[:, i] = result[0].flatten(), result[1].flatten()
167-
168-
else:
169-
l0 = len(average_chunks(arr[:, 0].flatten(), wts[:, 0].flatten(), chunksize[0])[0])
170-
l1 = len(average_chunks(arr[0, :].flatten(), wts[0, :].flatten(), chunksize[1])[0])
171-
172-
tempchunks = numpy.zeros([arr.shape[0], l1], dtype=arr.dtype)
173-
tempwt = numpy.zeros([arr.shape[0], l1])
174-
175-
tempchunks *= tempwt
176-
for i in range(arr.shape[0]):
177-
result = average_chunks(arr[i, :], wts[i, :], chunksize[1])
178-
tempchunks[i, :], tempwt[i, :] = result[0].flatten(), result[1].flatten()
179-
180-
chunks = numpy.zeros([l0, l1], dtype=arr.dtype)
181-
weights = numpy.zeros([l0, l1])
182-
183-
for i in range(l1):
184-
result = average_chunks(tempchunks[:, i], tempwt[:, i], chunksize[0])
185-
chunks[:, i], weights[:, i] = result[0].flatten(), result[1].flatten()
186-
187-
188-
# Original
189-
# l0 = len(average_chunks(arr[:, 0].flatten(), wts[:, 0].flatten(), chunksize[0])[0])
190-
# l1 = len(average_chunks(arr[0, :].flatten(), wts[0, :].flatten(), chunksize[1])[0])
191-
#
192-
# tempchunks = numpy.zeros([arr.shape[0], l1], dtype=arr.dtype)
193-
# tempwt = numpy.zeros([arr.shape[0], l1])
194-
#
195-
# tempchunks *= tempwt
196-
# for i in range(arr.shape[0]):
197-
# result = average_chunks(arr[i, :], wts[i, :], chunksize[1])
198-
# tempchunks[i, :], tempwt[i, :] = result[0].flatten(), result[1].flatten()
199-
#
200-
# chunks = numpy.zeros([l0, l1], dtype=arr.dtype)
201-
# weights = numpy.zeros([l0, l1])
202-
#
203-
# for i in range(l1):
204-
# result = average_chunks(tempchunks[:, i], tempwt[:, i], chunksize[0])
205-
# chunks[:, i], weights[:, i] = result[0].flatten(), result[1].flatten()
206-
99+
100+
l0 = len(average_chunks(arr[:, 0], wts[:, 0], chunksize[0])[0])
101+
l1 = len(average_chunks(arr[0, :], wts[0, :], chunksize[1])[0])
102+
103+
tempchunks = numpy.zeros([arr.shape[0], l1], dtype=arr.dtype)
104+
tempwt = numpy.zeros([arr.shape[0], l1])
105+
106+
tempchunks *= tempwt
107+
for i in range(arr.shape[0]):
108+
result = average_chunks(arr[i, :], wts[i, :], chunksize[1])
109+
tempchunks[i, :], tempwt[i, :] = result[0].flatten(), result[1].flatten()
110+
111+
chunks = numpy.zeros([l0, l1], dtype=arr.dtype)
112+
weights = numpy.zeros([l0, l1])
113+
114+
for i in range(l1):
115+
result = average_chunks(tempchunks[:, i], tempwt[:, i], chunksize[0])
116+
chunks[:, i], weights[:, i] = result[0].flatten(), result[1].flatten()
117+
207118
return chunks, weights
208119

209120

‎wrappers/arlexecute/execution_support/arlexecute.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,5 +3,4 @@
33
"""
44

55
from wrappers.arlexecute.execution_support.arlexecutebase import ARLExecuteBase
6-
76
arlexecute = ARLExecuteBase(use_dask=True)

0 commit comments

Comments
 (0)