started unit tests, added trust-remote-code param detection

Author: elizjo

ads/aqua/modeldeployment/deployment.py

@@ -1291,7 +1291,7 @@ def recommend_shape(self, **kwargs) -> Union[Table, ShapeRecommendationReport]:
         """
         deployment_config = self.get_deployment_config(model_id=kwargs.get("model_id"))
         kwargs["deployment_config"] = deployment_config
-        print(deployment_config)
+
         try:
             request = RequestRecommend(**kwargs)
         except ValidationError as e:

ads/aqua/shaperecommend/constants.py

@@ -88,6 +88,7 @@
 VLLM_PARAMS = {
     "max_model_len": "--max-model-len",
     "in_flight_quant": "--quantization bitsandbytes --load-format bitsandbytes",
+    "trust_remote_code": "--trust-remote-code"
 }
 
 DEFAULT_WEIGHT_SIZE = "bfloat16"

ads/aqua/shaperecommend/estimator.py

@@ -131,6 +131,10 @@ def construct_deployment_params(self) -> str:
             # vLLM only supports 4bit in-flight quantization
             params.append(VLLM_PARAMS["in_flight_quant"])
 
+        # add trust-remote-code if custom modules are specified
+        if c.trust_remote_code:
+            params.append(VLLM_PARAMS["trust_remote_code"])
+
         params = " ".join(params) if params else ""
         return params
 

ads/aqua/shaperecommend/llm_config.py

@@ -75,7 +75,9 @@ class LLMConfig(BaseModel):
         None, description="For MoE architectures, size of the MLP activation layer."
     )
 
-    tie_word_embeddings: Optional[bool] = Field(None)
+    tie_word_embeddings: Optional[bool] = Field(True, description="if True, input and output embedding matrices share the same parameters in memory.")
+
+    trust_remote_code: Optional[bool] = Field(False, description="if True, the model requires custom code to operate.")
 
     @property
     def bytes_per_parameter(self) -> float:
@@ -208,6 +210,17 @@ def validate_model_support(cls, raw: dict) -> ValueError:
                 "Encoder-decoder models (ex. T5, Gemma) and encoder-only (BERT) are not supported at this time."
             )
 
+    @staticmethod
+    def get_bool(raw, key, default=False):
+        val = raw.get(key)
+        if val is None:
+            return default
+        if isinstance(val, bool):
+            return val
+        if isinstance(val, str):
+            return val.lower() == "true"
+        return bool(val)
+
     @classmethod
     def from_raw_config(cls, raw: dict) -> "LLMConfig":
         """
@@ -258,6 +271,10 @@ def from_raw_config(cls, raw: dict) -> "LLMConfig":
             "intermediate_size"
         )
 
+        tie_word_embeddings = LLMConfig.get_bool(raw, "tie_word_embeddings", True)
+
+        trust_remote_code = "auto_map" in raw # trust-remote-code is always needed when this key is present
+
         # Type safety: minimal assertion
         if None in [
             num_hidden_layers,
@@ -281,4 +298,6 @@ def from_raw_config(cls, raw: dict) -> "LLMConfig":
             max_seq_len=int(max_seq_len),
             num_local_experts=num_local_experts,
             intermediate_size=intermediate_size,
+            tie_word_embeddings=tie_word_embeddings,
+            trust_remote_code=trust_remote_code
         )

ads/aqua/shaperecommend/shape_report.py

