Setting the file. One moment.
Invoke Endpoint · Hf Cloud Sagemaker Production Defaults · huggingface/skills · Skills Docs
Back to the top of the page scripts/ invoke_endpoint.py
Python · 180 lines · 6 KB
16
python invoke_endpoint.py --endpoint-name NAME --payload-file body.json
17 python invoke_endpoint.py --endpoint-name NAME --payload-file img.json \\
18 --content-type application/json --region eu-west-1
19
20 Inference-component endpoints (including scale-to-zero deployments) need the
21 component name, and a cold endpoint needs a wait:
22
23 python invoke_endpoint.py --endpoint-name NAME \\
24 --inference-component-name COMP --payload '{"prompt": "hi"}' \\
25 --wait-for-capacity 900
26
27 A component scaled to zero copies answers the first request with a 400
28 ValidationError ("has no capacity"). That error is also the wake signal: it
29 publishes NoCapacityInvocationFailures, which triggers the step-scaling policy.
30 With --wait-for-capacity the helper keeps retrying until a copy is in service.
31
32 The endpoint response body is printed to stdout.
33 """
34
35 from __future__ import annotations
36
37 import argparse
38 import json
39 import os
40 import shutil
41 import subprocess
42 import sys
43 import tempfile
44 import time
45
46 NO_CAPACITY_MARKERS = ( "has no capacity" , "NoCapacity" )
47
48
49 def log (msg: str ) -> None :
50 print ( f "[invoke] { msg } " , file = sys.stderr, flush = True )
51
52
53 def aws_bin () -> str :
54 exe = shutil.which( "aws" )
55 if not exe:
56 log( "ERROR: the 'aws' CLI was not found on PATH. Install AWS CLI v2." )
57 sys.exit( 2 )
58 return exe
59
60
61 def resolve_region (arg_region: str | None ) -> str :
62 if arg_region:
63 return arg_region
64 for var in ( "AWS_REGION" , "AWS_DEFAULT_REGION" ):
65 if os.environ.get(var):
66 return os.environ[var]
67 proc = subprocess.run(
68 [aws_bin(), "configure" , "get" , "region" ], capture_output = True , text = True
69 )
70 return proc.stdout.strip() if proc.returncode == 0 else ""
71
72
73 def load_payload (args: argparse.Namespace) -> str :
74 """Return the raw payload text, BOM stripped. Validates JSON when applicable."""
75 if args.payload_file:
76 # utf-8-sig decodes and discards a leading BOM if the file has one.
77 raw = open (args.payload_file, "r" , encoding = "utf-8-sig" ).read()
78 else :
79 raw = args.payload
80
81 # If it's meant to be JSON, validate and re-serialize so the body is clean.
82 if "json" in args.content_type:
83 try :
84 return json.dumps(json.loads(raw))
85 except json.JSONDecodeError as e:
86 log( f "ERROR: payload is not valid JSON: { e } " )
87 sys.exit( 1 )
88 return raw
89
90
91 def main () -> int :
92 parser = argparse.ArgumentParser( description = __doc__ )
93 parser.add_argument( "--endpoint-name" , required = True )
94 parser.add_argument( "--region" )
95 parser.add_argument( "--content-type" , default = "application/json" )
96 parser.add_argument(
97 "--inference-component-name" ,
98 help = "Required when the endpoint hosts inference components" ,
99 )
100 parser.add_argument(
101 "--wait-for-capacity" , type = int , default = 0 , metavar = "SECONDS" ,
102 help = (
103 "Retry while the component reports no capacity (scale-to-zero wake). "
104 "0 fails immediately."
105 ),
106 )
107 src = parser.add_mutually_exclusive_group( required = True )
108 src.add_argument( "--payload" , help = "Inline request body (e.g. a JSON string)" )
109 src.add_argument( "--payload-file" , help = "Path to a file containing the request body" )
110 args = parser.parse_args()
111
112 region = resolve_region(args.region)
113 if not region:
114 log( "ERROR: no AWS region. Pass --region or set AWS_REGION." )
115 return 1
116
117 body = load_payload(args)
118
119 # Write the request body as BOM-free UTF-8. NamedTemporaryFile + explicit
120 # utf-8 encoding guarantees no BOM regardless of platform.
121 body_path = None
122 out_path = None
123 try :
124 with tempfile.NamedTemporaryFile(
125 "w" , suffix = ".json" , encoding = "utf-8" , delete = False
126 ) as f:
127 f.write(body)
128 body_path = f.name
129 out_fd, out_path = tempfile.mkstemp( suffix = ".out" )
130 os.close(out_fd)
131
132 cmd = [
133 aws_bin(), "sagemaker-runtime" , "invoke-endpoint" ,
134 "--endpoint-name" , args.endpoint_name,
135 "--content-type" , args.content_type,
136 "--body" , f "fileb:// { body_path } " ,
137 "--region" , region,
138 ]
139 if args.inference_component_name:
140 cmd += [ "--inference-component-name" , args.inference_component_name]
141 cmd.append(out_path)
142
143 log( f "Invoking { args.endpoint_name } in { region } ..." )
144 deadline = time.time() + args.wait_for_capacity
145 attempt = 0
146 while True :
147 attempt += 1
148 proc = subprocess.run(cmd, capture_output = True , text = True )
149 if proc.returncode == 0 :
150 break
151
152 err = proc.stderr.strip()
153 no_capacity = any (marker in err for marker in NO_CAPACITY_MARKERS )
154 if no_capacity and time.time() < deadline:
155 if attempt == 1 :
156 log( "No capacity: the component is at zero copies. "
157 "This request triggered the wake alarm; retrying." )
158 log( f " attempt { attempt } : still no capacity, waiting 30s "
159 f "( { int (deadline - time.time()) } s left)" )
160 time.sleep( 30 )
161 continue
162
163 log( "Invocation failed:" )
164 log(err)
165 if no_capacity:
166 log( "HINT: the component has no copies in service. Pass "
167 "--wait-for-capacity 900 to wait for the wake-from-zero policy." )
168 return 1
169
170 # The response body landed in out_path; print it to stdout.
171 print ( open (out_path, "r" , encoding = "utf-8-sig" ).read())
172 return 0
173 finally :
174 for p in (body_path, out_path):
175 if p and os.path.exists(p):
176 os.remove(p)
177
178
179 if __name__ == "__main__" :
180 sys.exit(main())