[ML][Pipelines] Support serverless compute in pipeline & add end-to-end tests (Azure#28919)

* skip try get compute operation for serverless

* test: add first E2E test for serverless

* update pipeline job YAML

* test: add sweep

* test: add pipeline component

* test: add AutoML

* test: add parallel and spark

* test: update test and skip when playback mode

* run black

* update spark test case

* Update skip test with work item

* update VM size

* add env YAML

* add back default compute

* add comment and explicitl specify VM size

Author: zhengfeiwang

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py

@@ -57,6 +57,7 @@
     GIT_PATH_PREFIX,
     LEVEL_ONE_NAMED_RESOURCE_ID_FORMAT,
     LOCAL_COMPUTE_TARGET,
+    SERVERLESS_COMPUTE,
     SHORT_URI_FORMAT,
     SWEEP_JOB_BEST_CHILD_RUN_ID_PROPERTY_NAME,
     TID_FMT,
@@ -379,6 +380,8 @@ def _try_get_compute_arm_id(self, compute: Union[Compute, str]):
 
             if is_data_binding_expression(compute_name):
                 return compute_name
+            if compute_name == SERVERLESS_COMPUTE:
+                return compute_name
 
             try:
                 return self._compute_operations.get(compute_name).id

sdk/ml/azure-ai-ml/tests/internal/e2etests/test_pipeline_job.py

@@ -112,6 +112,7 @@ def pipeline_func():
             json.dumps(node_rest_dict, indent=2), json.dumps(mismatched_runsettings, indent=2)
         )
 
