Use Azure ML Environments and add code upload option

docs/source/03_quickstart.rst

@@ -19,7 +19,6 @@ created in Azure and have their **names** ready to input to the plugin:
 -  Azure ML Compute Cluster
 -  Azure Storage Account and Storage Container
 -  Azure Storage Key (will be used to execute the pipeline)
--  Azure Container Registry
 
 1. Make sure that you're logged into Azure (``az login``).
 2. Prepare new virtual environment with Python >=3.8. Install the
@@ -47,30 +46,68 @@ created in Azure and have their **names** ready to input to the plugin:
    - The project's python package in /Users/marcin/Dev/tmp/kedro-azureml-demo/src/kedro_azureml_demo
 
 3. Go to the project's directory: ``cd kedro-azureml-demo``
-4. Add ``kedro-azureml`` to ``src/requriements.txt``
+4. Add ``kedro-azureml`` to ``src/requirements.txt``
 5. (optional) Remove ``kedro-telemetry`` from ``src/requirements.txt``
    or set appopriate settings
    (`https://github.com/kedro-org/kedro-plugins/tree/main/kedro-telemetry <https://github.com/kedro-org/kedro-plugins/tree/main/kedro-telemetry>`__).
 6. Install the requirements ``pip install -r src/requirements.txt``
 7. Initialize Kedro Azure ML plugin, it requires the Azure resource
    names as stated above. Experiment name can be anything you like (as
-   long as it's allowed by Azure ML).
+   long as it's allowed by Azure ML). The environment name is the name
+   of the Azure ML Environment to be created in the next step. You can
+   use the syntax <environment_name>@latest for the latest version or
+   <environment-name>:<version> for a specific version.
 
 .. code:: console
 
    #Usage: kedro azureml init [OPTIONS] RESOURCE_GROUP WORKSPACE_NAME
    #                          EXPERIMENT_NAME CLUSTER_NAME STORAGE_ACCOUNT_NAME
-   #                          STORAGE_CONTAINER
-   kedro azureml init <resource-group-name> <workspace-name> <experiment-name> <compute-cluster-name> <storage-account-name> <storage-container-name> --acr <azure-container-registry-name>
+   #                          STORAGE_CONTAINER ENVIRONMENT_NAME
+   kedro azureml init <resource-group-name> <workspace-name> <experiment-name> <compute-cluster-name> <storage-account-name> <storage-container-name> <environment-name>
+
+
+8. Create an Azure ML Environment for the project:
+
+   For the project's code to run on Azure ML it needs to have an environment
+   with the necessary dependencies. Here is it shown how to do this from a
+   local Docker build context. Please refer to the
+   `Azure ML CLI documentation <https://learn.microsoft.com/en-us/azure/machine-learning/how-to-manage-environments-v2#create-an-environment>`
+   for more options.
+
+   Start by executing the following command:
+
+.. code:: console
+
+   kedro docker init
+
+   This command creates a several files, including ``Dockerfile`` and 
+   ``.dockerignore``. These can be adjusted to match the workflow for
+   your project.
+
+   Depending on whether you want to use code upload when submitting an
+   experiment or not, you would need to add the code and any possible input
+   data to the Docker image.
+
+   - If using code upload:
+      Everything apart from the section "install project requirements"
+      can be removed from the ``Dockerfile``. You can add a
+      ``.amlignore`` file to specify which files should be uploaded.
+
+      Set ``code_directory: "."`` (or a subdirectory containing the code to upload)
+      in the ``azureml.yml`` config file.
+
+   - If not using code upload:
+      Keep the sections in the ``Dockerfile`` and adjust the ``.dockerignore``
+      file to add any other files to be added to the Docker image,
+      such as ``!data/01_raw`` for the raw data files.
+
+   Create or update an Azure ML Environment by running the following command:
 
 .. code:: console
 
-   Configuration generated in /Users/marcin/Dev/tmp/kedro-azureml-demo/conf/base/azureml.yml
-   It's recommended to set Lifecycle management rule for storage container kedro-azure-storage to avoid costs of long-term storage of the temporary data.
-   Temporary data will be stored under abfs://kedro-azure-storage/kedro-azureml-temp path
-   See https://docs.microsoft.com/en-us/azure/storage/blobs/lifecycle-management-policy-configure?tabs=azure-portal
+   az ml environment create --name <environment-name> --version <version> --build-context . --dockerfile-path Dockerfile
 
-8. Adjust the Data Catalog - the default one stores all data locally,
+9. Adjust the Data Catalog - the default one stores all data locally,
    whereas the plugin will automatically use Azure Blob Storage. Only
    input data is required to be read locally. Final
    ``conf/base/catalog.yml`` should look like this:
@@ -92,37 +129,7 @@ created in Azure and have their **names** ready to input to the plugin:
      filepath: data/01_raw/shuttles.xlsx
      layer: raw
 
-8. Build docker image for the project:
-
-.. code:: console
-
-   kedro docker init
-
-This command creates a several files, including ``.dockerignore``. This
-file ensures that transient files are not included in the docker image
-and it requires small adjustment. Open it in your favourite text editor
-and extend the section ``# except the following`` by adding there:
-
-.. code:: console
-
-   !data/01_raw
-
-Invoke docker build
-
-.. code:: console
-
-   kedro docker build --docker-args "--build-arg=BASE_IMAGE=python:3.9" --image=<image tag from conf/base/azureml.yml>
-
-Once finished, push the image:
-
-.. code:: console
-
-   docker push <image tag from conf/base/azureml.yml>
-
-(you will need to authorize to the ACR first, e.g. by
-``az acr login --name <acr repo name>`` )
-
-9. Run the pipeline on Azure ML Pipelines. Here, the *Azure Subscription
+10. Run the pipeline on Azure ML Pipelines. Here, the *Azure Subscription
    ID* and *Storage Account Key* will be used:
 
 .. code:: console
@@ -138,10 +145,9 @@ You will most likely see the following prompt:
 
 Input the storage account key and press [ENTER] (input will be hidden).
 
-10. Plugin will verify the configuration (e.g. the existence of the
+11. Plugin will verify the configuration (e.g. the existence of the
     compute cluster) and then it will create a *Job* in the Azure ML.
     The URL to view the job will be displayed in the console output.
-11.
 
 12. (optional) You can also use
     ``kedro azureml run -s <azure-subscription-id> --wait-for-completion``

kedro_azureml/cli.py

@@ -15,7 +15,6 @@
 from kedro_azureml.config import CONFIG_TEMPLATE_YAML
 from kedro_azureml.constants import (
     AZURE_SUBSCRIPTION_ID,
-    FILL_IN_DOCKER_IMAGE,
     KEDRO_AZURE_BLOB_TEMP_DIR_NAME,
 )
 from kedro_azureml.distributed.utils import is_distributed_master_node
@@ -56,7 +55,7 @@ def azureml_group(ctx, metadata: ProjectMetadata, env):
 @click.argument("cluster_name")
 @click.argument("storage_account_name")
 @click.argument("storage_container")
-@click.option("--acr", help="Azure Container Registry repo name", default="")
+@click.argument("environment_name")
 @click.pass_obj
 def init(
     ctx: CliContext,
@@ -66,49 +65,36 @@ def init(
     cluster_name,
     storage_account_name,
     storage_container,
-    acr,
+    environment_name,
 ):
     """
     Creates basic configuration for Kedro AzureML plugin
     """
-    with KedroContextManager(ctx.metadata.package_name, ctx.env) as mgr:
-        target_path = Path.cwd().joinpath("conf/base/azureml.yml")
-        cfg = CONFIG_TEMPLATE_YAML.format(
-            **{
-                "resource_group": resource_group,
-                "workspace_name": workspace_name,
-                "experiment_name": experiment_name,
-                "cluster_name": cluster_name,
-                "docker_image": (
-                    f"{acr}.azurecr.io/{mgr.context.project_path.name}:latest"
-                    if acr
-                    else FILL_IN_DOCKER_IMAGE
-                ),
-                "storage_container": storage_container,
-                "storage_account_name": storage_account_name,
-            }
-        )
-        target_path.write_text(cfg)
+    target_path = Path.cwd().joinpath("conf/base/azureml.yml")
+    cfg = CONFIG_TEMPLATE_YAML.format(
+        **{
+            "resource_group": resource_group,
+            "workspace_name": workspace_name,
+            "experiment_name": experiment_name,
+            "cluster_name": cluster_name,
+            "environment_name": environment_name,
+            "storage_container": storage_container,
+            "storage_account_name": storage_account_name,
+        }
+    )
+    target_path.write_text(cfg)
 
-        click.echo(f"Configuration generated in {target_path}")
+    click.echo(f"Configuration generated in {target_path}")
 
-        click.echo(
-            click.style(
-                f"It's recommended to set Lifecycle management rule for storage container {storage_container} "
-                f"to avoid costs of long-term storage of the temporary data."
-                f"\nTemporary data will be stored under abfs://{storage_container}/{KEDRO_AZURE_BLOB_TEMP_DIR_NAME} path"  # noqa
-                f"\nSee https://docs.microsoft.com/en-us/azure/storage/blobs/lifecycle-management-policy-configure?tabs=azure-portal",  # noqa
-                fg="green",
-            )
+    click.echo(
+        click.style(
+            f"It's recommended to set Lifecycle management rule for storage container {storage_container} "
+            f"to avoid costs of long-term storage of the temporary data."
+            f"\nTemporary data will be stored under abfs://{storage_container}/{KEDRO_AZURE_BLOB_TEMP_DIR_NAME} path"  # noqa
+            f"\nSee https://docs.microsoft.com/en-us/azure/storage/blobs/lifecycle-management-policy-configure?tabs=azure-portal",  # noqa
+            fg="green",
         )
-
-        if not acr:
-            click.echo(
-                click.style(
-                    "Please fill in docker image name before running the pipeline",
-                    fg="yellow",
-                )
-            )
+    )
 
 
 @azureml_group.command()
@@ -120,10 +106,11 @@ def init(
     type=str,
 )
 @click.option(
-    "-i",
-    "--image",
+    "--azureml_environment",
+    "--aml_env",
+    "aml_env",
     type=str,
-    help="Docker image to use for pipeline execution.",
+    help="Azure ML Environment to use for pipeline execution.",
 )
 @click.option(
     "-p",
@@ -146,7 +133,7 @@ def run(
     click_context: click.Context,
     ctx: CliContext,
     subscription_id: str,
-    image: Optional[str],
+    aml_env: Optional[str],
     pipeline: str,
     params: str,
     wait_for_completion: bool,
@@ -159,11 +146,11 @@ def run(
         subscription_id
     ), f"Please provide Azure Subscription ID or set `{AZURE_SUBSCRIPTION_ID}` env"
 
-    if image:
-        click.echo(f"Overriding image for run to: {image}")
+    if aml_env:
+        click.echo(f"Overriding Azure ML Environment for run by: {aml_env}")
 
     mgr: KedroContextManager
-    with get_context_and_pipeline(ctx, image, pipeline, params) as (mgr, az_pipeline):
+    with get_context_and_pipeline(ctx, aml_env, pipeline, params) as (mgr, az_pipeline):
         az_client = AzureMLPipelinesClient(az_pipeline, subscription_id)
 
         is_ok = az_client.run(
@@ -193,10 +180,11 @@ def run(
 
 @azureml_group.command()
 @click.option(
-    "-i",
-    "--image",
+    "--azureml_environment",
+    "--aml_env",
+    "aml_env",
     type=str,
-    help="Docker image to use for pipeline execution.",
+    help="Azure ML Environment to use for pipeline execution.",
 )
 @click.option(
     "-p",
@@ -221,11 +209,11 @@ def run(
 )
 @click.pass_obj
 def compile(
-    ctx: CliContext, image: Optional[str], pipeline: str, params: list, output: str
+    ctx: CliContext, aml_env: Optional[str], pipeline: str, params: list, output: str
 ):
     """Compiles the pipeline into YAML format"""
     params = json.dumps(p) if (p := parse_extra_params(params)) else ""
-    with get_context_and_pipeline(ctx, image, pipeline, params) as (_, az_pipeline):
+    with get_context_and_pipeline(ctx, aml_env, pipeline, params) as (_, az_pipeline):
         Path(output).write_text(str(az_pipeline))
         click.echo(f"Compiled pipeline to {output}")
 

kedro_azureml/cli_functions.py

@@ -12,7 +12,7 @@
 
 
 @contextmanager
-def get_context_and_pipeline(ctx: CliContext, image: str, pipeline: str, params: str):
+def get_context_and_pipeline(ctx: CliContext, aml_env: str, pipeline: str, params: str):
     with KedroContextManager(
         ctx.metadata.package_name, ctx.env, parse_extra_params(params, True)
     ) as mgr:
@@ -35,7 +35,7 @@ def get_context_and_pipeline(ctx: CliContext, image: str, pipeline: str, params:
             ctx.env,
             mgr.plugin_config,
             mgr.context.params,
-            image,
+            aml_env,
             params,
             storage_account_key,
         )

kedro_azureml/config.py

@@ -12,10 +12,6 @@ def __getitem__(self, key):
         return defaults.copy(update=this.dict(exclude_none=True)) if defaults else this
 
 
-class DockerConfig(BaseModel):
-    image: str
-
-
 class AzureTempStorageConfig(BaseModel):
     account_name: str
     container: str
@@ -44,11 +40,13 @@ def _validate_compute(cls, value):
     resource_group: str
     temporary_storage: AzureTempStorageConfig
     compute: Optional[Dict[str, ComputeConfig]]
+    environment_name: str
+    code_directory: Optional[str]
+    working_directory: Optional[str]
 
 
 class KedroAzureMLConfig(BaseModel):
     azure: AzureMLConfig
-    docker: DockerConfig
 
 
 class KedroAzureRunnerConfig(BaseModel):
@@ -66,6 +64,13 @@ class KedroAzureRunnerConfig(BaseModel):
   resource_group: "{resource_group}"
   # Azure ML Workspace name
   workspace_name: "{workspace_name}"
+  # Azure ML Environment to use during pipeline execution
+  environment_name: "{environment_name}"
+  # Path to directory to upload, or null to disable code upload
+  code_directory: null
+  # Path to the directory in the Docker image to run the code from
+  # Ignored when code_directory is set
+  working_directory: /home/kedro
 
   # Temporary storage settings - this is used to pass some data between steps
   # if the data is not specified in the catalog directly
@@ -86,9 +91,6 @@ class KedroAzureRunnerConfig(BaseModel):
       cluster_name: "{cluster_name}"
     # <your_node_tag>:
     #   cluster_name: "<your_cluster_name>"
-docker:
-  # Docker image to use during pipeline execution
-  image: "{docker_image}"
 """.strip()
 
 # This auto-validates the template above during import

kedro_azureml/constants.py

@@ -2,6 +2,5 @@
 KEDRO_AZURE_RUNNER_CONFIG = "KEDRO_AZURE_RUNNER_CONFIG"
 KEDRO_AZURE_RUNNER_DATASET_TIMEOUT = "KEDRO_AZURE_RUNNER_DATASET_TIMEOUT"
 AZURE_SUBSCRIPTION_ID = "AZURE_SUBSCRIPTION_ID"
-FILL_IN_DOCKER_IMAGE = "<fill in docker image>"
 DISTRIBUTED_CONFIG_FIELD = "__kedroazureml_distributed_config__"
 PARAMS_PREFIX = "params:"