Skip to content

fix: Add pre-validation for Bedrock and MLflow role permissions - #6162

Open
rsareddy0329 wants to merge 2 commits into
aws:masterfrom
rsareddy0329:master-eval-role-log-fix
Open

fix: Add pre-validation for Bedrock and MLflow role permissions#6162
rsareddy0329 wants to merge 2 commits into
aws:masterfrom
rsareddy0329:master-eval-role-log-fix

Conversation

@rsareddy0329

Copy link
Copy Markdown
Contributor

Add pre-validation for Bedrock and MLflow role permissions

Validate execution role permissions before submitting LLM-as-Judge evaluation jobs to fail fast with actionable guidance instead of failing during execution.

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Validate execution role permissions before submitting evaluation jobs
to fail fast with actionable guidance instead of failing during execution.

Reuses existing _simulate_denied_actions and _role_trusts_service from
iam_role_resolver.py. Applies to all evaluators via BaseEvaluator.
@github-actions

github-actions Bot commented Aug 6, 2026

Copy link
Copy Markdown

Claude finished @rsareddy0329's task in 5m 52s —— View job


Code Review: PR #6162 — Pre-validation for Bedrock/MLflow role permissions

  • Read the PR diff
  • Inspect referenced helpers (iam_role_resolver, base_evaluator, iam_policies)
  • Analyze correctness, security, tests
  • Post review feedback

The goal — fail fast with actionable guidance instead of 30 minutes into a job — is worthwhile, and the graceful-degradation and error-message handling are well done. However, there are a few correctness issues that can hard-fail jobs that would otherwise succeed, so I'd hold on merging until these are addressed. (Note: my environment doesn't have the inline-comment tool, so findings are collected below with file/line references.)


🔴 High — validation runs for all evaluators, but Bedrock trust applies only to some

base_evaluator.py:756 — the call is added to _get_aws_execution_context(), which lives on BaseEvaluator and is invoked by every evaluator:

  • llm_as_judge_evaluator.py:815
  • inspect_ai_evaluator.py:727
  • benchmark_evaluator.py:673
  • custom_scorer_evaluator.py:453
  • multi_turn_rl_evaluator.py:583

The PR is scoped to "LLM-as-Judge evaluation jobs," but as written it also gates InspectAI (existing/create-endpoint modes), Benchmark, CustomScorer, and MultiTurnRL (which uses Bedrock AgentCore, not bedrock:CreateEvaluationJob). Those backends run as SageMaker training jobs — the execution role is assumed by sagemaker.amazonaws.com, not bedrock.amazonaws.com. _check_trust() (role_permission_validator.py:145) will then hit trusts_bedrock is False and raise ValueError, breaking jobs that don't need Bedrock at all. The validation needs to be gated to the LLM-as-Judge / Bedrock-evaluation path (e.g. called from LLMAsJudgeEvaluator.evaluate, or guarded by a flag), not applied unconditionally in the base class.

🔴 High — MLflow (and likely Bedrock) actions are resource-scoped → false denials

