Initial

Author: aocenas

.gitignore

@@ -0,0 +1,5 @@
+.cache
+.idea
+.python-version
+*.egg-info
+*.pyc

LICENSE

@@ -0,0 +1,7 @@
+Copyright 2017 Andrej Ocenas
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -0,0 +1,14 @@
+version: "2"
+services:
+  elastic:
+    image: elasticsearch
+    ports:
+      - 9200:9200
+      - 9300:9300
+
+  kibana:
+    image: kibana
+    ports:
+      - 5601:5601
+    links:
+      - elastic:elasticsearch

docker-compose.yml

@@ -0,0 +1,152 @@
+from __future__ import print_function
+from __future__ import division
+from __future__ import unicode_literals
+
+from datetime import datetime
+from keras.callbacks import Callback
+from elasticsearch import Elasticsearch
+from elasticsearch.helpers import bulk
+import numpy as np
+
+
+class ElasticCallback(Callback):
+    """
+    Sends all data to elasticsearch.
+    """
+    def __init__(
+            self,
+            run_name,
+            index_name='keras',
+            event_data=None,
+            url=None,
+            es_client=None,
+            buffer_length=1,
+    ):
+        """
+        :param str run_name: Name of the run that can be searched for.
+        :param str index_name: Name of index in elasticsearch. Must be lowercase
+        and cannot contain special characters.
+        :param Dict event_data: Additional data that will be merged each each
+        event.
+        :param str url: Elasticsearch url <host>[:<port>]
+        :param object es_client: Elasticsearch client
+        :param int buffer_length: Length of event buffer, before events are sent
+        to Elasticsearch. If length is 1, event is sent after and before
+        each batch, which can slow down learning. If buffer_length is 0, events
+        are flushed only on epoch end. If buffer_length is > 1, number of events
+        will be buffered before flushed, but still everything is flushed on
+        epoch end.
+        """
+        super(ElasticCallback, self).__init__()
+        if es_client:
+            self.es = es_client
+        else:
+            self.es = Elasticsearch([url])
+        self.index = index_name
+        self.run_name = run_name
+        self.data = event_data or {}
+        self.batch_start_time = None
+        self.epoch_start_time = None
+        self.event_buffer = []
+        self.buffer_length = buffer_length
+
+    def on_train_begin(self, logs={}):
+        self._index(
+            'train_begin',
+            logs,
+        )
+
+    def on_epoch_begin(self, epoch, logs={}):
+        self.epoch_start_time = datetime.now()
+        self._index(
+            'epoch_begin',
+            logs,
+            epoch=epoch,
+        )
+
+    def on_batch_begin(self, batch, logs={}):
+        self.batch_start_time = datetime.now()
+        self._add_to_queue(
+            'batch_begin',
+            logs,
+        )
+
+    def on_batch_end(self, batch, logs={}):
+        self._add_to_queue(
+            'batch_end',
+            logs,
+            duration=(datetime.now() - self.batch_start_time).seconds,
+        )
+
+    def on_epoch_end(self, epoch, logs={}):
+        self._flush_queue()
+        self._index(
+            'epoch_end',
+            logs,
+            duration=(datetime.now() - self.epoch_start_time).seconds,
+            epoch=epoch,
+        )
+
+    def on_train_end(self, logs={}):
+        self._index('train_end', logs)
+
+    def _add_to_queue(self, doc_type, logs, **kw):
+        if self.buffer_length == 1:
+            self._index(
+                doc_type,
+                logs,
+                **kw
+            )
+        else:
+            self.event_buffer.append((
+                doc_type,
+                self._create_event_body(doc_type, logs, **kw)
+            ))
+
+            if len(self.event_buffer) == self.buffer_length:
+                self._flush_queue()
+
+    def _flush_queue(self):
+        if len(self.event_buffer):
+            bulk(self.es, self._map_actions(self.event_buffer))
+            self.event_buffer = []
+
+    def _map_actions(self, events):
+        def mapper(event):
+            return {
+                '_op_type': 'index',
+                '_index': self.index,
+                '_type': event[0],
+                '_source': event[1],
+            }
+        return map(mapper, events)
+
+    def _index(self, doc_type, logs, **kw):
+        self.es.index(
+            index=self.index,
+            doc_type=doc_type,
+            body=self._create_event_body(doc_type, logs, **kw),
+        )
+
+    def _create_event_body(self, doc_type, logs, **kw):
+        body = dict(
+            timestamp=datetime.utcnow().isoformat(),
+            event=doc_type,
+            run_name=self.run_name,
+            **self._convert_np_arrays(logs)
+        )
+        body.update(kw)
+        body.update(self.data)
+        return body
+
+    @staticmethod
+    def _convert_np_arrays(data):
+        """
+        Convert numpy ndarrays in a dictionary to list, so it can be serialized
+        to JSON.
+        """
+        return {
+            k: v.tolist() if type(v) == np.ndarray else v
+            for k, v
+            in data.items()
+        }

