Fix a flaky test. Fix a bug in environment creation with private wheels (#169)

CHANGELOG.md

@@ -10,6 +10,20 @@ For each Pull Request, the affected code parts should be briefly described and a
 release. In the first PR after a release has been made, a section for the upcoming release should be added, by copying
 the section headers (Added/Changed/...) and incrementing the package version.
 
+## 0.1.13
+
+### Added
+
+### Changed
+
+### Fixed
+- ([#169](https://github.com/microsoft/hi-ml/pull/169)) Fix a test that was failing occasionally
+
+### Removed
+
+### Deprecated
+
+
 ## 0.1.12
 
 ### Added

docs/source/conf.py

@@ -28,7 +28,6 @@ project = 'hi-ml'
 copyright = '2021, InnerEye'
 author = 'InnerEye'
 
-
 # -- General configuration ---------------------------------------------------
 
 # Add any Sphinx extension module names here, as strings. They can be
@@ -42,8 +41,7 @@ extensions = [
     'sphinx_automodapi.automodapi',
     'sphinx_autodoc_typehints',
     'sphinx.ext.viewcode',
-    ]
-
+]
 
 numpydoc_show_class_members = False
 
@@ -60,7 +58,6 @@ templates_path = ['_templates']
 # This pattern also affects html_static_path and html_extra_path.
 exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
 
-
 # -- Options for HTML output -------------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for

docs/source/examples.md

@@ -371,4 +371,20 @@ this case, the DataFrame will contain a string representing the path to the arti
 
     |                | 0                                       | 1                                     |
     |----------------|-----------------------------------------|---------------------------------------|
-    | accuracy_plot  | aml://artifactId/ExperimentRun/dcid.... | aml://artifactId/ExperimentRun/dcid...|
+    | accuracy_plot  | aml://artifactId/ExperimentRun/dcid.... | aml://artifactId/ExperimentRun/dcid...|
+
+
+## Modifying checkpoints stored in an AzureML run
+
+The script in [examples/modify_checkpoint/modify_checkpoint.py](examples/modify_checkpoint/modify_checkpoint.rst)
+shows how checkpoints can be downloaded from an AzureML run, modified, and the uploaded back to a newly created run.
+
+This can be helpful for example if networks architecture changed, but you do not want to re-train the stored models
+with the new code.
+
+The essential bits are:
+* Download files from a run via `download_files_from_run_id`
+* Modify the checkpoints
+* Create a new run via `create_aml_run_object`
+* Then use `Run.upload_folder` to upload all modified checkpoints to that new run. From there, they can be consumed
+  in a follow-up training run again via `download_files_from_run_id`

docs/source/examples/10/sample.py

@@ -12,11 +12,10 @@ from pathlib import Path
 import numpy as np
 from azureml.core.run import Run
 
-from health_azure import submit_to_azure_if_needed, create_crossval_hyperdrive_config, get_workspace
+from health_azure import create_crossval_hyperdrive_config, submit_to_azure_if_needed
 
 
 def main() -> None:
-
     num_cross_validation_splits = 2
     metric_name = "val/loss"
     hyperdrive_config = create_crossval_hyperdrive_config(num_cross_validation_splits,
@@ -34,7 +33,7 @@ def main() -> None:
         tags=tags,
         hyperdrive_config=hyperdrive_config,
         submit_to_azureml=True
-        )
+    )
 
     if run_info.run is None:
         raise ValueError("run_info.run is None")
@@ -48,8 +47,7 @@ def main() -> None:
                         help='Penalty parameter of the error term')
     parser.add_argument('--cross_validation_split_index', help="An index denoting which split of the dataset this"
                                                                "run represents in k-fold cross-validation")
-    parser.add_argument("--num_splits", help="The total number of splits being used for k-fold"
-                                                              "cross validation")
+    parser.add_argument("--num_splits", help="The total number of splits being used for k-fol cross validation")
 
     args = parser.parse_args()
     run.log('Kernel type', args.kernel)

docs/source/examples/modify_checkpoint/modify_checkpoint.py

