1717 _get_fine_tuning_options_and_model_arn ,
1818 _validate_and_resolve_model_package_group ,
1919 _extract_evaluator_arn ,
20+ _get_lambda_arn_from_evaluator_arn ,
2021 _is_lambda_arn ,
2122 _is_nova_model ,
2223 _resolve_model_and_name ,
@@ -104,9 +105,9 @@ class RLVRTrainer(BaseTrainer):
104105 ARN, or ModelPackageGroup object. Required when model is not a ModelPackage.
105106 custom_reward_function (Optional[Union[str, Evaluator]]):
106107 The custom reward function evaluator. Can be an evaluator ARN string, a Lambda
107- function ARN string, or an Evaluator object. If a Lambda ARN is provided
108+ function ARN string, or an Evaluator object. If a Lambda ARN is provided for serverless training
108109 (e.g., "arn:aws:lambda:us-east-1:123456789012:function:my-reward"), an Evaluator
109- will be automatically created in the AI Registry and used for training .
110+ will be automatically created in the AI Registry.
110111 Required for RLVR training to provide reward signals.
111112 mlflow_resource_arn (Optional[Union[str, MlflowTrackingServer]]):
112113 The MLflow tracking server ARN for experiment tracking.
@@ -132,6 +133,9 @@ class RLVRTrainer(BaseTrainer):
132133 is_multimodal (Optional[bool]):
133134 Whether the training dataset contains multimodal data. If None (default),
134135 auto-detected from the training dataset at train time.
136+ skip_reward_validation (bool):
137+ If True, skips the reward function verification step before submitting
138+ the training job. Defaults to False.
135139 """
136140
137141 _customization_technique = CustomizationTechnique .RLVR .value
@@ -156,6 +160,7 @@ def __init__(
156160 recipe : Optional [str ] = None ,
157161 overrides : Optional [dict ] = None ,
158162 is_multimodal : Optional [bool ] = None ,
163+ skip_reward_validation : bool = False ,
159164 base_model_name : Optional [str ] = None ,
160165 disable_output_compression : Optional [bool ] = False ,
161166 ** kwargs ,
@@ -195,6 +200,7 @@ def __init__(
195200 self ._recipe_resolver = None
196201 self ._resolved_recipe_cache = None
197202 self .is_multimodal = is_multimodal
203+ self .skip_reward_validation = skip_reward_validation
198204
199205 # Initialize fine-tuning options with beta session fallback
200206 self .hyperparameters , self ._model_arn , is_gated_model = _get_fine_tuning_options_and_model_arn (self ._model_name ,
@@ -279,19 +285,13 @@ def _verify_reward_function(
279285 )
280286 elif isinstance (reward_function , str ) and not _is_lambda_arn (reward_function ):
281287 # It's a string but not a Lambda ARN — treat as an evaluator ARN,
282- # fetch the Evaluator object and extract the Lambda ARN from it .
283- # Evaluator Lambdas always use OSS format (statusCode/body envelope)
288+ # fetch the Lambda ARN directly from the hub content entry .
289+ # This bypasses AIRHub so it works for any hub name in the ARN.
284290 try :
285- # Parse evaluator name from the ARN
286- # ARN format: arn:aws:sagemaker:region:account:hub-content/hub/type/name/version
287- evaluator_name = reward_function .split ("/" )[- 2 ]
288- evaluator = Evaluator .get (evaluator_name , sagemaker_session = TrainDefaults .get_sagemaker_session (sagemaker_session = self .sagemaker_session ))
289- reward_function = evaluator .reference
290- if not reward_function :
291- raise ValueError (
292- f"Evaluator '{ evaluator_name } ' does not have a Lambda ARN reference. "
293- "Verification requires a Lambda ARN."
294- )
291+ reward_function = _get_lambda_arn_from_evaluator_arn (
292+ reward_function ,
293+ sagemaker_session = self .sagemaker_session ,
294+ )
295295 logger .info (f"Resolved evaluator ARN to Lambda ARN: { reward_function } " )
296296 except ValueError :
297297 raise
@@ -430,7 +430,21 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None,
430430 )
431431
432432 # Extract and validate evaluator ARN
433- evaluator_arn = _extract_evaluator_arn (self .custom_reward_function ) if self .custom_reward_function else None
433+ # If custom_reward_function is a Lambda ARN, create an Evaluator object first
434+ if self .custom_reward_function and isinstance (self .custom_reward_function , str ) and _is_lambda_arn (self .custom_reward_function ):
435+ lambda_arn = self .custom_reward_function
436+ evaluator_name = _get_unique_name (f"rlvr-reward-{ self ._model_name } " )
437+ logger .info (f"Creating Evaluator from Lambda ARN: { lambda_arn } " )
438+ evaluator_obj = Evaluator .create (
439+ name = evaluator_name ,
440+ type = "RewardFunction" ,
441+ source = lambda_arn ,
442+ sagemaker_session = sagemaker_session ,
443+ )
444+ evaluator_arn = _extract_evaluator_arn (evaluator_obj )
445+ logger .info (f"Created Evaluator with ARN: { evaluator_arn } " )
446+ else :
447+ evaluator_arn = _extract_evaluator_arn (self .custom_reward_function ) if self .custom_reward_function else None
434448 serverless_config = _create_serverless_config (model_arn = self ._model_arn ,
435449 customization_technique = CustomizationTechnique .RLVR .value ,
436450 training_type = self .training_type ,
@@ -465,7 +479,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None,
465479 )
466480
467481 # Verify reward function before submitting training job
468- if self .custom_reward_function :
482+ if self .custom_reward_function and not self . skip_reward_validation :
469483 effective_dataset = training_dataset or self .training_dataset
470484 if effective_dataset is not None :
471485 logger .info ("Verifying reward function before submitting training job..." )
0 commit comments