Add a few args to submit function

src/health/azure/aml.py

@@ -7,13 +7,74 @@
 Wrapper functions for running local Python scripts on Azure ML.
 """
 import logging
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import Optional
+from attr import dataclass
+
+from azureml.core import Workspace
 
 
-logger = logging.getLogger('hi_ml')
+logger = logging.getLogger('health.azure')
 logger.setLevel(logging.DEBUG)
 
 
-def submit_to_azure_if_needed() -> None:
+@dataclass
+class WorkspaceConfig:
+    workspace_name: str = ""
+    subscription_id: str = ""
+    resource_group: str = ""
+
+
+def submit_to_azure_if_needed(
+        workspace_config: Optional[WorkspaceConfig],
+        workspace_config_path: Optional[Path]) -> None:
     """
     Submit a folder to Azure, if needed and run it.
+
+    :param workspace_config: Optional workspace config.
+    :param workspace_config_file: Optional path to workspace config file.
+    :return: None.
     """
+    workspace: Workspace = None
+    if workspace_config is not None:
+        workspace = Workspace.get(
+            name=workspace_config.workspace_name,
+            subscription_id=workspace_config.subscription_id,
+            resource_group=workspace_config.resource_group)
+    else:
+        workspace = Workspace.from_config(path=workspace_config_path)
+
+    if workspace is None:
+        raise ValueError("Unable to get workspace.")
+
+    print(f"Loaded: {workspace.name}")
+
+
+def main() -> None:
+    """
+    Handle submit_to_azure if called from the command line.
+    """
+    parser = ArgumentParser()
+    parser.add_argument("-w", "--workspace_name", type=str, required=False,
+                        help="Azure ML workspace name")
+    parser.add_argument("-s", "--subscription_id", type=str, required=False,
+                        help="AzureML subscription id")
+    parser.add_argument("-r", "--resource_group", type=str, required=False,
+                        help="AzureML resource group")
+    parser.add_argument("-p", "--workspace_config_path", type=str, required=False,
+                        help="AzureML workspace config file")
+
+    args = parser.parse_args()
+    config = WorkspaceConfig(
+        workspace_name=args.workspace_name,
+        subscription_id=args.subscription_id,
+        resource_group=args.resource_group)
+
+    submit_to_azure_if_needed(
+        config,
+        args.workspace_config_path)
+
+
+if __name__ == "__main__":
+    main()

test_requirements.txt

@@ -1 +1,2 @@
+azureml-core==1.32.0
 pytest==6.2.2

tests/health/azure/test_aml.py

@@ -8,14 +8,16 @@ Tests for hi-ml.
 """
 
 import logging
+import pytest
+
 
 try:
-    from hi_ml.aml import submit_to_azure_if_needed  # type: ignore
+    from health.azure.aml import submit_to_azure_if_needed  # type: ignore
 except ImportError:
     logging.info("using local src")
-    from src.hi_ml.aml import submit_to_azure_if_needed  # type: ignore
+    from src.health.azure.aml import submit_to_azure_if_needed  # type: ignore
 
-logger = logging.getLogger('test_hi_ml')
+logger = logging.getLogger('test.health.azure')
 logger.setLevel(logging.DEBUG)
 
 
@@ -23,5 +25,8 @@ def test_submit_to_azure_if_needed() -> None:
     """
     Test that submit_to_azure_if_needed can be called.
     """
-    submit_to_azure_if_needed()
-
+    with pytest.raises(Exception) as ex:
+        submit_to_azure_if_needed(
+            workspace_config=None,
+            workspace_config_path=None)
+    assert "We could not find config.json in:" in str(ex)