def test_valid_api_lookups(self, mock_hook, mock_requests, valid_job_params):
"""test api lookups are correct"""
queue_mock = Mock(
**{
'json.return_value': {
'results': [
{'id': 2, 'public_name': 'queue2'},
{'id': 3, 'public_name': 'queue3'},
{'id': 1, 'public_name': 'queue'},
]
}
}
)
billing_mock = Mock(**{'json.return_value': {'results': [{'id': 1}]}})
app_mock = Mock(
**{
'json.return_value': {
'results': [
{'name': 'pipeline2', 'id': 2, 'is_service': True},
{'name': 'pipeline3', 'id': 3, 'is_service': False},
{'name': 'pipeline', 'id': 1, 'is_service': False},
]
}
}
)
mock_requests.get.side_effect = [queue_mock, billing_mock, app_mock]
plexus_job_operator = PlexusJobOperator(**valid_job_params)
assert plexus_job_operator._api_lookup('queue', mock_hook) == 1
assert plexus_job_operator._api_lookup('billing_account_id', mock_hook) == 1
assert plexus_job_operator.is_service is None
assert plexus_job_operator._api_lookup('app', mock_hook) == 1
assert plexus_job_operator.is_service is not None
def test_invalid_api_lookups(self, mock_hook, mock_requests,
invalid_lookups):
"""test api lookups that are not found"""
queue_mock = Mock(**{
'json.return_value': {
'results': [{
'id': 2,
'public_name': 'queue2'
}]
}
})
billing_mock = Mock(
**{'json.return_value': {
'results': [{
'id': None
}]
}})
app_mock = Mock(
**{
'json.return_value': {
'results': [{
'name': 'pipeline2',
'id': 2,
'is_service': True
}]
}
})
mock_requests.get.side_effect = [queue_mock, billing_mock, app_mock]
plexus_job_operator = PlexusJobOperator(**invalid_lookups)
with pytest.raises(AirflowException):
plexus_job_operator._api_lookup('queue', mock_hook)
plexus_job_operator._api_lookup('app', mock_hook)
plexus_job_operator._api_lookup('billing_account_id', mock_hook)
def test_valid_job_params(self, valid_job_params):
"""test params are initialized correctly"""
expected_job_params = {
'name': 'test',
'app': 'pipeline',
'queue': 'queue',
'num_nodes': 1,
'num_cores': 1,
'billing_account_id': None,
}
plexus_job_operator = PlexusJobOperator(**valid_job_params)
assert plexus_job_operator.job_params == expected_job_params
assert plexus_job_operator.is_service is None
def test_invalid_job_params(self, mock_hook, invalid_job_params):
"""test required job params are supplied"""
plexus_job_operator = PlexusJobOperator(**invalid_job_params)
with pytest.raises(AirflowException):
plexus_job_operator.construct_job_params(mock_hook)
def test_create_job(self, mock_hook, mock_construct_params, mock_request,
mock_time, valid_job_params):
"""test submit/get job"""
mock_construct_params.return_value = {
'name': 'test',
'app': 1,
'queue': 2,
'num_nodes': 1,
'num_cores': 1,
'billing_account_id': 3,
}
# assert Error for is_service is None
plexus_job_operator = PlexusJobOperator(**valid_job_params)
with pytest.raises(AirflowException):
plexus_job_operator.execute(None)
# assert create_job timeout
plexus_job_operator.is_service = False
mock_request.post.side_effect = Timeout
with pytest.raises(Timeout):
plexus_job_operator.execute(None)
# assert create_job error
mock_request.post.side_effect = None
mock_request.post.return_value = Mock(ok=False,
status_code=400,
reason='Bad',
text='Bad')
with pytest.raises(AirflowException):
plexus_job_operator.execute(None)
# assert get_job error
mock_request.post.return_value = Mock(**{
'ok': True,
'json.return_value': {
'id': 1,
'last_state': 'Pending'
}
})
mock_request.get.return_value = Mock(ok=False)
with pytest.raises(AirflowException):
plexus_job_operator.execute(None)
# assert get_job state error
mock_request.get.return_value = Mock(**{
'ok': True,
'json.return_value': {
'last_state': 'Failed'
}
})
with pytest.raises(AirflowException):
plexus_job_operator.execute(None)
# assert get_job state error
mock_request.get.return_value = Mock(**{
'ok': True,
'json.return_value': {
'last_state': 'Cancelled'
}
})
with pytest.raises(AirflowException):
plexus_job_operator.execute(None)
# assert get_job timeout
mock_request.get.side_effect = Timeout
with pytest.raises(Timeout):
plexus_job_operator.execute(None)
HOME = '/home/acc'
T3_PRERUN_SCRIPT = 'cp {home}/imdb/run_scripts/mlflow.sh {home}/ && chmod +x mlflow.sh'.format(
home=HOME)
args = {'owner': 'core scientific', 'retries': 1}
dag = DAG(
'test',
default_args=args,
description='testing plexus operator',
start_date=days_ago(1),
schedule_interval='@once',
catchup=False,
)
t1 = PlexusJobOperator(
task_id='test',
job_params={
'name': 'test',
'app': 'MLFlow Pipeline 01',
'queue': 'DGX-2 (gpu:Tesla V100-SXM3-32GB)',
'num_nodes': 1,
'num_cores': 1,
'prerun_script': T3_PRERUN_SCRIPT,
},
dag=dag,
)
t1
'depends_on_past': False,
'retries': 1,
}
dag = DAG('imdb',
default_args=args,
description='loads imdb dataset and trains a NN model',
schedule_interval='@once',
catchup=False)
t1 = PlexusJobOperator(
task_id='process_data',
job_params={
'name': 'process_data',
'app': 'MLFlow Pipeline 01',
'queue': 'DGX-2 (gpu:Tesla V100-SXM3-32GB)',
'num_nodes': 1,
'num_cores': 1,
'run_script': T1_RUN_SCRIPT,
},
dag=dag,
)
t2 = PlexusJobOperator(
task_id='train_model',
job_params={
'name': 'train_model',
'app': 'TensorFlow Pipeline 01',
'queue': 'DGX-2 (gpu:Tesla V100-SXM3-32GB)',
'num_nodes': 1,
'num_cores': 1,
'run_script': T2_RUN_SCRIPT,
