Setting the file. One moment.
Check Role · Hf Cloud Sagemaker IAM Preflight · huggingface/skills · Skills Docs
Back to the top of the page Reference
scripts/ check_role.py
Python · 184 lines · 7 KB
16 get-caller-identity` already succeeds.
17
18 Usage:
19 python check_role.py # discover
20 python check_role.py <role-name-or-arn> # validate a specific role
21
22 (On Windows the launcher is usually `python`; on macOS/Linux it's `python3`.)
23
24 Exit: 0 = usable role found (ARN printed to stdout)
25 1 = no usable role
26 2 = AWS CLI / credentials error
27 """
28
29 from __future__ import annotations
30
31 import fnmatch
32 import json
33 import shutil
34 import subprocess
35 import sys
36
37 ROLE_NAME_PATTERNS = [
38 "AmazonSageMaker-ExecutionRole-*" ,
39 "SageMakerExecutionRole*" ,
40 "*SageMaker*Execution*" ,
41 "*sagemaker*execution*" ,
42 ]
43
44
45 def log (msg: str ) -> None :
46 print ( f "[check_role] { msg } " , file = sys.stderr, flush = True )
47
48
49 def aws_bin () -> str :
50 """Resolve the `aws` executable, honoring PATHEXT (finds aws.exe/aws.cmd on Windows)."""
51 exe = shutil.which( "aws" )
52 if not exe:
53 log( "ERROR: the 'aws' CLI was not found on PATH." )
54 log( "Install AWS CLI v2 and confirm `aws sts get-caller-identity` works in this shell." )
55 sys.exit( 2 )
56 return exe
57
58
59 def run_aws (args: list[ str ]) -> subprocess.CompletedProcess:
60 """Run an aws CLI command, capturing output. Never raises on nonzero exit."""
61 return subprocess.run(
62 [aws_bin(), * args],
63 capture_output = True ,
64 text = True ,
65 )
66
67
68 def get_role (role_name: str ) -> dict | None :
69 """Return the Role dict from `aws iam get-role`, or None if it can't be read."""
70 proc = run_aws([ "iam" , "get-role" , "--role-name" , role_name, "--output" , "json" ])
71 if proc.returncode != 0 :
72 return None
73 try :
74 return json.loads(proc.stdout)[ "Role" ]
75 except (json.JSONDecodeError, KeyError ):
76 return None
77
78
79 def trust_allows_sagemaker (role: dict ) -> bool :
80 """True if the trust policy lets sagemaker.amazonaws.com assume the role."""
81 trust = role.get( "AssumeRolePolicyDocument" , {})
82 # Substring check over the serialized doc — matches the original shell helper
83 # and is robust to single/list Principal.Service shapes.
84 return "sagemaker.amazonaws.com" in json.dumps(trust)
85
86
87 def validate_role (role_name: str , quiet: bool = False ) -> str | None :
88 """Return the role ARN if it exists and trusts SageMaker, else None."""
89 role = get_role(role_name)
90 if role is None :
91 if not quiet:
92 log( f "Role ' { role_name } ' does not exist or you cannot describe it." )
93 return None
94 if not trust_allows_sagemaker(role):
95 if not quiet:
96 log( f "Role ' { role_name } ': trust policy does not allow sagemaker.amazonaws.com" )
97 return None
98 return role.get( "Arn" )
99
100
101 def main () -> int :
102 if run_aws([ "sts" , "get-caller-identity" ]).returncode != 0 :
103 log( "ERROR: 'aws sts get-caller-identity' failed. Run hf-cloud-aws-context-discovery first." )
104 log( "Run this helper from the shell where the AWS CLI is configured (e.g. PowerShell on Windows)." )
105 return 2
106
107 # Path 1: validate a user-supplied role
108 if len (sys.argv) >= 2 :
109 supplied = sys.argv[ 1 ]
110 role_name = supplied.rsplit( "/" , 1 )[ - 1 ] if supplied.startswith( "arn:aws:iam::" ) else supplied
111 log( f "Validating: { role_name } " )
112 arn = validate_role(role_name)
113 if arn:
114 log( f "OK: { arn } " )
115 print (arn)
116 return 0
117 return 1
118
119 # Path 2: discover candidates
120 log( "Searching for SageMaker execution roles in the account..." )
121 proc = run_aws([ "iam" , "list-roles" , "--query" , "Roles[*].RoleName" , "--output" , "json" ])
122 if proc.returncode != 0 :
123 log( "Could not list roles (caller likely lacks iam:ListRoles)." )
124 log( "Ask the user for an existing role ARN, or have someone with IAM access run this." )
125 return 1
126 try :
127 all_roles = json.loads(proc.stdout) or []
128 except json.JSONDecodeError:
129 all_roles = []
130
131 candidates = [
132 name
133 for name in all_roles
134 if any (fnmatch.fnmatchcase(name, pat) for pat in ROLE_NAME_PATTERNS )
135 ]
136
137 if not candidates:
138 log( "No matching roles. Options:" )
139 log( " 1. Ask the user for an ARN (role might have an unusual name)" )
140 log( " 2. Create one (see create_role.py, requires iam:CreateRole)" )
141 return 1
142
143 log( f "Found { len (candidates) } candidate(s): { ' ' .join(candidates) } " )
144
145 # Rank by RoleLastUsed (most recent first) — alphabetical order rarely picks
146 # the actively-maintained role in accounts with multiple SageMaker roles.
147 # Fetch each role once and reuse the result for both ranking and validation.
148 log( "Ranking by last-used date (fallback: creation date)..." )
149 roles_by_name: dict[ str , dict ] = {}
150 ranked: list[tuple[ str , str , str ]] = [] # (last_used, create_date, role_name)
151 for name in candidates:
152 role = get_role(name)
153 if role is None :
154 continue
155 roles_by_name[name] = role
156 last_used = (role.get( "RoleLastUsed" ) or {}).get( "LastUsedDate" ) or ""
157 # IAM often reports no RoleLastUsed at all (tracking only covers recent
158 # activity); when every candidate ties at "", newest CreateDate wins.
159 create_date = str (role.get( "CreateDate" ) or "" )
160 # ISO-8601 timestamps sort chronologically as strings; "" (never used)
161 # sorts before any timestamp, so descending order puts it last.
162 ranked.append((last_used, create_date, name))
163
164 ranked.sort( reverse = True )
165
166 log( "Ranking (most recent first):" )
167 for last_used, create_date, name in ranked:
168 when = f "last used { last_used } " if last_used else f "never used, created { create_date or '?' } "
169 log( f " - { name } ( { when } )" )
170
171 for _, _, name in ranked:
172 if trust_allows_sagemaker(roles_by_name[name]):
173 arn = roles_by_name[name].get( "Arn" )
174 log( f "Using: { arn } " )
175 print (arn)
176 return 0
177
178 log( "No candidate passed validation (they exist but lack correct trust policy)." )
179 log( "Fix the trust policy or create a new role (see create_role.py)." )
180 return 1
181
182
183 if __name__ == "__main__" :
184 sys.exit(main())