Skip to content

Commit abbe4db

Browse files
CheerfulUserCheerfulUser
committed
Switch from prefer='threads' to backend='multiprocessing' (fork-based)
prefer='threads' caused high system CPU time on OzSTAR because all threads share one address space: simultaneous malloc/free calls from numpy temporary arrays contend on the glibc allocator lock, and cross-socket memory access on NUMA hardware requires OS arbitration — both appearing as system rather than user CPU time. backend='multiprocessing' uses Linux fork() directly. Each worker gets its own address space via copy-on-write, eliminating allocator contention and NUMA cross-talk. Unlike the loky backend (which uses a forkserver that was failing silently on OzSTAR), the multiprocessing backend uses a direct fork() call that works reliably in SLURM environments.
1 parent 1cfb087 commit abbe4db

7 files changed

Lines changed: 34 additions & 34 deletions

File tree

‎tessreduce/adaptive_background.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ def _compute_dev(cw, data_metric=data_metric, segments=segments, gw=gw):
236236
dev[s:e] = median_filter(diff, size=(dw, 1, 1), mode='reflect')
237237
return dev
238238

239-
all_devs = Parallel(n_jobs=n_jobs, prefer="threads")(delayed(_compute_dev)(cw) for cw in scales)
239+
all_devs = Parallel(n_jobs=n_jobs, backend="multiprocessing")(delayed(_compute_dev)(cw) for cw in scales)
240240

