Skip to content

Commit 7e7b6f1

Browse files
mujtaba1747Syed Jafri
andauthored
rlvr validation bugfix (#6000)
* fix: removed reward verifier is_lambda + compute check * chore: add log to test to print job arn * fix: use hyp override in nova test as workaround for recipe resolver error * fix: added comments + fixed evaluator test --------- Co-authored-by: Syed Jafri <syedjfr@amazon.com>
1 parent 0fdbbdc commit 7e7b6f1

8 files changed

Lines changed: 159 additions & 116 deletions

File tree

‎sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -479,6 +479,67 @@ def _extract_evaluator_arn(evaluator, param_name: str = "custom_reward_function"
479479
return evaluator.arn
480480

481481

482+
def _get_lambda_arn_from_evaluator_arn(evaluator_arn: str, sagemaker_session=None) -> str:
483+
"""Extract the Lambda ARN from an evaluator's hub content entry.
484+
485+
Parses the hub name, content type, content name, and version directly from the
486+
evaluator ARN, then calls describe_hub_content on the SageMaker client to retrieve
487+
the stored Lambda function reference. This bypasses AIRHub (which hardcodes a
488+
single hub name) so it works for evaluators stored in any hub.
489+
490+
Args:
491+
evaluator_arn: Full evaluator ARN
492+
(e.g., arn:aws:sagemaker:us-west-2:123456789012:hub-content/MyHub/JsonDoc/my-eval/0.0.1)
493+
sagemaker_session: Optional SageMaker session
494+
495+
Returns:
496+
The Lambda function ARN stored as the evaluator's reference
497+
498+
Raises:
499+
ValueError: If the ARN cannot be parsed or the evaluator has no Lambda reference
500+
"""
501+
# Parse: arn:aws:sagemaker:region:account:hub-content/HubName/ContentType/ContentName/Version
502+
try:
503+
resource_part = evaluator_arn.split(":", 5)[5] # "hub-content/HubName/Type/Name/Version"
504+
parts = resource_part.split("/")
505+
# parts[0] = "hub-content", parts[1] = HubName, parts[2] = ContentType,
506+
# parts[3] = ContentName, parts[4] = Version (optional)
507+
hub_name = parts[1]
508+
hub_content_type = parts[2]
509+
hub_content_name = parts[3]
510+
hub_content_version = parts[4] if len(parts) > 4 else None
511+
except (IndexError, ValueError) as e:
512+
raise ValueError(
513+
f"Failed to parse evaluator ARN '{evaluator_arn}': {str(e)}"
514+
)
515+
516+
sagemaker_session = TrainDefaults.get_sagemaker_session(sagemaker_session=sagemaker_session)
517+
client = sagemaker_session.sagemaker_client
518+
519+
request = {
520+
"HubName": hub_name,
521+
"HubContentType": hub_content_type,
522+
"HubContentName": hub_content_name,
523+
}
524+
if hub_content_version:
525+
request["HubContentVersion"] = hub_content_version
526+
527+
response = client.describe_hub_content(**request)
528+
529+
# Parse the document: HubContentDocument -> JsonContent -> Reference
530+
doc = json.loads(response["HubContentDocument"])
531+
json_content = json.loads(doc.get("JsonContent", "{}"))
532+
reference = json_content.get("Reference", "")
533+
534+
if not reference:
535+
raise ValueError(
536+
f"Evaluator '{hub_content_name}' in hub '{hub_name}' does not have a "
537+
"Lambda ARN reference."
538+
)
539+
540+
return reference
541+
542+
482543
def _resolve_model_name(model_package) -> str:
483544
"""Resolve model_name from model_package if needed."""
484545
if model_package:

‎sagemaker-train/src/sagemaker/train/common_utils/rlvr_reward_verifier.py‎

Lines changed: 24 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -181,32 +181,34 @@ def verify_reward_function(
181181
# Determine if it's a Lambda ARN or local file
182182
is_lambda = bool(LAMBDA_ARN_REGEX.match(reward_function))
183183

184-
# Compute is required for Lambda ARNs with Nova models
185-
if is_lambda and is_nova and compute is None:
186-
raise ValueError(
187-
"The 'compute' parameter is required for Nova models when using a Lambda ARN. "
188-
"Please specify compute as a TrainingJobCompute or HyperPodCompute instance."
189-
)
190-
191184
# Validate Lambda ARN format for HyperPod compute
192-
if is_nova and isinstance(compute, HyperPodCompute) and is_lambda:
193-
# Extract function name from ARN: arn:aws[-*]:lambda:region:account:function:function-name
194-
function_name_match = re.search(
195-
r"arn:aws[a-zA-Z-]*:lambda:[^:]+:[^:]+:function:([^:]+)", reward_function
196-
)
197-
if function_name_match:
198-
function_name = function_name_match.group(1)
199-
# Check if function name contains 'SageMaker' (case-insensitive)
200-
if not re.search(r"sagemaker", function_name, re.IGNORECASE):
185+
if is_nova and isinstance(compute, HyperPodCompute):
186+
if is_lambda:
187+
# Extract function name from ARN: arn:aws[-*]:lambda:region:account:function:function-name
188+
function_name_match = re.search(
189+
r"arn:aws[a-zA-Z-]*:lambda:[^:]+:[^:]+:function:([^:]+)", reward_function
190+
)
191+
if function_name_match:
192+
function_name = function_name_match.group(1)
193+
# Check if function name contains 'SageMaker' (case-insensitive)
194+
if not re.search(r"sagemaker", function_name, re.IGNORECASE):
195+
raise ValueError(
196+
f"Lambda ARN for HyperPod compute must contain 'SageMaker' in the function name for Nova models. "
197+
f"Current function name: '{function_name}'. "
198+
f"Expected format: 'arn:aws:lambda:*:*:function:*SageMaker*'"
199+
)
200+
else:
201201
raise ValueError(
202-
f"Lambda ARN for HyperPod compute must contain 'SageMaker' in the function name for Nova models. "
203-
f"Current function name: '{function_name}'. "
204-
f"Expected format: 'arn:aws:lambda:*:*:function:*SageMaker*'"
202+
f"Invalid Lambda ARN format: {reward_function}. "
203+
f"Expected format: 'arn:aws:lambda:region:account:function:function-name'"
205204
)
206205
else:
207-
raise ValueError(
208-
f"Invalid Lambda ARN format: {reward_function}. "
209-
f"Expected format: 'arn:aws:lambda:region:account:function:function-name'"
206+
logger.warning(
207+
f"Skipping Lambda function name validation because a local Python file is being used: "
208+
f"'{reward_function}'. "
209+
"For Nova RLVR jobs on HyperPod, ensure your reward function Lambda name "
210+
"contains 'SageMaker' (case-insensitive). "
211+
"Expected format: 'arn:aws:lambda:*:*:function:*SageMaker*'"
210212
)
211213

212214
results = []

‎sagemaker-train/src/sagemaker/train/rlvr_trainer.py‎

Lines changed: 30 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
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...")

‎sagemaker-train/tests/integ/train/code/nova_reward_fn.py‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,9 @@ def lambda_handler(event, context):
2626
# print(sample)
2727
if not "id" in sample:
2828
print(f"ID is None/empty for sample: {sample}")
29-
else:
30-
idx = sample["id"]
29+
continue
30+
31+
idx = sample["id"]
3132

3233
ro = RewardOutput(id=idx, aggregate_reward_score=0.0)
3334

@@ -38,24 +39,24 @@ def lambda_handler(event, context):
3839

3940
# Extract answer from ground truth dict
4041
if ground_truth is None:
41-
print(f"No answer found in ground truth for id: {idx}")
42-
# scores.append(RewardOutput(id="0", aggregate_reward_score=0.0))
42+
print(f"Warning: No answer found in ground truth (reference_answer) for id: {idx}")
43+
scores.append(RewardOutput(id=idx, aggregate_reward_score=0.0))
4344
continue
4445

4546
# Get completion from last message (assistant message)
4647
last_message = sample["messages"][-1]
4748
# completion_text = last_message["content"]
48-
49-
if last_message["role"] != "assistant":
50-
print(f"Last message is not from assistant for id: {idx}")
51-
# scores.append(RewardOutput(id="0", aggregate_reward_score=0.0))
52-
continue
5349

5450
if not "content" in last_message:
5551
print(f"Completion text is empty for id: {idx}")
5652
# scores.append(RewardOutput(id="0", aggregate_reward_score=0.0))
5753
continue
5854

55+
if last_message["role"] != "assistant":
56+
print(f"Last message is not from assistant for id: {idx}")
57+
scores.append(RewardOutput(id=idx, aggregate_reward_score=0.0))
58+
continue
59+
5960
random_score = random.uniform(0.0, 1.0)
6061
ro = RewardOutput(id=idx, aggregate_reward_score=random_score)
6162

‎sagemaker-train/tests/integ/train/test_extract_evaluator_arn_integration.py‎

Lines changed: 0 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -94,34 +94,3 @@ def test_extract_evaluator_arn_with_evaluator_string(sagemaker_session, evaluato
9494
# Should return the ARN string unchanged
9595
assert result == evaluator.arn
9696

97-
98-
def test_extract_evaluator_arn_with_lambda_arn_string(sagemaker_session, lambda_arn):
99-
"""Test _extract_evaluator_arn with a Lambda ARN string.
100-
101-
Verifies that passing a Lambda ARN triggers auto-creation of an Evaluator
102-
in the AI Registry and returns the newly created evaluator's hub-content ARN.
103-
"""
104-
result = _extract_evaluator_arn(lambda_arn, "custom_reward_function")
105-
106-
# The result should be a SageMaker hub-content ARN (not the Lambda ARN)
107-
assert result != lambda_arn
108-
assert result.startswith("arn:aws:sagemaker:")
109-
assert "/JsonDoc/" in result
110-
# The evaluator name should be derived from the Lambda function name (sanitized)
111-
assert SANITIZED_LAMBDA_FUNCTION_NAME in result
112-
113-
114-
def test_extract_evaluator_arn_lambda_creates_retrievable_evaluator(sagemaker_session, lambda_arn):
115-
"""Test that the Evaluator auto-created from a Lambda ARN is retrievable.
116-
117-
Verifies that after _extract_evaluator_arn creates an Evaluator from a Lambda ARN,
118-
the evaluator can be retrieved from the AI Registry by name.
119-
"""
120-
result_arn = _extract_evaluator_arn(lambda_arn, "custom_reward_function")
121-
122-
# Verify the evaluator was actually created and is retrievable (use sanitized name)
123-
evaluator = Evaluator.get(SANITIZED_LAMBDA_FUNCTION_NAME, sagemaker_session=sagemaker_session)
124-
assert evaluator is not None
125-
assert evaluator.arn == result_arn
126-
assert evaluator.type == REWARD_FUNCTION
127-
assert evaluator.reference == lambda_arn

‎sagemaker-train/tests/integ/train/test_reward_verifier_integration.py‎

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -414,16 +414,16 @@ def test_lambda_nova_without_format_validation(self, nova_sample_data, nova_lamb
414414
assert result["total_samples"] == 2
415415
assert result["successful_samples"] == 2
416416

417-
def test_lambda_nova_requires_compute(self, nova_sample_data, nova_lambda_arn):
418-
"""Test that Nova Lambda ARN requires compute parameter."""
419-
with pytest.raises(ValueError, match="compute.*required"):
420-
verify_reward_function(
421-
reward_function=nova_lambda_arn,
422-
sample_data=nova_sample_data,
423-
validate_format=True,
424-
is_nova=True,
425-
)
426-
417+
def test_lambda_nova_serverless_compute(self, nova_sample_data, nova_lambda_arn):
418+
result = verify_reward_function(
419+
reward_function=nova_lambda_arn,
420+
sample_data=nova_sample_data,
421+
validate_format=True,
422+
is_nova=True,
423+
)
424+
assert result["success"] is True
425+
assert result["total_samples"] == 2
426+
assert result["successful_samples"] == 2
427427

428428
# ---------------------------------------------------------------------------
429429
# Test class: OSS Remote Lambda (is_nova=False)
@@ -533,16 +533,6 @@ def test_invalid_lambda_arn_format(self, nova_sample_data):
533533
is_nova=True,
534534
)
535535

536-
def test_lambda_arn_without_compute_raises_error(self, nova_sample_data, nova_lambda_arn):
537-
"""Test that Lambda ARN without compute raises ValueError."""
538-
with pytest.raises(ValueError, match="compute.*required"):
539-
verify_reward_function(
540-
reward_function=nova_lambda_arn,
541-
sample_data=nova_sample_data,
542-
validate_format=True,
543-
compute=None,
544-
is_nova=True,
545-
)
546536

547537
def test_hyperpod_compute_requires_sagemaker_in_function_name(
548538
self, nova_sample_data, nova_lambda_arn

0 commit comments

Comments
 (0)