server.py

import sqlite3
import xmltodict
from datetime import datetime
from flask import g, Flask, request

app = Flask(__name__)
DATABASE = 'data.db'

# Run the following command on the hosts:
# nvidia-smi -q -x -l 60 | curl -X POST -d @- http://eos12:5000/$HOSTNAME

def connect_to_database():
    conn = sqlite3.connect(DATABASE)
    return conn


def init_db():
    with app.app_context():
        db = get_db()
        with app.open_resource('schema.sql', mode='r') as f:
            db.cursor().executescript(f.read())
        db.commit()


def get_db():
    db = getattr(g, '_database', None)
    if db is None:
        db = g._database = connect_to_database()
        db.row_factory = sqlite3.Row
    return db


def query_db(query, args=(), one=False):
    cur = get_db().execute(query, args)
    rv = cur.fetchall()
    cur.close()
    return (rv[0] if rv else None) if one else rv


@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, '_database', None)
    if db is not None:
        db.close()


@app.route('/<host>', methods=['POST'])
def index(host):
    data = xmltodict.parse(request.get_data())['nvidia_smi_log']
    date = datetime.strptime(data['timestamp'], '%a %b %d %X %Y')
    timestamp = int(date.strftime('%s'))

    rows = []
    for gpu in data['gpu']:
        utilization = gpu['utilization']['gpu_util']
        if utilization == 'N/A':
            utilization = -1
        else:
            utilization = float(utilization[:-1])
        rows.append((host, gpu['@id'], gpu['product_name'], timestamp,
                     utilization))

    db = get_db()
    with db:
        print 'inserting', rows, db
        db.executemany('INSERT INTO data VALUES (?, ?, ?, ?, ?)', rows)

    return 'OK'

if __name__ == "__main__":
    app.debug = True
    app.run(host='0.0.0.0')