def test_encode():
"""
Test encoding of a message into 'RecordIO' format.
"""
try:
encoder = recordio.Encoder(lambda s: bytes(json.dumps(s), "UTF-8"))
except Exception as exception:
raise MesosException(
"Error instantiating 'RecordIO' encoder: {error}".format(
error=exception))
try:
message = {
"type": "ATTACH_CONTAINER_OUTPUT",
"containerId": "123456789"
}
encoded = encoder.encode(message)
except Exception as exception:
raise MesosException(
"Error encoding 'RecordIO' message: {error}".format(
error=exception))
string = json.dumps(message)
assert encoded == bytes(str(len(string)) + "\n" + string, "UTF-8")
def decode(self, data):
"""
Decode a 'RecordIO' formatted message to its original type.
:param data: an array of 'UTF-8' encoded bytes that make up a
partial 'RecordIO' message. Subsequent calls to this
function maintain state to build up a full 'RecordIO'
message and decode it
:type data: bytes
:returns: a list of deserialized messages
:rtype: list
"""
if not isinstance(data, bytes):
raise MesosException("Parameter 'data' must of of type 'bytes'")
if self.state == self.FAILED:
raise MesosException("Decoder is in a FAILED state")
records = []
for c in data:
if self.state == self.HEADER:
if c != ord('\n'):
self.buffer += bytes([c])
continue
try:
self.length = int(self.buffer.decode("UTF-8"))
except Exception as exception:
self.state = self.FAILED
raise MesosException("Failed to decode length"
"'{buffer}': {error}".format(
buffer=self.buffer,
error=exception))
self.buffer = bytes("", "UTF-8")
self.state = self.RECORD
# Note that for 0 length records, we immediately decode.
if self.length <= 0:
records.append(self.deserialize(self.buffer))
self.state = self.HEADER
elif self.state == self.RECORD:
assert self.length
assert len(self.buffer) < self.length
self.buffer += bytes([c])
if len(self.buffer) == self.length:
records.append(self.deserialize(self.buffer))
self.buffer = bytes("", "UTF-8")
self.state = self.HEADER
return records
def test_encode_decode():
"""
Test encoding/decoding of a message and records into 'RecordIO' format.
"""
total_messages = 10
try:
encoder = recordio.Encoder(lambda s: bytes(json.dumps(s), "UTF-8"))
except Exception as exception:
raise MesosException(
"Error instantiating 'RecordIO' encoder: {error}".format(
error=exception))
try:
decoder = recordio.Decoder(lambda s: json.loads(s.decode("UTF-8")))
except Exception as exception:
raise MesosException(
"Error instantiating 'RecordIO' decoder: {error}".format(
error=exception))
try:
message = {
"type": "ATTACH_CONTAINER_OUTPUT",
"containerId": "123456789"
}
encoded = b""
for _ in range(total_messages):
encoded += encoder.encode(message)
except Exception as exception:
raise MesosException(
"Error encoding 'RecordIO' message: {error}".format(
error=exception))
try:
all_records = []
offset = 0
chunk_size = 5
while offset < len(encoded):
records = decoder.decode(encoded[offset:offset + chunk_size])
all_records.extend(records)
offset += chunk_size
assert len(all_records) == total_messages
for record in all_records:
assert record == message
except Exception as exception:
raise MesosException(
"Error decoding 'RecordIO' messages: {error}".format(
error=exception))
def request_json(self,
method,
timeout=None,
auth=None,
payload=None,
decoder=None,
params=None,
**kwargs):
"""
Make an HTTP request and deserialize the response as JSON. Optionally
decode the deserialized json dict into a decoded object.
:param method: request method
:type method: str
:param timeout: timeout in seconds
:type timeout: float
:param auth: auth scheme for the request
:type auth: requests.auth.AuthBase
:param payload: json payload in the request
:type payload: dict[str, T] | str
:param decoder: decoder for json response
:type decoder: (dict) -> T
:param params: additional params to include in the request
:type params: str | dict[str, T]
:param kwargs: additional arguments to pass to requests.request
:type kwargs: dict[str, T]
:return: JSON response
:rtype: dict[str, T]
"""
resp = self.request(method=method,
timeout=timeout,
auth=auth,
json=payload,
additional_headers=REQUEST_JSON_HEADERS,
params=params,
**kwargs)
try:
json_dict = ujson.loads(resp.text)
except ValueError as exception:
raise MesosException(
'could not load JSON from "{data}"'.format(data=resp.text),
exception)
if decoder is not None:
return decoder(json_dict)
return json_dict
def encode(self, message):
"""
Encode a message into 'RecordIO' format.
:param message: a message to serialize and then wrap in
a 'RecordIO' frame.
:type message: object
:returns: a serialized message wrapped in a 'RecordIO' frame
:rtype: bytes
"""
s = self.serialize(message)
if not isinstance(s, bytes):
raise MesosException("Calling 'serialize(message)' must"
" return a 'bytes' object")
return bytes(str(len(s)) + "\n", "UTF-8") + s
def request(self,
method,
additional_headers=None,
retry=True,
timeout=None,
auth=None,
use_gzip_encoding=None,
params=None,
max_attempts=None,
**kwargs):
"""
Make an HTTP request by calling self._request with backoff retry.
:param method: request method
:type method: str
:param additional_headers: additional headers to include in the request
:type additional_headers: dict[str, str]
:param retry: boolean indicating whether to retry if the request fails
:type retry: boolean
:param timeout: timeout in seconds, overrides default_timeout_secs
:type timeout: float
:param timeout: timeout in seconds
:type timeout: float
:param auth: auth scheme for the request
:type auth: requests.auth.AuthBase
:param use_gzip_encoding: boolean indicating whether to pass gzip
encoding in the request headers or not
:type use_gzip_encoding: boolean | None
:param params: additional params to include in the request
:type params: str | dict[str, T] | None
:param max_attempts: maximum number of attempts to try for any request
:type max_attempts: int
:param kwargs: additional arguments to pass to requests.request
:type kwargs: dict[str, T]
:return: HTTP response
:rtype: requests.Response
"""
request = self._request
if retry:
if max_attempts is None:
max_attempts = self.default_max_attempts
# We retry only when it makes sense: either due to a network
# partition (e.g. connection errors) or if the request failed
# due to a server error such as 500s, timeouts, and so on.
request = tenacity.retry(
stop=tenacity.stop_after_attempt(max_attempts),
wait=tenacity.wait_exponential(),
retry=tenacity.retry_if_exception_type((
requests.exceptions.Timeout,
requests.exceptions.ConnectionError,
MesosServiceUnavailableException,
MesosInternalServerErrorException,
)),
reraise=True,
)(request)
try:
return request(method=method,
additional_headers=additional_headers,
timeout=timeout,
auth=auth,
use_gzip_encoding=use_gzip_encoding,
params=params,
**kwargs)
# If the request itself failed, an exception subclassed from
# RequestException will be raised. Catch this and reraise as
# MesosException since we want the caller to be able to catch
# and handle this.
except requests.exceptions.RequestException as err:
raise MesosException('Request failed', err)