Updating Nova validation (#338)

* Updating recipe validation models for Nova

* Skip job validation for nova-2 jobs due to RFT requirements

Author: niphaded-amzn

src/hyperpod_cli/validators/job_validator.py

@@ -292,6 +292,17 @@ def validate_scheduler_related_fields(
                 return False
     return True
 
+def skippable_recipe(model_type: str, recipe_path: str):
+    '''
+
+    :param model_type: Could be [nova-1, nova-2] for NovaModels and open source model types for others.
+    :param recipe_path: can contain possible values from sft, rft, cpt, eval,
+    :return: True if skippable, False if non-skippable
+    '''
+    return "nova-2" in model_type and any(x in recipe_path.lower() for x in ["sft", "rft", "cpt"])
+
+
+
 def validate_recipe_file(recipe: str):
     recipe_path = os.path.join(RECIPES_DIR, f"{recipe}.yaml")
 
@@ -314,6 +325,9 @@ def validate_recipe_file(recipe: str):
                 elif "nova" in model_type and "evaluation" in recipe_data:
                     NovaEvaluationRecipeSchema(**recipe_data)
                 elif "nova" in model_type:
+                    # Skip recipe validation for nova-2 models for beta
+                    if skippable_recipe(model_type, recipe_path):
+                        return True
                     NovaRecipeSchema(**recipe_data)
                 else:
                     raise Exception("Unsupported model_type {model_type}. Make sure the recipe exists in src/hyperpod_cli/sagemaker_hyperpod_recipes/recipes_collection/recipes")

src/hyperpod_cli/validators/recipe_models/nova/model.py

@@ -11,6 +11,7 @@ class RunConfig(BaseModel):
     replicas: Optional[int|str] = None
     data_s3_path: Optional[str] = None
     output_s3_path: Optional[str] = None
+    validation_data_s3_path: Optional[str] = None
 
     # PPO-specific replica configurations
     actor_train_replicas: Optional[int|str] = None
@@ -20,11 +21,20 @@ class RunConfig(BaseModel):
     am_replicas: Optional[int|str] = None
 
 
+    # MLFlow optional parameters
+    mlflow_tracking_uri: Optional[str] = None
+    mlflow_experiment_name: Optional[str] = None
+    mlflow_run_name: Optional[str] = None
+
+
 class TrainerConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
     max_epochs: Optional[int|str] = None
     num_nodes: Optional[int|str] = None
+    max_steps: Optional[int|str] = None
+    val_check_interval: Optional[int|float|str] = None
+    limit_val_batches: Optional[int|float|str] = None
 
 
 class SchedulerConfig(BaseModel):
@@ -36,7 +46,7 @@ class SchedulerConfig(BaseModel):
 
 
 class OptimizerConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="allow")
 
     name: Optional[str] = None
     lr: Optional[float] = None
@@ -45,6 +55,8 @@ class OptimizerConfig(BaseModel):
     weight_decay: Optional[float] = None
     betas: Optional[List[float]] = None
     sched: Optional[SchedulerConfig] = None
+    adam_beta1: Optional[float] = None
+    adam_beta2: Optional[float] = None
 
 
 class DpoConfig(BaseModel):
@@ -59,6 +71,7 @@ class LoraTuningConfig(BaseModel):
     loraplus_lr_ratio: Optional[float] = None
     alpha: Optional[float] = None
     adapter_dropout: Optional[float] = None
+    lora_plus_lr_ratio: Optional[float] = None
 
 
 class PeftConfig(BaseModel):
@@ -84,13 +97,22 @@ class ModelConfig(BaseModel):
     kl_reward_penalty_coeff: Optional[float] = None
 
 
+class ModelImportanceScore(BaseModel):
+    fine_tuned_model: Optional[float] = None
+
+
 class TrainingConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="allow")
 
     max_length: Optional[int|str] = None
     global_batch_size: Optional[int|str] = None
     trainer: Optional[TrainerConfig] = None
     model: Optional[ModelConfig] = None
+    max_steps: Optional[int|str] = None
+    save_steps: Optional[int | str] = None
+    save_top_k: Optional[int | str] = None
+    reasoning_enabled: Optional[int | str] = None
+    lr_scheduler: Optional[SchedulerConfig] = None
 
     # Distillation-specific fields
     distillation_data: Optional[str] = None
@@ -105,6 +127,14 @@ class TrainingConfig(BaseModel):
     top_p: Optional[str] = None
     customer_bucket: Optional[str] = None
     kms_key: Optional[str] = None
+    task_type: Optional[str] = None
+    optim: Optional[OptimizerConfig] = None
+
+    optim_config: Optional[OptimizerConfig] = None
+    peft: Optional[PeftConfig] = None
+
+    # RAI vector merge
+    model_importance_score: Optional[ModelImportanceScore] = None
 
 
 class PpoRewardConfig(BaseModel):
@@ -153,12 +183,20 @@ class PpoActorTrainConfig(BaseModel):
 class NovaRecipeSchema(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
+    display_name: Optional[str] = None
+    versions: Optional[list] = None
+    instance_types: Optional[list] = None
+
     # Common configurations
     run: RunConfig
 
     # Training and fine-tuning specific configurations
     training_config: Optional[TrainingConfig] = None
 
+    # Enable skipping recipe validation in the container
+    # This is controlled by an allowlist in the container
+    skip_recipe_validation: Optional[bool] = None
+
     # PPO-specific configurations
     ppo_reward: Optional[PpoRewardConfig] = None
     ppo_critic: Optional[PpoCriticConfig] = None