def main(argv=None):
parser = create_parser()
args = parser.parse_args(argv)
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_sagemaker_client(args.region, args.endpoint_url)
logging.info('Submitting Processing Job to SageMaker...')
job_name = _utils.create_processing_job(client, vars(args))
logging.info('Job request submitted. Waiting for completion...')
try:
_utils.wait_for_processing_job(client, job_name)
except:
raise
finally:
cw_client = _utils.get_cloudwatch_client(args.region)
_utils.print_logs_for_job(cw_client, '/aws/sagemaker/ProcessingJobs',
job_name)
outputs = _utils.get_processing_job_outputs(client, job_name)
with open('/tmp/job_name.txt', 'w') as f:
f.write(job_name)
with open('/tmp/output_artifacts.txt', 'w') as f:
f.write(json.dumps(outputs))
logging.info('Job completed.')
def test_wait_for_failed_job(self):
mock_client = MagicMock()
mock_client.describe_processing_job.side_effect = [{
"ProcessingJobStatus":
"Starting"
}, {
"ProcessingJobStatus":
"InProgress"
}, {
"ProcessingJobStatus":
"Downloading"
}, {
"ProcessingJobStatus":
"Failed",
"FailureReason":
"Something broke lol"
}, {
"ProcessingJobStatus":
"Should not be called"
}]
with self.assertRaises(Exception):
_utils.wait_for_processing_job(mock_client, 'processing-job', 0)
self.assertEqual(mock_client.describe_processing_job.call_count, 4)
def test_wait_for_processing_job(self):
mock_client = MagicMock()
mock_client.describe_processing_job.side_effect = [
{"ProcessingJobStatus": "Starting"},
{"ProcessingJobStatus": "InProgress"},
{"ProcessingJobStatus": "Downloading"},
{"ProcessingJobStatus": "Completed"},
{"ProcessingJobStatus": "Should not be called"}
]
_utils.wait_for_processing_job(mock_client, 'processing-job', 0)
self.assertEqual(mock_client.describe_processing_job.call_count, 4)
def main(argv=None):
parser = create_parser()
args = parser.parse_args(argv)
logging.getLogger().setLevel(logging.INFO)
client = _utils.get_sagemaker_client(args.region,
args.endpoint_url,
assume_role_arn=args.assume_role)
logging.info('Submitting Processing Job to SageMaker...')
job_name = _utils.create_processing_job(client, vars(args))
def signal_term_handler(signalNumber, frame):
logging.info(f"Stopping Processing Job: {job_name}")
_utils.stop_processing_job(client, job_name)
logging.info(f"Processing Job: {job_name} request submitted to Stop")
signal.signal(signal.SIGTERM, signal_term_handler)
logging.info('Job request submitted. Waiting for completion...')
try:
_utils.wait_for_processing_job(client, job_name)
except:
raise
finally:
cw_client = _utils.get_cloudwatch_client(
args.region, assume_role_arn=args.assume_role)
_utils.print_logs_for_job(cw_client, '/aws/sagemaker/ProcessingJobs',
job_name)
outputs = _utils.get_processing_job_outputs(client, job_name)
_utils.write_output(args.job_name_output_path, job_name)
_utils.write_output(args.output_artifacts_output_path,
outputs,
json_encode=True)
logging.info('Job completed.')