@@ -250,55 +250,83 @@ class ShapeRecommendationReport(BaseModel):
     @classmethod
     def from_deployment_config(cls, deployment_config: AquaDeploymentConfig, model_name: str, valid_shapes: List[ComputeShapeSummary]) -> "ShapeRecommendationReport":
         """
-        For service models, pre-set deployment configurations (AquaDeploymentConfig) are available.
-        Derives ShapeRecommendationReport from AquaDeploymentConfig (if service model & available)
+        Creates a ShapeRecommendationReport from an AquaDeploymentConfig, extracting recommended
+        model configurations for each valid compute shape.
+
+        Parameters
+        ----------
+        deployment_config : AquaDeploymentConfig
+            The object containing per-shape deployment configurations.
+        model_name : str
+            The name of the model for which to generate recommendations.
+        valid_shapes : list of ComputeShapeSummary
+            List of compute shapes to evaluate and recommend deployment configurations for.
+
+        Returns
+        -------
+        ShapeRecommendationReport
+            Report containing recommendations for each valid compute shape.
+
+        Notes
+        -----
+        For service models, this method interprets pre-set deployment configurations to derive
+        recommendations for each allowed compute shape, including environment variables, quantization,
+        and maximum model length parameters.
         """
 
         recs = []
-        # may need to sort?
         for shape in valid_shapes:
             current_config = deployment_config.configuration.get(shape.name)
-            if current_config:
-                quantization = None
-                max_model_len = None
-                recommendation = ""
-                current_params = current_config.parameters.get(VLLM_PARAMS_KEY)
-                current_env = current_config.env.get(VLLM_ENV_KEY)
+            if not current_config:
+                continue
 
-                if current_params:
-                    param_list = current_params.split()
+            quantization = None
+            max_model_len = None
+            recommendation = ""
+            current_params = current_config.parameters.get(VLLM_PARAMS_KEY)
+            current_env = current_config.env.get(VLLM_ENV_KEY)
 
-                    if QUANT_FLAG in param_list and (idx := param_list.index(QUANT_FLAG)) + 1 < len(param_list):
+            if current_params:
+                param_list = current_params.split()
+
+                if QUANT_FLAG in param_list:
+                    idx = param_list.index(QUANT_FLAG)
+                    if idx + 1 < len(param_list):
                         quantization = param_list[idx + 1]
 
-                    if MAX_MODEL_LEN_FLAG in param_list and (idx := param_list.index(MAX_MODEL_LEN_FLAG)) + 1 < len(param_list):
-                        max_model_len = param_list[idx + 1]
-                        max_model_len = int(max_model_len)
+                if MAX_MODEL_LEN_FLAG in param_list:
+                    idx = param_list.index(MAX_MODEL_LEN_FLAG)
+                    if idx + 1 < len(param_list):
+                        try:
+                            max_model_len = int(param_list[idx + 1])
+                        except ValueError:
+                            max_model_len = None
 
-                if current_env:
-                    recommendation += f"ENV: {json.dumps(current_env)}\n\n"
+            if current_env:
+                recommendation += f"ENV: {json.dumps(current_env)}\n\n"
 
-                recommendation += "Model fits well within the allowed compute shape."
+            if not current_params and not current_env: # model works with default params and no extra env variables
+                recommendation += "No override PARAMS and ENV variables needed. \n\n"
 
-                deployment_params = DeploymentParams(
-                    quantization=quantization if quantization else DEFAULT_WEIGHT_SIZE,
-                    max_model_len=max_model_len,
-                    params=current_params if current_params else "",
-                )
+            recommendation += "Model fits well within the allowed compute shape."
 
-                # TODO: calculate memory footprint based on params??
-                # TODO: add --env vars not just params, current_config.env
-                # are there multiple configurations in the SMM configs per shape??
-                configuration = [ModelConfig(
-                    deployment_params=deployment_params,
-                    recommendation=recommendation,
-                )]
-
-                recs.append(ShapeReport(
-                    shape_details=shape,
-                    configurations=configuration
-                )
-                )
+            deployment_params = DeploymentParams(
+                quantization=quantization if quantization else DEFAULT_WEIGHT_SIZE,
+                max_model_len=max_model_len,
+                params=current_params if current_params else "",
+            )
+
+            # need to adjust for multiple configs per shape
+            configuration = [ModelConfig(
+                deployment_params=deployment_params,
+                recommendation=recommendation,
+            )]
+
+            recs.append(ShapeReport(
+                shape_details=shape,
+                configurations=configuration
+            )
+            )
 
         return ShapeRecommendationReport(
             display_name=model_name,