<s> import contextlib <EOL> import gc <EOL> import os <EOL> import pickle <EOL> import time <EOL> from helpers import unittest <EOL> import luigi <EOL> import mock <EOL> import psutil <EOL> from luigi . worker import Worker <EOL> from luigi . task_status import UNKNOWN <EOL> def running_children ( ) : <EOL> children = set ( ) <EOL> process = psutil . Process ( os . getpid ( ) ) <EOL> for child in process . children ( ) : <EOL> if child . is_running ( ) : <EOL> children . add ( child . pid ) <EOL> return children <EOL> @ contextlib . contextmanager <EOL> def pause_gc ( ) : <EOL> if not gc . isenabled ( ) : <EOL> yield <EOL> try : <EOL> gc . disable ( ) <EOL> yield <EOL> finally : <EOL> gc . enable ( ) <EOL> class SlowCompleteWrapper ( luigi . WrapperTask ) : <EOL> def requires ( self ) : <EOL> return [ SlowCompleteTask ( i ) for i in range ( <NUM_LIT:4> ) ] <EOL> class SlowCompleteTask ( luigi . Task ) : <EOL> n = luigi . IntParameter ( ) <EOL> def complete ( self ) : <EOL> time . sleep ( <NUM_LIT:0.1> ) <EOL> return True <EOL> class OverlappingSelfDependenciesTask ( luigi . Task ) : <EOL> n = luigi . IntParameter ( ) <EOL> k = luigi . IntParameter ( ) <EOL> def complete ( self ) : <EOL> return self . n < self . k or self . k == <NUM_LIT:0> <EOL> def requires ( self ) : <EOL> return [ OverlappingSelfDependenciesTask ( self . n - <NUM_LIT:1> , k ) for k in range ( self . k + <NUM_LIT:1> ) ] <EOL> class ExceptionCompleteTask ( luigi . Task ) : <EOL> def complete ( self ) : <EOL> assert False <EOL> class ExceptionRequiresTask ( luigi . Task ) : <EOL> def requires ( self ) : <EOL> assert False <EOL> class UnpicklableExceptionTask ( luigi . Task ) : <EOL> def complete ( self ) : <EOL> class UnpicklableException ( Exception ) : <EOL> pass <EOL> raise UnpicklableException ( ) <EOL> class ParallelSchedulingTest ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> self . sch = mock . Mock ( ) <EOL> self . w = Worker ( scheduler = self . sch , worker_id = '<STR_LIT:x>' ) <EOL> def added_tasks ( self , status ) : <EOL> return [ kw [ '<STR_LIT>' ] for args , kw in self . sch . add_task . call_args_list if kw [ '<STR_LIT:status>' ] == status ] <EOL> def test_children_terminated ( self ) : <EOL> before_children = running_children ( ) <EOL> with pause_gc ( ) : <EOL> self . w . add ( <EOL> OverlappingSelfDependenciesTask ( <NUM_LIT:5> , <NUM_LIT:2> ) , <EOL> multiprocess = True , <EOL> ) <EOL> self . assertLessEqual ( running_children ( ) , before_children ) <EOL> def test_multiprocess_scheduling_with_overlapping_dependencies ( self ) : <EOL> self . w . add ( OverlappingSelfDependenciesTask ( <NUM_LIT:5> , <NUM_LIT:2> ) , True ) <EOL> self . assertEqual ( <NUM_LIT:15> , self . sch . add_task . call_count ) <EOL> self . assertEqual ( set ( ( <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:1> , k = <NUM_LIT:1> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:2> , k = <NUM_LIT:1> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:2> , k = <NUM_LIT:2> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:3> , k = <NUM_LIT:1> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:3> , k = <NUM_LIT:2> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:4> , k = <NUM_LIT:1> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:4> , k = <NUM_LIT:2> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:5> , k = <NUM_LIT:2> ) . task_id , <EOL> ) ) , set ( self . added_tasks ( '<STR_LIT>' ) ) ) <EOL> self . assertEqual ( set ( ( <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:0> , k = <NUM_LIT:0> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:0> , k = <NUM_LIT:1> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:1> , k = <NUM_LIT:0> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:1> , k = <NUM_LIT:2> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:2> , k = <NUM_LIT:0> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:3> , k = <NUM_LIT:0> ) . task_id , <EOL> OverlappingSelfDependenciesTask ( n = <NUM_LIT:4> , k = <NUM_LIT:0> ) . task_id , <EOL> ) ) , set ( self . added_tasks ( '<STR_LIT>' ) ) ) <EOL> @ mock . patch ( '<STR_LIT>' ) <EOL> def test_raise_exception_in_complete ( self , send ) : <EOL> self . w . add ( ExceptionCompleteTask ( ) , multiprocess = True ) <EOL> send . check_called_once ( ) <EOL> self . assertEqual ( UNKNOWN , self . sch . add_task . call_args [ <NUM_LIT:1> ] [ '<STR_LIT:status>' ] ) <EOL> self . assertFalse ( self . sch . add_task . call_args [ <NUM_LIT:1> ] [ '<STR_LIT>' ] ) <EOL> self . assertTrue ( '<STR_LIT>' in send . call_args [ <NUM_LIT:0> ] [ <NUM_LIT:1> ] ) <EOL> @ mock . patch ( '<STR_LIT>' ) <EOL> def test_raise_unpicklable_exception_in_complete ( self , send ) : <EOL> self . assertRaises ( Exception , UnpicklableExceptionTask ( ) . complete ) <EOL> try : <EOL> UnpicklableExceptionTask ( ) . complete ( ) <EOL> except Exception as e : <EOL> ex = e <EOL> self . assertRaises ( ( pickle . PicklingError , AttributeError ) , pickle . dumps , ex ) <EOL> self . w . add ( UnpicklableExceptionTask ( ) , multiprocess = True ) <EOL> send . check_called_once ( ) <EOL> self . assertEqual ( UNKNOWN , self . sch . add_task . call_args [ <NUM_LIT:1> ] [ '<STR_LIT:status>' ] ) <EOL> self . assertFalse ( self . sch . add_task . call_args [ <NUM_LIT:1> ] [ '<STR_LIT>' ] ) <EOL> self . assertTrue ( '<STR_LIT>' in send . call_args [ <NUM_LIT:0> ] [ <NUM_LIT:1> ] ) <EOL> @ mock . patch ( '<STR_LIT>' ) <EOL> def test_raise_exception_in_requires ( self , send ) : <EOL> self . w . add ( ExceptionRequiresTask ( ) , multiprocess = True ) <EOL> send . check_called_once ( ) <EOL> self . assertEqual ( UNKNOWN , self . sch . add_task . call_args [ <NUM_LIT:1> ] [ '<STR_LIT:status>' ] ) <EOL> self . assertFalse ( self . sch . add_task . call_args [ <NUM_LIT:1> ] [ '<STR_LIT>' ] ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
