def dask_scheduler_create_command(args): from tethys_compute.models.dask.dask_scheduler import DaskScheduler name = args.name host = args.endpoint timeout = args.timeout heartbeat_interval = args.heartbeat_interval dashboard = args.dashboard existing_scheduler = DaskScheduler.objects.filter(name=name).first() if existing_scheduler: with pretty_output(FG_YELLOW) as p: p.write('A Dask Scheduler with name "{}" already exists. Command aborted.'.format(name)) exit(0) scheduler = DaskScheduler( name=name, host=host, timeout=timeout, heartbeat_interval=heartbeat_interval, dashboard=dashboard, ) scheduler.save() with pretty_output(FG_GREEN) as p: p.write('Dask Scheduler created successfully!') exit(0)
def dask_scheduler_create_command(args): load_apps() from tethys_compute.models.dask.dask_scheduler import DaskScheduler name = args.name host = args.endpoint timeout = args.timeout heartbeat_interval = args.heartbeat_interval dashboard = args.dashboard existing_scheduler = DaskScheduler.objects.filter(name=name).first() if existing_scheduler: with pretty_output(FG_YELLOW) as p: p.write('A Dask Scheduler with name "{}" already exists. Command aborted.'.format(name)) exit(0) scheduler = DaskScheduler( name=name, host=host, timeout=timeout, heartbeat_interval=heartbeat_interval, dashboard=dashboard, ) scheduler.save() with pretty_output(FG_GREEN) as p: p.write('Dask Scheduler created successfully!') exit(0)
def set_up(self): self.user = User.objects.create_user('tethys_super', '*****@*****.**', 'pass') self.scheduler = DaskScheduler( name='test_dask_scheduler', host='127.0.0.1:8000', timeout=10, heartbeat_interval=5, dashboard='test_dashboard', ) self.scheduler.save()
class DaskJobMockedResultsPropertyTests(TethysTestCase): def set_up(self): self.user = User.objects.create_user('tethys_super', '*****@*****.**', 'pass') self.scheduler = DaskScheduler( name='test_dask_scheduler', host='127.0.0.1:8000', timeout=10, heartbeat_interval=5, dashboard='test_dashboard', ) self.scheduler.save() def tear_down(self): pass @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.log') def test_process_result_serialize_exception(self, mock_log, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # NOTE: To mock the "result" property, we must mock it on the type object, not the instance. # Unfortunately, this will persist for any test in the same test case that runs after this test. # That's why this test is pulled out in a separate test case (the other tests on "result" won't work # after this one runs). type(djob).result = mock.PropertyMock( side_effect=[Exception, 'foo', 'foo']) djob._process_results() # check the result mock_client.gather.assert_called_with(mock_future) mock_function.assert_called_with(mock_client.gather()) mock_log.exception.assert_called_with('Results Serialization Error') mock_re_lock.assert_called() self.assertEqual('ERR', djob._status)
class DaskJobMockedResultsPropertyTests(TethysTestCase): def set_up(self): self.user = User.objects.create_user('tethys_super', '*****@*****.**', 'pass') self.scheduler = DaskScheduler( name='test_dask_scheduler', host='127.0.0.1:8000', timeout=10, heartbeat_interval=5, dashboard='test_dashboard', ) self.scheduler.save() def tear_down(self): pass @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.log') def test_process_result_serialize_exception(self, mock_log, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function' ) # NOTE: To mock the "result" property, we must mock it on the type object, not the instance. # Unfortunately, this will persist for any test in the same test case that runs after this test. # That's why this test is pulled out in a separate test case (the other tests on "result" won't work # after this one runs). type(djob).result = mock.PropertyMock(side_effect=[Exception, 'foo', 'foo']) djob._process_results() # check the result mock_client.gather.assert_called_with(mock_future) mock_function.assert_called_with(mock_client.gather()) mock_log.exception.assert_called_with('Results Serialization Error') mock_re_lock.assert_called() self.assertEqual('ERR', djob._status)
def create_dask_scheduler(name, host, timeout=None, heartbeat_interval=None, dashboard=None): """ Creates a new dask scheduler Args: name (str): The name of the scheduler host (str): The hostname or IP address of the scheduler timeout (str, optional): heartbeat_interval (str, optional): dashboard (str, optional): Returns: The newly created dask scheduler Note: The newly created dask scheduler object is not committed to the database. """ dask_scheduler = DaskScheduler(name, host, timeout=timeout, heartbeat_interval=heartbeat_interval, dashboard=dashboard) return dask_scheduler
class DaskJobTest(TethysTestCase): def set_up(self): self.user = User.objects.create_user('tethys_super', '*****@*****.**', 'pass') self.scheduler = DaskScheduler( name='test_dask_scheduler', host='127.0.0.1:8000', timeout=10, heartbeat_interval=5, dashboard='test_dashboard', ) self.scheduler.save() def tear_down(self): self.scheduler.delete() @mock.patch('tethys_compute.models.dask.dask_job.Client') def test_client_prop_with_invalid_scheduler(self, mock_client): mock_client.return_value = 'test_client' djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=None) # Execute ret = djob.client # Check result self.assertEqual('test_client', ret) mock_client.assert_called() @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_client_prop_with_valid_scheduler(self, mock_client): mock_client.return_value = 'test_client' dask_scheduler = Scheduler.objects.get_subclass( name='test_dask_scheduler') djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=dask_scheduler) # Execute ret = djob.client # Check result self.assertEqual('test_client', ret) mock_client.assert_called_with(address='127.0.0.1:8000', heartbeat_interval=5, timeout=10) @mock.patch('tethys_compute.models.dask.dask_job.Client') def test_client_no_scheduler_prop(self, mock_client): mock_client.return_value = 'test_default_client' # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label') # Execute ret = djob.client # Check result self.assertEqual('test_default_client', ret) mock_client.assert_called_with() @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') @mock.patch('tethys_compute.models.dask.dask_job.Future') def test_future_prop(self, mock_future, mock_client): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run rando function with a future handler future = client.submit(inc, 1) # Get the key from future handler and assign it to DaskJob key to keep track of this inc function djob.key = future.key # Use DaskJob future property to get back the inc function ret = djob.future # Check result mock_future.assert_called_with(key='test_key', client=mock_client_ret) self.assertEqual(mock_future(), ret) @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_future_prop_no_key(self, mock_client): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run inc function with a future handler client.submit(inc, 1) # Use DaskJob future property to get back the inc function ret = djob.future # Check result self.assertIsNone(ret) @mock.patch('tethys_compute.models.dask.dask_job.log') @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') @mock.patch('tethys_compute.models.dask.dask_job.Future') def test_future_prop_exception(self, mock_future, mock_client, mock_log): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') mock_future.side_effect = Exception('exception in creating future') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run inc function with a future handler future = client.submit(inc, 1) # Get the key from future handler and assign it to DaskJob key to keep track of this inc function djob.key = future.key # Use DaskJob future property to get back the inc function ret = djob.future # Check result self.assertIsNone(ret) mock_log.exception.assert_called_with('Dask Future Init Error') @mock.patch('tethys_compute.models.dask.dask_job.fire_and_forget') @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_execute_delayed(self, mock_client, mock_save, mock_ff): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_future = mock.MagicMock(key='test_key') mock_client_ret.compute.return_value = mock_future # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Delayed option delayed = dask.delayed(inc)(1) # _Execute djob._execute(delayed) # Check result mock_client_ret.compute.assert_called_with(delayed) self.assertEqual('test_key', djob.key) mock_save.assert_called() mock_ff.assert_called_with(mock_future) @mock.patch('tethys_compute.models.dask.dask_job.isinstance') @mock.patch('tethys_compute.models.dask.dask_job.fire_and_forget') @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_execute_future(self, mock_client, mock_save, mock_ff, mock_isinstance): mock_client.return_value = mock.MagicMock() mock_isinstance.side_effect = [True, False] # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # get client from DaskJob client = djob.client # Future option future = client.submit(inc, 2) # _Execute djob._execute(future) # Check result self.assertEqual(future.key, djob.key) mock_save.assert_called() mock_ff.assert_called_with(future) def test_execute_not_future_delayed(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # _Execute self.assertRaises(ValueError, djob._execute, 1) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_update_status(self, mock_future, mock_save, mock_client): mock_future.status = 'finished' # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the function djob._update_status() # check the results mock_client.close.assert_called() mock_save.assert_called() def test_update_status_with_no_future(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # check the results self.assertIsNone(djob._update_status()) @mock.patch('tethys_compute.models.dask.dask_job.log') @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_update_status_exception(self, mock_future, mock_save, mock_log): # Invalid status key mock_future.status = 'foo' # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the function djob._update_status() # check the results mock_log.error.assert_called_with('Unknown Dask Status: "foo"') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') def test_process_result_with_failed_lock(self, mock_re_lock, mock_apl): mock_apl.return_value = False # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # call the function self.assertIsNone(djob._process_results()) # check the result mock_re_lock.assert_not_called() @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future', new_callable=mock.PropertyMock(return_value=None)) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') def test_process_result_no_future(self, mock_apl, _): mock_apl.return_value = True # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # call the function self.assertIsNone(djob._process_results()) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future', new_callable=mock.PropertyMock()) def test_process_result_forget(self, _, mock_client): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, forget=True) # call the function ret = djob._process_results() # check the result mock_client.close.assert_called() self.assertIsNone(ret) @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') def test_process_result_with_result_function(self, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): fake_key = 'sum_faef' mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') djob.key = fake_key # call the function djob._process_results() # check the result mock_client.close.assert_called() mock_client.gather.assert_called_with(mock_future) mock_function.assert_called_with(mock_client.gather()) mock_client.set_metadata.assert_called_with(fake_key, False) self.assertEqual('', djob.key) mock_re_lock.assert_called() @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.log') def test_process_result_with_client_gather_exception( self, mock_logger, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True gather_exception = Exception('Fake exception') mock_client.gather.side_effect = gather_exception # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # call the function djob._process_results() # check the result mock_client.gather.assert_called_with(mock_future) mock_logger.warning.assert_called() mock_function.assert_called_with(gather_exception) mock_re_lock.assert_called() @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_job.log') @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') def test_process_result_with_result_function_with_exception( self, mock_re_lock, mock_apl, _, mock_client, mock_tfe, mock_log, mock_save): mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock() mock_function.side_effect = Exception mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # call the function djob._process_results() # check the result mock_log.exception.assert_called_with('Process Results Function Error') self.assertEqual('ERR', djob._status) mock_save.assert_called() mock_re_lock.assert_called() @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_stop(self, mock_future): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the stop function djob.stop() # Check result mock_future.cancel.assert_called() def test_pause(self): djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=self.scheduler) # Execute and heck result self.assertRaises(NotImplementedError, djob.pause) def test_resume(self): djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=self.scheduler) # Execute and heck result self.assertRaises(NotImplementedError, djob.resume) def test_result(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # need to convert to string because it will convert to string when saving to the database djob.result = 'serialized_results' # call the function ret = djob.result # Check result self.assertEqual('serialized_results', ret) def test_result_none(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) djob.result = None # call the function ret = djob.result # Check result self.assertIsNone(ret) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_done(self, mock_future): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the done function ret = djob.done() # Check result mock_future.done.assert_called() self.assertEqual(mock_future.done(), ret) def test_done_with_no_future(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Check result self.assertIsNone(djob.done()) def test_update_status_interval_prop(self): from datetime import timedelta # Create DaskJob djob = DaskJob(name='test_daskjob', user=self.user, label='label') djob.save() ret = DaskJob.objects.get(name='test_daskjob').update_status_interval # Check result self.assertIsInstance(ret, timedelta) self.assertEqual(timedelta(0, 0), ret) djob.delete() @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_retry(self, mock_future): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the done function djob.retry() # Check result mock_future.retry.assert_called() def test_retry_no_future(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the done function self.assertIsNone(djob.retry()) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test__resubmit(self, mock_future): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the done function djob._resubmit() # Check result mock_future.retry.assert_called() @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_get_logs(self, mock_client): mock_get_log = mock.MagicMock() mock_get_log.get_scheduler_logs.return_value = (( 'INFO', 'dask_scheduler_log1'), ('INFO', 'dask_scheduler_log2')) mock_get_log.get_worker_logs.return_value = { 'worker1': (('INFO', 'dask_worker1_log1'), ('INFO', 'dask_worker1_log2')) } mock_client.return_value = mock_get_log # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) expected_ret = { 'Scheduler': '"INFO", "dask_scheduler_log1"\n"INFO", "dask_scheduler_log2"', 'Workers': { 'Worker-1': '"INFO", "dask_worker1_log1"\n"INFO", "dask_worker1_log2"' } } ret = djob._get_logs() self.assertEqual(expected_ret, ret) @mock.patch('tethys_compute.models.dask.dask_job.log') def test_fail_acquire_pr_lock(self, mock_log): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) djob.extended_properties['processing_results'] = True self.assertFalse(djob._acquire_pr_lock()) mock_log.warning.assert_called_with( 'Unable to aquire lock. Processing results already occurring. Skipping...' ) @mock.patch('django.db.models.base.Model.save') def test_fail_release_pr_lock(self, mock_save): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) djob.extended_properties['processing_results'] = True djob._release_pr_lock() self.assertFalse(djob.extended_properties['processing_results']) mock_save.assert_called()
class DaskJobTest(TethysTestCase): def set_up(self): self.user = User.objects.create_user('tethys_super', '*****@*****.**', 'pass') self.scheduler = DaskScheduler( name='test_dask_scheduler', host='127.0.0.1:8000', timeout=10, heartbeat_interval=5, dashboard='test_dashboard', ) self.scheduler.save() def tear_down(self): self.scheduler.delete() @mock.patch('tethys_compute.models.dask.dask_job.Client') def test_client_prop_with_invalid_scheduler(self, mock_client): mock_client.return_value = 'test_client' djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=None) # Execute ret = djob.client # Check result self.assertEqual('test_client', ret) mock_client.assert_called() @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_client_prop_with_valid_scheduler(self, mock_client): mock_client.return_value = 'test_client' dask_scheduler = Scheduler.objects.get_subclass(name='test_dask_scheduler') djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=dask_scheduler) # Execute ret = djob.client # Check result self.assertEqual('test_client', ret) mock_client.assert_called_with(address='127.0.0.1:8000', heartbeat_interval=5, timeout=10) @mock.patch('tethys_compute.models.dask.dask_job.Client') def test_client_no_scheduler_prop(self, mock_client): mock_client.return_value = 'test_default_client' # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label') # Execute ret = djob.client # Check result self.assertEqual('test_default_client', ret) mock_client.assert_called_with() @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') @mock.patch('tethys_compute.models.dask.dask_job.Future') def test_future_prop(self, mock_future, mock_client): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run rando function with a future handler future = client.submit(inc, 1) # Get the key from future handler and assign it to DaskJob key to keep track of this inc function djob.key = future.key # Use DaskJob future property to get back the inc function ret = djob.future # Check result mock_future.assert_called_with(key='test_key', client=mock_client_ret) self.assertEqual(mock_future(), ret) @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_future_prop_no_key(self, mock_client): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run inc function with a future handler client.submit(inc, 1) # Use DaskJob future property to get back the inc function ret = djob.future # Check result self.assertIsNone(ret) @mock.patch('tethys_compute.models.dask.dask_job.log') @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') @mock.patch('tethys_compute.models.dask.dask_job.Future') def test_future_prop_exception(self, mock_future, mock_client, mock_log): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') mock_future.side_effect = Exception('exception in creating future') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run inc function with a future handler future = client.submit(inc, 1) # Get the key from future handler and assign it to DaskJob key to keep track of this inc function djob.key = future.key # Use DaskJob future property to get back the inc function ret = djob.future # Check result self.assertIsNone(ret) mock_log.exception.assert_called_with('Dask Future Init Error') @mock.patch('tethys_compute.models.dask.dask_job.fire_and_forget') @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_execute_delayed(self, mock_client, mock_save, mock_ff): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_future = mock.MagicMock(key='test_key') mock_client_ret.compute.return_value = mock_future # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Delayed option delayed = dask.delayed(inc)(1) # _Execute djob._execute(delayed) # Check result mock_client_ret.compute.assert_called_with(delayed) self.assertEqual('test_key', djob.key) mock_save.assert_called() mock_ff.assert_called_with(mock_future) @mock.patch('tethys_compute.models.dask.dask_job.isinstance') @mock.patch('tethys_compute.models.dask.dask_job.fire_and_forget') @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_scheduler.Client') def test_execute_future(self, mock_client, mock_save, mock_ff, mock_isinstance): mock_client.return_value = mock.MagicMock() mock_isinstance.side_effect = [True, False] # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler ) # get client from DaskJob client = djob.client # Future option future = client.submit(inc, 2) # _Execute djob._execute(future) # Check result self.assertEqual(future.key, djob.key) mock_save.assert_called() mock_ff.assert_called_with(future) def test_execute_not_future_delayed(self): # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler ) # _Execute self.assertRaises(ValueError, djob._execute, 1) @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_update_status(self, mock_future, mock_save): mock_future.status = 'finished' # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler ) # call the function djob._update_status() # check the results mock_save.assert_called() def test_update_status_with_no_future(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # check the results self.assertIsNone(djob._update_status()) @mock.patch('tethys_compute.models.dask.dask_job.log') @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_update_status_exception(self, mock_future, mock_save, mock_log): # Invalid status key mock_future.status = 'foo' # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler ) # call the function djob._update_status() # check the results mock_log.error.assert_called_with('Unknown Dask Status: "foo"') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') def test_process_result_with_failed_lock(self, mock_re_lock, mock_apl): mock_apl.return_value = False # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # call the function self.assertIsNone(djob._process_results()) # check the result mock_re_lock.assert_not_called() @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future', new_callable=mock.PropertyMock(return_value=None)) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') def test_process_result_no_future(self, mock_apl, _): mock_apl.return_value = True # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # call the function self.assertIsNone(djob._process_results()) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future', new_callable=mock.PropertyMock()) def test_process_result_forget(self, _): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, forget=True) # call the function ret = djob._process_results() # check the result self.assertIsNone(ret) @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') def test_process_result_with_result_function(self, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): fake_key = 'sum_faef' mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function' ) djob.key = fake_key # call the function djob._process_results() # check the result mock_client.gather.assert_called_with(mock_future) mock_function.assert_called_with(mock_client.gather()) mock_client.set_metadata.assert_called_with(fake_key, False) self.assertEqual('', djob.key) mock_re_lock.assert_called() @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.log') def test_process_result_with_client_gather_exception(self, mock_logger, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True gather_exception = Exception('Fake exception') mock_client.gather.side_effect = gather_exception # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function' ) # call the function djob._process_results() # check the result mock_client.gather.assert_called_with(mock_future) mock_logger.warning.assert_called() mock_function.assert_called_with(gather_exception) mock_re_lock.assert_called() @mock.patch('django.db.models.base.Model.save') @mock.patch('tethys_compute.models.dask.dask_job.log') @mock.patch('tethys_compute.models.tethys_job.TethysFunctionExtractor') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.client') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._acquire_pr_lock') @mock.patch('tethys_compute.models.dask.dask_job.DaskJob._release_pr_lock') def test_process_result_with_result_function_with_exception(self, mock_re_lock, mock_apl, _, mock_client, mock_tfe, mock_log, mock_save): mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock() mock_function.side_effect = Exception mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') # call the function djob._process_results() # check the result mock_log.exception.assert_called_with('Process Results Function Error') self.assertEquals('ERR', djob._status) mock_save.assert_called() mock_re_lock.assert_called() @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_stop(self, mock_future): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the stop function djob.stop() # Check result mock_future.cancel.assert_called() def test_pause(self): djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=self.scheduler) # Execute and heck result self.assertRaises(NotImplementedError, djob.pause) def test_resume(self): djob = DaskJob(name='test_dj', user=self.user, key='test_key', label='label', scheduler=self.scheduler) # Execute and heck result self.assertRaises(NotImplementedError, djob.resume) def test_result(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # need to convert to string because it will convert to string when saving to the database djob.result = 'serialized_results' # call the function ret = djob.result # Check result self.assertEquals('serialized_results', ret) def test_result_none(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) djob.result = None # call the function ret = djob.result # Check result self.assertIsNone(ret) @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_done(self, mock_future): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the done function ret = djob.done() # Check result mock_future.done.assert_called() self.assertEqual(mock_future.done(), ret) def test_done_with_no_future(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Check result self.assertIsNone(djob.done()) def test_update_status_interval_prop(self): from datetime import timedelta # Create DaskJob djob = DaskJob(name='test_daskjob', user=self.user, label='label') djob.save() ret = DaskJob.objects.get(name='test_daskjob').update_status_interval # Check result self.assertIsInstance(ret, timedelta) self.assertEqual(timedelta(0, 0), ret) djob.delete() @mock.patch('tethys_compute.models.dask.dask_job.DaskJob.future') def test_retry(self, mock_future): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the done function djob.retry() # Check result mock_future.retry.assert_called() def test_retry_no_future(self): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # call the done function self.assertIsNone(djob.retry()) @mock.patch('tethys_compute.models.dask.dask_job.log') def test_fail_acquire_pr_lock(self, mock_log): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) djob.extended_properties['processing_results'] = True self.assertFalse(djob._acquire_pr_lock()) mock_log.warning.assert_called_with('Unable to aquire lock. Processing results already occurring. Skipping...') @mock.patch('django.db.models.base.Model.save') def test_fail_release_pr_lock(self, mock_save): # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) djob.extended_properties['processing_results'] = True djob._release_pr_lock() self.assertFalse(djob.extended_properties['processing_results']) mock_save.assert_called()