Adding AquaFineTuningConfig

Author: kumar-shivam-ranjan

ads/aqua/finetuning/entities.py

@@ -3,7 +3,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import json
-from typing import List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from pydantic import Field, model_validator
 
@@ -55,6 +55,28 @@ def validate_restricted_fields(cls, data: dict):
         return data
 
 
+class AquaFineTuningConfig(Serializable):
+    """Represents model's shape list and detailed configuration for fine-tuning.
+
+    Attributes:
+        shape (List[str], optional): A list of shape names (e.g., BM.GPU.A10.4).
+        configuration (Dict[str, Any], optional): Configuration details of fine-tuning.
+    """
+
+    shape: Optional[Dict[str, Any]] = Field(
+        default_factory=dict, description="List of supported shapes for the model."
+    )
+    finetuning_params: Optional[str] = Field(
+        default_factory=str, description="Fine tuning parameters."
+    )
+    configuration: Optional[Dict[str, Any]] = Field(
+        default_factory=dict, description="Configuration details keyed by shape."
+    )
+
+    class Config:
+        extra = "allow"
+
+
 class AquaFineTuningSummary(Serializable):
     """Represents a summary of Aqua Finetuning job."""
 

ads/aqua/finetuning/finetuning.py

@@ -38,6 +38,7 @@
     FineTuneCustomMetadata,
 )
 from ads.aqua.finetuning.entities import (
+    AquaFineTuningConfig,
     AquaFineTuningParams,
     AquaFineTuningSummary,
     CreateFineTuningDetails,
@@ -371,11 +372,11 @@ def create(
             is_custom_container = True
 
         ft_parameters.batch_size = ft_parameters.batch_size or (
-            ft_config.get("shape", UNKNOWN_DICT)
+            (ft_config.shape if ft_config else UNKNOWN_DICT)
             .get(create_fine_tuning_details.shape_name, UNKNOWN_DICT)
             .get("batch_size", DEFAULT_FT_BATCH_SIZE)
         )
-        finetuning_params = ft_config.get("finetuning_params")
+        finetuning_params = ft_config.finetuning_params if ft_config else UNKNOWN
 
         ft_job.with_runtime(
             self._build_fine_tuning_runtime(
@@ -626,7 +627,7 @@ def _build_oci_launch_cmd(
     @telemetry(
         entry_point="plugin=finetuning&action=get_finetuning_config", name="aqua"
     )
-    def get_finetuning_config(self, model_id: str) -> Dict:
+    def get_finetuning_config(self, model_id: str) -> AquaFineTuningConfig:
         """Gets the finetuning config for given Aqua model.
 
         Parameters
@@ -641,12 +642,12 @@ def get_finetuning_config(self, model_id: str) -> Dict:
         """
         config = self.get_config_from_metadata(
             model_id, AquaModelMetadataKeys.FINE_TUNING_CONFIGURATION
-        )
+        ).config
         if config:
             logger.info(
                 f"Fetched {AquaModelMetadataKeys.FINE_TUNING_CONFIGURATION} from defined metadata for model: {model_id}."
             )
-            return config
+            return AquaFineTuningConfig(**(config or UNKNOWN_DICT))
         config = self.get_config(
             model_id,
             DEFINED_METADATA_TO_FILE_MAP.get(
@@ -657,7 +658,7 @@ def get_finetuning_config(self, model_id: str) -> Dict:
             logger.debug(
                 f"Fine-tuning config for custom model: {model_id} is not available. Use defaults."
             )
-        return config
+        return AquaFineTuningConfig(**(config or UNKNOWN_DICT))
 
     @telemetry(
         entry_point="plugin=finetuning&action=get_finetuning_default_params",
@@ -680,7 +681,9 @@ def get_finetuning_default_params(self, model_id: str) -> Dict:
         """
         default_params = {"params": {}}
         finetuning_config = self.get_finetuning_config(model_id)
-        config_parameters = finetuning_config.get("configuration", UNKNOWN_DICT)
+        config_parameters = (
+            finetuning_config.configuration if finetuning_config else UNKNOWN_DICT
+        )
         dataclass_fields = self._get_finetuning_params(
             config_parameters, validate=False
         ).to_dict()

tests/unitary/with_extras/aqua/test_data/finetuning/ft_config.json

@@ -10,6 +10,7 @@
     "sequence_len": 2048,
     "val_set_size": 0.1
   },
+  "finetuning_params": "",
   "shape": {
     "VM.GPU.A10.2": {
       "batch_size": 2,

tests/unitary/with_extras/aqua/test_finetuning.py

@@ -21,7 +21,7 @@
 from ads.aqua.common.errors import AquaValueError
 from ads.aqua.finetuning import AquaFineTuningApp
 from ads.aqua.finetuning.constants import FineTuneCustomMetadata
-from ads.aqua.finetuning.entities import AquaFineTuningParams
+from ads.aqua.finetuning.entities import AquaFineTuningParams, AquaFineTuningConfig
 from ads.aqua.model.entities import AquaFineTuneModel
 from ads.jobs.ads_job import Job
 from ads.jobs.builders.infrastructure.dsc_job import DataScienceJobRun
@@ -113,11 +113,13 @@ def test_create_fine_tuning(
         ft_model.time_created = "test_time_created"
         mock_ds_model_create.return_value = ft_model
 
-        mock_get_finetuning_config.return_value = {
-            "shape": {
-                "VM.GPU.A10.1": {"batch_size": 1, "replica": 1},
+        mock_get_finetuning_config.return_value = AquaFineTuningConfig(
+            **{
+                "shape": {
+                    "VM.GPU.A10.1": {"batch_size": 1, "replica": 1},
+                }
             }
-        }
+        )
         mock_get_container_image.return_value = "test_container_image"
 
         mock_job_id.return_value = "test_ft_job_id"
@@ -279,11 +281,13 @@ def test_get_finetuning_config(self):
         config_json = os.path.join(self.curr_dir, "test_data/finetuning/ft_config.json")
         with open(config_json, "r") as _file:
             config = json.load(_file)
-        self.app.get_config_from_metadata = MagicMock(return_value={})
+        self.app.get_config_from_metadata = MagicMock(
+            return_value=ModelConfigResult(config=config, model_details=None)
+        )
         self.app.get_config = MagicMock(
             return_value=ModelConfigResult(config=config, model_details=None)
         )
-        result = self.app.get_finetuning_config(model_id="test-model-id")
+        result = self.app.get_finetuning_config(model_id="test-model-id").to_dict()
         assert result == config
 
     def test_get_finetuning_default_params(self):
@@ -306,12 +310,16 @@ def test_get_finetuning_default_params(self):
         with open(config_json, "r") as _file:
             config = json.load(_file)
 
-        self.app.get_finetuning_config = MagicMock(return_value=config)
+        self.app.get_finetuning_config = MagicMock(
+            return_value=AquaFineTuningConfig(**config)
+        )
         result = self.app.get_finetuning_default_params(model_id="test_model_id")
         assert result == params_dict
 
         # check when config json is not available
-        self.app.get_finetuning_config = MagicMock(return_value={})
+        self.app.get_finetuning_config = MagicMock(
+            return_value=AquaFineTuningConfig(**{})
+        )
         result = self.app.get_finetuning_default_params(model_id="test_model_id")
         assert result == {"params": {}}