241241
_frame_mean = data_metric.mean(axis=(1, 2))
242242
_clipped = _frame_mean[np.isfinite(_frame_mean)]
@@ -280,7 +280,7 @@ def _compute_norm(i_dev, bright_mask=bright_mask, segments=segments,
280280
norm_scale = norm_scale * bright_mask
281281
return norm_scale
282282

283-
scale_norms = Parallel(n_jobs=n_jobs, prefer="threads")(
283+
scale_norms = Parallel(n_jobs=n_jobs, backend="multiprocessing")(
284284
delayed(_compute_norm)((i, dev)) for i, dev in enumerate(all_devs)
285285
)
286286

@@ -377,7 +377,7 @@ def _norm01(x):
377377
def _smooth_seg(w, seg=seg_data):
378378
return w, median_filter(seg, size=(w, 1, 1), mode='reflect')
379379

380-
for w, smoothed_w in Parallel(n_jobs=n_jobs, prefer="threads")(delayed(_smooth_seg)(w) for w in seg_levels):
380+
for w, smoothed_w in Parallel(n_jobs=n_jobs, backend="multiprocessing")(delayed(_smooth_seg)(w) for w in seg_levels):
381381
result[s:e][seg_wins == w] = smoothed_w[seg_wins == w]
382382

383383
if sigma_clip is not None:

‎tessreduce/background.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ def _smooth_wrapper(self):
8080

8181
if self.parallel:
8282
num_cores = multiprocessing.cpu_count()
83-
bkg_smth = Parallel(n_jobs=num_cores, prefer="threads")(
83+
bkg_smth = Parallel(n_jobs=num_cores, backend="multiprocessing")(
8484
delayed(Smooth_bkg)(frame) for frame in flux * m)
8585
else:
8686
bkg_smth = np.zeros_like(flux) * np.nan

‎tessreduce/helpers.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ def parallel_background2d(cube, box_size=5, filter_size=3, sigma=3, maxiters=5,
197197
from astropy.stats import SigmaClip
198198
sc = SigmaClip(sigma=sigma, maxiters=maxiters)
199199
estimator = MedianBackground()
200-
return np.array(Parallel(n_jobs=n_jobs, prefer="threads")(
200+
return np.array(Parallel(n_jobs=n_jobs, backend="multiprocessing")(
201201
delayed(_background2d_frame)(frame, box_size, filter_size, sc, estimator, mask)
202202
for frame in cube))
203203

@@ -1130,7 +1130,7 @@ def regional_stats_mask(image, size=90, sigma=3, iters=10, n_jobs=1):
11301130
region_pixels = [np.where(regions == i) for i in range(max_reg + 1)]
11311131

11321132
clip = np.zeros_like(image)
1133-
results = Parallel(n_jobs=n_jobs, prefer="threads")(
1133+
results = Parallel(n_jobs=n_jobs, backend="multiprocessing")(
11341134
delayed(_clip_region)(image, rx, ry, sigma, iters)
11351135
for rx, ry in region_pixels)
11361136
for rx_cut, ry_cut in results:
@@ -1355,7 +1355,7 @@ def parallel_strap_fit(frame, frame_bkg, frame_err, mask, repeats=3, tol=3, n_jo
13551355
qe = np.ones_like(frame)
13561356
if len(sind) == 0:
13571357
return qe
1358-
results = Parallel(n_jobs=n_jobs, prefer="threads")(
1358+
results = Parallel(n_jobs=n_jobs, backend="multiprocessing")(
13591359
delayed(_strap_fit_col)(frame[:, i], norm[:, i])
13601360
for i in sind)
13611361
for col, q in zip(sind, results):
@@ -1844,7 +1844,7 @@ def blend_dynamic_background(bkg_new, bkg_prev, flux, sigma=2.0, sharp_masks=Non
18441844

18451845
sharp_list = [sharp_masks[i] if sharp_masks is not None else None for i in range(T)]
18461846

1847-
results = Parallel(n_jobs=n_jobs, prefer="threads")(
1847+
results = Parallel(n_jobs=n_jobs, backend="multiprocessing")(
18481848
delayed(_blend_frame)(bkg_new[i], bkg_prev[i], delta[i], resid_prev[i],
18491849
sigma, sharp_list[i], gauss_kernel)
18501850
for i in range(T))

‎tessreduce/lastpercent.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ def _find_bkg_cor(tess,cores):
5555
coord = np.c_[y,x]
5656
cors = np.zeros_like(tess.ref)
5757

58-
cor = Parallel(n_jobs=cores, prefer="threads")(delayed(_parallel_correlation)
58+
cor = Parallel(n_jobs=cores, backend="multiprocessing")(delayed(_parallel_correlation)
5959
(tess.flux[:,coord[i,0],coord[i,1]],
6060
tess.bkg[:,coord[i,0],coord[i,1]],
6161
cors,coord[i],30) for i in range(len(coord)))
@@ -156,7 +156,7 @@ def multi_correlation_cor(tess, limit=0.8, cores=7):
156156
if len(y) == 0:
157157
return flux, bkg
158158

159-
results = Parallel(n_jobs=cores, prefer="threads")(
159+
results = Parallel(n_jobs=cores, backend="multiprocessing")(
160160
delayed(_correct_pixel_correlation)(
161161
tess.flux[:, y[i], x[i]],
162162
tess.bkg[:, y[i], x[i]],

‎tessreduce/rescale_straps.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ def correct_straps(Image,mask,av_size=5,parallel=True):
123123
if parallel:
124124
num_cores = multiprocessing.cpu_count()
125125
x = np.arange(0,len(breaks),dtype=int)
126-
qe = np.array(Parallel(n_jobs=num_cores, prefer="threads")(delayed(calc_strap_factor)(i,breaks,size,av_size,normals,data) for i in x))
126+
qe = np.array(Parallel(n_jobs=num_cores, backend="multiprocessing")(delayed(calc_strap_factor)(i,breaks,size,av_size,normals,data) for i in x))
127127
qe = np.nanmedian(qe,axis=0)
128128
qe[np.isnan(qe)] = 1
129129
else:

‎tessreduce/sep_aligner.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -703,7 +703,7 @@ def run(self, time: Optional[np.ndarray] = None,
703703
self._sub_ref, cores, positions, np.asarray(weights), p['core_half'])
704704

705705
results = Parallel(n_jobs=self.n_jobs, verbose=verbose,
706-
prefer="threads")(
706+
backend="multiprocessing")(
707707
delayed(_align_one_frame)(
708708
t, self.flux[t],
709709
ref_comp, w_cols,

‎tessreduce/tessreduce.py‎

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ def _fit_frame(residual):
127127

128128
n_jobs = _available_cores()
129129
residuals = flux - bkg
130-
corrections = Parallel(n_jobs=n_jobs, prefer="threads")(delayed(_fit_frame)(residuals[i]) for i in range(flux.shape[0]))
130+
corrections = Parallel(n_jobs=n_jobs, backend="multiprocessing")(delayed(_fit_frame)(residuals[i]) for i in range(flux.shape[0]))
131131
bkg += np.array(corrections)
132132

133133
return bkg
@@ -651,7 +651,7 @@ def psf_source_mask(self,sigma=5):
651651
data = (self._flux_aligned - self.ref) #* mask
652652
if self.parallel:
653653
try:
654-
m = Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_source_mask)(frame,self.prf,sigma) for frame in data)
654+
m = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_source_mask)(frame,self.prf,sigma) for frame in data)
655655
m = np.array(m)
656656
except:
657657
m = np.ones_like(data)
@@ -778,7 +778,7 @@ def background(self,gauss_smooth=None,calc_qe=True,strap_iso=True,source_hunt=Fa
778778
bkg_smth = np.zeros_like(flux) * np.nan
779779
if self.parallel:
780780
_t = time.perf_counter()
781-
bkg_smth = Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(Smooth_bkg)(frame,0,interpolate) for frame in flux*m)
781+
bkg_smth = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(Smooth_bkg)(frame,0,interpolate) for frame in flux*m)
782782
_times['initial smooth background'] = time.perf_counter() - _t
783783
if rerun_negative:
784784
_t = time.perf_counter()
@@ -800,7 +800,7 @@ def background(self,gauss_smooth=None,calc_qe=True,strap_iso=True,source_hunt=Fa
800800
else:
801801
m[over_sub] = 1
802802
self._bkgmask = m
803-
bkg_smth = Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(Smooth_bkg)(frame,gauss_smooth,interpolate) for frame in flux*m)
803+
bkg_smth = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(Smooth_bkg)(frame,gauss_smooth,interpolate) for frame in flux*m)
804804
_times['negative over-subtraction rerun'] = time.perf_counter() - _t
805805

806806

@@ -853,7 +853,7 @@ def background(self,gauss_smooth=None,calc_qe=True,strap_iso=True,source_hunt=Fa
853853
new_mask = abs(new_mask - 1)
854854
self._bkgmask = new_mask
855855
bkg_s1 = np.array(bkg_smth)
856-
bkg_smth = Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(Smooth_bkg)(frame,0,interpolate) for frame in flux*new_mask)
856+
bkg_smth = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(Smooth_bkg)(frame,0,interpolate) for frame in flux*new_mask)
857857
if blend_dynamic:
858858
bkg_smth = blend_dynamic_background(bkg_smth, bkg_s1, flux, n_jobs=self.num_cores)
859859
_times['residual surface rerun'] = time.perf_counter() - _t
@@ -996,7 +996,7 @@ def _bkg_round_3(self,iters=5):
996996
kern = np.ones((1,3,3))
997997
dist_mask = convolve(dist_mask,kern) > 0
998998
if self.parallel:
999-
bkg_3 = Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(parallel_bkg3)(self.bkg[i],dist_mask[i])
999+
bkg_3 = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(parallel_bkg3)(self.bkg[i],dist_mask[i])
10001000
for i in np.arange(len(dist_mask)))
10011001
else:
10021002
bkg_3 = np.zeros_like(self.bkg)
@@ -1020,7 +1020,7 @@ def _clip_background(self,sigma=5,ideal_size=90):
10201020
"""
10211021

10221022
if self.parallel:
1023-
bkg_clip = Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(clip_background)(self.bkg[i],self.mask,sigma,ideal_size)
1023+
bkg_clip = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(clip_background)(self.bkg[i],self.mask,sigma,ideal_size)
10241024
for i in np.arange(len(self.bkg)))
10251025
else:
10261026
bkg_clip = np.zeros_like(self.bkg)
@@ -1046,7 +1046,7 @@ def _grad_bkg_clip(self,sigma=3,max_size=1000):
10461046
"""
10471047

10481048
if self.parallel:
1049-
bkg_clip = Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(grad_clip_fill_bkg)(self.bkg[i],sigma,max_size)
1049+
bkg_clip = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(grad_clip_fill_bkg)(self.bkg[i],sigma,max_size)
10501050
for i in np.arange(len(self.bkg)))
10511051
else:
10521052
bkg_clip = np.zeros_like(self.bkg)
@@ -1402,7 +1402,7 @@ def centroids_shifts_starfind(self,plot=None,savename=None):
14021402
self._dat_sources = s.to_pandas()
14031403

14041404
if self.parallel:
1405-
shifts = Parallel(n_jobs=self.num_cores, prefer="threads")(
1405+
shifts = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(
14061406
delayed(Calculate_shifts)(frame,mx,my,finder) for frame in f)
14071407
shifts = np.array(shifts)
14081408
else:
@@ -1474,7 +1474,7 @@ def fit_shift(self,smooth=True,plot=None,savename=None):
14741474

14751475
if self.parallel:
14761476
ind = np.arange(len(f))
1477-
shifts = Parallel(n_jobs=self.num_cores, prefer="threads")(
1477+
shifts = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(
14781478
#delayed(difference_shifts)(f[i],m,self.eflux[i],eref) for i in ind)
14791479
delayed(difference_shifts)(f[i],m) for i in ind)
14801480
shifts = np.array(shifts)
@@ -1548,7 +1548,7 @@ def shift_images(self,median=False):
15481548
shifted[nans] = 0.
15491549
if median:
15501550
if self.parallel:
1551-
result = Parallel(n_jobs=self.num_cores, prefer="threads")(
1551+
result = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(
15521552
delayed(_shift_ref_one)(self.ref, shifted[i], self.shift[i])
15531553
for i in range(len(shifted)))
15541554
shifted = np.array(result)
@@ -1559,7 +1559,7 @@ def shift_images(self,median=False):
15591559

15601560
else:
15611561
if self.parallel:
1562-
result = Parallel(n_jobs=self.num_cores, prefer="threads")(
1562+
result = Parallel(n_jobs=self.num_cores, backend="multiprocessing")(
15631563
delayed(_shift_one)(shifted[i], self.shift[i])
15641564
for i in range(len(shifted)))
15651565
self.flux = np.array(result)
@@ -2267,7 +2267,7 @@ def moving_psf_photometry(self,xpos,ypos,size=5,time_ind=None,xlim=2,ylim=2):
22672267
raise ValueError(m)
22682268
inds = np.arange(0,len(xpos))
22692269
if self.parallel:
2270-
prfs, cutouts, ecutouts = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_initialise)(self.flux,self.tpf.camera,self.tpf.ccd,
2270+
prfs, cutouts, ecutouts = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_initialise)(self.flux,self.tpf.camera,self.tpf.ccd,
22712271
self.tpf.sector,self.tpf.column,self.tpf.row,
22722272
size,[xpos[i],ypos[i]],time_ind) for i in inds))
22732273
else:
@@ -2280,7 +2280,7 @@ def moving_psf_photometry(self,xpos,ypos,size=5,time_ind=None,xlim=2,ylim=2):
22802280
cutouts = np.array(cutouts)
22812281
print('made cutouts')
22822282
if self.parallel:
2283-
flux, pos = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_full)(cutouts[i],prfs[i],self.shift[i],xlim,ylim) for i in inds))
2283+
flux, pos = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_full)(cutouts[i],prfs[i],self.shift[i],xlim,ylim) for i in inds))
22842284
else:
22852285
flux = []
22862286
pos = []
@@ -2366,14 +2366,14 @@ def psf_photutils(self,xPix=None,yPix=None,size=5,local_bkg=False,epsf=None,
23662366
eflux = np.zeros(len(self.flux)) * np.nan
23672367
psfphot2 = PSFPhotometry(epsf, fit_shape, finder=None,aperture_radius=1.5,
23682368
xy_bounds=(0.05),localbkg_estimator=localbkg_estimator)
2369-
f,ef = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(parallel_photutils)(cutouts[i],ecutouts[i],psfphot2,init) for i in np.arange(len(cutouts))))
2369+
f,ef = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(parallel_photutils)(cutouts[i],ecutouts[i],psfphot2,init) for i in np.arange(len(cutouts))))
23702370
f = np.array(f).flatten()
23712371
ef = np.array(ef).flatten()
23722372
phot = phot.to_pandas()
23732373
pos = phot[['x_fit','y_fit']].values + np.array([xPix,yPix]) - size//2
23742374
epos = phot[['x_err','y_err']].values
23752375
else:
2376-
f,ef,pos,epos = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(parallel_photutils)(cutouts[i],ecutouts[i],psfphot,init,True) for i in np.arange(len(cutouts))))
2376+
f,ef,pos,epos = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(parallel_photutils)(cutouts[i],ecutouts[i],psfphot,init,True) for i in np.arange(len(cutouts))))
23772377
pos['x_fit'] += xPix - size//2
23782378
pos['y_fit'] += xPix - size//2
23792379

