clarifai 11.1.6__py3-none-any.whl → 11.1.6rc1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (112) hide show
  1. clarifai/__init__.py +1 -1
  2. clarifai/__pycache__/__init__.cpython-310.pyc +0 -0
  3. clarifai/__pycache__/errors.cpython-310.pyc +0 -0
  4. clarifai/__pycache__/versions.cpython-310.pyc +0 -0
  5. clarifai/cli/__main__.py~ +4 -0
  6. clarifai/cli/__pycache__/__init__.cpython-310.pyc +0 -0
  7. clarifai/cli/__pycache__/__main__.cpython-310.pyc +0 -0
  8. clarifai/cli/__pycache__/base.cpython-310.pyc +0 -0
  9. clarifai/cli/__pycache__/compute_cluster.cpython-310.pyc +0 -0
  10. clarifai/cli/__pycache__/deployment.cpython-310.pyc +0 -0
  11. clarifai/cli/__pycache__/model.cpython-310.pyc +0 -0
  12. clarifai/cli/__pycache__/nodepool.cpython-310.pyc +0 -0
  13. clarifai/client/#model_client.py# +430 -0
  14. clarifai/client/__pycache__/__init__.cpython-310.pyc +0 -0
  15. clarifai/client/__pycache__/app.cpython-310.pyc +0 -0
  16. clarifai/client/__pycache__/base.cpython-310.pyc +0 -0
  17. clarifai/client/__pycache__/dataset.cpython-310.pyc +0 -0
  18. clarifai/client/__pycache__/input.cpython-310.pyc +0 -0
  19. clarifai/client/__pycache__/lister.cpython-310.pyc +0 -0
  20. clarifai/client/__pycache__/model.cpython-310.pyc +0 -0
  21. clarifai/client/__pycache__/module.cpython-310.pyc +0 -0
  22. clarifai/client/__pycache__/runner.cpython-310.pyc +0 -0
  23. clarifai/client/__pycache__/search.cpython-310.pyc +0 -0
  24. clarifai/client/__pycache__/user.cpython-310.pyc +0 -0
  25. clarifai/client/__pycache__/workflow.cpython-310.pyc +0 -0
  26. clarifai/client/auth/__pycache__/__init__.cpython-310.pyc +0 -0
  27. clarifai/client/auth/__pycache__/helper.cpython-310.pyc +0 -0
  28. clarifai/client/auth/__pycache__/register.cpython-310.pyc +0 -0
  29. clarifai/client/auth/__pycache__/stub.cpython-310.pyc +0 -0
  30. clarifai/client/model_client.py +447 -0
  31. clarifai/constants/__pycache__/dataset.cpython-310.pyc +0 -0
  32. clarifai/constants/__pycache__/model.cpython-310.pyc +0 -0
  33. clarifai/constants/__pycache__/search.cpython-310.pyc +0 -0
  34. clarifai/datasets/__pycache__/__init__.cpython-310.pyc +0 -0
  35. clarifai/datasets/export/__pycache__/__init__.cpython-310.pyc +0 -0
  36. clarifai/datasets/export/__pycache__/inputs_annotations.cpython-310.pyc +0 -0
  37. clarifai/datasets/upload/__pycache__/__init__.cpython-310.pyc +0 -0
  38. clarifai/datasets/upload/__pycache__/base.cpython-310.pyc +0 -0
  39. clarifai/datasets/upload/__pycache__/features.cpython-310.pyc +0 -0
  40. clarifai/datasets/upload/__pycache__/image.cpython-310.pyc +0 -0
  41. clarifai/datasets/upload/__pycache__/text.cpython-310.pyc +0 -0
  42. clarifai/datasets/upload/__pycache__/utils.cpython-310.pyc +0 -0
  43. clarifai/datasets/upload/loaders/__pycache__/__init__.cpython-310.pyc +0 -0
  44. clarifai/datasets/upload/loaders/__pycache__/coco_detection.cpython-310.pyc +0 -0
  45. clarifai/models/__pycache__/__init__.cpython-310.pyc +0 -0
  46. clarifai/models/model_serving/__pycache__/__init__.cpython-310.pyc +0 -0
  47. clarifai/models/model_serving/__pycache__/constants.cpython-310.pyc +0 -0
  48. clarifai/models/model_serving/cli/__pycache__/__init__.cpython-310.pyc +0 -0
  49. clarifai/models/model_serving/cli/__pycache__/_utils.cpython-310.pyc +0 -0
  50. clarifai/models/model_serving/cli/__pycache__/base.cpython-310.pyc +0 -0
  51. clarifai/models/model_serving/cli/__pycache__/build.cpython-310.pyc +0 -0
  52. clarifai/models/model_serving/cli/__pycache__/create.cpython-310.pyc +0 -0
  53. clarifai/models/model_serving/model_config/__pycache__/__init__.cpython-310.pyc +0 -0
  54. clarifai/models/model_serving/model_config/__pycache__/base.cpython-310.pyc +0 -0
  55. clarifai/models/model_serving/model_config/__pycache__/config.cpython-310.pyc +0 -0
  56. clarifai/models/model_serving/model_config/__pycache__/inference_parameter.cpython-310.pyc +0 -0
  57. clarifai/models/model_serving/model_config/__pycache__/output.cpython-310.pyc +0 -0
  58. clarifai/models/model_serving/model_config/triton/__pycache__/__init__.cpython-310.pyc +0 -0
  59. clarifai/models/model_serving/model_config/triton/__pycache__/serializer.cpython-310.pyc +0 -0
  60. clarifai/models/model_serving/model_config/triton/__pycache__/triton_config.cpython-310.pyc +0 -0
  61. clarifai/models/model_serving/model_config/triton/__pycache__/wrappers.cpython-310.pyc +0 -0
  62. clarifai/models/model_serving/repo_build/__pycache__/__init__.cpython-310.pyc +0 -0
  63. clarifai/models/model_serving/repo_build/__pycache__/build.cpython-310.pyc +0 -0
  64. clarifai/models/model_serving/repo_build/static_files/__pycache__/base_test.cpython-310-pytest-7.2.0.pyc +0 -0
  65. clarifai/rag/__pycache__/__init__.cpython-310.pyc +0 -0
  66. clarifai/rag/__pycache__/rag.cpython-310.pyc +0 -0
  67. clarifai/rag/__pycache__/utils.cpython-310.pyc +0 -0
  68. clarifai/runners/__pycache__/__init__.cpython-310.pyc +0 -0
  69. clarifai/runners/__pycache__/server.cpython-310.pyc +0 -0
  70. clarifai/runners/dockerfile_template/Dockerfile.debug +11 -0
  71. clarifai/runners/dockerfile_template/Dockerfile.debug~ +9 -0
  72. clarifai/runners/models/__pycache__/__init__.cpython-310.pyc +0 -0
  73. clarifai/runners/models/__pycache__/base_typed_model.cpython-310.pyc +0 -0
  74. clarifai/runners/models/__pycache__/model_builder.cpython-310.pyc +0 -0
  75. clarifai/runners/models/__pycache__/model_class.cpython-310.pyc +0 -0
  76. clarifai/runners/models/__pycache__/model_run_locally.cpython-310.pyc +0 -0
  77. clarifai/runners/models/__pycache__/model_runner.cpython-310.pyc +0 -0
  78. clarifai/runners/models/__pycache__/model_servicer.cpython-310.pyc +0 -0
  79. clarifai/runners/models/__pycache__/model_upload.cpython-310.pyc +0 -0
  80. clarifai/runners/utils/__pycache__/__init__.cpython-310.pyc +0 -0
  81. clarifai/runners/utils/__pycache__/const.cpython-310.pyc +0 -0
  82. clarifai/runners/utils/__pycache__/data_handler.cpython-310.pyc +0 -0
  83. clarifai/runners/utils/__pycache__/data_types.cpython-310.pyc +0 -0
  84. clarifai/runners/utils/__pycache__/data_utils.cpython-310.pyc +0 -0
  85. clarifai/runners/utils/__pycache__/loader.cpython-310.pyc +0 -0
  86. clarifai/runners/utils/__pycache__/logging.cpython-310.pyc +0 -0
  87. clarifai/runners/utils/__pycache__/method_signatures.cpython-310.pyc +0 -0
  88. clarifai/runners/utils/__pycache__/serializers.cpython-310.pyc +0 -0
  89. clarifai/runners/utils/__pycache__/url_fetcher.cpython-310.pyc +0 -0
  90. clarifai/runners/utils/data_types.py +386 -0
  91. clarifai/runners/utils/method_signatures.py +435 -0
  92. clarifai/runners/utils/serializers.py +208 -0
  93. clarifai/schema/__pycache__/search.cpython-310.pyc +0 -0
  94. clarifai/urls/__pycache__/helper.cpython-310.pyc +0 -0
  95. clarifai/utils/__pycache__/__init__.cpython-310.pyc +0 -0
  96. clarifai/utils/__pycache__/logging.cpython-310.pyc +0 -0
  97. clarifai/utils/__pycache__/misc.cpython-310.pyc +0 -0
  98. clarifai/utils/__pycache__/model_train.cpython-310.pyc +0 -0
  99. clarifai/utils/evaluation/__pycache__/__init__.cpython-310.pyc +0 -0
  100. clarifai/utils/evaluation/__pycache__/helpers.cpython-310.pyc +0 -0
  101. clarifai/utils/evaluation/__pycache__/main.cpython-310.pyc +0 -0
  102. clarifai/workflows/__pycache__/__init__.cpython-310.pyc +0 -0
  103. clarifai/workflows/__pycache__/export.cpython-310.pyc +0 -0
  104. clarifai/workflows/__pycache__/utils.cpython-310.pyc +0 -0
  105. clarifai/workflows/__pycache__/validate.cpython-310.pyc +0 -0
  106. {clarifai-11.1.6.dist-info → clarifai-11.1.6rc1.dist-info}/METADATA +15 -26
  107. clarifai-11.1.6rc1.dist-info/RECORD +205 -0
  108. {clarifai-11.1.6.dist-info → clarifai-11.1.6rc1.dist-info}/WHEEL +1 -1
  109. clarifai-11.1.6.dist-info/RECORD +0 -101
  110. {clarifai-11.1.6.dist-info → clarifai-11.1.6rc1.dist-info}/LICENSE +0 -0
  111. {clarifai-11.1.6.dist-info → clarifai-11.1.6rc1.dist-info}/entry_points.txt +0 -0
  112. {clarifai-11.1.6.dist-info → clarifai-11.1.6rc1.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,447 @@
1
+ import inspect
2
+ import json
3
+ import time
4
+ from typing import Any, Dict, Iterator, List
5
+
6
+ from clarifai_grpc.grpc.api import resources_pb2, service_pb2
7
+ from clarifai_grpc.grpc.api.status import status_code_pb2
8
+
9
+ from clarifai.constants.model import MAX_MODEL_PREDICT_INPUTS
10
+ from clarifai.errors import UserError
11
+ from clarifai.runners.utils.method_signatures import (CompatibilitySerializer, deserialize,
12
+ get_stream_from_signature, serialize,
13
+ signatures_from_json)
14
+ from clarifai.utils.logging import logger
15
+ from clarifai.utils.misc import BackoffIterator, status_is_retryable
16
+
17
+
18
+ class ModelClient:
19
+ '''
20
+ Client for calling model predict, generate, and stream methods.
21
+ '''
22
+
23
+ def __init__(self, stub, request_template: service_pb2.PostModelOutputsRequest = None):
24
+ '''
25
+ Initialize the model client.
26
+
27
+ Args:
28
+ stub: The gRPC stub for the model.
29
+ request_template: The template for the request to send to the model, including
30
+ common fields like model_id, model_version, cluster, etc.
31
+ '''
32
+ self.STUB = stub
33
+ self.request_template = request_template or service_pb2.PostModelOutputsRequest()
34
+ self._method_signatures = None
35
+ self._defined = False
36
+
37
+ def fetch(self):
38
+ '''
39
+ Fetch function signature definitions from the model and define the functions in the client
40
+ '''
41
+ if self._defined:
42
+ return
43
+ try:
44
+ self._fetch_signatures()
45
+ self._define_functions()
46
+ finally:
47
+ self._defined = True
48
+
49
+ def __getattr__(self, name):
50
+ if not self._defined:
51
+ self.fetch()
52
+ return self.__getattribute__(name)
53
+
54
+ def _fetch_signatures(self):
55
+ '''
56
+ Fetch the method signatures from the model.
57
+
58
+ Returns:
59
+ Dict: The method signatures.
60
+ '''
61
+ #request = resources_pb2.GetModelSignaturesRequest()
62
+ #response = self.stub.GetModelSignatures(request)
63
+ #self._method_signatures = json.loads(response.signatures) # or define protos
64
+ # TODO this could use a new endpoint to get the signatures
65
+ # for local grpc models, we'll also have to add the endpoint to the model servicer
66
+ # for now we'll just use the predict endpoint with a special method name
67
+
68
+ request = service_pb2.PostModelOutputsRequest()
69
+ request.CopyFrom(self.request_template)
70
+ # request.model.model_version.output_info.params['_method_name'] = '_GET_SIGNATURES'
71
+ inp = request.inputs.add() # empty input for this method
72
+ inp.data.parts.add() # empty part for this input
73
+ inp.data.metadata['_method_name'] = '_GET_SIGNATURES'
74
+ start_time = time.time()
75
+ backoff_iterator = BackoffIterator(10)
76
+ while True:
77
+ response = self.STUB.PostModelOutputs(request)
78
+ if status_is_retryable(
79
+ response.status.code) and time.time() - start_time < 60 * 10: # 10 minutes
80
+ logger.info(f"Retrying model info fetch with response {response.status!r}")
81
+ time.sleep(next(backoff_iterator))
82
+ continue
83
+ break
84
+ if response.status.code == status_code_pb2.INPUT_UNSUPPORTED_FORMAT:
85
+ # return code from older models that don't support _GET_SIGNATURES
86
+ self._method_signatures = {}
87
+ self._define_compatability_functions()
88
+ return
89
+ if response.status.code != status_code_pb2.SUCCESS:
90
+ raise Exception(f"Model failed with response {response!r}")
91
+ self._method_signatures = signatures_from_json(response.outputs[0].data.text.raw)
92
+
93
+ def _define_functions(self):
94
+ '''
95
+ Define the functions based on the method signatures.
96
+ '''
97
+ for method_name, method_signature in self._method_signatures.items():
98
+ # define the function in this client instance
99
+ if method_signature.method_type == 'predict':
100
+ call_func = self._predict
101
+ elif method_signature.method_type == 'generate':
102
+ call_func = self._generate
103
+ elif method_signature.method_type == 'stream':
104
+ call_func = self._stream
105
+ else:
106
+ raise ValueError(f"Unknown method type {method_signature.method_type}")
107
+
108
+ # method argnames, in order, collapsing nested keys to corresponding user function args
109
+ method_argnames = []
110
+ for var in method_signature.inputs:
111
+ outer = var.name.split('.', 1)[0]
112
+ if outer in method_argnames:
113
+ continue
114
+ method_argnames.append(outer)
115
+
116
+ def bind_f(method_name, method_argnames, call_func):
117
+
118
+ def f(*args, **kwargs):
119
+ if len(args) > len(method_argnames):
120
+ raise TypeError(
121
+ f"{method_name}() takes {len(method_argnames)} positional arguments but {len(args)} were given"
122
+ )
123
+ for name, arg in zip(method_argnames, args): # handle positional with zip shortest
124
+ if name in kwargs:
125
+ raise TypeError(f"Multiple values for argument {name}")
126
+ kwargs[name] = arg
127
+ return call_func(kwargs, method_name)
128
+
129
+ return f
130
+
131
+ # need to bind method_name to the value, not the mutating loop variable
132
+ f = bind_f(method_name, method_argnames, call_func)
133
+
134
+ # set names, annotations and docstrings
135
+ f.__name__ = method_name
136
+ f.__qualname__ = f'{self.__class__.__name__}.{method_name}'
137
+ f.__doc__ = method_signature.docstring
138
+ input_annotations = json.loads(method_signature.annotations_json)
139
+ return_annotation = input_annotations.pop('return', None)
140
+ sig = inspect.signature(f).replace(
141
+ parameters=[
142
+ inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=v)
143
+ for k, v in input_annotations.items()
144
+ ],
145
+ return_annotation=return_annotation,
146
+ )
147
+ f.__signature__ = sig
148
+ f.__doc__ = method_signature.docstring
149
+ setattr(self, method_name, f)
150
+
151
+ def _define_compatability_functions(self):
152
+
153
+ serializer = CompatibilitySerializer()
154
+
155
+ def predict(input: Any) -> Any:
156
+ proto = resources_pb2.Input()
157
+ serializer.serialize(proto.data, input)
158
+ # always use text.raw for compat
159
+ if proto.data.string_value:
160
+ proto.data.text.raw = proto.data.string_value
161
+ proto.data.string_value = ''
162
+ response = self._predict_by_proto([proto])
163
+ if response.status.code != status_code_pb2.SUCCESS:
164
+ raise Exception(f"Model predict failed with response {response!r}")
165
+ response_data = response.outputs[0].data
166
+ if response_data.text.raw:
167
+ response_data.string_value = response_data.text.raw
168
+ response_data.text.raw = ''
169
+ return serializer.deserialize(response_data)
170
+
171
+ self.predict = predict
172
+
173
+ def _predict(
174
+ self,
175
+ inputs, # TODO set up functions according to fetched signatures?
176
+ method_name: str = 'predict',
177
+ ) -> Any:
178
+ input_signature = self._method_signatures[method_name].inputs
179
+ output_signature = self._method_signatures[method_name].outputs
180
+
181
+ batch_input = True
182
+ if isinstance(inputs, dict):
183
+ inputs = [inputs]
184
+ batch_input = False
185
+
186
+ proto_inputs = []
187
+ for input in inputs:
188
+ proto = resources_pb2.Input()
189
+ serialize(input, input_signature, proto.data)
190
+ proto_inputs.append(proto)
191
+
192
+ response = self._predict_by_proto(proto_inputs, method_name)
193
+ #print(response)
194
+
195
+ outputs = []
196
+ for output in response.outputs:
197
+ outputs.append(deserialize(output.data, output_signature, is_output=True))
198
+ if batch_input:
199
+ return outputs
200
+ return outputs[0]
201
+
202
+ def _predict_by_proto(
203
+ self,
204
+ inputs: List[resources_pb2.Input],
205
+ method_name: str = None,
206
+ inference_params: Dict = None,
207
+ output_config: Dict = None,
208
+ ) -> service_pb2.MultiOutputResponse:
209
+ """Predicts the model based on the given inputs.
210
+
211
+ Args:
212
+ inputs (List[resources_pb2.Input]): The inputs to predict.
213
+ method_name (str): The remote method name to call.
214
+ inference_params (Dict): Inference parameters to override.
215
+ output_config (Dict): Output configuration to override.
216
+
217
+ Returns:
218
+ service_pb2.MultiOutputResponse: The prediction response(s).
219
+ """
220
+ if not isinstance(inputs, list):
221
+ raise UserError('Invalid inputs, inputs must be a list of Input objects.')
222
+ if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
223
+ raise UserError(f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}.")
224
+
225
+ request = service_pb2.PostModelOutputsRequest()
226
+ request.CopyFrom(self.request_template)
227
+
228
+ request.inputs.extend(inputs)
229
+
230
+ if method_name:
231
+ # TODO put in new proto field?
232
+ for inp in request.inputs:
233
+ inp.data.metadata['_method_name'] = method_name
234
+ if inference_params:
235
+ request.model.model_version.output_info.params.update(inference_params)
236
+ if output_config:
237
+ request.model.model_version.output_info.output_config.MergeFrom(
238
+ resources_pb2.OutputConfig(**output_config))
239
+
240
+ start_time = time.time()
241
+ backoff_iterator = BackoffIterator(10)
242
+ while True:
243
+ response = self.STUB.PostModelOutputs(request)
244
+ if status_is_retryable(
245
+ response.status.code) and time.time() - start_time < 60 * 10: # 10 minutes
246
+ logger.info(f"Model predict failed with response {response!r}")
247
+ time.sleep(next(backoff_iterator))
248
+ continue
249
+
250
+ if response.status.code != status_code_pb2.SUCCESS:
251
+ raise Exception(f"Model predict failed with response {response!r}")
252
+ break
253
+
254
+ return response
255
+
256
+ def _generate(
257
+ self,
258
+ inputs, # TODO set up functions according to fetched signatures?
259
+ method_name: str = 'generate',
260
+ ) -> Any:
261
+ input_signature = self._method_signatures[method_name].inputs
262
+ output_signature = self._method_signatures[method_name].outputs
263
+
264
+ batch_input = True
265
+ if isinstance(inputs, dict):
266
+ inputs = [inputs]
267
+ batch_input = False
268
+
269
+ proto_inputs = []
270
+ for input in inputs:
271
+ proto = resources_pb2.Input()
272
+ serialize(input, input_signature, proto.data)
273
+ proto_inputs.append(proto)
274
+
275
+ response_stream = self._generate_by_proto(proto_inputs, method_name)
276
+ #print(response)
277
+
278
+ for response in response_stream:
279
+ outputs = []
280
+ for output in response.outputs:
281
+ outputs.append(deserialize(output.data, output_signature, is_output=True))
282
+ if batch_input:
283
+ yield outputs
284
+ yield outputs[0]
285
+
286
+ def _generate_by_proto(
287
+ self,
288
+ inputs: List[resources_pb2.Input],
289
+ method_name: str = None,
290
+ inference_params: Dict = {},
291
+ output_config: Dict = {},
292
+ ):
293
+ """Generate the stream output on model based on the given inputs.
294
+
295
+ Args:
296
+ inputs (list[Input]): The inputs to generate, must be less than 128.
297
+ method_name (str): The remote method name to call.
298
+ inference_params (dict): The inference params to override.
299
+ output_config (dict): The output config to override.
300
+ """
301
+ if not isinstance(inputs, list):
302
+ raise UserError('Invalid inputs, inputs must be a list of Input objects.')
303
+ if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
304
+ raise UserError(f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}."
305
+ ) # TODO Use Chunker for inputs len > 128
306
+
307
+ request = service_pb2.PostModelOutputsRequest()
308
+ request.CopyFrom(self.request_template)
309
+
310
+ request.inputs.extend(inputs)
311
+
312
+ if method_name:
313
+ # TODO put in new proto field?
314
+ for inp in request.inputs:
315
+ inp.data.metadata['_method_name'] = method_name
316
+ if inference_params:
317
+ request.model.model_version.output_info.params.update(inference_params)
318
+ if output_config:
319
+ request.model.model_version.output_info.output_config.MergeFromDict(output_config)
320
+
321
+ start_time = time.time()
322
+ backoff_iterator = BackoffIterator(10)
323
+ started = False
324
+ while not started:
325
+ stream_response = self.STUB.GenerateModelOutputs(request)
326
+ try:
327
+ response = next(stream_response) # get the first response
328
+ except StopIteration:
329
+ raise Exception("Model Generate failed with no response")
330
+ if status_is_retryable(response.status.code) and \
331
+ time.time() - start_time < 60 * 10:
332
+ logger.info("Model is still deploying, please wait...")
333
+ time.sleep(next(backoff_iterator))
334
+ continue
335
+ if response.status.code != status_code_pb2.SUCCESS:
336
+ raise Exception(f"Model Generate failed with response {response.status!r}")
337
+ started = True
338
+
339
+ yield response # yield the first response
340
+
341
+ for response in stream_response:
342
+ if response.status.code != status_code_pb2.SUCCESS:
343
+ raise Exception(f"Model Generate failed with response {response.status!r}")
344
+ yield response
345
+
346
+ def _stream(
347
+ self,
348
+ inputs,
349
+ method_name: str = 'stream',
350
+ ) -> Any:
351
+ input_signature = self._method_signatures[method_name].inputs
352
+ output_signature = self._method_signatures[method_name].outputs
353
+
354
+ if isinstance(inputs, list):
355
+ assert len(inputs) == 1, 'streaming methods do not support batched calls'
356
+ inputs = inputs[0]
357
+ assert isinstance(inputs, dict)
358
+ kwargs = inputs
359
+
360
+ # find the streaming vars in the input signature, and the streaming input python param
361
+ stream_sig = get_stream_from_signature(input_signature)
362
+ if stream_sig is None:
363
+ raise ValueError("Streaming method must have a Stream input")
364
+ stream_argname = stream_sig.name
365
+
366
+ # get the streaming input generator from the user-provided function arg values
367
+ user_inputs_generator = kwargs.pop(stream_argname)
368
+
369
+ def _input_proto_stream():
370
+ # first item contains all the inputs and the first stream item
371
+ proto = resources_pb2.Input()
372
+ try:
373
+ item = next(user_inputs_generator)
374
+ except StopIteration:
375
+ return # no items to stream
376
+ kwargs[stream_argname] = item
377
+ serialize(kwargs, input_signature, proto.data)
378
+
379
+ yield proto
380
+
381
+ # subsequent items are just the stream items
382
+ for item in user_inputs_generator:
383
+ proto = resources_pb2.Input()
384
+ serialize({stream_argname: item}, [stream_sig], proto.data)
385
+ yield proto
386
+
387
+ response_stream = self._stream_by_proto(_input_proto_stream(), method_name)
388
+ #print(response)
389
+
390
+ for response in response_stream:
391
+ assert len(response.outputs) == 1, 'streaming methods must have exactly one output'
392
+ yield deserialize(response.outputs[0].data, output_signature, is_output=True)
393
+
394
+ def _req_iterator(self,
395
+ input_iterator: Iterator[List[resources_pb2.Input]],
396
+ method_name: str = None,
397
+ inference_params: Dict = {},
398
+ output_config: Dict = {}):
399
+ request = service_pb2.PostModelOutputsRequest()
400
+ request.CopyFrom(self.request_template)
401
+ if inference_params:
402
+ request.model.model_version.output_info.params.update(inference_params)
403
+ if output_config:
404
+ request.model.model_version.output_info.output_config.MergeFromDict(output_config)
405
+ for inputs in input_iterator:
406
+ req = service_pb2.PostModelOutputsRequest()
407
+ req.CopyFrom(request)
408
+ if isinstance(inputs, list):
409
+ req.inputs.extend(inputs)
410
+ else:
411
+ req.inputs.append(inputs)
412
+ # TODO: put into new proto field?
413
+ for inp in req.inputs:
414
+ inp.data.metadata['_method_name'] = method_name
415
+ yield req
416
+
417
+ def _stream_by_proto(self,
418
+ inputs: Iterator[List[resources_pb2.Input]],
419
+ method_name: str = None,
420
+ inference_params: Dict = {},
421
+ output_config: Dict = {}):
422
+ """Generate the stream output on model based on the given stream of inputs.
423
+ """
424
+ # if not isinstance(inputs, Iterator[List[Input]]):
425
+ # raise UserError('Invalid inputs, inputs must be a iterator of list of Input objects.')
426
+
427
+ request = self._req_iterator(inputs, method_name, inference_params, output_config)
428
+
429
+ start_time = time.time()
430
+ backoff_iterator = BackoffIterator(10)
431
+ generation_started = False
432
+ while True:
433
+ if generation_started:
434
+ break
435
+ stream_response = self.STUB.StreamModelOutputs(request)
436
+ for response in stream_response:
437
+ if status_is_retryable(response.status.code) and \
438
+ time.time() - start_time < 60 * 10:
439
+ logger.info("Model is still deploying, please wait...")
440
+ time.sleep(next(backoff_iterator))
441
+ break
442
+ if response.status.code != status_code_pb2.SUCCESS:
443
+ raise Exception(f"Model Predict failed with response {response.status!r}")
444
+ else:
445
+ if not generation_started:
446
+ generation_started = True
447
+ yield response
@@ -0,0 +1,11 @@
1
+ FROM --platform=$TARGETPLATFORM public.ecr.aws/docker/library/python:3.12-slim-bookworm as pybase
2
+
3
+ #############################
4
+ # Final runtime image
5
+ #############################
6
+ FROM --platform=$TARGETPLATFORM ${RUNTIME_IMAGE} as final
7
+
8
+ COPY --from=pybase --link=true /usr/bin/ls /usr/bin/cat /usr/bin/which /usr/bin/bash /usr/bin/sort /usr/bin/du /usr/bin/
9
+ COPY --from=pybase --link=true /bin/rbash /bin/sh /bin/rm /bin/
10
+ COPY --from=pybase --link=true /lib/*-linux-gnu/libselinux.so.1 /lib/*-linux-gnu/libpcre2-8.so.0 /lib/x86_64-linux-gnu/
11
+ COPY --from=pybase --link=true /lib/*-linux-gnu/libselinux.so.1 /lib/*-linux-gnu/libpcre2-8.so.0 /lib/aarch64-linux-gnu/
@@ -0,0 +1,9 @@
1
+ FROM --platform=$TARGETPLATFORM public.ecr.aws/docker/library/python:3.12-slim-bookworm as pybase
2
+
3
+ #############################
4
+ # Final runtime image
5
+ #############################
6
+ FROM --platform=$TARGETPLATFORM ${RUNTIME_IMAGE} as final
7
+
8
+ COPY --from=pybase --link=true /usr/bin/ls /usr/bin/cat /usr/bin/which /usr/bin/bash /usr/bin/sort /usr/bin/du /usr/bin/
9
+ COPY --from=pybase --link=true /bin/rbash /bin/sh /bin/rm /bin/