Updating max threads inference

Author: harsh97

ads/aqua/common/entities.py

@@ -14,3 +14,4 @@ class ContainerSpec:
     HEALTH_CHECK_PORT = "healthCheckPort"
     ENV_VARS = "envVars"
     RESTRICTED_PARAMS = "restrictedParams"
+    EVALUATION_CONFIGURATION = "evaluationConfiguration"

ads/aqua/evaluation/entities.py

@@ -12,7 +12,6 @@
 from dataclasses import dataclass, field
 from typing import List, Optional, Union
 
-from ads.aqua.constants import EVALUATION_INFERENCE_DEFAULT_THREADS
 from ads.aqua.data import AquaResourceIdentifier
 from ads.common.serializer import DataClassSerializable
 
@@ -88,7 +87,7 @@ class CreateAquaEvaluationDetails(DataClassSerializable):
     log_id: Optional[str] = None
     metrics: Optional[List] = None
     force_overwrite: Optional[bool] = False
-    inference_max_threads: Optional[int] = EVALUATION_INFERENCE_DEFAULT_THREADS
+    inference_max_threads: Optional[int] = None
 
 
 @dataclass(repr=False)

ads/aqua/evaluation/evaluation.py

@@ -13,15 +13,9 @@
 from threading import Lock
 from typing import Any, Dict, List, Union
 
-import oci
 from cachetools import TTLCache
-from oci.data_science.models import (
-    JobRun,
-    Metadata,
-    UpdateModelDetails,
-    UpdateModelProvenanceDetails,
-)
 
+import oci
 from ads.aqua import logger
 from ads.aqua.app import AquaApp
 from ads.aqua.common import utils
@@ -47,6 +41,7 @@
 )
 from ads.aqua.constants import (
     CONSOLE_LINK_RESOURCE_TYPE_MAPPING,
+    EVALUATION_INFERENCE_DEFAULT_THREADS,
     EVALUATION_REPORT,
     EVALUATION_REPORT_JSON,
     EVALUATION_REPORT_MD,
@@ -76,6 +71,7 @@
     ModelParams,
 )
 from ads.aqua.evaluation.errors import EVALUATION_JOB_EXIT_CODE_MESSAGE
+from ads.aqua.ui import AquaContainerConfig
 from ads.common.auth import default_signer
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.utils import get_console_link, get_files, get_log_links
@@ -90,7 +86,9 @@
 from ads.jobs.builders.runtimes.base import Runtime
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
 from ads.model.datascience_model import DataScienceModel
+from ads.model.deployment import ModelDeploymentContainerRuntime
 from ads.model.deployment.model_deployment import ModelDeployment
+from ads.model.generic_model import ModelDeploymentRuntimeType
 from ads.model.model_metadata import (
     MetadataTaxonomyKeys,
     ModelCustomMetadata,
@@ -99,6 +97,12 @@
 )
 from ads.model.model_version_set import ModelVersionSet
 from ads.telemetry import telemetry
+from oci.data_science.models import (
+    JobRun,
+    Metadata,
+    UpdateModelDetails,
+    UpdateModelProvenanceDetails,
+)
 
 
 class AquaEvaluationApp(AquaApp):
@@ -166,7 +170,6 @@ def create(
                 f"Invalid evaluation source {create_aqua_evaluation_details.evaluation_source_id}. "
                 "Specify either a model or model deployment id."
             )
