Passing Early stopping params to SMC container for finetuning (#970)

Author: kumar-shivam-ranjan

ads/aqua/extension/finetune_handler.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
@@ -9,8 +8,8 @@
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
-from ads.aqua.extension.errors import Errors
 from ads.aqua.extension.base_handler import AquaAPIhandler
+from ads.aqua.extension.errors import Errors
 from ads.aqua.extension.utils import validate_function_parameters
 from ads.aqua.finetuning import AquaFineTuningApp
 from ads.aqua.finetuning.entities import CreateFineTuningDetails

ads/aqua/finetuning/entities.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 from dataclasses import dataclass, field
@@ -14,16 +13,18 @@ class AquaFineTuningParams(DataClassSerializable):
     epochs: int
     learning_rate: Optional[float] = None
     sample_packing: Optional[bool] = "auto"
-    batch_size: Optional[
-        int
-    ] = None  # make it batch_size for user, but internally this is micro_batch_size
+    batch_size: Optional[int] = (
+        None  # make it batch_size for user, but internally this is micro_batch_size
+    )
     sequence_len: Optional[int] = None
     pad_to_sequence_len: Optional[bool] = None
     lora_r: Optional[int] = None
     lora_alpha: Optional[int] = None
     lora_dropout: Optional[float] = None
     lora_target_linear: Optional[bool] = None
     lora_target_modules: Optional[List] = None
+    early_stopping_patience: Optional[int] = None
+    early_stopping_threshold: Optional[float] = None
 
 
 @dataclass(repr=False)

ads/aqua/finetuning/finetuning.py

@@ -1,11 +1,10 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import json
 import os
-from dataclasses import asdict, fields, MISSING
+from dataclasses import MISSING, asdict, fields
 from typing import Dict
 
 from oci.data_science.models import (
@@ -14,14 +13,15 @@
     UpdateModelProvenanceDetails,
 )
 
-from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, logger
+from ads.aqua import logger
 from ads.aqua.app import AquaApp
 from ads.aqua.common.enums import Resource, Tags
 from ads.aqua.common.errors import AquaFileExistsError, AquaValueError
 from ads.aqua.common.utils import (
     get_container_image,
     upload_local_to_os,
 )
+from ads.aqua.config.config import get_finetuning_config_defaults
 from ads.aqua.constants import (
     DEFAULT_FT_BATCH_SIZE,
     DEFAULT_FT_BLOCK_STORAGE_SIZE,
@@ -31,7 +31,6 @@
     UNKNOWN,
     UNKNOWN_DICT,
 )
-from ads.aqua.config.config import get_finetuning_config_defaults
 from ads.aqua.data import AquaResourceIdentifier
 from ads.aqua.finetuning.constants import *
 from ads.aqua.finetuning.entities import *
@@ -132,7 +131,7 @@ def create(
             or create_fine_tuning_details.validation_set_size >= 1
         ):
             raise AquaValueError(
-                f"Fine tuning validation set size should be a float number in between [0, 1)."
+                "Fine tuning validation set size should be a float number in between [0, 1)."
             )
 
         if create_fine_tuning_details.replica < DEFAULT_FT_REPLICA:
@@ -394,7 +393,7 @@ def create(
         )
         # track shapes that were used for fine-tune creation
         self.telemetry.record_event_async(
-            category=f"aqua/service/finetune/create/shape/",
+            category="aqua/service/finetune/create/shape/",
             action=f"{create_fine_tuning_details.shape_name}x{create_fine_tuning_details.replica}",
             **telemetry_kwargs,
         )
@@ -533,6 +532,12 @@ def _build_oci_launch_cmd(
                     oci_launch_cmd += f"--num_{key} {value} "
                 elif key == "lora_target_modules":
                     oci_launch_cmd += f"--{key} {','.join(str(k) for k in value)} "
+                elif key == "early_stopping_patience":
+                    if value != 0:
+                        oci_launch_cmd += f"--{key} {value} "
+                elif key == "early_stopping_threshold":
+                    if "early_stopping_patience" in oci_launch_cmd:
+                        oci_launch_cmd += f"--{key} {value} "
                 else:
                     oci_launch_cmd += f"--{key} {value} "