[WIP] Add e2e test for tune api with LLM hyperparameter optimization

.github/workflows/e2e-test-tune-api.yaml

@@ -22,10 +22,15 @@ jobs:
         with:
           kubernetes-version: ${{ matrix.kubernetes-version }}
 
+      - name: Install Training Operator SDK
+        shell: bash
+        run: pip install kubeflow-training[huggingface]
+
       - name: Run e2e test with tune API
         uses: ./.github/workflows/template-e2e-test
         with:
           tune-api: true
+          training-operator: true
 
     strategy:
       fail-fast: false

sdk/python/v1beta1/kubeflow/katib/api/katib_client.py

@@ -16,12 +16,13 @@
 import logging
 import multiprocessing
 import time
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING, Union
 
 import grpc
 import kubeflow.katib.katib_api_pb2 as katib_api_pb2
 import kubeflow.katib.katib_api_pb2_grpc as katib_api_pb2_grpc
 from kubeflow.katib import models
+from kubeflow.katib import types
 from kubeflow.katib.api_client import ApiClient
 from kubeflow.katib.constants import constants
 from kubeflow.katib.types.trainer_resources import TrainerResources
@@ -30,6 +31,12 @@
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from kubeflow.storage_initializer.hugging_face import HuggingFaceDatasetParams
+    from kubeflow.storage_initializer.hugging_face import HuggingFaceModelParams
+    from kubeflow.storage_initializer.hugging_face import HuggingFaceTrainerParams
+    from kubeflow.storage_initializer.s3 import S3DatasetParams
+
 
 class KatibClient(object):
     def __init__(
@@ -338,6 +345,7 @@ class name in this argument.
                 to the base image packages. These packages are installed before
                 executing the objective function.
             pip_index_url: The PyPI url from which to install Python packages.
+            metrics_collector_config: Specify the config of metrics collector,
             metrics_collector_config: Specify the config of metrics collector,
                 for example, `metrics_collector_config = {"kind": "Push"}`.
                 Currently, we only support `StdOut` and `Push` metrics collector.

sdk/python/v1beta1/kubeflow/katib/constants/constants.py

@@ -33,6 +33,7 @@
 
 
 DEFAULT_PRIMARY_CONTAINER_NAME = "training-container"
+PYTORCHJOB_PRIMARY_CONTAINER_NAME = "pytorch"
 
 # Label to identify Experiment's resources.
 EXPERIMENT_LABEL = "katib.kubeflow.org/experiment"

sdk/python/v1beta1/kubeflow/katib/types/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import absolute_import
+
+# Import types into type package.
+from kubeflow.katib.types.trainer_resources import TrainerResources
+
+# Import Kubernetes models.
+from kubernetes.client import *

test/e2e/v1beta1/scripts/gh-actions/run-e2e-tune-api.py

@@ -1,8 +1,15 @@
 import argparse
 import logging
 
-from kubeflow.katib import KatibClient, search
+import transformers
+from kubeflow.katib import KatibClient, search, types
+from kubeflow.storage_initializer.hugging_face import (
+    HuggingFaceDatasetParams,
+    HuggingFaceModelParams,
+    HuggingFaceTrainerParams,
+)
 from kubernetes import client
+from peft import LoraConfig
 from verify import verify_experiment_results
 
 # Experiment timeout is 40 min.
@@ -12,7 +19,8 @@
 logging.basicConfig(level=logging.INFO)
 
 
-def run_e2e_experiment_create_by_tune(
+# Test for Experiment created with custom objective.
+def run_e2e_experiment_create_by_tune_with_custom_objective(
     katib_client: KatibClient,
     exp_name: str,
     exp_namespace: str,
@@ -57,6 +65,70 @@ def objective(parameters):
     logging.debug(katib_client.get_experiment(exp_name, exp_namespace))
     logging.debug(katib_client.get_suggestion(exp_name, exp_namespace))
 
+# Test for Experiment created with external models and datasets.
+def run_e2e_experiment_create_by_tune_with_external_model(
+    katib_client: KatibClient,
+    exp_name: str,
+    exp_namespace: str,
+):
+    # Create Katib Experiment and wait until it is finished.
+    logging.debug("Creating Experiment: {}/{}".format(exp_namespace, exp_name))
+
+    # Use the test case from fine-tuning API tutorial.
+    # https://www.kubeflow.org/docs/components/training/user-guides/fine-tuning/
+    # Create Katib Experiment.
+    # And Wait until Experiment reaches Succeeded condition.
+    katib_client.tune(
+        name=exp_name,
+        namespace=exp_namespace,
+        # BERT model URI and type of Transformer to train it.
+        model_provider_parameters=HuggingFaceModelParams(
+            model_uri="hf://google-bert/bert-base-cased",
+            transformer_type=transformers.AutoModelForSequenceClassification,
+            num_labels=5,
+        ),
+        # In order to save test time, use 8 samples from Yelp dataset.
+        dataset_provider_parameters=HuggingFaceDatasetParams(
+            repo_id="yelp_review_full",
+            split="train[:8]",
+        ),
+        # Specify HuggingFace Trainer parameters.
+        trainer_parameters=HuggingFaceTrainerParams(
+            training_parameters=transformers.TrainingArguments(
+                output_dir="test_tune_api",
+                save_strategy="no",
+                learning_rate = search.double(min=1e-05, max=5e-05),
+                num_train_epochs=1,
+            ),
+            # Set LoRA config to reduce number of trainable model parameters.
+            lora_config=LoraConfig(
+                r = search.int(min=8, max=32),
+                lora_alpha=8,
+                lora_dropout=0.1,
+                bias="none",
+            ),
+        ),
+        objective_metric_name = "train_loss", 
+        objective_type = "minimize", 
+        algorithm_name = "random",
+        max_trial_count = 1,
+        parallel_trial_count = 1,
+        resources_per_trial=types.TrainerResources(
+            num_workers=1,
+            num_procs_per_worker=1,
+            resources_per_worker={"cpu": "2", "memory": "10G",},
+        ),
+    )
+    experiment = katib_client.wait_for_experiment_condition(
+        exp_name, exp_namespace, timeout=EXPERIMENT_TIMEOUT
+    )
+
+    # Verify the Experiment results.
+    verify_experiment_results(katib_client, experiment, exp_name, exp_namespace)
+
+    # Print the Experiment and Suggestion.
+    logging.debug(katib_client.get_experiment(exp_name, exp_namespace))
+    logging.debug(katib_client.get_suggestion(exp_name, exp_namespace))
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -82,15 +154,29 @@ def objective(parameters):
     exp_name = "tune-example"
     exp_namespace = args.namespace
     try:
-        run_e2e_experiment_create_by_tune(katib_client, exp_name, exp_namespace)
+        run_e2e_experiment_create_by_tune_with_custom_objective(katib_client, f"{exp_name}", exp_namespace)
+        logging.info("---------------------------------------------------------------")
+        logging.info(f"E2E is succeeded for Experiment created by tune: {exp_namespace}/{f"{exp_name}"}")
+    except Exception as e:
+        logging.info("---------------------------------------------------------------")
+        logging.info(f"E2E is failed for Experiment created by tune: {exp_namespace}/{f"{exp_name}"}")
+        raise e
+    finally:
+        # Delete the Experiment.
+        logging.info("---------------------------------------------------------------")
+        logging.info("---------------------------------------------------------------")
+        katib_client.delete_experiment(f"{exp_name}", exp_namespace)
+
+    try:
+        run_e2e_experiment_create_by_tune_with_external_model(katib_client, f"{exp_name}", exp_namespace)
         logging.info("---------------------------------------------------------------")
-        logging.info(f"E2E is succeeded for Experiment created by tune: {exp_namespace}/{exp_name}")
+        logging.info(f"E2E is succeeded for Experiment created by tune: {exp_namespace}/{f"{exp_name}"}")
     except Exception as e:
         logging.info("---------------------------------------------------------------")
-        logging.info(f"E2E is failed for Experiment created by tune: {exp_namespace}/{exp_name}")
+        logging.info(f"E2E is failed for Experiment created by tune: {exp_namespace}/{f"{exp_name}"}")
         raise e
     finally:
         # Delete the Experiment.
         logging.info("---------------------------------------------------------------")
         logging.info("---------------------------------------------------------------")
-        katib_client.delete_experiment(exp_name, exp_namespace)
+        katib_client.delete_experiment(f"{exp_name}", exp_namespace)