-
         evaluation_source = None
         if (
             DataScienceResource.MODEL_DEPLOYMENT
@@ -175,6 +178,33 @@ def create(
             evaluation_source = ModelDeployment.from_id(
                 create_aqua_evaluation_details.evaluation_source_id
             )
+            if evaluation_source.runtime.type == ModelDeploymentRuntimeType.CONTAINER:
+                runtime = ModelDeploymentContainerRuntime.from_dict(
+                    evaluation_source.runtime.to_dict()
+                )
+                container_config = AquaContainerConfig.from_container_index_json(
+                    enable_spec=True
+                )
+                for container in container_config.inference.values():
+                    if container.name == runtime.image.split(":")[0]:
+                        max_threads = container.spec.evaluation_configuration.evaluation_max_threads
+                        if (
+                            max_threads
+                            and create_aqua_evaluation_details.inference_max_threads
+                            and max_threads
+                            < create_aqua_evaluation_details.inference_max_threads
+                        ):
+                            raise AquaValueError(
+                                f"Invalid inference max threads. The maximum allowed value for {runtime.image} is {max_threads}."
+                            )
+                        if not create_aqua_evaluation_details.inference_max_threads:
+                            create_aqua_evaluation_details.inference_max_threads = container.spec.evaluation_configuration.evaluation_default_threads
+                        break
+                if not create_aqua_evaluation_details.inference_max_threads:
+                    create_aqua_evaluation_details.inference_max_threads = (
+                        EVALUATION_INFERENCE_DEFAULT_THREADS
+                    )
+
         elif (
             DataScienceResource.MODEL
             in create_aqua_evaluation_details.evaluation_source_id
@@ -1197,7 +1227,7 @@ def _delete_job_and_model(job, model):
                 f"Exception message: {ex}"
             )
 
-    def load_evaluation_config(self, eval_id):
+    def load_evaluation_config(self, _):
         """Loads evaluation config."""
         return {
             "model_params": {

ads/aqua/modeldeployment/deployment.py

@@ -167,7 +167,6 @@ def create(
             Tags.AQUA_SERVICE_MODEL_TAG,
             Tags.AQUA_FINE_TUNED_MODEL_TAG,
             Tags.AQUA_TAG,
-            Tags.MODEL_FORMAT,
         ]:
             if tag in aqua_model.freeform_tags:
                 tags[tag] = aqua_model.freeform_tags[tag]

ads/aqua/ui.py

@@ -9,15 +9,14 @@
 from typing import Dict, List, Optional
 
 from cachetools import TTLCache
-from oci.exceptions import ServiceError
-from oci.identity.models import Compartment
 
 from ads.aqua import logger
 from ads.aqua.app import AquaApp
 from ads.aqua.common.entities import ContainerSpec
 from ads.aqua.common.enums import Tags
 from ads.aqua.common.errors import AquaResourceAccessError, AquaValueError
 from ads.aqua.common.utils import get_container_config, load_config, sanitize_response
+from ads.aqua.constants import EVALUATION_INFERENCE_DEFAULT_THREADS
 from ads.common import oci_client as oc
 from ads.common.auth import default_signer
 from ads.common.object_storage_details import ObjectStorageDetails
@@ -30,6 +29,8 @@
     TENANCY_OCID,
 )
 from ads.telemetry import telemetry
+from oci.exceptions import ServiceError
+from oci.identity.models import Compartment
 
 
 class ModelFormat(Enum):
@@ -45,13 +46,37 @@ def to_dict(self):
 #   within ads.aqua.common.entities. In that case, check for circular imports due to usage of get_container_config.
 
 
+@dataclass(repr=False)
+class AquaContainerEvaluationConfiguration(DataClassSerializable):
+    """
+    Represents the evaluation configuration for the container.
+    """
+
+    evaluation_max_threads: Optional[int] = None
+    evaluation_default_threads: int = field(
+        default=EVALUATION_INFERENCE_DEFAULT_THREADS
+    )
+
+    @classmethod
+    def from_config(cls, config: dict) -> "AquaContainerEvaluationConfiguration":
+        return cls(
+            evaluation_max_threads=config.get("MAX_THREADS"),
+            evaluation_default_threads=config.get(
+                "DEFAULT_THREADS", EVALUATION_INFERENCE_DEFAULT_THREADS
+            ),
+        )
+
+
 @dataclass(repr=False)
 class AquaContainerConfigSpec(DataClassSerializable):
     cli_param: str = None
     server_port: str = None
     health_check_port: str = None
     env_vars: List[dict] = None
     restricted_params: List[str] = None
+    evaluation_configuration: AquaContainerEvaluationConfiguration = field(
+        default_factory=AquaContainerEvaluationConfiguration
+    )
 
 
 @dataclass(repr=False)
@@ -161,6 +186,11 @@ def from_container_index_json(
                             restricted_params=container_spec.get(
                                 ContainerSpec.RESTRICTED_PARAMS, []
                             ),
+                            evaluation_configuration=AquaContainerEvaluationConfiguration.from_config(
+                                container_spec.get(
+                                    ContainerSpec.EVALUATION_CONFIGURATION, {}
+                                )
+                            ),
                         )
                         if container_spec
                         else None,