@@ -2437,7 +2437,7 @@ def psf_photometry(self,xPix,yPix,size=7,snap='brightest',ext_shift=True,plot=Fa
24372437
prf, cutouts, ecutouts = self._psf_initialise(size,(xPix,yPix),ref=(not diff)) # gather base PRF and the array of cutouts data
24382438
inds = np.arange(len(cutouts))
24392439
base = create_psf(prf,size)
2440-
flux, eflux, pos = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_full)(cutouts[i],base,self.shift[i]) for i in inds))
2440+
flux, eflux, pos = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_full)(cutouts[i],base,self.shift[i]) for i in inds))
24412441

24422442
#prf, cutouts = self._psf_initialise(size,(xPix,yPix)) # gather base PRF and the array of cutouts data
24432443
#xShifts = []
@@ -2483,9 +2483,9 @@ def psf_photometry(self,xPix,yPix,size=7,snap='brightest',ext_shift=True,plot=Fa
24832483
if self.parallel:
24842484
inds = np.arange(len(cutouts))
24852485
if self.delta_kernel is not None:
2486-
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order,self.delta_kernel[i]) for i in inds))
2486+
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order,self.delta_kernel[i]) for i in inds))
24872487
else:
2488-
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order) for i in inds))
2488+
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order) for i in inds))
24892489
else:
24902490
for i in range(len(cutouts)):
24912491
flux += [par_psf_flux(cutouts[i],ecutouts[i],base,self.shift[i])]
@@ -2501,9 +2501,9 @@ def psf_photometry(self,xPix,yPix,size=7,snap='brightest',ext_shift=True,plot=Fa
25012501
if self.parallel:
25022502
inds = np.arange(len(cutouts))
25032503
if self.delta_kernel is not None:
2504-
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order,self.delta_kernel[i]) for i in inds))
2504+
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order,self.delta_kernel[i]) for i in inds))
25052505
else:
2506-
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order) for i in inds))
2506+
flux, eflux = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(par_psf_flux)(cutouts[i],ecutouts[i],base,self.shift[i],bkg_poly_order) for i in inds))
25072507
else:
25082508
for i in range(len(cutouts)):
25092509
flux += [par_psf_flux(cutouts[i],ecutouts[i],base,self.shift[i])]
@@ -2548,7 +2548,7 @@ def kernel_matching(self,size=7,diff=True):
25482548
mask = self.mask == 1
25492549

25502550
if self.parallel:
2551-
d, kernel = zip(*Parallel(n_jobs=self.num_cores, prefer="threads")(delayed(parallel_delta_diff)(frame,self.ref,mask,size) for frame in flux))
2551+
d, kernel = zip(*Parallel(n_jobs=self.num_cores, backend="multiprocessing")(delayed(parallel_delta_diff)(frame,self.ref,mask,size) for frame in flux))
25522552
else:
25532553
d = []
25542554
kernel = []

0 commit comments

Comments
 (0)