class DataServer:
def __init__(self, config):
self._cache = Cache(config)
self._host = config['host']
self._port = config['port']
self._encodingType = config['encoding']
self._encoding = Encoding(self._encodingType)
self._adapter = ZMQServers(self._port, self._createReply, config)
self.notAvailable = self._encoding.encode(
self._createError('notAvailable', 'taskId notAvailable'))
def listen(self):
log.info('discovery serving on {host}:{port} with {encoding} encoding',
host=self._host,
port=self._port,
encoding=self._encodingType)
self._adapter.listen()
@timing
def _createReply(self, message):
try:
decoded = self._encoding.decode(value=message, plainEncode=True)
tasks = decoded.get('tasks')
resultsAsTuple = self.getDataByTaskId(tasks)
except Exception as e:
result = self._createError('unknown', str(e))
header, encoded = self._encoding.encode(result)
return [header, encoded]
parts = []
for header, content in resultsAsTuple:
parts.append(header)
parts.append(content)
return parts
def getDataByTaskId(self, tasks):
results = []
for task in tasks:
if (task not in self._cache):
result = self.notAvailable
else:
result = self._cache.getWithHeader(task)
results.append(result)
return results
def setSendingState(self, taskId, header, encoded, size):
return self._cache.update(taskId, encoded, size=size, header=header)
def _createError(self, code, message):
return {'hkube_error': {'code': code, 'message': message}}
def isLocal(self, host, port):
return host == self._host and port == self._port
def isServing(self):
return self._adapter.isServing()
def shutDown(self):
self._adapter.close()
def __init__(self, options, receiverNode):
remoteAddress = options['remoteAddress']
encodingType = options['encoding']
self._encoding = Encoding(encodingType)
self.adapater = ZMQListener(remoteAddress, self.onMessage,
self._encoding, receiverNode)
self.messageOriginNodeName = options['messageOriginNodeName']
self.messageListeners = []
def __init__(self, options, dataServer=None):
self._dataServer = dataServer
self._storageCache = Cache(config.storage)
self._encoding = Encoding(options.storage['encoding'])
self._storageManager = StorageManager(options.storage)
self._requestEncoding = options.storage['encoding']
self._requestTimeout = options.discovery['timeout']
self._networkTimeout = options.discovery['networkTimeout']
self._maxWorkers = min(32, (multiprocessing.cpu_count() or 1) + 4)
log.info('using {workers} workers for DataAdapter', workers=self._maxWorkers)
def __init__(self, config):
self._cache = Cache(config)
self._host = config['host']
self._port = config['port']
self._encodingType = config['encoding']
self._encoding = Encoding(self._encodingType)
self._adapter = ZMQServers(self._port, self._createReply, config)
self.notAvailable = self._encoding.encode(
self._createError('notAvailable', 'taskId notAvailable'))
def __init__(self, reqDetails):
encoding = reqDetails.get('encoding')
address = reqDetails.get('address')
timeout = reqDetails.get('timeout')
networkTimeout = reqDetails.get('networkTimeout')
tasks = reqDetails.get('tasks')
options = {u'tasks': tasks}
self.tasks = tasks
self.encoding = Encoding(encoding)
content = self.encoding.encode(options, plainEncode=True)
self.request = dict()
self.request.update(address)
self.request.update({
"content": content,
"timeout": timeout,
"networkTimeout": networkTimeout
})
def __init__(self, encoding, server):
sm = StorageManager(config.storage)
storageEncoding = Encoding(config.storage.get('encoding'))
self._server = server
self._server.set_fn_new_client(self.handleConnected)
self._server.set_fn_client_left(self.handleDisconnected)
self._server.set_fn_message_received(self.handleMessage)
self._encoding = Encoding(encoding)
def getAlgorithmResult(request):
execId = request.get('execId')
storage= request.get('storage')
storageInput= request.get('storageInput')
input= request.get('input')
if (not storageInput):
response=input
else:
oneInput = storageInput[0]
key = oneInput[2:]
storageInfo=storage.get(key).get('storageInfo')
(header,encoded)=sm.storage.get(storageInfo)
decoded=storageEncoding.decode(header=header, value=encoded)
response=[decoded]
return {
'execId': execId,
'storage': storage,
'response': response
}
self._commands = {
"initialized": {
'command': "start",
'data': lambda x: x
},
"startAlgorithmExecution": {
'command': "algorithmExecutionDone",
'data': getAlgorithmResult
},
"startStoredSubPipeline": {
'command': "subPipelineDone",
'data': lambda x: {
'subPipelineId': x.get('subPipelineId'),
'response': x.get('subPipeline').get('flowInput')
}
}
}
class DataRequest:
def __init__(self, reqDetails):
encoding = reqDetails.get('encoding')
address = reqDetails.get('address')
timeout = reqDetails.get('timeout')
networkTimeout = reqDetails.get('networkTimeout')
tasks = reqDetails.get('tasks')
options = {u'tasks': tasks}
self.tasks = tasks
self.encoding = Encoding(encoding)
content = self.encoding.encode(options, plainEncode=True)
self.request = dict()
self.request.update(address)
self.request.update({
"content": content,
"timeout": timeout,
"networkTimeout": networkTimeout
})
@timing
def invoke(self):
try:
log.info('tcp://{host}:{port}',
host=self.request['host'],
port=self.request['port'])
adapter = ZMQRequest(self.request)
responseFrames = adapter.invokeAdapter()
results = []
for i in range(0, int(len(responseFrames) / 2)):
header = responseFrames[i * 2]
content = responseFrames[i * 2 + 1]
decoded = self.encoding.decode(header=header, value=content)
results.append((len(content), decoded))
return results
except Exception as e:
results = []
for _ in self.tasks:
results.append((0, self._createError('unknown', str(e))))
return results
finally:
adapter.close()
def _createError(self, code, message):
return {'hkube_error': {'code': code, 'message': message}}
def __init__(self, options, consumerNodes, nodeName):
self.nodeNames = consumerNodes
port = options['port']
maxMemorySize = options['messagesMemoryBuff'] * 1024 * 1024
encodingType = options['encoding']
statisticsInterval = options['statisticsInterval']
self._encoding = Encoding(encodingType)
self.adapter = ZMQProducer(port,
maxMemorySize,
self.responseAccumulator,
self.queueTimeAccumulator,
consumerTypes=self.nodeNames,
encoding=self._encoding,
nodeName=nodeName)
self.durationsCache = {}
self.grossDurationCache = {}
self.queueTimeCache = {}
self.responseCount = {}
self.active = True
self.printStatistics = 0
for consumer in consumerNodes:
self.durationsCache[consumer] = FifoArray(RESPONSE_CACHE)
self.grossDurationCache[consumer] = FifoArray(RESPONSE_CACHE)
self.queueTimeCache[consumer] = FifoArray(RESPONSE_CACHE)
self.responseCount[consumer] = 0
self.listeners = []
def sendStatisticsEvery(interval):
while (self.active):
self.sendStatistics()
time.sleep(interval)
self.sendStatistics()
if (self.nodeNames):
runThread = Thread(name="Statistics",
target=sendStatisticsEvery,
args=[statisticsInterval])
runThread.daemon = True
runThread.start()
DaemonThread.__init__(self, "MessageProducer")
class MessageListener():
def __init__(self, options, receiverNode):
remoteAddress = options['remoteAddress']
encodingType = options['encoding']
self._encoding = Encoding(encodingType)
self.adapater = ZMQListener(remoteAddress, self.onMessage,
self._encoding, receiverNode)
self.messageOriginNodeName = options['messageOriginNodeName']
self.messageListeners = []
def registerMessageListener(self, listener):
self.messageListeners.append(listener)
def onMessage(self, messageFlowPattern, header, msg):
start = time.time()
decodedMsg = self._encoding.decode(header=header, value=msg)
for listener in self.messageListeners:
try:
listener(messageFlowPattern, decodedMsg,
self.messageOriginNodeName)
except Exception as e:
log.error('Error during MessageListener onMessage {e}',
e=str(e))
end = time.time()
duration = float((end - start) * 1000)
return self._encoding.encode({'duration': round(duration, 4)},
plainEncode=True)
def fetch(self):
self.adapater.fetch()
def close(self, force=True):
closed = False
try:
closed = self.adapater.close(force)
except Exception as e:
log.error('Exception in adapater.close {e}', e=str(e))
return closed
import time
from hkube_python_wrapper.storage.storage_manager import StorageManager
from hkube_python_wrapper.util.encoding import Encoding
from hkube_python_wrapper import Algorunner
from tests.configs import config
from hkube_python_wrapper.tracing.tracer import Tracer
algorithmName = 'eval-alg'
subpipelineName = 'simple'
encoding = Encoding(config.storage['encoding'])
inp1 = {'nodeName': 'node1','result':'originalRelsult'}
inp2 = {'storageInfo': {'path': 'a'}}
storageMock = {
'a': [{'nodeName': 'node1','info':{'path':'b','isBigData':True}}],
'b':'stam',
}
outp2 = [{'nodeName': 'node1','info':{'path':'b','isBigData':True},'result':'stam'}]
def start(args, hkubeApi=None):
waiter1 = hkubeApi.start_algorithm(algorithmName, [inp1])
waiter2 = hkubeApi.start_stored_subpipeline(subpipelineName, inp2)
res = [waiter1, waiter2]
return res
def test_callback():
sm = StorageManager(config.storage)
algorunner = Algorunner()
algorunner.loadAlgorithmCallbacks(start, options=config)
class MessageProducer(DaemonThread):
def __init__(self, options, consumerNodes, nodeName):
self.nodeNames = consumerNodes
port = options['port']
maxMemorySize = options['messagesMemoryBuff'] * 1024 * 1024
encodingType = options['encoding']
statisticsInterval = options['statisticsInterval']
self._encoding = Encoding(encodingType)
self.adapter = ZMQProducer(port,
maxMemorySize,
self.responseAccumulator,
self.queueTimeAccumulator,
consumerTypes=self.nodeNames,
encoding=self._encoding,
nodeName=nodeName)
self.durationsCache = {}
self.grossDurationCache = {}
self.queueTimeCache = {}
self.responseCount = {}
self.active = True
self.printStatistics = 0
for consumer in consumerNodes:
self.durationsCache[consumer] = FifoArray(RESPONSE_CACHE)
self.grossDurationCache[consumer] = FifoArray(RESPONSE_CACHE)
self.queueTimeCache[consumer] = FifoArray(RESPONSE_CACHE)
self.responseCount[consumer] = 0
self.listeners = []
def sendStatisticsEvery(interval):
while (self.active):
self.sendStatistics()
time.sleep(interval)
self.sendStatistics()
if (self.nodeNames):
runThread = Thread(name="Statistics",
target=sendStatisticsEvery,
args=[statisticsInterval])
runThread.daemon = True
runThread.start()
DaemonThread.__init__(self, "MessageProducer")
def produce(self, messageFlowPattern, obj):
header, encodedMessage = self._encoding.encode(obj)
self.adapter.produce(header,
encodedMessage,
messageFlowPattern=messageFlowPattern)
def responseAccumulator(self, response, consumerType, grossDuration):
decodedResponse = self._encoding.decode(value=response,
plainEncode=True)
duration = decodedResponse['duration']
self.durationsCache[consumerType].append(float(duration))
self.grossDurationCache[consumerType].append(grossDuration)
self.responseCount[consumerType] += 1
def queueTimeAccumulator(self, consumerType, queueTime):
self.queueTimeCache[consumerType].append(queueTime)
def resetDurationsCache(self, consumerType):
durationPerNode = self.durationsCache[consumerType].getAsArray()
self.durationsCache[consumerType].reset()
return durationPerNode
def resetGrossDurationsCache(self, consumerType):
durationPerNode = self.grossDurationCache[consumerType].getAsArray()
self.grossDurationCache[consumerType].reset()
return durationPerNode
def resetQueueDurationsCache(self, consumerType):
durationPerNode = self.queueTimeCache[consumerType].getAsArray()
self.queueTimeCache[consumerType].reset()
return durationPerNode
def getResponseCount(self, consumerType):
return self.responseCount[consumerType]
def registerStatisticsListener(self, listener):
self.listeners.append(listener)
def sendStatistics(self):
statistics = []
for nodeName in self.nodeNames:
queueSize = self.adapter.queueSize(nodeName)
sent = self.adapter.sent(nodeName)
singleNodeStatistics = {
"nodeName": nodeName,
"sent": sent,
"queueSize": queueSize,
"netDurations": self.resetDurationsCache(nodeName),
"durations": self.resetGrossDurationsCache(nodeName),
"queueDurations": self.resetQueueDurationsCache(nodeName),
"responses": self.getResponseCount(nodeName),
"dropped": self.adapter.messageQueue.lostMessages[nodeName]
}
statistics.append(singleNodeStatistics)
for listener in self.listeners:
listener(statistics)
if (self.printStatistics % 30 == 0):
log.debug("statistics {stats}", stats=str(statistics))
self.printStatistics += 1
def run(self):
self.adapter.start()
def close(self, force=True):
if not (self.active):
log.warning("Attempting to close inactive MessageProducer")
else:
self.adapter.close(force)
self.active = False
def test_json_encoding():
encoding = Encoding('json')
data = createObjectJson(size)
encoded = encoding.encode(data, plainEncode=True)
decoded = encoding.decode(value=encoded, plainEncode=True)
assert data == decoded
import os
import shutil
import pytest
from hkube_python_wrapper.storage.storage_manager import StorageManager
from hkube_python_wrapper.util.encoding import Encoding
from tests.configs import config
def ensure_dir(dirName):
if not os.path.exists(dirName):
os.makedirs(dirName)
config = config.storage
baseDirectory = config["fs"]["baseDirectory"]
rootDirectory = baseDirectory.split('/')[0]
sm = StorageManager(config)
encoding = Encoding(config['encoding'])
dir1 = 'dir1'
dir2 = 'dir2'
raw = {"data": 'all_my_data'}
(header, payload) = encoding.encode({"data": 'all_my_data'})
@pytest.fixture(scope="session", autouse=True)
def beforeall(request):
ensure_dir(baseDirectory)
def afterall():
shutil.rmtree(rootDirectory, ignore_errors=True)
def __init__(self, config, encoding):
self.basePath = config['baseDirectory']
self.encoding = Encoding(encoding)
class DataAdapter:
def __init__(self, options, dataServer=None):
self._dataServer = dataServer
self._storageCache = Cache(config.storage)
self._encoding = Encoding(options.storage['encoding'])
self._storageManager = StorageManager(options.storage)
self._requestEncoding = options.storage['encoding']
self._requestTimeout = options.discovery['timeout']
self._networkTimeout = options.discovery['networkTimeout']
self._maxWorkers = min(32, (multiprocessing.cpu_count() or 1) + 4)
log.info('using {workers} workers for DataAdapter', workers=self._maxWorkers)
def encode(self, value):
return self._encoding.encode(value)
def decode(self, header=None, value=None):
return self._encoding.decode(header=header, value=value)
@trace()
def getData(self, options):
jobId = options.get('jobId')
inputArgs = options.get('input')
flatInput = options.get('flatInput')
storage = options.get('storage')
if (not flatInput):
return inputArgs
for k, v in flatInput.items():
if self._isStorage(v):
key = v[2:]
link = storage.get(key, None)
if (link is None):
raise Exception('unable to find storage key')
if (typeCheck.isList(link)):
data = self.batchRequest(link, jobId)
else:
data = self.tryGetDataFromPeerOrStorage(link)
setPath(inputArgs, k, data)
return inputArgs
def _isStorage(self, value):
return typeCheck.isString(value) and value.startswith('$$')
def setAlgorithmStorage(self, jobId, input):
storage = {}
mappedInput = []
for item in input:
taskId = uid(8)
(header, data) = self.encode(item)
storageInfo = self.setData({'jobId': jobId, 'taskId': taskId, 'header': header, 'data': data})
storage[taskId] = {'storageInfo': storageInfo}
mappedInput.append('$${taskId}'.format(taskId=taskId))
return (storage, mappedInput)
@trace()
def setData(self, options):
jobId = options.get('jobId')
taskId = options.get('taskId')
header = options.get('header')
data = options.get('data')
result = self._storageManager.hkube.put(jobId, taskId, header=header, value=data)
return result
@timing
def batchRequest(self, options, jobId):
batchResponse = []
for d in options:
d.update({"jobId": jobId})
with concurrent.futures.ThreadPoolExecutor(max_workers=self._maxWorkers) as executor:
for out in executor.map(self._batchRequest, options):
batchResponse += out
return batchResponse
def _batchRequest(self, options):
batchResponse = []
jobId = options.get('jobId')
tasks = options.get('tasks')
dataPath = options.get('path')
storageInfo = options.get('storageInfo')
if (storageInfo):
storageResult = self._getFromCacheOrStorage(storageInfo, dataPath, storageInfo.get("path"))
batchResponse.append(storageResult)
return batchResponse
tasksNotInCache, batchResponse = self._storageCache.getAll(tasks)
if (tasksNotInCache):
options['tasks'] = tasksNotInCache
results = self._getFromPeer(options)
for i, item in enumerate(results):
size, content = item
peerError = self._getPeerError(content)
taskId = tasksNotInCache[i]
if (peerError):
storageData = self._getDataForTask(jobId, taskId, dataPath)
batchResponse.append(storageData)
else:
self._storageCache.update(taskId, content, size)
content = self._getPath(content, dataPath)
batchResponse.append(content)
return batchResponse
def _getDataForTask(self, jobId, taskId, dataPath):
path = self._storageManager.hkube.createPath(jobId, taskId)
return self._getFromCacheOrStorage({'path': path}, dataPath, taskId)
def tryGetDataFromPeerOrStorage(self, options):
dataPath = options.get('path')
storageInfo = options.get('storageInfo')
discovery = options.get('discovery')
data = None
hasResponse = False
if (discovery):
cacheId = options.get('taskId')
else:
cacheId = storageInfo.get('path')
data = self._getFromCache(cacheId, dataPath)
if not (data):
if (discovery):
size, data = self._getFromPeer(options)[0]
peerError = self._getPeerError(data)
hasResponse = not peerError
data = None if peerError else data
if (hasResponse):
self._setToCache(cacheId, data, size)
data = self._getPath(data, dataPath)
if (not hasResponse and storageInfo):
data = self._getFromCacheOrStorage(storageInfo, dataPath, cacheId)
return data
@trace(name='getFromPeer')
@timing
def _getFromPeer(self, options):
taskId = options.get('taskId')
tasks = [taskId] if taskId else options.get('tasks')
discovery = options.get('discovery')
port = discovery.get('port')
host = discovery.get('host')
if (self._dataServer and self._dataServer.isLocal(host, port)):
dataList = self._dataServer.getDataByTaskId(tasks)
responses = []
for header, payload in dataList:
responses.append((len(payload), self.decode(header=header, value=payload)))
else:
request = {
'address': {
'port': port,
'host': host
},
'tasks': tasks,
'encoding': self._requestEncoding,
'timeout': self._requestTimeout,
'networkTimeout': self._networkTimeout
}
dataRequest = DataRequest(request)
responses = dataRequest.invoke()
return responses
def _getPeerError(self, options):
error = None
if (typeCheck.isDict(options)):
error = options.get('hkube_error')
return error
def _getFromCacheOrStorage(self, options, dataPath, cacheID):
data = self._getFromCache(cacheID, dataPath)
if (data is None):
size, data = self._getFromStorage(options)
self._setToCache(cacheID, data, size)
data = self._getPath(data, dataPath)
return data
@trace(name='getFromCache')
@timing
def _getFromCache(self, cacheId, dataPath):
data = self._storageCache.get(cacheId)
data = self._getPath(data, dataPath)
return data
def _setToCache(self, cacheId, data, size):
self._storageCache.update(cacheId, data, size)
@trace(name='getFromStorage')
@timing
def _getFromStorage(self, options):
(header, payload) = self._storageManager.storage.get(options)
decoded = self.decode(header=header, value=payload)
size = len(payload)
return (size, decoded)
def createStorageInfo(self, options):
jobId = options.get('jobId')
taskId = options.get('taskId')
encodedData = options.get('encodedData')
path = self._storageManager.hkube.createPath(jobId, taskId)
metadata = self.createMetadata(options)
storageInfo = {
'storageInfo': {
'path': path,
'size': len(encodedData) if encodedData else 0
},
'metadata': metadata
}
return storageInfo
def createMetadata(self, options):
nodeName = options.get('nodeName')
data = options.get('data')
savePaths = options.get('savePaths', [])
metadata = dict()
objData = dict()
objData[nodeName] = data
for path in savePaths:
try:
value = getPath(objData, path)
if (value != 'DEFAULT'):
meta = self._getMetadata(value)
metadata[path] = meta
except Exception:
pass
return metadata
def _getMetadata(self, value):
if (typeCheck.isDict(value)):
meta = {'type': 'object'}
elif (typeCheck.isList(value)):
meta = {'type': 'array', 'size': len(value)}
else:
meta = {'type': str(type(value).__name__)}
return meta
def _getPath(self, data, dataPath):
if (data and dataPath):
newData = getPath(data, dataPath)
if (newData == 'DEFAULT'):
newData = None
else:
newData = data
return newData
def test_bson_encoding():
encoding = Encoding('bson')
data = createObject(size, size)
(header, payload) = encoding.encode(data)
decoded = encoding.decode(header=header, value=payload)
assert data == decoded
def test_encoding_header_in_payload_object():
encoding = Encoding('msgpack')
data = createObject(size, size)
(header, payload) = encoding.encode(data)
decoded = encoding.decode(header=None, value=header + payload)
assert data == decoded
import time
from threading import Thread
from hkube_python_wrapper.communication.zmq.streaming.producer.ZMQProducer import ZMQProducer
from hkube_python_wrapper.communication.zmq.streaming.ZMQListener import ZMQListener
from hkube_python_wrapper.util.encoding import Encoding
PROTOCOL_TYPE_MSGPACK = 0x03
DATA_TYPE_RAW = 0x01
encoding = Encoding('msgpack')
header = encoding.createHeader(DATA_TYPE_RAW, PROTOCOL_TYPE_MSGPACK)
def xtest_zmq_streaming():
def doNothing(header, msg, dur):
return b'5'
def queueTimeAccumulator(msg, dur):
return b'5'
count = [0, 0]
producer = ZMQProducer(port=5559,
maxMemorySize=5000,
responseAccumulator=doNothing,
queueTimeAccumulator=queueTimeAccumulator,
consumerTypes=['b', 'c'],
encoding=encoding,
nodeName='a')
runThread = Thread(name="Producer", target=producer.start)
runThread.start()
class WebSocketServerClass:
def __init__(self, encoding, server):
sm = StorageManager(config.storage)
storageEncoding = Encoding(config.storage.get('encoding'))
self._server = server
self._server.set_fn_new_client(self.handleConnected)
self._server.set_fn_client_left(self.handleDisconnected)
self._server.set_fn_message_received(self.handleMessage)
self._encoding = Encoding(encoding)
def getAlgorithmResult(request):
execId = request.get('execId')
storage= request.get('storage')
storageInput= request.get('storageInput')
input= request.get('input')
if (not storageInput):
response=input
else:
oneInput = storageInput[0]
key = oneInput[2:]
storageInfo=storage.get(key).get('storageInfo')
(header,encoded)=sm.storage.get(storageInfo)
decoded=storageEncoding.decode(header=header, value=encoded)
response=[decoded]
return {
'execId': execId,
'storage': storage,
'response': response
}
self._commands = {
"initialized": {
'command': "start",
'data': lambda x: x
},
"startAlgorithmExecution": {
'command': "algorithmExecutionDone",
'data': getAlgorithmResult
},
"startStoredSubPipeline": {
'command': "subPipelineDone",
'data': lambda x: {
'subPipelineId': x.get('subPipelineId'),
'response': x.get('subPipeline').get('flowInput')
}
}
}
def handleMessage(self, client, server, message):
decoded = self._encoding.decode(value=message, plainEncode=True)
command = decoded["command"]
data = decoded.get("data", None)
commandBack = self._commands.get(command)
if(commandBack):
msgBack = {
"command": commandBack["command"],
"data": commandBack["data"](data)
}
self.sendMsgToClient(client, msgBack)
def handleConnected(self, client, server):
# print('ws connected')
self.sendMsgToClient(client, {'command': 'initialize', 'data': mockdata.initData})
def handleDisconnected(self, client, server):
# print('ws disconnected')
pass
def sendMsgToClient(self, client, data):
self._server.send_message(client, self._encoding.encode(data, plainEncode=True))
def test_msgpack_encoding():
encoding = Encoding('msgpack')
data = create_bytearray(size)
(header, payload) = encoding.encode(data)
decoded = encoding.decode(header=header, value=payload)
assert data == decoded
def test_encoding_no_header_bytes():
encoding = Encoding('msgpack')
data = create_bytearray(size)
(_, payload) = encoding.encode(data)
decoded = encoding.decode(header=None, value=payload)
assert data == decoded
"bytesData": bytearray(b'\xdd'*(sizeBytes)),
"anotherBytesData": bytearray(sizeBytes),
"randomString": randomString(sizeRandom),
"randomIntArray": randomInt(sizeRandom),
"dataString": randomString(sizeRandom),
"bool": False,
"anotherBool": False,
"nestedObj": {
"dataString": randomString(sizeRandom),
"randomIntArray": randomInt(sizeRandom)
}
}
return obj
encoding = Encoding('msgpack')
storageManager = StorageManager(config.storage)
@timing
def create_bytearray(sizeBytes):
return bytearray(b'\xdd'*(sizeBytes))
def test_msgpack_bytearray_storage(sizeBytes):
data = create_bytearray(sizeBytes)
storageManager.storage.put({"path": "tests/file2", "data": data})
(header, payload) = storageManager.storage.get({"path": "tests/file2"})
from hkube_python_wrapper.communication.DataRequest import DataRequest
from hkube_python_wrapper.communication.DataServer import DataServer
from hkube_python_wrapper.util.encoding import Encoding
from tests.mocks import mockdata
from tests.configs import config
discovery = dict(config.discovery)
discovery.update({"port": 9024})
data1 = mockdata.dataTask1
data2 = mockdata.dataTask2
data2_original = mockdata.dataTask2_original
taskId1 = mockdata.taskId1
taskId2 = mockdata.taskId2
address1 = {'port': "9024", 'host': discovery['host']}
encoding = Encoding(discovery['encoding'])
timeout = 5000
networkTimeout = 1000
ds = DataServer(discovery)
ds.setSendingState(mockdata.taskId1, data1)
ds.setSendingState(mockdata.taskId2, data2)
ds.setSendingState(mockdata.taskId1, mockdata.dataTask1)
ds.setSendingState(mockdata.taskId2, mockdata.dataTask2)
thread = threading.Thread(target=ds.listen)
thread.start()
gevent.sleep(2)
def test_get_data_bytes():
class FSAdapter:
def __init__(self, config, encoding):
self.basePath = config['baseDirectory']
self.encoding = Encoding(encoding)
def init(self, options):
path = options["path"]
return self.ensure_dir(path)
def put(self, options):
filePath = self.getPath(self.basePath, options['path'])
self.ensure_dir(filePath)
data = options['data']
header = options.get('header')
with open(filePath, 'wb') as f:
if (header):
f.write(header)
f.write(data)
return {'path': options['path']}
def get(self, options):
filePath = self.getPath(self.basePath, options['path'])
header = None
payload = None
headerLength = self.encoding.headerLength()
with open(filePath, 'rb') as f:
header = f.read(headerLength)
if (not self.encoding.isHeader(header)):
f.seek(0)
payload = f.read()
return (header, payload)
def list(self, options):
filePath = self.basePath + os.path.sep + options['path']
if not (os.path.exists(filePath)):
return None
recursive_files_in_dir = []
for r, d, f in os.walk(filePath): # pylint: disable=unused-variable
files_in_dir = []
relativePath = r.replace(self.basePath, '')
for fname in f:
files_in_dir.append(
{'path': relativePath + os.path.sep + fname})
recursive_files_in_dir = recursive_files_in_dir + files_in_dir
return recursive_files_in_dir
def delete(self, options):
filePath = self.getPath(self.basePath, options['path'])
os.remove(filePath)
def listPrefix(self, options):
filePath = self.basePath + os.path.sep + options['path']
if not (os.path.exists(filePath)):
return None
for r, d, f in os.walk(filePath): # pylint: disable=unused-variable
files_in_dir = []
relativePath = r.replace(self.basePath, '')
for fname in f:
files_in_dir.append(relativePath + os.path.sep + fname)
break
return files_in_dir
@staticmethod
def ensure_dir(dirName):
return FSAdapter.ensure_dir_py3(
dirName) if PY3 else FSAdapter.ensure_dir_py27(dirName)
@staticmethod
def ensure_dir_py3(dirName):
d = os.path.dirname(dirName)
# pylint: disable=unexpected-keyword-arg
os.makedirs(d, exist_ok=True)
return os.path.exists(dirName)
@staticmethod
def ensure_dir_py27(dirName):
d = os.path.dirname(dirName)
try:
os.makedirs(d)
except OSError as e:
if e.errno != errno.EEXIST:
raise
return os.path.exists(dirName)
@staticmethod
def getPath(base, dirName):
return base + os.path.sep + dirName
import os
import pytest
from hkube_python_wrapper.util.encoding import Encoding
from hkube_python_wrapper.storage.storage_manager import StorageManager
from tests.configs import config
config = config.storage.copy()
config['type'] = 's3'
bucket = 'local-hkube'
encoding = Encoding(config['encoding'])
raw = {"data": 'all_my_data'}
(header, payload) = encoding.encode(raw)
sm = StorageManager(config)
@pytest.fixture(scope="session", autouse=True)
def beforeall():
sm.storage.adapter.init({'bucket': bucket})
def test_put_get():
options = {
'path': bucket + os.path.sep + 'key1',
'header': header,
'data': payload
}
sm.storage.put(options)
(header1, payload1) = sm.storage.get(options)
assert encoding.decode(header=header1, value=payload1) == raw
def test_none_encoding():
encoding = Encoding('msgpack')
decoded = encoding.decode(header=None, value=None)
assert decoded is None