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.')