def test_download_checkpoints_from_aml(test_output_dirs: OutputFolderForTests) -> None:
"""
Check that we can download checkpoint files from an AzureML run, if they are not available on disk.
"""
run = get_most_recent_run(fallback_run_id_for_local_execution=FALLBACK_SINGLE_RUN)
checkpoint_folder = f"{DEFAULT_AML_UPLOAD_DIR}/{CHECKPOINT_FOLDER}/"
temp_folder = download_folder_from_run_to_temp_folder(folder=checkpoint_folder,
run=run,
workspace=get_default_workspace())
files = list(temp_folder.glob("*"))
assert len(files) == 1
assert (temp_folder / LAST_CHECKPOINT_FILE_NAME_WITH_SUFFIX).is_file()
# Test if what's in the folder are really files, not directories
for file in files:
assert file.is_file()
# Now test if that is correctly integrated into the checkpoint finder. To avoid downloading a second time,
# now mock the call to the actual downloader.
with mock.patch("InnerEye.ML.utils.checkpoint_handling.is_running_in_azure_ml", return_value=True):
with mock.patch("InnerEye.ML.utils.checkpoint_handling.download_folder_from_run_to_temp_folder",
return_value=temp_folder) as download:
# Call the checkpoint finder with a temp folder that does not contain any files, so it should try to
# download
result = find_recovery_checkpoint_on_disk_or_cloud(test_output_dirs.root_dir)
download.assert_called_once_with(folder=checkpoint_folder)
assert result is not None
assert result.name == LAST_CHECKPOINT_FILE_NAME_WITH_SUFFIX
def test_is_completed_single_run() -> None:
"""
Test if we can correctly check run status for a single run.
:return:
"""
logging_to_stdout()
workspace = get_default_workspace()
get_run_and_check(get_most_recent_run_id(), True, workspace)
def get_most_recent_model(fallback_run_id_for_local_execution: str = FALLBACK_SINGLE_RUN) -> Model:
"""
Gets the string name of the most recently executed AzureML run, extracts which model that run had registered,
and return the instantiated model object.
:param fallback_run_id_for_local_execution: A hardcoded AzureML run ID that is used when executing this code
on a local box, outside of Azure build agents.
"""
model_id = get_most_recent_model_id(fallback_run_id_for_local_execution=fallback_run_id_for_local_execution)
return Model(workspace=get_default_workspace(), id=model_id)
def test_is_completed_ensemble_run() -> None:
"""
Test if we can correctly check run status and status of child runs for an ensemble run.
:return:
"""
logging_to_stdout()
workspace = get_default_workspace()
run_id = get_most_recent_run_id(
fallback_run_id_for_local_execution=FALLBACK_ENSEMBLE_RUN)
get_run_and_check(run_id, True, workspace)
def test_get_checkpoints_from_model_ensemble_run(test_output_dirs: OutputFolderForTests) -> None:
model_id = get_most_recent_model_id(fallback_run_id_for_local_execution=FALLBACK_ENSEMBLE_RUN)
downloaded_checkpoints = CheckpointHandler.get_checkpoints_from_model(model_id=model_id,
workspace=get_default_workspace(),
download_path=test_output_dirs.root_dir)
# Check that all the ensemble checkpoints have been downloaded
expected_model_root = test_output_dirs.root_dir / FINAL_ENSEMBLE_MODEL_FOLDER
assert expected_model_root.is_dir()
model_inference_config = read_model_inference_config(expected_model_root / MODEL_INFERENCE_JSON_FILE_NAME)
expected_paths = [expected_model_root / x for x in model_inference_config.checkpoint_paths]
assert len(expected_paths) == len(downloaded_checkpoints)
assert set(expected_paths) == set(downloaded_checkpoints)
for expected_path in expected_paths:
assert expected_path.is_file()
def test_get_checkpoints_from_model_single_run(test_output_dirs: OutputFolderForTests) -> None:
model_id = get_most_recent_model_id(fallback_run_id_for_local_execution=FALLBACK_SINGLE_RUN)
downloaded_checkpoints = CheckpointHandler.get_checkpoints_from_model(model_id=model_id,
workspace=get_default_workspace(),
download_path=test_output_dirs.root_dir)
# Check a single checkpoint has been downloaded
expected_model_root = test_output_dirs.root_dir / FINAL_MODEL_FOLDER
assert expected_model_root.is_dir()
model_inference_config = read_model_inference_config(expected_model_root / MODEL_INFERENCE_JSON_FILE_NAME)
expected_paths = [expected_model_root / x for x in model_inference_config.checkpoint_paths]
assert len(expected_paths) == 1 # A registered model for a non-ensemble run should contain only one checkpoint
assert len(downloaded_checkpoints) == 1
assert expected_paths[0] == downloaded_checkpoints[0]
assert expected_paths[0].is_file()
def test_is_cross_validation_child_run(is_ensemble: bool,
is_numeric: bool) -> None:
"""
Test that cross validation child runs are identified correctly.
"""
if is_ensemble:
rid = DEFAULT_ENSEMBLE_RUN_RECOVERY_ID_NUMERIC if is_numeric else DEFAULT_ENSEMBLE_RUN_RECOVERY_ID
else:
rid = DEFAULT_RUN_RECOVERY_ID_NUMERIC if is_numeric else DEFAULT_RUN_RECOVERY_ID
run = fetch_run(workspace=get_default_workspace(), run_recovery_id=rid)
# check for offline run
assert not is_cross_validation_child_run(Run.get_context())
# check for online runs
assert not is_cross_validation_child_run(run)
if is_ensemble:
assert all(
[is_cross_validation_child_run(x) for x in fetch_child_runs(run)])
def test_get_cross_validation_split_index(is_ensemble: bool) -> None:
"""
Test that retrieved cross validation split index is as expected, for single runs and ensembles.
"""
run = fetch_run(workspace=get_default_workspace(),
run_recovery_id=DEFAULT_ENSEMBLE_RUN_RECOVERY_ID
if is_ensemble else DEFAULT_RUN_RECOVERY_ID)
# check for offline run
assert get_cross_validation_split_index(
Run.get_context()) == DEFAULT_CROSS_VALIDATION_SPLIT_INDEX
# check for online runs
assert get_cross_validation_split_index(
run) == DEFAULT_CROSS_VALIDATION_SPLIT_INDEX
if is_ensemble:
assert all([
get_cross_validation_split_index(x) >
DEFAULT_CROSS_VALIDATION_SPLIT_INDEX for x in fetch_child_runs(run)
])
def test_is_completed() -> None:
"""
Test if we can correctly check run status and status of child runs.
:return:
"""
logging_to_stdout()
workspace = get_default_workspace()
def get_run_and_check(run_id: str, expected: bool) -> None:
run = fetch_run(workspace, run_id)
status = is_run_and_child_runs_completed(run)
assert status == expected
get_run_and_check(DEFAULT_RUN_RECOVERY_ID, True)
get_run_and_check(DEFAULT_ENSEMBLE_RUN_RECOVERY_ID, True)
# This Hyperdrive run has 1 failing child run, the parent run completed successfully.
get_run_and_check(
"refs_pull_326_merge:HD_d123f042-ca58-4e35-9a64-48d71c5f63a7", False)
def test_register_model_invalid() -> None:
ws = get_default_workspace()
config = get_model_loader().create_model_config_from_name("Lung")
with pytest.raises(Exception):
ml_runner = MLRunner(config, None)
ml_runner.register_segmentation_model(
run=Run.get_context(),
workspace=ws,
best_epoch=0,
best_epoch_dice=0,
checkpoint_paths=checkpoint_paths,
model_proc=ModelProcessing.DEFAULT)
with pytest.raises(Exception):
ml_runner = MLRunner(config, get_default_azure_config())
ml_runner.register_segmentation_model(
best_epoch=0,
best_epoch_dice=0,
checkpoint_paths=checkpoint_paths,
model_proc=ModelProcessing.DEFAULT)
def test_register_and_score_model(is_ensemble: bool,
dataset_expected_spacing_xyz: Any,
model_outside_package: bool,
test_output_dirs: OutputFolderForTests) -> None:
"""
End-to-end test which ensures the scoring pipeline is functioning as expected by performing the following:
1) Registering a pre-trained model to AML
2) Checking that a model zip from the registered model can be created successfully
3) Calling the scoring pipeline to check inference can be run from the published model successfully
"""
ws = get_default_workspace()
# Get an existing config as template
loader = get_model_loader("Tests.ML.configs" if model_outside_package else None)
config: SegmentationModelBase = loader.create_model_config_from_name(
model_name="BasicModel2EpochsOutsidePackage" if model_outside_package else "BasicModel2Epochs"
)
config.dataset_expected_spacing_xyz = dataset_expected_spacing_xyz
config.set_output_to(test_output_dirs.root_dir)
# copy checkpoints into the outputs (simulating a run)
stored_checkpoints = full_ml_test_data_path(os.path.join("train_and_test_data", "checkpoints"))
shutil.copytree(str(stored_checkpoints), str(config.checkpoint_folder))
paths = [config.checkpoint_folder / "1_checkpoint.pth.tar"]
checkpoints = paths * 2 if is_ensemble else paths
model = None
model_path = None
# Mocking to get the source from the current directory
# the score.py and python_wrapper.py cannot be moved inside the InnerEye package, which will be the
# only code running (if these tests are run on the package).
with mock.patch('InnerEye.Common.fixed_paths.repository_root_directory',
return_value=tests_root_directory().parent):
try:
tags = {"model_name": config.model_name}
azure_config = get_default_azure_config()
if model_outside_package:
azure_config.extra_code_directory = "Tests" # contains DummyModel
deployment_hook = lambda cfg, azure_cfg, mdl, is_ens: (Path(cfg.model_name), azure_cfg.docker_shm_size)
ml_runner = MLRunner(config, azure_config, model_deployment_hook=deployment_hook)
model, deployment_path, deployment_details = ml_runner.register_segmentation_model(
workspace=ws,
tags=tags,
best_epoch=0,
best_epoch_dice=0,
checkpoint_paths=checkpoints,
model_proc=ModelProcessing.DEFAULT)
assert model is not None
model_path = Path(model.get_model_path(model.name, model.version, ws))
assert (model_path / fixed_paths.ENVIRONMENT_YAML_FILE_NAME).exists()
assert (model_path / Path("InnerEye/ML/runner.py")).exists()
assert deployment_path == Path(config.model_name)
assert deployment_details == azure_config.docker_shm_size
# move test data into the data folder to simulate an actual run
train_and_test_data_dir = full_ml_test_data_path("train_and_test_data")
img_channel_1_name = "id1_channel1.nii.gz"
img_channel_1_path = train_and_test_data_dir / img_channel_1_name
img_channel_2_name = "id1_channel2.nii.gz"
img_channel_2_path = train_and_test_data_dir / img_channel_2_name
# download the registered model and test that we can run the score pipeline on it
model_root = Path(model.download(str(test_output_dirs.root_dir)))
# create a dummy datastore to store model checkpoints and image data
# this simulates the code shapshot being executed in a real run
test_datastore = test_output_dirs.root_dir / "test_datastore"
shutil.move(
str(model_root / "test_outputs"),
str(test_datastore / RELATIVE_TEST_OUTPUTS_PATH)
)
data_root = test_datastore / DEFAULT_DATA_FOLDER
os.makedirs(data_root)
shutil.copy(str(img_channel_1_path), data_root)
shutil.copy(str(img_channel_2_path), data_root)
# run score pipeline as a separate process using the python_wrapper.py code to simulate a real run
return_code = SubprocessConfig(process="python", args=[
str(model_root / "python_wrapper.py"),
"--spawnprocess=python",
str(model_root / "score.py"),
f"--data-folder={str(test_datastore)}",
f"--test_image_channels={img_channel_1_name},{img_channel_2_name}",
"--use_gpu=False"
]).spawn_and_monitor_subprocess()
# check that the process completed as expected
assert return_code == 0
expected_segmentation_path = Path(model_root) / DEFAULT_RESULT_IMAGE_NAME
assert expected_segmentation_path.exists()
# sanity check the resulting segmentation
expected_shape = get_nifti_shape(img_channel_1_path)
image_header = get_unit_image_header()
assert_nifti_content(str(expected_segmentation_path), expected_shape, image_header, [0], np.ubyte)
finally:
# delete the registered model, and any downloaded artifacts
shutil.rmtree(test_output_dirs.root_dir)
if model and model_path:
model.delete()
shutil.rmtree(model_path)