@@ -0,0 +1,55 @@
+#  ------------------------------------------------------------------------------------------
+#  Copyright (c) Microsoft Corporation. All rights reserved.
+#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+#  ------------------------------------------------------------------------------------------
+
+# This script shows how to download files from a run, modify a checkpoint, and upload to a new run.
+# From that new run, the modified checkpoint can be easily consumed in other training runs, both inside and
+# outside AzureML.
+
+from pathlib import Path
+
+import torch
+
+from health_azure import download_files_from_run_id, create_aml_run_object
+
+if __name__ == "__main__":
+
+    root_folder = Path.cwd()
+
+    # This is the run from which we want to download checkpoints
+    experiment_name = "SSLCheckpoint"
+    old_run = f"{experiment_name}:SSLCheckpoint_1629396871_2263a0ec"
+
+    # Specify where your AML workspace config.json file lives. If you set that to None, the code will try to find a file
+    # called config.json in the current folder
+    workspace_config_json = root_folder / "myworkspace.config.json"
+
+    download_folder = Path(root_folder / "old_run")
+    download_folder.mkdir(exist_ok=True)
+
+    # Download all checkpoints in the run
+    checkpoint_folder = "outputs/checkpoints"
+    download_files_from_run_id(run_id=old_run, workspace_config_path=workspace_config_json,
+                               output_folder=download_folder, prefix=checkpoint_folder)
+
+    for file in download_folder.rglob("*.ckpt"):
+        checkpoint = torch.load(file)
+        state_dict = checkpoint['state_dict']
+        # Here we modify the checkpoints: They reference weights from an older version of the code, delete any
+        # such weights
+        linear_head_states = [name for name in state_dict.keys() if name.startswith("non_linear_evaluator")]
+        print(linear_head_states)
+        if linear_head_states:
+            print(f"Removing linear head from {file}")
+            for state in linear_head_states:
+                del checkpoint['state_dict'][state]
+            torch.save(checkpoint, file)
+
+    # Create a new AzureML run in the same experiment. The run will get a new unique ID
+    new_run = create_aml_run_object(experiment_name=experiment_name, workspace_config_path=workspace_config_json)
+    new_run.upload_folder(name=checkpoint_folder, path=str(download_folder / checkpoint_folder))
+    new_run.complete()
+
+    print(f"Uploaded the modified checkpoints to this run: {new_run.get_portal_url()}")
+    print(f"Use this RunID to download the modified checkpoints: {new_run.id}")

docs/source/examples/modify_checkpoint/modify_checkpoint.rst

@@ -0,0 +1,5 @@
+Example script: Modifying and uploading checkpoints
+===================================================
+
+.. literalinclude:: modify_checkpoint.py
+   :emphasize-lines: 21-22,33-34,50-52

hi-ml-azure/src/health_azure/himl.py