+    @pytest.mark.skip(reason="Task 2262955: skip to unblock March release.")
     @pytest.mark.parametrize(
         "test_case",
         PARAMETERS_TO_TEST_WITH_OUTPUT,

sdk/ml/azure-ai-ml/tests/pipeline_job/e2etests/test_pipeline_job.py

@@ -1900,6 +1900,22 @@ def check_name_version_and_register_succeed(output, asset_name):
         assert pipeline_job.jobs["compare"].outputs.best_model.name == "best_model"
         assert pipeline_job.jobs["compare"].outputs.best_model.version == random_version
 
+    @pytest.mark.skipif(condition=not is_live(), reason="Task 2177353: component version changes across tests.")
+    @pytest.mark.parametrize(
+        "yaml_path",
+        [
+            "./tests/test_configs/pipeline_jobs/serverless_compute/all_types/command/pipeline.yml",
+            "./tests/test_configs/pipeline_jobs/serverless_compute/all_types/sweep/pipeline.yml",
+            "./tests/test_configs/pipeline_jobs/serverless_compute/all_types/pipeline/pipeline.yml",
+            "./tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/pipeline.yml",
+            "./tests/test_configs/pipeline_jobs/serverless_compute/all_types/parallel/pipeline.yml",
+            "./tests/test_configs/pipeline_jobs/serverless_compute/all_types/spark/pipeline.yml",
+        ],
+    )
+    def test_serverless_compute_in_pipeline(self, client: MLClient, yaml_path: str) -> None:
+        pipeline_job = load_job(yaml_path)
+        assert_job_cancel(pipeline_job, client)
+
 
 @pytest.mark.usefixtures("enable_pipeline_private_preview_features")
 @pytest.mark.e2etest

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/components/component_preprocessing.yaml

@@ -0,0 +1,28 @@
+$schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
+name: data_preprocessing
+display_name: Data Preprocessing
+version: 1
+type: command
+inputs:
+  train_data:
+    type: mltable
+  validation_data:
+    type: mltable
+outputs:
+  preprocessed_train_data:
+    type: mltable
+  preprocessed_validation_data:
+    type: mltable
+code: ./src/
+environment:
+  name: pipeline-custom-environment
+  image: mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04
+  conda_file: ./preprocessing_env.yaml
+command: >-
+  python preprocess.py
+  --train_data ${{inputs.train_data}}
+  --validation_data ${{inputs.validation_data}}
+  --preprocessed_train_data ${{outputs.preprocessed_train_data}}
+  --preprocessed_validation_data ${{outputs.preprocessed_validation_data}}
+resources:
+  instance_type: Standard_D3_v2

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/components/component_register_model.yaml

@@ -0,0 +1,18 @@
+$schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
+name: register_model
+display_name: Register Model
+version: 1
+type: command
+inputs:
+  model_input_path:
+    type: uri_folder
+  model_base_name:
+    type: string
+code: ./src/
+environment: azureml:AzureML-sklearn-1.0-ubuntu20.04-py38-cpu:1
+command: >-
+  python register.py
+  --model_input_path ${{inputs.model_input_path}}
+  --model_base_name ${{inputs.model_base_name}}
+resources:
+  instance_type: Standard_D3_v2

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/components/preprocessing_env.yaml

@@ -0,0 +1,16 @@
+channels:
+  - conda-forge
+dependencies:
+  - python=3.8
+  - numpy=1.21.2
+  - pip=21.2.4
+  - scikit-learn=0.24.2
+  - scipy=1.7.1
+  - 'pandas>=1.1,<1.2'
+  - pip:
+      - 'inference-schema[numpy-support]==1.3.0'
+      - xlrd==2.0.1
+      - azureml-mlflow==1.42.0
+      - mltable==0.1.0b4
+      - pyyaml==6.0
+name: model-env

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/components/src/preprocess.py

@@ -0,0 +1,87 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import os
+import argparse
+import datetime
+from pathlib import Path
+import yaml
+from mltable import load
+
+
+def parse_args():
+    # setup arg parser
+    parser = argparse.ArgumentParser()
+
+    # add arguments
+    parser.add_argument("--train_data", type=str)
+    parser.add_argument("--validation_data", type=str)
+    parser.add_argument("--preprocessed_train_data", type=str)
+    parser.add_argument("--preprocessed_validation_data", type=str)
+    # parse args
+    args = parser.parse_args()
+    print("args received ", args)
+    # return args
+    return args
+
+
+def get_preprocessed_data(text):
+    """
+    Do preprocessing as needed
+    Currently we are just passing text file as it is
+    """
+    return text
+
+
+def main(args):
+    """
+    Preprocessing of training/validation data
+    """
+    train_data_path = os.path.join(args.train_data, "train.txt")
+    with open(train_data_path) as f:
+        train_data = f.read()
+    preprocessed_train_data = get_preprocessed_data(train_data)
+
+    # write preprocessed train txt file
+    preprocessed_train_data_path = os.path.join(args.preprocessed_train_data, "train.txt")
+    with open(preprocessed_train_data_path, "w") as f:
+        f.write(preprocessed_train_data)
+
+    validation_data_path = os.path.join(args.validation_data, "valid.txt")
+    with open(validation_data_path) as f:
+        validation_data = f.read()
+    preprocessed_validation_data = get_preprocessed_data(validation_data)
+
+    # write preprocessed validation txt file
+    preprocessed_validation_data_path = os.path.join(args.preprocessed_validation_data, "valid.txt")
+    with open(preprocessed_validation_data_path, "w") as f:
+        f.write(preprocessed_validation_data)
+
+    # Write MLTable yaml file as well in output folder
+    # Since in this example we are not doing any preprocessing, we are just copying same yaml file from input,change it if needed
+
+    # read and write MLModel yaml file for train data
+    train_data_mltable_path = os.path.join(args.train_data, "MLTable")
+    preprocessed_train_data_mltable_path = os.path.join(args.preprocessed_train_data, "MLTable")
+    with open(train_data_mltable_path, "r") as file:
+        yaml_file = yaml.safe_load(file)
+    with open(preprocessed_train_data_mltable_path, "w") as file:
+        yaml.dump(yaml_file, file)
+
+    # read and write MLModel yaml file for validation data
+    validation_data_mltable_path = os.path.join(args.validation_data, "MLTable")
+    preprocessed_validation_data_mltable_path = os.path.join(args.preprocessed_validation_data, "MLTable")
+    with open(validation_data_mltable_path, "r") as file:
+        yaml_file = yaml.safe_load(file)
+    with open(preprocessed_validation_data_mltable_path, "w") as file:
+        yaml.dump(yaml_file, file)
+
+
+# run script
+if __name__ == "__main__":
+    # parse args
+    args = parse_args()
+
+    # run main function
+    main(args)

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/components/src/register.py

@@ -0,0 +1,73 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import argparse
+import json
+import os
+import time
+
+
+from azureml.core import Run
+
+import mlflow
+import mlflow.sklearn
+
+# Based on example:
+# https://docs.microsoft.com/en-us/azure/machine-learning/how-to-train-cli
+# which references
+# https://github.com/Azure/azureml-examples/tree/main/cli/jobs/train/lightgbm/iris
+
+
+def parse_args():
+    # setup arg parser
+    parser = argparse.ArgumentParser()
+
+    # add arguments
+    parser.add_argument("--model_input_path", type=str, help="Path to input model")
+    parser.add_argument("--model_base_name", type=str, help="Name of the registered model")
+
+    # parse args
+    args = parser.parse_args()
+    print("Path: " + args.model_input_path)
+    # return args
+    return args
+
+
+def main(args):
+    """
+    Register Model Example
+    """
+    # Set Tracking URI
+    current_experiment = Run.get_context().experiment
+    tracking_uri = current_experiment.workspace.get_mlflow_tracking_uri()
+    print("tracking_uri: {0}".format(tracking_uri))
+    mlflow.set_tracking_uri(tracking_uri)
+    mlflow.set_experiment(current_experiment.name)
+
+    # Get Run ID from model path
+    print("Getting model path")
+    mlmodel_path = os.path.join(args.model_input_path, "MLmodel")
+    runid = ""
+    with open(mlmodel_path, "r") as modelfile:
+        for line in modelfile:
+            if "run_id" in line:
+                runid = line.split(":")[1].strip()
+
+    # Construct Model URI from run ID extract previously
+    model_uri = "runs:/{}/outputs/".format(runid)
+    print("Model URI: " + model_uri)
+
+    # Register the model with Model URI and Name of choice
+    registered_name = args.model_base_name
+    print(f"Registering model as {registered_name}")
+    mlflow.register_model(model_uri, registered_name)
+
+
+# run script
+if __name__ == "__main__":
+    # parse args
+    args = parse_args()
+
+    # run main function
+    main(args)

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/pipeline.yml

@@ -0,0 +1,47 @@
+$schema: https://azuremlschemas.azureedge.net/latest/pipelineJob.schema.json
+type: pipeline
+display_name: Serverless in Pipeline - AutoML
+settings:
+  default_compute: azureml:serverless
+inputs:
+  text_ner_training_data:
+    type: mltable
+    path: ./training-mltable-folder
+  text_ner_validation_data:
+    type: mltable
+    path: ./validation-mltable-folder
+jobs:
+  preprocessing_node:
+    type: command
+    component: file:./components/component_preprocessing.yaml
+    inputs:
+      train_data: ${{parent.inputs.text_ner_training_data}}
+      validation_data: ${{parent.inputs.text_ner_validation_data}}
+    outputs:
+      preprocessed_train_data:
+        type: mltable
+      preprocessed_validation_data:
+        type: mltable
+  text_ner_node:
+    type: automl
+    task: text_ner
+    log_verbosity: info
+    primary_metric: accuracy
+    limits:
+      max_trials: 1
+      timeout_minutes: 60
+    target_column_name: label
+    training_data: ${{parent.jobs.preprocessing_node.outputs.preprocessed_train_data}}
+    validation_data: ${{parent.jobs.preprocessing_node.outputs.preprocessed_validation_data}}
+    outputs:
+      best_model:
+        type: mlflow_model
+    # BUG 2246569: resources should be optional, explicitly specify here with NC/ND-family for dnn tasks.
+    resources:
+      instance_type: Standard_NC6
+  register_model_node:
+    type: command
+    component: file:./components/component_register_model.yaml
+    inputs:
+      model_input_path: ${{parent.jobs.text_ner_node.outputs.best_model}}
+      model_base_name: paper_categorization_model

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/serverless_compute/all_types/automl/training-mltable-folder/MLTable

@@ -0,0 +1,4 @@
+paths:
+  - file: ./train.txt
+transformations:
+  - take: 1