role_permission_validator.py:107-125 simulates the actions via _simulate_denied_actions, which calls iam:SimulatePrincipalPolicy without ResourceArns. The MLflow actions checked (sagemaker-mlflow:CreateRun, LogBatch, …) are scoped in the real managed policy to arn:aws:sagemaker:*:*:mlflow-app/* and …:mlflow-tracking-server/* (see iam_policies.py:252-268) — not "*". When you simulate a resource-scoped action without a ResourceArns, IAM evaluates it against an implicit * and returns implicitDeny even for a perfectly valid role, producing a false positive that hard-fails the job.

This is exactly the trap the existing code documents and deliberately avoids — see iam_role_resolver.py:109-134 (_get_smoke_test_actions), which only gates on *-resource actions for this reason. This new validator reintroduces the false-positive that the smoke-test design was written to prevent. bedrock:CreateEvaluationJob/GetEvaluationJob may have the same issue if a customer scopes them by resource. Recommendation: either pass appropriate ResourceArns, or only gate (raise) on actions whose policy resource is * and warn (don't block) otherwise — matching the established pattern.

🟡 Medium — hard-raise diverges from the codebase's "warn, don't block" convention

resolve_and_validate_role (the sibling that already runs at base_evaluator.py:741) only blocks on definitive denials of *-resource actions and otherwise warns and proceeds. This validator raises ValueError on any denied action. Combined with the two issues above, the blast radius of a false positive is a blocked job rather than a warning. Please align with the warn-unless-definitive-*-denial behavior.

🟡 Medium — depends on private cross-package helpers

role_permission_validator.py:24-27 imports _simulate_denied_actions and _role_trusts_service — both underscore-private — from sagemaker.core.helper.iam_role_resolver. This is a fragile cross-package coupling to a private API that can change without notice. Consider exposing a supported public entry point in iam_role_resolver (it already has the machinery) rather than importing private symbols.

Minor upside: passing "bedrock" as the 3rd arg to _role_trusts_service happens to work because "bedrock" is a valid role_type key in IAM_POLICY_CONFIG (resolving to bedrock.amazonaws.com). It reads like a service string but is actually a role-type lookup — worth a comment, since it's coincidental.

🟢 Low

  • Redundant AWS round-trips / latency: _get_iam_client (role_permission_validator.py:181) builds a second IAM client and issues additional SimulatePrincipalPolicy + GetRole calls right after resolve_and_validate_role already simulated + checked trust. Consider folding the Bedrock/MLflow checks into the existing single resolution pass. Also, the fallback boto3.Session() bypasses the SDK's config/region defaults that _get_boto_session provides in iam_role_resolver.py:235.
  • Blank-line-only change: llm_as_judge_evaluator.py:755 adds a stray blank line with no functional effect — drop it to keep the diff clean.
  • Test coverage gap: test_bedrock_role_validation.py tests the helper in isolation well, but there's no test asserting the base-class integration — specifically that non-LLM-as-Judge evaluators (InspectAI/Benchmark/MTRL) are not broken by the new gate. Given the High issue above, a regression test there would be valuable. There's also no test for the resource-scoped false-denial scenario.

Summary: the graceful-degradation paths and error messaging are good, but the unconditional placement in BaseEvaluator plus simulating resource-scoped actions without ResourceArns can turn this fail-fast guard into a source of false failures for jobs that would otherwise run. I'd recommend scoping it to the LLM-as-Judge path and gating raises to *-resource actions (or supplying ResourceArns), consistent with iam_role_resolver.py.
· master-eval-role-log-fix

@github-actions

github-actions Bot commented Aug 6, 2026

Copy link
Copy Markdown

Claude finished @rsareddy0329's task in 3m 26s —— View job


Code Review: Pre-validation for Bedrock/MLflow role permissions

I reviewed the diff (/tmp/pr.diff) against origin/master and the surrounding code in iam_role_resolver.py and the evaluator classes. The evaluation role-type addition is well-formed and consistent with the existing config schema, and the new test file is thorough for the resolver behavior it targets. However, there are a few correctness concerns worth addressing before merge.

  • Read PR diff and changed files
  • Analyze iam_role_resolver / iam_policies changes
  • Trace _get_aws_execution_context callers
  • Check existing test impact

🔴 1. Existing tests will now fail — test_base_evaluator.py still asserts role_type="training"

base_evaluator.py now passes role_type="evaluation", but the existing tests still assert the old value and were not updated in this PR:

  • sagemaker-train/tests/unit/train/evaluate/test_base_evaluator.py:737role_type="training"
  • sagemaker-train/tests/unit/train/evaluate/test_base_evaluator.py:764role_type="training"

Both test_get_aws_execution_context and test_get_aws_execution_context_with_explicit_role do mock_role.assert_called_once_with(..., role_type="training", ...), which will break. These need to be updated to "evaluation".

🔴 2. Change is in the shared base class — it forces Bedrock trust/permissions on all evaluators

_get_aws_execution_context() in base_evaluator.py:738 is the common path for every evaluator, not just LLM-as-Judge:

  • benchmark_evaluator.py:673
  • custom_scorer_evaluator.py:453
  • inspect_ai_evaluator.py:727
  • multi_turn_rl_evaluator.py:583
  • llm_as_judge_evaluator.py:815

By switching this shared method to role_type="evaluation", the validation now requires Bedrock permissions and Bedrock trust for jobs that never touch Bedrock (e.g. a custom-scorer or inspect-ai eval). Since only the bedrock:* actions have Resource: "*", they are exactly the ones _get_smoke_test_actions gates on — so a role without Bedrock permissions will hit verdict is False and raise RoleValidationError, blocking a perfectly valid non-Bedrock eval job. Given the PR title/description scope this to LLM-as-Judge, consider parameterizing the role type per-backend rather than flipping the shared base.

🟠 3. Requiring bedrock.amazonaws.com in the trust policy may be wrong for the SMTJ/serverless path

_role_trusts_service uses expected.issubset(trusted) (line 468), so the evaluation type now requires the role to trust both sagemaker.amazonaws.com and bedrock.amazonaws.com. A standard SageMaker execution role trusts only sagemaker.amazonaws.com.

For the serverless/SMTJ backend the eval container runs as the SageMaker execution role and calls Bedrock via InvokeModel using that role's own credentials — that needs the bedrock:InvokeModel permission, but not bedrock.amazonaws.com in the trust policy (trust for the Bedrock service principal is only needed when Bedrock itself assumes the role, as in native Bedrock evaluation jobs). As written, existing execution roles that work today will be rejected with "Could not assume role" guidance that doesn't match the actual failure mode. Please confirm the backend architecture requires Bedrock trust; if not, drop bedrock.amazonaws.com from the trust policy (or make it backend-specific).

Note: in Studio/notebook contexts the caller often lacks iam:SimulatePrincipalPolicy/iam:GetRole, so both checks degrade to None and won't block — which means this regression will surface inconsistently (only for callers who can introspect), making it easy to miss in testing.

🟡 4. Minor

  • llm_as_judge_evaluator.py:753 (diff line 120) adds a stray extra blank line above the @_telemetry_emitter decorator — unrelated noise.
  • New test file is named test_bedrock_role_validation.py and lives under sagemaker-train/tests/..., but it exclusively exercises sagemaker.core.helper.iam_role_resolver. Consider placing it alongside the core resolver tests (and a name like test_evaluation_role_type.py) for discoverability.
  • The evaluation role type has no KMS policy, unlike bedrock. If eval S3 artifacts can be SSE-KMS encrypted, kms:Decrypt/kms:GenerateDataKey may be needed for the container to read/write them. Worth confirming.

Summary: the resolver-side additions look good, but (1) breaks existing unit tests, (2) widens the requirement to all evaluators, and (3) may require Bedrock trust that standard execution roles don't have. I'd recommend scoping the evaluation role type to the LLM-as-Judge/Bedrock path and updating the base-evaluator tests.

def test_evaluation_trust_includes_bedrock(self):
"""Evaluation role type should require bedrock.amazonaws.com trust."""
expected = _expected_trust_services("evaluation")
assert "bedrock.amazonaws.com" in expected
def test_evaluation_trust_includes_sagemaker(self):
"""Evaluation role type should require sagemaker.amazonaws.com trust."""
expected = _expected_trust_services("evaluation")
assert "sagemaker.amazonaws.com" in expected
Add 'evaluation' role type to IAM_POLICY_CONFIG with Bedrock and MLflow
permissions. BaseEvaluator now calls resolve_and_validate_role with
role_type='evaluation' which validates permissions and trust using the
existing iam_role_resolver infrastructure.

Removes the standalone role_permission_validator utility in favor of
the existing centralized approach.
@github-actions

github-actions Bot commented Aug 6, 2026

Copy link
Copy Markdown

Claude finished @rsareddy0329's task in 1m 42s —— View job


Code Review: Pre-validation for Bedrock/MLflow role permissions

  • Read the PR diff (/tmp/pr.diff)
  • Trace _get_aws_execution_context callers and iam_role_resolver internals
  • Check existing test impact
  • Post review

Note: the inline-comment tool isn't available in this environment, so findings are collected below with file/line references.

The evaluation role-type addition is well-formed and consistent with the config schema, and — importantly — the smoke-test design already protects against the resource-scoped false-denial trap: _get_smoke_test_actions (iam_role_resolver.py:109) only gates on Resource: "*" actions, so the MLflow (mlflow-app/*) and sagemaker:CreateTrainingJob (training-job/*) statements you added are correctly excluded from blocking. Good.

That leaves the placement and trust concerns below.


🔴 High — the change is in the shared base class; it forces Bedrock permission + trust on every evaluator

base_evaluator.py:740 switches _get_aws_execution_context() from role_type="training" to role_type="evaluation". That method is the common path for all evaluators, not just LLM-as-Judge:

  • llm_as_judge_evaluator.py:815
  • benchmark_evaluator.py:673
  • custom_scorer_evaluator.py:453
  • inspect_ai_evaluator.py:727
  • multi_turn_rl_evaluator.py:583

Your bedrock_evaluation_policy uses Resource: "*", so bedrock:CreateEvaluationJob / GetEvaluationJob / InvokeModel* are exactly the actions the smoke test does gate on. A role that runs a Benchmark, CustomScorer, InspectAI, or MultiTurnRL (Bedrock AgentCore, not bedrock:CreateEvaluationJob) eval — none of which need bedrock:CreateEvaluationJob — will hit verdict is False at resolve_and_validate_role (iam_role_resolver.py:572) and raise RoleValidationError, hard-failing a job that would otherwise succeed. Given the PR is scoped to LLM-as-Judge, please parameterize the role type per-backend (e.g. resolve "evaluation" only on the LLM-as-Judge/Bedrock path and keep "training" for the others) rather than flipping the shared base method.

🔴 High — existing base-evaluator tests will fail

test_base_evaluator.py:737 and :764 still assert mock_role.assert_called_once_with(..., role_type="training", ...) in test_get_aws_execution_context and test_get_aws_execution_context_with_explicit_role. They were not updated in this PR and will break. (Whatever the resolution to the High issue above, these assertions need to match.)

🟠 Medium — requiring bedrock.amazonaws.com in the trust policy is likely wrong for the SMTJ/serverless backend

_role_trusts_service uses expected.issubset(trusted) (iam_role_resolver.py:468), so the evaluation type now requires the role to trust both sagemaker.amazonaws.com and bedrock.amazonaws.com. A standard SageMaker execution role trusts only sagemaker.amazonaws.com.

For the serverless/SMTJ path the eval container runs as the execution role and calls Bedrock via InvokeModel using that role's own credentials — that needs the bedrock:InvokeModel permission, but not bedrock.amazonaws.com in the trust policy. Bedrock-service trust is only needed when Bedrock itself assumes the role (native CreateEvaluationJob). As written, working execution roles will be rejected with "Could not assume role" guidance that doesn't match the real failure mode. Please confirm whether the backend actually requires Bedrock to assume the role; if not, drop bedrock.amazonaws.com from the trust policy (or make it backend-specific).

This regression surfaces inconsistently: in Studio/notebook contexts the caller often lacks iam:SimulatePrincipalPolicy/iam:GetRole, so both checks degrade to None and won't block — easy to miss in testing.

🟡 Low

  • llm_as_judge_evaluator.py:755 — stray extra blank line above the @_telemetry_emitter decorator; drop it to keep the diff clean.
  • Test file placement/name: test_bedrock_role_validation.py lives under sagemaker-train/tests/... but exclusively exercises sagemaker.core.helper.iam_role_resolver. Consider moving it next to the core resolver tests (e.g. test_evaluation_role_type.py) for discoverability.
  • CodeQL warnings (lines 28, 33): the in-substring assertions on service principals trip "incomplete URL substring sanitization." Assert equality against the returned set instead (assert "bedrock.amazonaws.com" in _expected_trust_services("evaluation") on a set is fine, but the flagged lines compare strings) — prefer set-membership on parsed values.
  • Missing regression test: no test asserts that non-LLM-as-Judge evaluators (Benchmark/InspectAI/CustomScorer/MTRL) are not broken by the new gate. Given the High issue, that regression test is the most valuable one to add.

Summary: the resolver-side additions and the smoke-test gating are sound, but flipping the shared _get_aws_execution_context() to "evaluation" widens Bedrock permission+trust requirements to evaluators that don't touch Bedrock (breaking valid jobs and existing tests), and requiring bedrock.amazonaws.com trust may reject standard execution roles on the serverless/SMTJ path. I'd scope the evaluation role type to the LLM-as-Judge/Bedrock path, reconsider the trust requirement, and update/extend the tests.
· master-eval-role-log-fix

@rsareddy0329
rsareddy0329 deployed to auto-approve August 6, 2026 15:33 — with GitHub Actions Active
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants