Aqua refactor: Move all enums under enums.py for centralized management (#848)

Author: mingkang111

ads/aqua/app.py

@@ -12,6 +12,7 @@
 
 from ads import set_auth
 from ads.aqua import logger
+from ads.aqua.common.enums import Tags
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     UNKNOWN,
@@ -20,7 +21,6 @@
     is_valid_ocid,
     load_config,
 )
-from ads.aqua.data import Tags
 from ads.common import oci_client as oc
 from ads.common.auth import default_signer
 from ads.common.utils import extract_region

ads/aqua/common/enums.py

@@ -11,8 +11,14 @@
 from ads.common.extended_enum import ExtendedEnumMeta
 
 
+class DataScienceResource(str, metaclass=ExtendedEnumMeta):
+    MODEL_DEPLOYMENT = "datasciencemodeldeployment"
+    MODEL = "datasciencemodel"
+
+
 class Resource(str, metaclass=ExtendedEnumMeta):
     JOB = "jobs"
+    JOBRUN = "jobruns"
     MODEL = "models"
     MODEL_DEPLOYMENT = "modeldeployments"
     MODEL_VERSION_SET = "model-version-sets"
@@ -21,3 +27,28 @@ class Resource(str, metaclass=ExtendedEnumMeta):
 class DataScienceResource(str, metaclass=ExtendedEnumMeta):
     MODEL_DEPLOYMENT = "datasciencemodeldeployment"
     MODEL = "datasciencemodel"
+
+
+class Tags(str, metaclass=ExtendedEnumMeta):
+    TASK = "task"
+    LICENSE = "license"
+    ORGANIZATION = "organization"
+    AQUA_TAG = "OCI_AQUA"
+    AQUA_SERVICE_MODEL_TAG = "aqua_service_model"
+    AQUA_FINE_TUNED_MODEL_TAG = "aqua_fine_tuned_model"
+    AQUA_MODEL_NAME_TAG = "aqua_model_name"
+    AQUA_EVALUATION = "aqua_evaluation"
+    AQUA_FINE_TUNING = "aqua_finetuning"
+    READY_TO_FINE_TUNE = "ready_to_fine_tune"
+    READY_TO_IMPORT = "ready_to_import"
+    BASE_MODEL_CUSTOM = "aqua_custom_base_model"
+
+
+class RqsAdditionalDetails(str, metaclass=ExtendedEnumMeta):
+    METADATA = "metadata"
+    CREATED_BY = "createdBy"
+    DESCRIPTION = "description"
+    MODEL_VERSION_SET_ID = "modelVersionSetId"
+    MODEL_VERSION_SET_NAME = "modelVersionSetName"
+    PROJECT_ID = "projectId"
+    VERSION_LABEL = "versionLabel"

ads/aqua/common/utils.py

@@ -21,15 +21,13 @@
 import oci
 from oci.data_science.models import JobRun, Model
 
+from ads.aqua.common.enums import RqsAdditionalDetails
 from ads.aqua.common.errors import (
     AquaFileNotFoundError,
     AquaRuntimeError,
     AquaValueError,
 )
-from ads.aqua.constants import (
-    SERVICE_MANAGED_CONTAINER_URI_SCHEME,
-    RqsAdditionalDetails,
-)
+from ads.aqua.constants import SERVICE_MANAGED_CONTAINER_URI_SCHEME
 from ads.aqua.data import AquaResourceIdentifier
 from ads.common.auth import AuthState, default_signer
 from ads.common.extended_enum import ExtendedEnumMeta

ads/aqua/constants.py

@@ -3,7 +3,6 @@
 # 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/
 """This module defines constants used in ads.aqua module."""
-from ads.common.extended_enum import ExtendedEnumMeta
 
 UNKNOWN = ""
 UNKNOWN_VALUE = ""
@@ -40,34 +39,3 @@
 VALIDATION_METRICS = "validation_metrics"
 
 SERVICE_MANAGED_CONTAINER_URI_SCHEME = "dsmc://"
-
-
-class RqsAdditionalDetails(str, metaclass=ExtendedEnumMeta):
-    METADATA = "metadata"
-    CREATED_BY = "createdBy"
-    DESCRIPTION = "description"
-    MODEL_VERSION_SET_ID = "modelVersionSetId"
-    MODEL_VERSION_SET_NAME = "modelVersionSetName"
-    PROJECT_ID = "projectId"
-    VERSION_LABEL = "versionLabel"
-
-
-class FineTuningDefinedMetadata(str, metaclass=ExtendedEnumMeta):
-    """Represents the defined metadata keys used in Fine Tuning."""
-
-    VAL_SET_SIZE = "val_set_size"
-    TRAINING_DATA = "training_data"
-
-
-class FineTuningCustomMetadata(str, metaclass=ExtendedEnumMeta):
-    """Represents the custom metadata keys used in Fine Tuning."""
-
-    FT_SOURCE = "fine_tune_source"
-    FT_SOURCE_NAME = "fine_tune_source_name"
-    FT_OUTPUT_PATH = "fine_tune_output_path"
-    FT_JOB_ID = "fine_tune_job_id"
-    FT_JOB_RUN_ID = "fine_tune_jobrun_id"
-    TRAINING_METRICS_FINAL = "train_metrics_final"
-    VALIDATION_METRICS_FINAL = "val_metrics_final"
-    TRAINING_METRICS_EPOCH = "train_metrics_epoch"
-    VALIDATION_METRICS_EPOCH = "val_metrics_epoch"

ads/aqua/data.py

@@ -5,7 +5,6 @@
 
 from dataclasses import dataclass
 
-from ads.common.extended_enum import ExtendedEnumMeta
 from ads.common.serializer import DataClassSerializable
 
 
@@ -14,46 +13,3 @@ class AquaResourceIdentifier(DataClassSerializable):
     id: str = ""
     name: str = ""
     url: str = ""
-
-
-class Resource(str, metaclass=ExtendedEnumMeta):
-    JOB = "jobs"
-    JOBRUN = "jobruns"
-    MODEL = "models"
-    MODEL_DEPLOYMENT = "modeldeployments"
-    MODEL_VERSION_SET = "model-version-sets"
-
-
-class DataScienceResource(str, metaclass=ExtendedEnumMeta):
-    MODEL_DEPLOYMENT = "datasciencemodeldeployment"
-    MODEL = "datasciencemodel"
-
-
-class Tags(str, metaclass=ExtendedEnumMeta):
-    TASK = "task"
-    LICENSE = "license"
-    ORGANIZATION = "organization"
-    AQUA_TAG = "OCI_AQUA"
-    AQUA_SERVICE_MODEL_TAG = "aqua_service_model"
-    AQUA_FINE_TUNED_MODEL_TAG = "aqua_fine_tuned_model"
-    AQUA_MODEL_NAME_TAG = "aqua_model_name"
-    AQUA_EVALUATION = "aqua_evaluation"
-    AQUA_FINE_TUNING = "aqua_finetuning"
-    READY_TO_FINE_TUNE = "ready_to_fine_tune"
-    READY_TO_IMPORT = "ready_to_import"
-    BASE_MODEL_CUSTOM = "aqua_custom_base_model"
-
-
-class InferenceContainerType(str, metaclass=ExtendedEnumMeta):
-    CONTAINER_TYPE_VLLM = "vllm"
-    CONTAINER_TYPE_TGI = "tgi"
-
-
-class InferenceContainerTypeKey(str, metaclass=ExtendedEnumMeta):
-    AQUA_VLLM_CONTAINER_KEY = "odsc-vllm-serving"
-    AQUA_TGI_CONTAINER_KEY = "odsc-tgi-serving"
-
-
-class InferenceContainerParamType(str, metaclass=ExtendedEnumMeta):
-    PARAM_TYPE_VLLM = "VLLM_PARAMS"
-    PARAM_TYPE_TGI = "TGI_PARAMS"

ads/aqua/evaluation/evaluation.py

@@ -26,7 +26,12 @@
 from ads.aqua import logger
 from ads.aqua.app import AquaApp
 from ads.aqua.common import utils
-from ads.aqua.common.enums import *
+from ads.aqua.common.enums import (
+    DataScienceResource,
+    Resource,
+    RqsAdditionalDetails,
+    Tags,
+)
 from ads.aqua.common.errors import (
     AquaFileExistsError,
     AquaFileNotFoundError,
@@ -44,7 +49,6 @@
     is_valid_ocid,
     upload_local_to_os,
 )
-from ads.aqua.data import Tags
 from ads.aqua.evaluation.constants import *
 from ads.aqua.evaluation.entities import *
 from ads.aqua.evaluation.errors import *

ads/aqua/extension/model_handler.py

@@ -179,7 +179,7 @@ def post(self, *args, **kwargs):
             )
 
         # Check pipeline_tag, it should be `text-generation`
-        if hf_model_info.pipeline_tag.lower() != ModelTask.TEXT_GENERATION.value:
+        if hf_model_info.pipeline_tag.lower() != ModelTask.TEXT_GENERATION:
             raise AquaRuntimeError(
                 f"Unsupported pipeline tag for the chosen model: '{hf_model_info.pipeline_tag}'. "
                 f"AQUA currently supports the following tasks only: {', '.join(ModelTask.values())}. "

ads/aqua/extension/ui_handler.py

@@ -9,7 +9,7 @@
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
-from ads.aqua.data import Tags
+from ads.aqua.common.enums import Tags
 from ads.aqua.extension.base_handler import AquaAPIhandler, Errors
 from ads.aqua.extension.utils import validate_function_parameters
 from ads.aqua.model.entities import ImportModelDetails

ads/aqua/finetuning/finetuning.py

@@ -16,6 +16,7 @@
 
 from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, 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 (
     DEFAULT_FT_BATCH_SIZE,
@@ -28,7 +29,7 @@
     get_container_image,
     upload_local_to_os,
 )
-from ads.aqua.data import AquaResourceIdentifier, Resource, Tags
+from ads.aqua.data import AquaResourceIdentifier
 from ads.aqua.finetuning.constants import *
 from ads.aqua.finetuning.entities import *
 from ads.common.auth import default_signer

ads/aqua/job.py

@@ -7,8 +7,9 @@
 
 import logging
 from dataclasses import dataclass, field
-from ads.common.serializer import DataClassSerializable
+
 from ads.aqua.data import AquaResourceIdentifier
+from ads.common.serializer import DataClassSerializable
 
 logger = logging.getLogger(__name__)