Repository: microsoft/hi-ml

Issue:
get_workspace fails when running on Azure DevOps build pipelines
It is calling is_running_on_azur_agent - but actually it should check if it is really running in AzureML


Discussion:

Resolving patch:
diff --git a/docs/source/conf.py b/docs/source/conf.py
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -18,8 +18,9 @@
 import os
 import sys
 
-root_dir = os.path.abspath('../src')
-sys.path.insert(0, root_dir)
+# Make the source code for both packages available here
+sys.path.insert(0, os.path.abspath('../../hi-ml/src'))
+sys.path.insert(0, os.path.abspath('../../hi-ml-azure/src'))
 
 # -- Project information -----------------------------------------------------
 
diff --git a/hi-ml-azure/src/health/azure/azure_util.py b/hi-ml-azure/src/health/azure/azure_util.py
--- a/hi-ml-azure/src/health/azure/azure_util.py
+++ b/hi-ml-azure/src/health/azure/azure_util.py
@@ -5,13 +5,13 @@
 """
 Utility functions for interacting with AzureML runs
 """
-from argparse import Namespace
-from enum import Enum
 import hashlib
-from itertools import islice
 import logging
 import os
 import re
+from argparse import Namespace
+from enum import Enum
+from itertools import islice
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -100,7 +100,7 @@ def get_workspace(aml_workspace: Optional[Workspace], workspace_config_path: Opt
         config
     :return: An AzureML workspace.
     """
-    if is_running_on_azure_agent():
+    if is_running_in_azure_ml(RUN_CONTEXT):
         return RUN_CONTEXT.experiment.workspace
 
     if aml_workspace:
@@ -172,14 +172,6 @@ def fetch_run(workspace: Workspace, run_recovery_id: str) -> Run:
     return run_to_recover
 
 
-def is_running_on_azure_agent() -> bool:
-    """
-    Returns True if the code appears to be running on an Azure build agent, and False otherwise.
-    """
-    # Guess by looking at the AGENT_OS variable, that all Azure hosted agents define.
-    return bool(os.environ.get("AGENT_OS", None))
-
-
 def fetch_run_for_experiment(experiment_to_recover: Experiment, run_id: str) -> Run:
     """
     Gets an AzureML Run object for a given run ID in an experiment.
@@ -905,3 +897,16 @@ def download_checkpoints_from_run_id(run_id: str, checkpoint_dir: str, output_fo
     workspace = get_workspace(aml_workspace=aml_workspace, workspace_config_path=workspace_config_path)
     download_files_from_run_id(run_id, output_folder, prefix=checkpoint_dir, workspace=workspace,
                                validate_checksum=True)
+
+
+def is_running_in_azure_ml(aml_run: Run = RUN_CONTEXT) -> bool:
+    """
+    Returns True if the given run is inside of an AzureML machine, or False if it is on a machine outside AzureML.
+    When called without arguments, this functions returns True if the present code is running in AzureML.
+    Note that in runs with "compute_target='local'" this function will also return True. Such runs execute outside
+    of AzureML, but are able to log all their metrics, etc to an AzureML run.
+
+    :param aml_run: The run to check. If omitted, use the default run in RUN_CONTEXT
+    :return: True if the given run is inside of an AzureML machine, or False if it is a machine outside AzureML.
+    """
+    return hasattr(aml_run, 'experiment')
diff --git a/hi-ml-azure/src/health/azure/himl.py b/hi-ml-azure/src/health/azure/himl.py
--- a/hi-ml-azure/src/health/azure/himl.py
+++ b/hi-ml-azure/src/health/azure/himl.py
@@ -27,7 +27,7 @@
 from azureml.train.hyperdrive import HyperDriveConfig
 
 from health.azure.azure_util import (create_python_environment, create_run_recovery_id, _find_file,
-                                     is_run_and_child_runs_completed, register_environment,
+                                     is_run_and_child_runs_completed, is_running_in_azure_ml, register_environment,
                                      run_duration_string_to_seconds,
                                      to_azure_friendly_string, RUN_CONTEXT, get_workspace)
 from health.azure.datasets import (DatasetConfig, StrOrDatasetConfig, _input_dataset_key, _output_dataset_key,
@@ -68,7 +68,7 @@ class AzureRunInfo:
     run: Optional[Run]
     """An AzureML Run object if the present script is executing inside AzureML, or None if outside of AzureML.
     The Run object has methods to log metrics, upload files, etc."""
-    is_running_in_azure: bool
+    is_running_in_azure_ml: bool
     """If True, the present script is executing inside AzureML. If False, outside AzureML."""
     output_folder: Path
     """The output folder into which all script outputs should be written, if they should be later available in the
@@ -78,17 +78,6 @@ class AzureRunInfo:
     folder will be uploaded to blob storage regularly during the script run."""
 
 
-def is_running_in_azure(aml_run: Run = RUN_CONTEXT) -> bool:
-    """
-    Returns True if the given run is inside of an AzureML machine, or False if it is a machine outside AzureML.
-    When called without arguments, this functions returns True if the present code is running in AzureML.
-
-    :param aml_run: The run to check. If omitted, use the default run in RUN_CONTEXT
-    :return: True if the given run is inside of an AzureML machine, or False if it is a machine outside AzureML.
-    """
-    return hasattr(aml_run, 'experiment')
-
-
 def create_run_configuration(workspace: Workspace,
                              compute_cluster_name: str,
                              conda_environment_file: Optional[Path] = None,
@@ -373,7 +362,7 @@ def submit_to_azure_if_needed(  # type: ignore
     # The present function will most likely be called from the script once it is running in AzureML.
     # The '--azureml' flag will not be present anymore, but we don't want to rely on that. From Run.get_context we
     # can infer if the present code is running in AzureML.
-    in_azure = is_running_in_azure()
+    in_azure = is_running_in_azure_ml(RUN_CONTEXT)
     if in_azure:
         return _generate_azure_datasets(cleaned_input_datasets, cleaned_output_datasets)
 
@@ -397,7 +386,7 @@ def submit_to_azure_if_needed(  # type: ignore
             input_datasets=[d.local_folder for d in cleaned_input_datasets],
             output_datasets=[d.local_folder for d in cleaned_output_datasets],
             run=None,
-            is_running_in_azure=False,
+            is_running_in_azure_ml=False,
             output_folder=output_folder,
             logs_folder=logs_folder
         )
@@ -523,7 +512,7 @@ def _generate_azure_datasets(
         input_datasets=returned_input_datasets,  # type: ignore
         output_datasets=returned_output_datasets,  # type: ignore
         run=RUN_CONTEXT,
-        is_running_in_azure=True,
+        is_running_in_azure_ml=True,
         output_folder=Path.cwd() / OUTPUT_FOLDER,
         logs_folder=Path.cwd() / LOGS_FOLDER)
 


Test patch:
diff --git a/hi-ml-azure/tests/testazure/test_azure_util.py b/hi-ml-azure/tests/testazure/test_azure_util.py
--- a/hi-ml-azure/tests/testazure/test_azure_util.py
+++ b/hi-ml-azure/tests/testazure/test_azure_util.py
@@ -62,6 +62,23 @@ def test_find_file(tmp_path: Path) -> None:
     os.chdir(where_are_we_now)
 
 
+def test_is_running_in_azureml() -> None:
+    """
+    Test if the code correctly recognizes that it is executed in AzureML
+    """
+    # These tests would always run outside of AzureML, on local box or Azure build agents. Function should return
+    # False in all those cases
+    assert not util.is_running_in_azure_ml()
+    assert not util.is_running_in_azure_ml(util.RUN_CONTEXT)
+    # When in AzureML, the Run has a field "experiment"
+    mock_workspace = "foo"
+    with patch("health.azure.azure_util.RUN_CONTEXT") as mock_run_context:
+        mock_run_context.experiment = MagicMock(workspace=mock_workspace)
+        # We can't try that with the default argument because of Python's handling of mutable default arguments
+        # (default argument value has been assigned already before mocking)
+        assert util.is_running_in_azure_ml(util.RUN_CONTEXT)
+
+
 @pytest.mark.fast
 @patch("health.azure.azure_util.Workspace.from_config")
 @patch("health.azure.azure_util.get_authentication")
@@ -73,16 +90,16 @@ def test_get_workspace(
         tmp_path: Path) -> None:
 
     # Test the case when running on AML
-    with patch("health.azure.azure_util.is_running_on_azure_agent") as mock_is_is_running_on_azure_agent:
-        mock_is_is_running_on_azure_agent.return_value = True
-        with patch("health.azure.azure_util.RUN_CONTEXT") as mock_run_context:
-            mock_run_context.experiment = MagicMock(workspace=mock_workspace)
-            workspace = util.get_workspace(None, None)
-            assert workspace == mock_workspace
+    with patch("health.azure.azure_util.RUN_CONTEXT") as mock_run_context:
+        mock_run_context.experiment = MagicMock(workspace=mock_workspace)
+        workspace = util.get_workspace(None, None)
+        assert workspace == mock_workspace
 
-    # Test the case when a workspace object is provided
-    workspace = util.get_workspace(mock_workspace, None)
-    assert workspace == mock_workspace
+    # Test the case when a workspace object is provided. The test always runs outside AzureML, and should return the
+    # workspace object unchanged
+    mock_workspace2 = "foo"
+    workspace = util.get_workspace(mock_workspace2, None)  # type: ignore
+    assert workspace == mock_workspace2
 
     # Test the case when a workspace config path is provided
     mock_get_authentication.return_value = "auth"
@@ -111,12 +128,6 @@ def test_create_run_recovery_id(mock_run: MagicMock) -> None:
     assert recovery_id == EXPERIMENT_NAME + util.EXPERIMENT_RUN_SEPARATOR + RUN_ID
 
 
-@pytest.mark.parametrize("on_azure", [True, False])
-def test_is_running_on_azure_agent(on_azure: bool) -> None:
-    with mock.patch.dict(os.environ, {"AGENT_OS": "LINUX" if on_azure else ""}):
-        assert on_azure == util.is_running_on_azure_agent()
-
-
 @patch("health.azure.azure_util.Workspace")
 @patch("health.azure.azure_util.Experiment")
 @patch("health.azure.azure_util.Run")
diff --git a/hi-ml-azure/tests/testazure/test_himl.py b/hi-ml-azure/tests/testazure/test_himl.py
--- a/hi-ml-azure/tests/testazure/test_himl.py
+++ b/hi-ml-azure/tests/testazure/test_himl.py
@@ -28,7 +28,8 @@
 from azureml.train.hyperdrive import HyperDriveConfig
 
 import health.azure.himl as himl
-from health.azure.azure_util import EXPERIMENT_RUN_SEPARATOR, get_most_recent_run, WORKSPACE_CONFIG_JSON, get_workspace
+from health.azure.azure_util import (EXPERIMENT_RUN_SEPARATOR, WORKSPACE_CONFIG_JSON, get_most_recent_run,
+                                     get_workspace, is_running_in_azure_ml)
 from health.azure.datasets import DatasetConfig, _input_dataset_key, _output_dataset_key, get_datastore
 from testazure.test_data.make_tests import render_environment_yaml, render_test_script
 from testazure.util import DEFAULT_DATASTORE, change_working_directory, check_config_json, repository_root
@@ -66,7 +67,7 @@ def test_submit_to_azure_if_needed_returns_immediately() -> None:
             compute_cluster_name="foo",
             conda_environment_file=Path("env.yml"))
         assert isinstance(result, himl.AzureRunInfo)
-        assert not result.is_running_in_azure
+        assert not result.is_running_in_azure_ml
         assert result.run is None
 
 
@@ -259,7 +260,7 @@ def test_get_script_params() -> None:
 
 
 @pytest.mark.fast
-@patch("health.azure.himl.is_running_in_azure")
+@patch("health.azure.azure_util.is_running_in_azure_ml")
 def test_get_workspace_no_config(
         mock_is_running_in_azure: mock.MagicMock,
         tmp_path: Path) -> None:
@@ -279,18 +280,26 @@ def test_get_workspace_no_config(
 @patch("health.azure.himl.Run")
 @patch("health.azure.himl.Workspace")
 @patch("health.azure.himl._generate_azure_datasets")
-@patch("health.azure.himl.is_running_in_azure")
+@patch("health.azure.himl.RUN_CONTEXT")
 def test_submit_to_azure_if_needed_azure_return(
-        mock_is_running_in_azure: mock.MagicMock,
+        mock_run_context: mock.MagicMock,
         mock_generate_azure_datasets: mock.MagicMock,
         mock_workspace: mock.MagicMock,
         mock_run: mock.MagicMock) -> None:
-    mock_is_running_in_azure.return_value = True
+    """
+    When running in AzureML, the call to submit_to_azure_if_needed should return immediately, without trying to
+    submit a new job.
+    """
+    # The presence of the "experiment" flag is the trigger to recognize an AzureML run.
+    mock_run_context.experiment = mock.MagicMock(workspace=mock_workspace)
+    # This import needs to be local, after mocking the RUN_CONTEXT
+    import health
+    assert is_running_in_azure_ml(health.azure.himl.RUN_CONTEXT)
     expected_run_info = himl.AzureRunInfo(
         run=mock_run,
         input_datasets=[],
         output_datasets=[],
-        is_running_in_azure=True,
+        is_running_in_azure_ml=True,
         output_folder=Path.cwd(),
         logs_folder=Path.cwd())
     mock_generate_azure_datasets.return_value = expected_run_info
@@ -317,7 +326,7 @@ def test_generate_azure_datasets(
     run_info = himl._generate_azure_datasets(
         cleaned_input_datasets=[mock_dataset_config] * 2,
         cleaned_output_datasets=[mock_dataset_config] * 3)
-    assert run_info.is_running_in_azure
+    assert run_info.is_running_in_azure_ml
     assert len(run_info.input_datasets) == 2
     assert len(run_info.output_datasets) == 3
     for i, d in enumerate(run_info.input_datasets):