in cdk-project/lib/images/codebuild-image/python/src/notebooks/cli/run_pr_notebooks.py [0:0]
def main():
args = parse_args(sys.argv[1:])
skip_args = {
"docker": args.skip_docker,
"local_mode": args.skip_local,
"fsx_esx": args.skip_filesystem
}
jobs = {}
session = ensure_session()
instance_type = args.instance or "ml.m5.xlarge"
for notebook in parse.pr_notebook_filenames(args.pr):
if parse.is_notebook_skipped(notebook, skip_args):
job_name = None
else:
image = kernels.kernel_image_for(notebook)
s3path = upload_notebook(notebook, session)
parameters = {"kms_key": kms_key()}
job_name = execute_notebook(
image=image,
input_path=s3path,
notebook=notebook,
role="SageMakerRole",
instance_type=instance_type,
session=session,
output_prefix=get_output_prefix(),
parameters=parameters,
)
time.sleep(1)
jobs[notebook] = job_name
failures = {}
while jobs:
for notebook in list(jobs):
job_name = jobs[notebook]
if not is_running(job_name, session):
if job_name:
status, failure_reason = wait_for_complete(
job_name, progress=False, session=session
)
else:
status, failure_reason = (
"Skipped",
"This notebook was skipped because it either uses Docker or Local Mode.",
)
basename = os.path.basename(notebook)
print("\n" * 2)
print(f"* {basename} " + "*" * (97 - len(basename)))
print("*")
print(f"* {'job name':>11}: {str(job_name):<11}")
print("*")
print(f"* {'kernel':>11}: {kernels.kernel_type_for(notebook):<11}")
print("*")
print(f"* {'status':>11}: {status:<11}")
print("*")
if status != "Completed":
print(failure_reason)
if status != "Skipped":
failures[notebook] = failure_reason
jobs.pop(notebook)
time.sleep(10)
print("\n" * 2)
print("-" * 100)
if failures:
raise Exception(
"One or more notebooks failed to execute. Please see above for error messages. "
"If you need more information, please see the CloudWatch logs for the corresponding Processing job."
)