@@ -54,16 +54,40 @@ def teardown_temp_folder():
5454with_temp_folder = with_setup (setup_temp_folder , teardown_temp_folder )
5555
5656
57- def double (input ):
58- """Dummy helper function to be executed in subprocesses"""
57+ def setup_if_has_dev_shm ():
58+ if not os .path .exists ('/dev/shm' ):
59+ raise SkipTest ("This test requires the /dev/shm shared memory fs." )
60+
61+
62+ with_dev_shm = with_setup (setup_if_has_dev_shm )
63+
64+
65+ def check_array (args ):
66+ """Dummy helper function to be executed in subprocesses
67+
68+ Check that the provided array has the expected values in the provided
69+ range.
70+
71+ """
5972 assert_array_equal = np .testing .assert_array_equal
73+ data , position , expected = args
74+ assert_equal (data [position ], expected )
6075
61- data , position , expected = input
62- if expected is not None :
63- assert_equal (data [position ], expected )
76+
77+ def inplace_double (args ):
78+ """Dummy helper function to be executed in subprocesses
79+
80+
81+ Check that the input array has the right values in the provided range
82+ and perform an inplace modification to double the values in the range by
83+ two.
84+
85+ """
86+ assert_array_equal = np .testing .assert_array_equal
87+ data , position , expected = args
88+ assert_equal (data [position ], expected )
6489 data [position ] *= 2
65- if expected is not None :
66- assert_array_equal (data [position ], 2 * expected )
90+ assert_equal (data [position ], 2 * expected )
6791
6892
6993@with_numpy
@@ -210,18 +234,18 @@ def test_pool_with_memmap():
210234 a = np .memmap (filename , dtype = np .float32 , shape = (3 , 5 ), mode = 'w+' )
211235 a .fill (1.0 )
212236
213- p .map (double , [(a , (i , j ), 1.0 )
214- for i in range (a .shape [0 ])
215- for j in range (a .shape [1 ])])
237+ p .map (inplace_double , [(a , (i , j ), 1.0 )
238+ for i in range (a .shape [0 ])
239+ for j in range (a .shape [1 ])])
216240
217241 assert_array_equal (a , 2 * np .ones (a .shape ))
218242
219243 # Open a copy-on-write view on the previous data
220244 b = np .memmap (filename , dtype = np .float32 , shape = (5 , 3 ), mode = 'c' )
221245
222- p .map (double , [(b , (i , j ), 2.0 )
223- for i in range (b .shape [0 ])
224- for j in range (b .shape [1 ])])
246+ p .map (inplace_double , [(b , (i , j ), 2.0 )
247+ for i in range (b .shape [0 ])
248+ for j in range (b .shape [1 ])])
225249
226250 # Passing memmap instances to the pool should not trigger the creation
227251 # of new files on the FS
@@ -235,12 +259,12 @@ def test_pool_with_memmap():
235259 c = np .memmap (filename , dtype = np .float32 , shape = (10 ,), mode = 'r' ,
236260 offset = 5 * 4 )
237261
238- assert_raises (AssertionError , p .map , double ,
262+ assert_raises (AssertionError , p .map , check_array ,
239263 [(c , i , 3.0 ) for i in range (c .shape [0 ])])
240264
241265 # depending on the version of numpy one can either get a RuntimeError
242266 # or a ValueError
243- assert_raises ((RuntimeError , ValueError ), p .map , double ,
267+ assert_raises ((RuntimeError , ValueError ), p .map , inplace_double ,
244268 [(c , i , 2.0 ) for i in range (c .shape [0 ])])
245269 finally :
246270 # Clean all filehandlers held by the pool
@@ -270,9 +294,9 @@ def test_pool_with_memmap_array_view():
270294 assert_false (isinstance (a_view , np .memmap ))
271295 assert_true (has_shareable_memory (a_view ))
272296
273- p .map (double , [(a_view , (i , j ), 1.0 )
274- for i in range (a .shape [0 ])
275- for j in range (a .shape [1 ])])
297+ p .map (inplace_double , [(a_view , (i , j ), 1.0 )
298+ for i in range (a .shape [0 ])
299+ for j in range (a .shape [1 ])])
276300
277301 # Both a and the a_view have been updated
278302 assert_array_equal (a , 2 * np .ones (a .shape ))
@@ -307,24 +331,17 @@ def test_memmaping_pool_for_large_arrays():
307331
308332 small = np .ones (5 , dtype = np .float32 )
309333 assert_equal (small .nbytes , 20 )
310- p .map (double , [(small , i , 1.0 ) for i in range (small .shape [0 ])])
334+ p .map (check_array , [(small , i , 1.0 ) for i in range (small .shape [0 ])])
311335
312336 # Memory has been copied, the pool filesystem folder is unused
313337 assert_equal (os .listdir (TEMP_FOLDER ), [])
314338
315339 # Try with a file larger than the memmap threshold of 40 bytes
316340 large = np .ones (100 , dtype = np .float64 )
317341 assert_equal (large .nbytes , 800 )
318- p .map (double , [(large , i , 1.0 ) for i in range (large .shape [0 ])])
342+ p .map (check_array , [(large , i , 1.0 ) for i in range (large .shape [0 ])])
319343
320- # By defaul, the mmap_mode is copy-on-write to make the pool
321- # process able to modify their view individually as if they would have
322- # received their own copy of the original array. The original array
323- # (which is not a shared memmap instance is untouched)
324- assert_false (has_shareable_memory (large ))
325- assert_array_equal (large , np .ones (100 ))
326-
327- # The data has been dump in a temp folder for subprocess to share it
344+ # The data has been dumped in a temp folder for subprocess to share it
328345 # without per-child memory copies
329346 assert_true (os .path .isdir (p ._temp_folder ))
330347 dumped_filenames = os .listdir (p ._temp_folder )
@@ -352,7 +369,7 @@ def test_memmaping_pool_for_large_arrays_disabled():
352369 # Try with a file largish than the memmap threshold of 40 bytes
353370 large = np .ones (100 , dtype = np .float64 )
354371 assert_equal (large .nbytes , 800 )
355- p .map (double , [(large , i , 1.0 ) for i in range (large .shape [0 ])])
372+ p .map (check_array , [(large , i , 1.0 ) for i in range (large .shape [0 ])])
356373
357374 # Check that the tempfolder is still empty
358375 assert_equal (os .listdir (TEMP_FOLDER ), [])
@@ -363,6 +380,43 @@ def test_memmaping_pool_for_large_arrays_disabled():
363380 del p
364381
365382
383+ @with_numpy
384+ @with_multiprocessing
385+ @with_dev_shm
386+ def test_memmaping_on_dev_shm ():
387+ """Check that large arrays memmaping can be disabled"""
388+ p = MemmapingPool (3 , max_nbytes = 10 )
389+ try :
390+ # Check that the pool has correctly detected the presence of the
391+ # shared memory filesystem.
392+ pool_temp_folder = p ._temp_folder
393+ folder_prefix = '/dev/shm/joblib_memmaping_pool_'
394+ assert_true (pool_temp_folder .startswith (folder_prefix ))
395+ assert_true (os .path .exists (pool_temp_folder ))
396+
397+ # Try with a file larger than the memmap threshold of 10 bytes
398+ a = np .ones (100 , dtype = np .float64 )
399+ assert_equal (a .nbytes , 800 )
400+ p .map (id , [a ] * 10 )
401+ # a should have been memmaped to the pool temp folder: the joblib
402+ # pickling procedure generate a .pkl and a .npy file:
403+ assert_equal (len (os .listdir (pool_temp_folder )), 2 )
404+
405+ b = np .ones (100 , dtype = np .float64 )
406+ assert_equal (b .nbytes , 800 )
407+ p .map (id , [b ] * 10 )
408+ # A copy of both a and b are not stored in the shared memory folder
409+ assert_equal (len (os .listdir (pool_temp_folder )), 4 )
410+
411+ finally :
412+ # Cleanup open file descriptors
413+ p .terminate ()
414+ del p
415+
416+ # The temp folder is cleaned up upon pool termination
417+ assert_false (os .path .exists (pool_temp_folder ))
418+
419+
366420@with_numpy
367421@with_multiprocessing
368422@with_temp_folder
0 commit comments