Skip to content

Commit 0316c38

Browse files
authored
fix(core): resolve default training role from sagemaker config (#6228)
Training (and feature_store) role resolution only inferred a role from the caller identity, so a caller authenticating as an IAM user or the account root - whose identity has no backing role - hit "No IAM role could be resolved from your caller identity" and was forced to pass role= on every call, even when a default execution role was configured in the SageMaker intelligent-defaults config. resolve_and_validate_role now consults the config default (SageMaker.TrainingJob.RoleArn / FeatureGroup.RoleArn) before falling back to caller-identity inference, matching the pattern already used by processing and model monitor. Resolution order is now: explicit role -> config default -> caller identity -> raise. Behavior for a bare IAM user with no configured default is unchanged (still raises the same error). --- X-AI-Prompt: Deep-dive RLVR trainer default-role resolution failing for IAM users; add a sagemaker-config default-role fallback X-AI-Tool: Kiro
1 parent 499ac2f commit 0316c38

2 files changed

Lines changed: 143 additions & 10 deletions

File tree

‎sagemaker-core/src/sagemaker/core/helper/iam_role_resolver.py‎

Lines changed: 76 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,59 @@ def _resolve_caller_role_arn(
306306
raise
307307

308308

309+
def _config_path_for_role_type(role_type: str) -> Optional[str]:
310+
"""Return the SageMaker config key path holding a default role ARN for a role type.
311+
312+
Returns None for role types the config schema has no dedicated role-ARN path
313+
for, in which case there is no config default to consult.
314+
"""
315+
try:
316+
from sagemaker.core.config.config_schema import (
317+
TRAINING_JOB_ROLE_ARN_PATH,
318+
FEATURE_GROUP_ROLE_ARN_PATH,
319+
)
320+
except Exception: # pragma: no cover - defensive against import/layout changes
321+
return None
322+
return {
323+
"training": TRAINING_JOB_ROLE_ARN_PATH,
324+
"feature_store": FEATURE_GROUP_ROLE_ARN_PATH,
325+
}.get(role_type)
326+
327+
328+
def _resolve_config_default_role(role_type: str, sagemaker_session=None) -> Optional[str]:
329+
"""Return a default role ARN from the SageMaker intelligent-defaults config, if set.
330+
331+
This lets a caller whose own identity has no backing role (an IAM user or the
332+
account root) configure a default execution role in the SageMaker config
333+
(e.g. ``SageMaker.TrainingJob.RoleArn``) instead of being forced to pass
334+
``role=`` on every call. Returns None when no config default is set, the role
335+
type has no config path, or the config cannot be read — in every such case the
336+
caller falls back to caller-identity resolution exactly as before.
337+
"""
338+
config_path = _config_path_for_role_type(role_type)
339+
if not config_path:
340+
return None
341+
try:
342+
from sagemaker.core.common_utils import resolve_value_from_config
343+
344+
config_role = resolve_value_from_config(
345+
direct_input=None,
346+
config_path=config_path,
347+
sagemaker_session=sagemaker_session,
348+
)
349+
except Exception as e: # pragma: no cover - defensive; treat as "no config default"
350+
logger.debug(
351+
"Could not read a default role from the SageMaker config for '%s': %s",
352+
role_type,
353+
e,
354+
)
355+
return None
356+
# Only trust a concrete string ARN/name; anything else means "not configured".
357+
if isinstance(config_role, str) and config_role:
358+
return config_role
359+
return None
360+
361+
309362
# ---------------------------------------------------------------------------
310363
# Read-only permission / trust validation
311364
# ---------------------------------------------------------------------------
@@ -530,9 +583,11 @@ def resolve_and_validate_role(
530583
) -> str:
531584
"""Resolve the role to use and validate it (read-only; does not mutate IAM).
532585
533-
Resolution:
586+
Resolution (first match wins):
534587
1. ``provided_role`` given → resolve it to an ARN (must exist).
535-
2. Otherwise → resolve the caller's own identity role.
588+
2. A default role set in the SageMaker config for this role type
589+
(e.g. ``SageMaker.TrainingJob.RoleArn``) → resolve it to an ARN.
590+
3. Otherwise → resolve the caller's own identity role.
536591
537592
The resolved role is then VALIDATED (read-only, via iam:SimulatePrincipalPolicy
538593
+ trust inspection):
@@ -564,14 +619,25 @@ def resolve_and_validate_role(
564619
if provided_role:
565620
role_arn = _resolve_explicit_role(provided_role, sagemaker_session)
566621
else:
567-
sts_client = boto_session.client("sts")
568-
caller_identity = sts_client.get_caller_identity()
569-
caller_arn = caller_identity["Arn"]
570-
account_id = caller_identity["Account"]
571-
partition = _partition_from_arn(caller_arn)
572-
role_arn = _resolve_caller_role_arn(iam_client, caller_arn, account_id, partition)
573-
if not role_arn:
574-
raise RoleValidationError(_build_validation_error_message(None, role_type))
622+
# Prefer a default role configured in the SageMaker config for this role
623+
# type (e.g. SageMaker.TrainingJob.RoleArn) before falling back to
624+
# caller-identity inference. This is what lets an IAM-user or root caller
625+
# (whose identity has no backing role) run without passing role= on every
626+
# call, as long as they have configured a default execution role.
627+
config_role = _resolve_config_default_role(role_type, sagemaker_session)
628+
if config_role:
629+
role_arn = _resolve_explicit_role(config_role, sagemaker_session)
630+
else:
631+
sts_client = boto_session.client("sts")
632+
caller_identity = sts_client.get_caller_identity()
633+
caller_arn = caller_identity["Arn"]
634+
account_id = caller_identity["Account"]
635+
partition = _partition_from_arn(caller_arn)
636+
role_arn = _resolve_caller_role_arn(
637+
iam_client, caller_arn, account_id, partition
638+
)
639+
if not role_arn:
640+
raise RoleValidationError(_build_validation_error_message(None, role_type))
575641

576642
# Permission check (definitive denial blocks; unverifiable warns).
577643
verdict, denied = _evaluate_permissions(iam_client, role_arn, role_type)

‎sagemaker-core/tests/unit/helper/test_iam_role_resolver.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -351,6 +351,73 @@ def test_no_resolvable_caller_role_raises(self):
351351
assert "No IAM role could be resolved" in str(exc.value)
352352
mock_iam.create_role.assert_not_called()
353353

354+
def test_config_default_role_used_when_caller_is_iam_user(self):
355+
"""An IAM user with a configured default training role uses it, not failing."""
356+
role_arn = "arn:aws:iam::123456789012:role/ConfiguredRole"
357+
mock_session, mock_iam, _ = _make_session(
358+
"arn:aws:iam::123456789012:user/dev-user"
359+
)
360+
mock_iam.get_role.return_value = {
361+
"Role": {"Arn": role_arn, "AssumeRolePolicyDocument": _trusted_doc()}
362+
}
363+
mock_iam.get_paginator.return_value = _paginator_allowing(["s3:GetObject"])
364+
365+
with patch(
366+
"sagemaker.core.common_utils.resolve_value_from_config",
367+
return_value=role_arn,
368+
) as mock_cfg:
369+
result = resolve_and_validate_role(
370+
provided_role=None,
371+
role_type="training",
372+
sagemaker_session=mock_session,
373+
)
374+
375+
assert result == role_arn
376+
mock_cfg.assert_called_once()
377+
# The config default is used instead of caller-identity inference.
378+
mock_iam.create_role.assert_not_called()
379+
380+
def test_config_default_role_takes_precedence_over_caller_role(self):
381+
"""A configured default role wins over the caller's own backing role."""
382+
config_role = "arn:aws:iam::123456789012:role/ConfiguredRole"
383+
mock_session, mock_iam, _ = _make_session(
384+
"arn:aws:sts::123456789012:assumed-role/CallerRole/sess"
385+
)
386+
mock_iam.get_role.return_value = {
387+
"Role": {"Arn": config_role, "AssumeRolePolicyDocument": _trusted_doc()}
388+
}
389+
mock_iam.get_paginator.return_value = _paginator_allowing(["s3:GetObject"])
390+
391+
with patch(
392+
"sagemaker.core.common_utils.resolve_value_from_config",
393+
return_value=config_role,
394+
):
395+
result = resolve_and_validate_role(
396+
provided_role=None,
397+
role_type="training",
398+
sagemaker_session=mock_session,
399+
)
400+
401+
assert result == config_role
402+
403+
def test_iam_user_without_config_default_still_raises(self):
404+
"""No configured default + IAM-user caller still raises (behavior preserved)."""
405+
mock_session, mock_iam, _ = _make_session(
406+
"arn:aws:iam::123456789012:user/dev-user"
407+
)
408+
with patch(
409+
"sagemaker.core.common_utils.resolve_value_from_config",
410+
return_value=None,
411+
):
412+
with pytest.raises(RoleValidationError) as exc:
413+
resolve_and_validate_role(
414+
provided_role=None,
415+
role_type="training",
416+
sagemaker_session=mock_session,
417+
)
418+
assert "No IAM role could be resolved" in str(exc.value)
419+
mock_iam.create_role.assert_not_called()
420+
354421
def test_invalid_role_type_raises(self):
355422
"""Invalid role_type raises ValueError."""
356423
with pytest.raises(ValueError, match="Invalid role_type"):

0 commit comments

Comments
 (0)