keras_elastic_callback/__init__.py

@@ -0,0 +1,2 @@
+[metadata]
+description-file = README.md

setup.cfg

@@ -0,0 +1,23 @@
+from setuptools import setup
+
+setup(
+    name='keras_elastic_callback',
+    version='0.1',
+    description='Keras callbacks that sends all events and data to elasticsearch',
+    url='http://github.com/aocenas/keras_elastic_callback',
+    author='Andrej Ocenas',
+    author_email='[email protected]',
+    license='MIT',
+    packages=['keras_elastic_callback'],
+    install_requires=[
+        'keras',
+        'numpy',
+        'elasticsearch',
+    ],
+    tests_require=[
+        'pytest',
+        'mock',
+    ],
+    keywords=['keras', 'elasticsearch', 'machine learning'],
+    zip_safe=False
+)

setup.py

@@ -0,0 +1,68 @@
+import sys
+print sys.path
+from keras_elastic_callback import ElasticCallback
+from elasticsearch import Elasticsearch
+from mock import NonCallableMock
+
+
+def test_data():
+    es_mock = NonCallableMock(spec=Elasticsearch)
+    callback = ElasticCallback(
+        'test_run',
+        'test_index',
+        event_data={'data_key': 'data_value'},
+        es_client=es_mock,
+    )
+
+    callback.on_epoch_begin(1)
+    es_mock.index.assert_called_once()
+    args = es_mock.index.call_args[1]
+    assert args['index'] == 'test_index'
+    assert args['doc_type'] == 'epoch_begin'
+    assert args['body']['event'] == 'epoch_begin'
+    assert args['body']['run_name'] == 'test_run'
+    assert args['body']['data_key'] == 'data_value'
+    assert args['body']['epoch'] == 1
+
+
+events_names = [
+    'train_begin',
+    'epoch_begin',
+    'batch_begin',
+    'batch_end',
+    'epoch_end',
+    'train_end',
+]
+
+
+def test_all_events():
+    es_mock = NonCallableMock(spec=Elasticsearch)
+    callback = ElasticCallback(
+        'test_run',
+        'test_index',
+        es_client=es_mock,
+    )
+
+    for event in events_names:
+
+        func = getattr(callback, 'on_' + event)
+        if 'batch' in event or 'epoch' in event:
+            func(1)
+        else:
+            func()
+        args = es_mock.index.call_args[1]
+        assert args['doc_type'] == event
+
+
+def test_buffer():
+    es_mock = NonCallableMock(spec=Elasticsearch)
+    callback = ElasticCallback(
+        'test_run',
+        'test_index',
+        es_client=es_mock,
+        buffer_length=10
+    )
+
+    callback.on_batch_begin(1)
+    callback.on_batch_end(1)
+    es_mock.index.assert_not_called()