@@ -129,14 +129,16 @@ def create_run_configuration(workspace: Workspace,
     if aml_environment_name:
         run_config.environment = Environment.get(workspace, aml_environment_name)
     elif conda_environment_file:
-        run_config.environment = create_python_environment(
+        # Create an AzureML environment, then check if it exists already. If it exists, use the registered
+        # environment, otherwise register the new environment.
+        new_environment = create_python_environment(
             conda_environment_file=conda_environment_file,
             pip_extra_index_url=pip_extra_index_url,
             workspace=workspace,
             private_pip_wheel_path=private_pip_wheel_path,
             docker_base_image=docker_base_image,
             environment_variables=environment_variables)
-        registered_env = register_environment(workspace, run_config.environment)
+        registered_env = register_environment(workspace, new_environment)
         run_config.environment = registered_env
     else:
         raise ValueError("One of the two arguments 'aml_environment_name' or 'conda_environment_file' must be given.")

hi-ml-azure/src/health_azure/utils.py

@@ -475,24 +475,21 @@ def is_private_field_name(name: str) -> bool:
 
 def determine_run_id_type(run_or_recovery_id: str) -> str:
     """
-    Determine whether a run id is of type "run id" or "run recovery id". This distinction is made
-    by checking for telltale patterns within the string. Run recovery ideas take the form "experiment_name:run_id"
-    whereas run_ids follow the pattern of a mixture of strings and decimals, separated by underscores. If the input
+    Determine whether a run id is of type "run id" or "run recovery id". Run recovery ideas take the form
+    "experiment_name:run_id". If the input
     string takes the format of a run recovery id, only the run id part will be returned. If it is a run id already,
-    it will be returned without transformation. If neither, a ValueError is raised.
+    it will be returned without transformation.
 
     :param run_or_recovery_id: The id to determine as either a run id or a run recovery id
     :return: A string representing the run id
     """
     if run_or_recovery_id is None:
         raise ValueError("Expected run_id or run_recovery_id but got None")
-    elif len(run_or_recovery_id.split(EXPERIMENT_RUN_SEPARATOR)) > 1:
+    parts = run_or_recovery_id.split(EXPERIMENT_RUN_SEPARATOR)
+    if len(parts) > 1:
         # return only the run_id, which comes after the colon
-        return run_or_recovery_id.split(EXPERIMENT_RUN_SEPARATOR)[1]
-    elif re.search(r"\d", run_or_recovery_id) and re.search('_', run_or_recovery_id):
-        return run_or_recovery_id
-    else:
-        raise ValueError("Unknown run type. Expected run_id or run_recovery id")
+        return parts[1]
+    return run_or_recovery_id
 
 
 def _find_file(file_name: str, stop_at_pythonpath: bool = True) -> Optional[Path]:
@@ -565,7 +562,7 @@ def get_workspace(aml_workspace: Optional[Workspace] = None, workspace_config_pa
 
 def create_run_recovery_id(run: Run) -> str:
     """
-   Creates an recovery id for a run so it's checkpoints could be recovered for training/testing
+    Creates a unique ID for a run, from which the experiment name and the run ID can be re-created
 
    :param run: an instantiated run.
    :return: recovery id for a given run in format: [experiment name]:[run id]
@@ -754,16 +751,16 @@ def create_python_environment(conda_environment_file: Path,
                               docker_base_image: str = "",
                               environment_variables: Optional[Dict[str, str]] = None) -> Environment:
     """
-    Creates a description for the Python execution environment in AzureML, based on the Conda environment
-    definition files that are specified in `source_config`. If such environment with this Conda environment already
-    exists, it is retrieved, otherwise created afresh.
+    Creates a description for the Python execution environment in AzureML, based on the arguments.
+    The environment will have a name that uniquely identifies it (it is based on hashing the contents of the
+    Conda file, the docker base image, environment variables and private wheels.
 
     :param environment_variables: The environment variables that should be set when running in AzureML.
     :param docker_base_image: The Docker base image that should be used when creating a new Docker image.
     :param pip_extra_index_url: If provided, use this PIP package index to find additional packages when building
         the Docker image.
     :param workspace: The AzureML workspace to work in, required if private_pip_wheel_path is supplied.
-    :param private_pip_wheel_path: If provided, add this wheel as a private package to the AzureML workspace.
+    :param private_pip_wheel_path: If provided, add this wheel as a private package to the AzureML environment.
     :param conda_environment_file: The file that contains the Conda environment definition.
     """
     conda_dependencies = CondaDependencies(conda_dependencies_file_path=conda_environment_file)
@@ -780,20 +777,31 @@ def create_python_environment(conda_environment_file: Path,
         **(environment_variables or {})
     }
     # See if this package as a whl exists, and if so, register it with AzureML environment.
-    if workspace is not None and private_pip_wheel_path is not None:
-        if private_pip_wheel_path.is_file():
-            whl_url = Environment.add_private_pip_wheel(workspace=workspace,
-                                                        file_path=str(private_pip_wheel_path),
-                                                        exist_ok=True)
-            conda_dependencies.add_pip_package(whl_url)
-            print(f"Added add_private_pip_wheel {private_pip_wheel_path} to AzureML environment.")
-        else:
-            raise FileNotFoundError(f"Cannot add add_private_pip_wheel: {private_pip_wheel_path}, it is not a file.")
+    if private_pip_wheel_path is not None:
+        if not private_pip_wheel_path.is_file():
+            raise FileNotFoundError(f"Cannot add private wheel: {private_pip_wheel_path} is not a file.")
+        if workspace is None:
+            raise ValueError("To use a private pip wheel, an AzureML workspace must be provided.")
+        whl_url = Environment.add_private_pip_wheel(workspace=workspace,
+                                                    file_path=str(private_pip_wheel_path),
+                                                    exist_ok=True)
+        conda_dependencies.add_pip_package(whl_url)
+        logging.info(f"Added add_private_pip_wheel {private_pip_wheel_path} to AzureML environment.")
     # Create a name for the environment that will likely uniquely identify it. AzureML does hashing on top of that,
     # and will re-use existing environments even if they don't have the same name.
     # Hashing should include everything that can reasonably change. Rely on hashlib here, because the built-in
-    # hash function gives different results for the same string in different python instances.
-    hash_string = "\n".join([yaml_contents, docker_base_image, str(environment_variables)])
+    hash_string = "\n".join([yaml_contents,
+                             docker_base_image,
+                             # Changing the index URL can lead to differences in package version resolution
+                             pip_extra_index_url,
+                             str(environment_variables),
+                             # Use the path of the private wheel as a proxy. This could lead to problems if
+                             # a new environment uses the same private wheel file name, but the wheel has different
+                             # contents. In hi-ml PR builds, the wheel file name is unique to the build, so it
+                             # should not occur there.
+                             str(private_pip_wheel_path)])
+    # Python's hash function gives different results for the same string in different python instances,
+    # hence need to use hashlib
     sha1 = hashlib.sha1(hash_string.encode("utf8"))
     overall_hash = sha1.hexdigest()[:32]
     unique_env_name = f"HealthML-{overall_hash}"

hi-ml-azure/testazure/testazure/test_azure_util.py

@@ -31,7 +31,6 @@ from azureml.core.environment import CondaDependencies
 from azureml.data.azure_storage_datastore import AzureBlobDatastore
 
 import health_azure.utils as util
-from health_azure import himl
 from health_azure.himl import AML_IGNORE_FILE, append_to_amlignore
 from testazure.test_himl import RunTarget, render_and_run_test_script
 from testazure.utils_testazure import (DEFAULT_IGNORE_FOLDERS, DEFAULT_WORKSPACE, MockRun, change_working_directory,
@@ -397,7 +396,6 @@ def test_create_python_environment(
         mock_workspace: mock.MagicMock,
         random_folder: Path,
 ) -> None:
-    just_conda_str_env_name = "HealthML-6555b3cfac0b3ee24349701f07dad394"
     conda_str = """name: simple-env
 dependencies:
   - pip=20.1.1
@@ -422,7 +420,8 @@ dependencies:
     assert "AZUREML_RUN_KILL_SIGNAL_TIMEOUT_SEC" in env.environment_variables
     assert "RSLEX_DIRECT_VOLUME_MOUNT" in env.environment_variables
     assert "RSLEX_DIRECT_VOLUME_MOUNT_MAX_CACHE_SIZE" in env.environment_variables
-    assert env.name == just_conda_str_env_name
+    # Just check that the environment has a reasonable name. Detailed checks for uniqueness of the name follow below.
+    assert env.name.startswith("HealthML")
 
     pip_extra_index_url = "https://where.great.packages.live/"
     docker_base_image = "viennaglobal.azurecr.io/azureml/azureml_a187a87cc7c31ac4d9f67496bc9c8239"
@@ -431,8 +430,9 @@ dependencies:
         pip_extra_index_url=pip_extra_index_url,
         docker_base_image=docker_base_image,
         environment_variables={"HELLO": "world"})
+    # Environment variables should be added to the default ones
     assert "HELLO" in env.environment_variables
-    assert env.name != just_conda_str_env_name
+    assert "RSLEX_DIRECT_VOLUME_MOUNT" in env.environment_variables
     assert env.docker.base_image == docker_base_image
 
     private_pip_wheel_url = "https://some.blob/private/wheel"
@@ -446,13 +446,80 @@ dependencies:
     assert "hi-ml-azure" in envs_pip_packages
     assert private_pip_wheel_url in envs_pip_packages
 
-    private_pip_wheel_path = Path("a_file_that_does_not.exist")
-    with pytest.raises(FileNotFoundError) as e:
-        _ = util.create_python_environment(
+
+def test_create_environment_unique_name(random_folder: Path) -> None:
+    """
+    Test if the name of the conda environment changes with each of the components
+    """
+    conda_str1 = """name: simple-env
+dependencies:
+  - pip=20.1.1
+  - python=3.7.3
+"""
+    conda_environment_file = random_folder / "environment.yml"
+    conda_environment_file.write_text(conda_str1)
+    env1 = util.create_python_environment(conda_environment_file=conda_environment_file)
+
+    # Changing the contents of the conda file should create a new environment names
+    conda_str2 = """name: simple-env
+dependencies:
+  - pip=20.1.1
+"""
+    assert conda_str1 != conda_str2
+    conda_environment_file.write_text(conda_str2)
+    env2 = util.create_python_environment(conda_environment_file=conda_environment_file)
+    assert env1.name != env2.name
+
+    # Using a different PIP index URL can lead to different package resolution, so this should change name too
+    env3 = util.create_python_environment(conda_environment_file=conda_environment_file,
+                                          pip_extra_index_url="foo")
+    assert env3.name != env2.name
+
+    # Environment variables
+    env4 = util.create_python_environment(conda_environment_file=conda_environment_file,
+                                          environment_variables={"foo": "bar"})
+    assert env4.name != env2.name
+
+    # Docker base image
+    env5 = util.create_python_environment(conda_environment_file=conda_environment_file,
+                                          docker_base_image="docker")
+    assert env5.name != env2.name
+
+    # PIP wheel
+    with mock.patch("health_azure.utils.Environment") as mock_environment:
+        mock_environment.add_private_pip_wheel.return_value = "private_pip_wheel_url"
+        env6 = util.create_python_environment(
             conda_environment_file=conda_environment_file,
-            workspace=mock_workspace,
-            private_pip_wheel_path=private_pip_wheel_path)
-    assert f"Cannot add add_private_pip_wheel: {private_pip_wheel_path}" in str(e.value)
+            workspace=DEFAULT_WORKSPACE.workspace,
+            private_pip_wheel_path=Path(__file__))
+        assert env6.name != env2.name
+
+    all_names = [env1.name, env2.name, env3.name, env4.name, env5.name, env6.name]
+    all_names_set = {*all_names}
+    assert len(all_names) == len(all_names_set), "Environment names are not unique"
+
+
+def test_create_environment_wheel_fails(random_folder: Path) -> None:
+    """
+    Test if all necessary checks are carried out when adding private wheels to an environment.
+    """
+    conda_str = """name: simple-env
+dependencies:
+  - pip=20.1.1
+  - python=3.7.3
+"""
+    conda_environment_file = random_folder / "environment.yml"
+    conda_environment_file.write_text(conda_str)
+    # Wheel file does not exist at all:
+    with pytest.raises(FileNotFoundError) as ex1:
+        util.create_python_environment(conda_environment_file=conda_environment_file,
+                                       private_pip_wheel_path=Path("does_not_exist"))
+        assert "Cannot add private wheel" in str(ex1)
+    # Wheel exists, but no workspace provided:
+    with pytest.raises(ValueError) as ex2:
+        util.create_python_environment(conda_environment_file=conda_environment_file,
+                                       private_pip_wheel_path=Path(__file__))
+        assert "AzureML workspace must be provided" in str(ex2)
 
 
 class MockEnvironment:
@@ -777,45 +844,56 @@ def test_download_file_from_run_remote(tmp_path: Path) -> None:
 
 
 def test_download_run_file_during_run(tmp_path: Path) -> None:
-    # This test will create a Run in your workspace (using only local compute)
+    """
+    Test if we can download files from a run, when executing inside AzureML. This should not require any additional
+    information about the workspace to use, but pick up the current workspace.
+    """
+    # Create a run that contains a simple txt file
+    experiment_name = "himl-tests"
+    run_to_download_from = util.create_aml_run_object(experiment_name=experiment_name,
+                                                      workspace=DEFAULT_WORKSPACE.workspace)
+    file_contents = "Hello World!"
+    file_name = "hello.txt"
+    full_file_path = tmp_path / file_name
+    full_file_path.write_text(file_contents)
+    run_to_download_from.upload_file(file_name, str(full_file_path))
+    run_to_download_from.complete()
+    run_id = run_to_download_from.id
 
-    expected_file_path = tmp_path / "azureml-logs"
-    # Check that at first the path to downloaded logs doesnt exist (will be created by the later test script)
-    assert not expected_file_path.exists()
+    # Test if we can retrieve the run directly from the workspace. This tests for a bug in an earlier version
+    # of the code where run IDs as those created from runs outside AML were not recognized
+    run_2 = util.get_aml_run_from_run_id(run_id, aml_workspace=DEFAULT_WORKSPACE.workspace)
+    assert run_2.id == run_id
 
-    ws = DEFAULT_WORKSPACE.workspace
+    # Now create an AzureML run with a simple script that uses that file. The script will download the file,
+    # where the download is should pick up the workspace from the current AML run.
+    script_body = ""
+    script_body += f"run_id = '{run_id}'\n"
+    script_body += f"    file_name = '{file_name}'\n"
+    script_body += f"    file_contents = '{file_contents}'\n"
+    script_body += """
+    output_path = Path("outputs")
+    output_path.mkdir(exist_ok=True)
 
-    # call the script here
+    download_files_from_run_id(run_id, output_path, prefix=file_name)
+    full_file_path = output_path / file_name
+    actual_contents = full_file_path.read_text().strip()
+    print(f"{actual_contents}")
+    assert actual_contents == file_contents
+"""
     extra_options = {
         "imports": """
 import sys
+from pathlib import Path
 from azureml.core import Run
-from health_azure.utils import _download_files_from_run""",
-        "args": """
-    parser.add_argument("--output_path", type=str, required=True)
-        """,
-        "body": """
-    output_path = Path(args.output_path)
-    output_path.mkdir(exist_ok=True)
-
-    run_ctx = Run.get_context()
-    available_files = run_ctx.get_file_names()
-    print(f"available files: {available_files}")
-    first_file_name = available_files[0]
-    output_file_path = output_path / first_file_name
-
-    _download_files_from_run(run_ctx, output_path, prefix=first_file_name)
-
-    print(f"Downloaded file {first_file_name} to location {output_file_path}")
-        """
+from health_azure.utils import download_files_from_run_id""",
+        "body": script_body
     }
-
-    extra_args = ["--output_path", 'outputs']
-    render_and_run_test_script(tmp_path, RunTarget.AZUREML, extra_options, extra_args, True)
-
-    run = util.get_most_recent_run(run_recovery_file=tmp_path / himl.RUN_RECOVERY_FILE,
-                                   workspace=ws)
-    assert run.status == "Completed"
+    # Run the script locally first, then in the cloud. In local runs, the workspace should be picked up from the
+    # config.json file, in AzureML runs it should be read off the run context.
+    render_and_run_test_script(tmp_path, RunTarget.LOCAL, extra_options, extra_args=[], expected_pass=True)
+    print("Local run finished")
+    render_and_run_test_script(tmp_path / "foo", RunTarget.AZUREML, extra_options, extra_args=[], expected_pass=True)
 
 
 def test_is_global_rank_zero() -> None:

hi-ml-azure/testazure/testazure/test_himl.py

@@ -607,6 +607,7 @@ def render_and_run_test_script(path: Path,
     :param suppress_config_creation: (Optional, defaults to False) do not create a config.json file if none exists
     :return: Either response from spawn_and_monitor_subprocess or run output if in AzureML.
     """
+    path.mkdir(exist_ok=True)
     # target hi-ml-azure package version, if specified in an environment variable.
     version = ""
     run_requirements = False
@@ -687,7 +688,7 @@ def render_and_run_test_script(path: Path,
 
         run = get_most_recent_run(run_recovery_file=path / himl.RUN_RECOVERY_FILE,
                                   workspace=workspace)
-        if run.status not in ["Failed", "Completed"]:
+        if run.status not in ["Failed", "Completed", "Cancelled"]:
             run.wait_for_completion()
         assert run.status == "Completed"
         log_root = path / "logs"
@@ -795,11 +796,8 @@ def test_invoking_hello_world_no_private_pip_fails(tmp_path: Path) -> None:
     extra_args: List[str] = []
     with mock.patch.dict(os.environ, {"HIML_AZURE_WHEEL_FILENAME": 'not_a_known_file.whl'}):
         output = render_and_run_test_script(tmp_path, RunTarget.AZUREML, extra_options, extra_args, False)
-    error_message_begin = "FileNotFoundError: Cannot add add_private_pip_wheel:"
-    error_message_end = "not_a_known_file.whl, it is not a file."
-
+    error_message_begin = "FileNotFoundError: Cannot add private wheel"
     assert error_message_begin in output
-    assert error_message_end in output
 
 
 @pytest.mark.parametrize("run_target", [RunTarget.LOCAL, RunTarget.AZUREML])

hi-ml-azure/testazure/testazure/utils_testazure.py

@@ -63,6 +63,13 @@ def change_working_directory(path_or_str: Path) -> Generator:
     os.chdir(old_path)
 
 
+def get_shared_config_json() -> Path:
+    """
+    Gets the path to the config.json file that should exist for running tests locally (outside github build agents).
+    """
+    return repository_root() / "hi-ml-azure" / "testazure" / WORKSPACE_CONFIG_JSON
+
+
 @contextmanager
 def check_config_json(script_folder: Path) -> Generator:
     """
@@ -70,20 +77,27 @@ def check_config_json(script_folder: Path) -> Generator:
     from the repository root folder (this should be the case when executing a test on a dev machine), or create
     it from environment variables (this should trigger in builds on the github agents).
     """
-    shared_config_json = repository_root() / WORKSPACE_CONFIG_JSON
+    shared_config_json = get_shared_config_json()
     target_config_json = script_folder / WORKSPACE_CONFIG_JSON
     if shared_config_json.exists():
         logging.info(f"Copying {WORKSPACE_CONFIG_JSON} from repository root to folder {script_folder}")
         shutil.copy(shared_config_json, target_config_json)
     else:
         logging.info(f"Creating {str(target_config_json)} from environment variables.")
-        with open(str(target_config_json), 'w', encoding="utf-8") as file:
-            config = {
-                "subscription_id": os.getenv(ENV_SUBSCRIPTION_ID, ""),
-                "resource_group": os.getenv(ENV_RESOURCE_GROUP, ""),
-                "workspace_name": os.getenv(ENV_WORKSPACE_NAME, "")
-            }
-            json.dump(config, file)
+        subscription_id = os.getenv(ENV_SUBSCRIPTION_ID, "")
+        resource_group = os.getenv(ENV_RESOURCE_GROUP, "")
+        workspace_name = os.getenv(ENV_WORKSPACE_NAME, "")
+        if subscription_id and resource_group and workspace_name:
+            with open(str(target_config_json), 'w', encoding="utf-8") as file:
+                config = {
+                    "subscription_id": os.getenv(ENV_SUBSCRIPTION_ID, ""),
+                    "resource_group": os.getenv(ENV_RESOURCE_GROUP, ""),
+                    "workspace_name": os.getenv(ENV_WORKSPACE_NAME, "")
+                }
+                json.dump(config, file)
+        else:
+            raise ValueError("Either a shared config.json must be present, or all 3 environment variables for "
+                             "workspace creation must exist.")
     try:
         yield
     finally: