Repository: microsoft/hi-ml

Issue:
`--azureml` flag not working as expected
It seems that the `--azureml` flag is not working, at least not in the the [first steps docs](https://hi-ml.readthedocs.io/en/latest/first_steps.html) explains it.

If I try running my script with the `--azureml` flag it runs locally, unless I specify `submit_to_azureml=True` in the call to `submit_to_azureml_if_needed()`. After having a look at the code in [himl.py](https://github.com/microsoft/hi-ml/blob/main/hi-ml-azure/src/health_azure/himl.py#L834) this behaviour makes sense, as the code is run locally `if not submit_to_azureml`, where `submit_to_azureml` is set to `None` even if the `--azureml` flag is used. 

Not sure if I'm missing something here, but it appears the logic for checking for `--azureml` in `sys.argv` is missing.


Discussion:
Expected behaviour on local box with `--azureml`:
* Script runs on the local box, with `--azureml` on the commandline
* Submit function will prepare datasets if needed
* Then create a snapshot of the current codebase and send it to AzureML
* Create a `job` or `Run` from the codebase that invokes the current script, send that to AzureML. Commandline arguments for that job in AzureML should be the same as the current script, possibly minus the `--azureml` flag
* Stop execution by calling `exit` - AzureML will now take over and run the script in the cloud

Expected behaviour in AzureML:
* `submit_to_azure_if_needed` function returns a `RunInfo` object that contains paths to mounted datasets

Expected behaviour on local box without `--azureml`:
* `submit_to_azure_if_needed` function returns a `RunInfo` object that contains paths to either mounted datasets or the `local_path`(?) field of a `DatasetConfig` object


The test `test_submit_to_azure_if_needed_azure_return()` in `testazure/testazure/test_himl.py:564` is coded to test for incorrect behaviour. if the `--azureml` flag is passed then we'd expect the program to exit entirely with no run info returned.

Resolving patch:
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
@@ -50,7 +50,7 @@
 logger.setLevel(logging.DEBUG)
 
 AML_IGNORE_FILE = ".amlignore"
-AZUREML_COMMANDLINE_FLAG = "--azureml"
+AZUREML_FLAG = "--azureml"
 CONDA_ENVIRONMENT_FILE = "environment.yml"
 LOGS_FOLDER = "logs"
 OUTPUT_FOLDER = "outputs"
@@ -338,18 +338,21 @@ def create_crossval_hyperparam_args_v2(num_splits: int,
                                           metric_name=metric_name)
 
 
-def create_script_run(snapshot_root_directory: Optional[Path] = None,
-                      entry_script: Optional[PathOrString] = None,
-                      script_params: Optional[List[str]] = None) -> ScriptRunConfig:
+def create_script_run(
+    script_params: List[str],
+    snapshot_root_directory: Optional[Path] = None,
+    entry_script: Optional[PathOrString] = None,
+) -> ScriptRunConfig:
     """
     Creates an AzureML ScriptRunConfig object, that holds the information about the snapshot, the entry script, and
     its arguments.
 
-    :param entry_script: The script that should be run in AzureML.
+    :param script_params: A list of parameter to pass on to the script as it runs in AzureML. Required arg. Script
+        parameters can be generated using the ``_get_script_params()`` function.
     :param snapshot_root_directory: The directory that contains all code that should be packaged and sent to AzureML.
         All Python code that the script uses must be copied over.
-    :param script_params: A list of parameter to pass on to the script as it runs in AzureML. If empty (or None, the
-        default) these will be copied over from sys.argv, omitting the --azureml flag.
+    :param entry_script: The script that should be run in AzureML. If None, the current main Python file will be
+        executed.
     :return:
     """
     if snapshot_root_directory is None:
@@ -371,7 +374,6 @@ def create_script_run(snapshot_root_directory: Optional[Path] = None,
                              f"Snapshot root: {snapshot_root_directory}, entry script: {entry_script}")
     else:
         entry_script_relative = str(entry_script)
-    script_params = _get_script_params(script_params)
     print(f"This command will be run in AzureML: {entry_script_relative} {' '.join(script_params)}")
     return ScriptRunConfig(
         source_directory=str(snapshot_root_directory),
@@ -831,6 +833,8 @@ def submit_to_azure_if_needed(  # type: ignore
             return _generate_v2_azure_datasets(cleaned_input_datasets, cleaned_output_datasets)
     # This codepath is reached when executing outside AzureML. Here we first check if a script submission to AzureML
     # is necessary. If not, return to the caller for local execution.
+    if submit_to_azureml is None:
+        submit_to_azureml = AZUREML_FLAG in sys.argv[1:]
     if not submit_to_azureml:
         # Set the environment variables for local execution.
         environment_variables = {
@@ -877,6 +881,7 @@ def submit_to_azure_if_needed(  # type: ignore
 
     amlignore_path = snapshot_root_directory / AML_IGNORE_FILE
     lines_to_append = [str(path) for path in (ignored_folders or [])]
+    script_params = _get_script_params(script_params)
 
     with append_to_amlignore(amlignore=amlignore_path, lines_to_append=lines_to_append):
         if strictly_aml_v1:
@@ -895,9 +900,12 @@ def submit_to_azure_if_needed(  # type: ignore
                 input_datasets=cleaned_input_datasets,
                 output_datasets=cleaned_output_datasets,
             )
-            script_run_config = create_script_run(snapshot_root_directory=snapshot_root_directory,
-                                                  entry_script=entry_script,
-                                                  script_params=script_params)
+
+            script_run_config = create_script_run(
+                script_params=script_params,
+                snapshot_root_directory=snapshot_root_directory,
+                entry_script=entry_script,
+            )
             script_run_config.run_config = run_config
 
             if hyperdrive_config:
@@ -921,7 +929,6 @@ def submit_to_azure_if_needed(  # type: ignore
             )
             if entry_script is None:
                 entry_script = Path(sys.argv[0])
-            script_params = script_params or sys.argv[1:]
 
             ml_client = get_ml_client(ml_client=ml_client, aml_workspace=workspace)
             registered_env = register_environment_v2(environment, ml_client)
@@ -1008,7 +1015,7 @@ def _get_script_params(script_params: Optional[List[str]] = None) -> List[str]:
     """
     if script_params:
         return script_params
-    return [p for p in sys.argv[1:] if p != AZUREML_COMMANDLINE_FLAG]
+    return [p for p in sys.argv[1:] if p != AZUREML_FLAG]
 
 
 def _generate_azure_datasets(


Test patch:
diff --git a/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt b/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt
--- a/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt
+++ b/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt
@@ -52,7 +52,7 @@ def main() -> None:
         wait_for_completion={{ wait_for_completion }},
         wait_for_completion_show_output={{ wait_for_completion_show_output }},
         tags={{ tags }},
-        submit_to_azureml={{ submit_to_azureml}},
+        submit_to_azureml={{ submit_to_azureml }},
         strictly_aml_v1={{ strictly_aml_v1 }})
 
     parser = ArgumentParser()
diff --git a/hi-ml-azure/testazure/testazure/test_himl.py b/hi-ml-azure/testazure/testazure/test_himl.py
--- a/hi-ml-azure/testazure/testazure/test_himl.py
+++ b/hi-ml-azure/testazure/testazure/test_himl.py
@@ -67,6 +67,7 @@
 INEXPENSIVE_TESTING_CLUSTER_NAME = "lite-testing-ds2"
 EXPECTED_QUEUED = "This command will be run in AzureML:"
 GITHUB_SHIBBOLETH = "GITHUB_RUN_ID"  # https://docs.github.com/en/actions/reference/environment-variables
+AZUREML_FLAG = himl.AZUREML_FLAG
 
 logger = logging.getLogger('test.health_azure')
 logger.setLevel(logging.DEBUG)
@@ -514,13 +515,15 @@ def test_invalid_entry_script(tmp_path: Path) -> None:
 
     with pytest.raises(ValueError) as e:
         himl.create_script_run(
+            script_params=[],
             entry_script=problem_entry_script,
             snapshot_root_directory=snapshot_dir,
-            script_params=[])
+        )
     assert "entry script must be inside of the snapshot root directory" in str(e)
 
     with mock.patch("sys.argv", ["foo"]):
-        script_run = himl.create_script_run()
+        script_params = himl._get_script_params()
+        script_run = himl.create_script_run(script_params)
         assert script_run.source_directory == str(Path.cwd())
         assert script_run.script == "foo"
         assert script_run.arguments == []
@@ -535,7 +538,7 @@ def test_invalid_entry_script(tmp_path: Path) -> None:
 def test_get_script_params() -> None:
     expected_params = ["a string"]
     assert expected_params == himl._get_script_params(expected_params)
-    with mock.patch("sys.argv", ["", "a string", "--azureml"]):
+    with mock.patch("sys.argv", ["", "a string", AZUREML_FLAG]):
         assert expected_params == himl._get_script_params()
     with mock.patch("sys.argv", ["", "a string"]):
         assert expected_params == himl._get_script_params()
@@ -582,7 +585,7 @@ def test_submit_to_azure_if_needed_azure_return(
         output_folder=Path.cwd() / himl.OUTPUT_FOLDER,
         logs_folder=Path.cwd() / himl.LOGS_FOLDER)
     mock_generate_azure_datasets.return_value = expected_run_info
-    with mock.patch("sys.argv", ["", "--azureml"]):
+    with mock.patch("sys.argv", ["", AZUREML_FLAG]):
         run_info = himl.submit_to_azure_if_needed(
             aml_workspace=mock_workspace,
             entry_script=Path(__file__),
@@ -1030,7 +1033,10 @@ def spawn() -> Tuple[int, List[str]]:
                                               f"got a return code {code}"
 
     if run_target == RunTarget.LOCAL or not expected_pass:
-        assert EXPECTED_QUEUED not in captured
+        if AZUREML_FLAG in extra_args:
+            assert EXPECTED_QUEUED in captured
+        else:
+            assert EXPECTED_QUEUED not in captured
         return captured
     else:
         assert EXPECTED_QUEUED in captured
@@ -1125,6 +1131,25 @@ def test_invoking_hello_world_config(run_target: RunTarget, use_package: bool, t
     assert expected_output in output
 
 
+def test_invoking_hello_world_using_azureml_flag(tmp_path: Path) -> None:
+    """
+    Test that invoking hello_world.py with the --azureml flag will submit to AzureML and not run locally.
+    :param tmp_path: PyTest test fixture for temporary path.
+    """
+
+    message_guid = uuid4().hex
+    parser_args = "parser.add_argument('-m', '--message', type=str, required=True, help='The message to print out')"
+    extra_options = {
+        'args': parser_args,
+        'body': 'print(f"The message was: {args.message}")',
+        'submit_to_azureml': None,
+    }
+    extra_args = [f"--message={message_guid}", AZUREML_FLAG]
+    output = render_and_run_test_script(tmp_path, RunTarget.LOCAL, extra_options, extra_args, True)
+    expected_output = f"The message was: {message_guid}"
+    assert expected_output in output
+
+
 @patch("health_azure.himl.submit_to_azure_if_needed")
 def test_calling_script_directly(mock_submit_to_azure_if_needed: mock.MagicMock) -> None:
     with mock.patch("sys.argv", ["",
@@ -1509,7 +1534,7 @@ def test_submit_to_azure_if_needed_with_hyperdrive(mock_sys_args: MagicMock,
     Test that himl.submit_to_azure_if_needed can be called, and returns immediately.
     """
     cross_validation_metric_name = cross_validation_metric_name or ""
-    mock_sys_args.return_value = ["", "--azureml"]
+    mock_sys_args.return_value = ["", AZUREML_FLAG]
     with patch("health_azure.himl.get_ml_client") as mock_get_ml_client:
         mock_ml_client = MagicMock()
         mock_get_ml_client.return_value = mock_ml_client