def __init__(self, opts):
self.db = StatDB(opts)
if opts.purge: self.db.purge()
self.data = None
self.last_query = {
'timestamp': 0.0,
'color': '',
'year': 0,
'month': 0
}
self.tasks = TaskManager(opts)
self.districts = None
self.districts_xs = []
self.districts_ys = []
self.districts_names = []
self.selected_type = 'Pickups'
self.selected_borough = 0
self.selected_color = 'yellow'
self.selected_year = 2016
self.selected_month = 1
self.refresh_ticks = 0
self.logger = logging.getLogger(self.__class__.__name__)
self.logger.setLevel(opts.verbose)
def start_multiprocess(opts):
def init():
_, idx = multiprocessing.current_process().name.split('-')
multiprocessing.current_process().name = 'mapper%02d' % int(idx)
db = StatDB(opts)
tasks = []
for start, end in TaskManager.cut(opts.start, opts.end, opts.nprocs):
opts_copy = copy.deepcopy(opts)
opts_copy.start, opts_copy.end = start, end
tasks.append(opts_copy)
try:
procs = multiprocessing.Pool(processes=opts.nprocs, initializer=init)
results = procs.map(start_process, tasks)
except Exception as e:
fatal(e)
finally:
procs.close()
procs.join()
master = results[0]
for res in results:
logger.info('%r => reducer' % res)
master += res
db.append(master)
if opts.report: master.report()
return True
def task_monitor(taskid):
print("SocketIO moinitor: " + taskid)
task = TaskManager.getInstance().get(taskid)
if not task:
emit('error', 'Download task "{}" do not exists'.format(taskid))
return
task.on('progress', lambda x: emit('progress', x))
task.on('complete', lambda x: emit('complete', x))
task.on('error', lambda x: emit('error', x))
task.run(app, request, request.namespace)
def task_monitor(taskid):
print("SocketIO moinitor: " + taskid)
task = TaskManager.getInstance().get(taskid)
if not task:
emit('error', 'Download task "{}" do not exists'.format(taskid))
return
task.on('progress', lambda x: emit('progress', x))
task.on('complete', lambda x: emit('complete', x))
task.on('error', lambda x: emit('error', x))
task.run(app, request, request.namespace)
def start_worker(opts):
task_manager = TaskManager(opts)
if not opts.debug: opts.nprocs = multiprocessing.cpu_count()
nth_task = 0
while True:
task = task_manager.retrieve_task(delete=False)
if task:
logger.info('task %d => start' % nth_task)
opts.color = task.color
opts.year = task.year
opts.month = task.month
opts.start = task.start
opts.end = task.end
if start_multiprocess(opts):
logger.info("task %r => succeeded" % task)
task_manager.delete_task(task)
nth_task += 1
else:
logger.info("no task, wait for %d seconds..." % opts.sleep)
time.sleep(opts.sleep)
async def start(self) -> bool:
"""Initialize the CS/ESPHome API."""
def _start_influxdb(config) -> bool:
success = False
if 'influxdb2' in config.keys():
self._influxdb_client = InfluxDB(config)
success = self._influxdb_client.start()
if not success:
self._influxdb_client = None
else:
_LOGGER.error(
"'influxdb2' options missing from YAML configuration file")
return success
_LOGGER.info(
f"CS/ESPHome energy collection utility {version.get_version()}, PID is {os.getpid()}"
)
config = self._config
if 'settings' in config.keys():
if 'watchdog' in config.settings.keys():
self._watchdog = config.settings.get(
'watchdog', CircuitSetup._DEFAULT_WATCHDOG)
if not _start_influxdb(config=config):
return False
self._esphome_api = ESPHomeApi(config=config)
if not await self._esphome_api.start():
return False
self._esphome_name = self._esphome_api.name()
self._task_manager = TaskManager(config=config,
influxdb_client=self._influxdb_client)
if not await self._task_manager.start(
by_location=self._esphome_api.sensors_by_location(),
by_integration=self._esphome_api.sensors_by_integration()):
return False
return True
from registry import Register
from discovery import ConsumerManager
from tasks import TaskManager
_consumer_manager = ConsumerManager()
_task_manager = TaskManager()
_register = Register()
def register(calendar_id, transformers):
_register.register(calendar_id, transformers)
_consumer_manager.autodiscover()
_task_manager.setup_tasks(_register.consumers)
from tasks import TaskManager
import logging
from flask import session, abort
# Disable flask logging
flaskLogger = logging.getLogger('werkzeug')
flaskLogger.setLevel(logging.ERROR)
logger = logging.getLogger("default")
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s')
file_handler = FileHandler("oddb.log")
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
logger.addHandler(StreamHandler(sys.stdout))
taskManager = TaskManager()
searchEngine = ElasticSearchEngine("od-database")
searchEngine.start_stats_scheduler()
db = Database("db.sqlite3")
# temporary hotfix...
sessionStore = dict()
def require_role(role: str):
if db.get_user_role(session.get("username", None)) != role:
abort(403)
def download_post():
task = DownloadTask(request.form)
TaskManager.getInstance().add(task)
return render_template('download_progress.html', taskid=task.id)
def download_post():
task = DownloadTask(request.form)
TaskManager.getInstance().add(task)
return render_template('download_progress.html', taskid=task.id)
class InteractivePlot:
def __init__(self, opts):
self.db = StatDB(opts)
if opts.purge: self.db.purge()
self.data = None
self.last_query = {
'timestamp': 0.0,
'color': '',
'year': 0,
'month': 0
}
self.tasks = TaskManager(opts)
self.districts = None
self.districts_xs = []
self.districts_ys = []
self.districts_names = []
self.selected_type = 'Pickups'
self.selected_borough = 0
self.selected_color = 'yellow'
self.selected_year = 2016
self.selected_month = 1
self.refresh_ticks = 0
self.logger = logging.getLogger(self.__class__.__name__)
self.logger.setLevel(opts.verbose)
def query(self, color, year, month):
if time.time() - self.last_query['timestamp'] < 2: return
year, mont = int(year), int(month)
self.data = self.db.get(color, year, month)
self.last_query['color'] = color
self.last_query['year'] = year
self.last_query['month'] = month
self.last_query['timestamp'] = time.time()
def hot_map_init(self, width=700, height=700, webgl=True):
self.districts = NYCGeoPolygon.load_districts()
rates = []
for district in self.districts:
x, y = district.xy()
self.districts_xs.append(x)
self.districts_ys.append(y)
self.districts_names.append(district.name)
rates.append(
self.data.pickups[district.index]) # default uses pickups
self.hot_map_source = ColumnDataSource(data=dict(
x=self.districts_xs,
y=self.districts_ys,
name=self.districts_names,
rate=rates,
))
palette.reverse()
color_mapper = LogColorMapper(palette=palette)
self.hot_map = figure(webgl=webgl,
plot_height=height,
plot_width=width,
tools='pan,wheel_zoom,box_zoom,reset,hover,save',
x_axis_location=None,
y_axis_location=None)
self.hot_map.grid.grid_line_color = None
self.hot_map.patches('x',
'y',
source=self.hot_map_source,
fill_color={
'field': 'rate',
'transform': color_mapper
},
fill_alpha=0.7,
line_color="white",
line_width=0.5)
self.hot_map.title.text = "%s %s/%s, %s" % \
(self.selected_type,
self.selected_year, self.selected_month,
NYCBorough.BOROUGHS[self.selected_borough])
hover = self.hot_map.select_one(HoverTool)
hover.point_policy = "follow_mouse"
hover.tooltips = [
("District", "@name"),
("Trips", "@rate"),
("Coordinates", "($x, $y)"),
]
def hot_map_update(self):
rates = []
for district in self.districts:
rate = 0
borough = self.selected_borough
if borough == 0 or borough == district.region:
if self.selected_type == 'Pickups':
rate = self.data.pickups[district.index]
else:
rate = self.data.dropoffs[district.index]
rates.append(rate)
self.hot_map_source.data = dict(
x=self.districts_xs,
y=self.districts_ys,
name=self.districts_names,
rate=rates,
)
self.hot_map.title.text = "%s %s/%s, %s" % \
(self.selected_type,
self.selected_year, self.selected_month,
NYCBorough.BOROUGHS[self.selected_borough])
def trip_hour_init(self, width=620, height=350, webgl=True):
self.trip_hour = figure(webgl=webgl,
toolbar_location=None,
width=width,
height=height,
title='Hour')
self.trip_hour_source = ColumnDataSource(
data=dict(x=range(24), hour=self.data.get_hour()))
vbar = self.trip_hour.vbar(width=0.6,
bottom=0,
x='x',
top='hour',
source=self.trip_hour_source,
fill_alpha=0.7,
line_color="white",
color='#D35400')
self.trip_hour.y_range.start = 0
self.trip_hour.xaxis.major_tick_line_color = None
self.trip_hour.xaxis.minor_tick_line_color = None
self.trip_hour.xaxis.ticker = FixedTicker(ticks=range(24))
self.trip_hour.select(dict(type=GlyphRenderer))
self.trip_hour.add_tools(
HoverTool(renderers=[vbar], tooltips=[("Trips", "@hour")]))
def trip_hour_update(self):
self.trip_hour_source.data = dict(x=range(24),
hour=self.data.get_hour())
def trip_distance_init(self, width=310, height=350, webgl=True):
def ticker():
labels = {
0: '0~1',
1: '1~2',
2: '2~5',
3: '5~10',
4: '10~20',
5: '>20'
}
return labels[tick]
self.trip_distance = figure(webgl=webgl,
toolbar_location=None,
width=width,
height=height,
title='Distance (miles)')
self.trip_distance_source = ColumnDataSource(
data=dict(x=range(6), dist=self.data.get_distance()))
vbar = self.trip_distance.vbar(width=1,
bottom=0,
x='x',
top='dist',
source=self.trip_distance_source,
fill_alpha=0.7,
line_color="white",
color='#588c7e')
self.trip_distance.y_range.start = 0
self.trip_distance.xaxis.major_tick_line_color = None
self.trip_distance.xaxis.minor_tick_line_color = None
self.trip_distance.xaxis.formatter = FuncTickFormatter.from_py_func(
ticker)
self.trip_distance.select(dict(type=GlyphRenderer))
self.trip_distance.add_tools(
HoverTool(renderers=[vbar], tooltips=[("Trips", "@dist")]))
def trip_distance_update(self):
self.trip_distance_source.data = dict(x=range(6),
dist=self.data.get_distance())
def trip_fare_init(self, width=310, height=350, webgl=True):
def ticker():
labels = {
0: '0~5',
1: '5~10',
2: '10~25',
3: '25~50',
4: '50~100',
5: '>100'
}
return labels[tick]
self.trip_fare = figure(webgl=webgl,
toolbar_location=None,
width=width,
height=height,
title='Fare (US dolloars)')
self.trip_fare_source = ColumnDataSource(
data=dict(x=range(6), fare=self.data.get_fare()))
vbar = self.trip_fare.vbar(width=1,
bottom=0,
x='x',
top='fare',
source=self.trip_fare_source,
fill_alpha=0.7,
line_color="white",
color='#ffcc5c')
self.trip_fare.y_range.start = 0
self.trip_fare.xaxis.major_tick_line_color = None
self.trip_fare.xaxis.minor_tick_line_color = None
self.trip_fare.xaxis.formatter = FuncTickFormatter.from_py_func(ticker)
self.trip_fare.select(dict(type=GlyphRenderer))
self.trip_fare.add_tools(
HoverTool(renderers=[vbar], tooltips=[("Trips", "@fare")]))
def trip_fare_update(self):
self.trip_fare_source.data = dict(x=range(6),
fare=self.data.get_fare())
def resource_usage_init(self, width=740, height=120):
data_len = 4
self.resource_usage_source = ColumnDataSource(data=dict(
x=[0, 1, 2, 3, 4], cpu=[2, 3, 5, 4, 10], mem=[20, 10, 40, 30, 15]))
self.resource_usage = figure(plot_width=width,
plot_height=height,
toolbar_location=None,
title=None,
x_axis_label='Elapsed (seconds)',
y_axis_label='%')
self.resource_usage.line(x='x',
y='cpu',
color='firebrick',
legend='CPU',
line_alpha=0.8,
line_width=2,
source=self.resource_usage_source)
self.resource_usage.line(x='x',
y='mem',
color='dodgerblue',
legend='MEM',
line_alpha=0.8,
line_width=2,
source=self.resource_usage_source)
self.resource_usage.xgrid.visible = False
self.resource_usage.ygrid.visible = False
self.resource_usage.x_range.start = 0
self.resource_usage.x_range.end = data_len * 1.07
self.resource_usage.y_range.start = 0
def tasks_stat_init(self, width=740, height=120):
self.tasks_stat_tick = 1
remain, retry = self.tasks.count_tasks()
self.tasks_stat_source = ColumnDataSource(data=dict(
x=range(self.tasks_stat_tick), remain=[remain], retry=[retry]))
self.tasks_stat = figure(plot_width=width,
plot_height=height,
title=None,
toolbar_location=None,
x_axis_label='elapsed (seconds)',
y_axis_label='tasks')
self.tasks_stat.line(x='x',
y='remain',
color='firebrick',
line_alpha=0.8,
line_width=2,
legend='Remain',
source=self.tasks_stat_source)
self.tasks_stat.line(x='x',
y='retry',
color='dodgerblue',
line_alpha=0.8,
line_width=2,
legend='Retry',
source=self.tasks_stat_source)
self.tasks_stat.legend.location = "bottom_left"
self.tasks_stat.xgrid.visible = False
self.tasks_stat.ygrid.visible = False
self.tasks_stat.x_range.start = 0
self.tasks_stat.y_range.start = 0
def tasks_stat_update(self):
self.tasks_stat_tick += 1
rm, re = self.tasks.count_tasks()
self.tasks_stat_source.data['remain'].append(rm)
self.tasks_stat_source.data['retry'].append(re)
self.tasks_stat_source.data = dict(
x=range(self.tasks_stat_tick),
remain=self.tasks_stat_source.data['remain'],
retry=self.tasks_stat_source.data['retry'])
def plot(self):
def update():
self.refresh_ticks += 1
self.query(self.selected_color, self.selected_year,
self.selected_month)
self.hot_map_update()
self.trip_hour_update()
self.trip_distance_update()
self.trip_fare_update()
self.tasks_stat_update()
# self.resource_usage_update()
def on_select():
BOROUGHS_CODE = {v: k for k, v in NYCBorough.BOROUGHS.items()}
self.selected_color = 'green' if color.active == 1 else 'yellow'
pickup.label = 'Pickups' if pickup.active else 'Dropoffs'
self.selected_type = pickup.label
self.selected_borough = BOROUGHS_CODE[borough.value]
borough.label = borough.value
self.selected_year = int(year.value)
self.selected_month = int(month.value)
def on_submit():
self.logger.debug('submit (%s, %s, %s, %s, %s)' % \
(self.selected_type,
NYCBorough.BOROUGHS[self.selected_borough],
self.selected_color,
self.selected_year, self.selected_month))
self.tasks.create_tasks(self.selected_color, self.selected_year,
self.selected_month)
cwd = os.path.dirname(__file__)
desc = Div(text=open(os.path.join(cwd, "description.html")).read(),
width=1000)
# Create input controls
color = RadioButtonGroup(labels=['Yellow', 'Green'], active=1)
color.on_change('active', lambda attr, old, new: on_select())
pickup = Toggle(label='Pickups', button_type="primary", active=True)
pickup.on_change('active', lambda attr, old, new: on_select())
# BUG: Dropdown menu value cannot be integer, i.e., ('Mahattan', '1')
borough_menu = [('All Boroughs', 'All Boroughs'), None,
('Manhattan', 'Manhattan'), ('Bronx', 'Bronx'),
('Brooklyn', 'Brooklyn'), ('Queens', 'Queens'),
('Staten Island', 'Staten Island')]
# https://github.com/bokeh/bokeh/issues/4915
borough = Dropdown(label="Boroughs",
button_type="warning",
menu=borough_menu,
value='All Boroughs')
borough.on_change('value', lambda attr, old, new: on_select())
year = Select(title="Year:",
value=str(self.selected_year),
options=[
str(y) for y in range(MIN_DATE['yellow'].year,
MAX_DATE['yellow'].year + 1)
])
year.on_change('value', lambda attr, old, new: on_select())
month = Select(title="Month:",
value=str(self.selected_month),
options=[str(m) for m in range(1, 13)])
month.on_change('value', lambda attr, old, new: on_select())
submit = Button(label="Submit", button_type="success")
submit.on_click(on_submit)
controls = [color, pickup, borough, year, month, submit]
self.query(self.selected_color, self.selected_year,
self.selected_month)
self.hot_map_init()
self.trip_hour_init()
self.trip_distance_init()
self.trip_fare_init()
self.tasks_stat_init()
self.resource_usage_init()
rightdown_row = row([self.trip_distance, self.trip_fare])
right_column = column([self.trip_hour, rightdown_row])
inputs = widgetbox(*controls, width=140, sizing_mode="fixed")
l = layout([
[desc],
[inputs, self.hot_map, right_column],
[self.tasks_stat, self.resource_usage],
],
sizing_mode="fixed")
curdoc().add_root(l)
curdoc().add_periodic_callback(update, 5000)
curdoc().title = "NYC Taxi Data Explorer"
class CircuitSetup():
"""Class to describe the CircuitSetup ESPHome API."""
_DEFAULT_WATCHDOG = 60
_WATCHDOG = 0
def __init__(self, config):
"""Create a new CircuitSetup object."""
self._config = config
self._task_manager = None
self._influxdb_client = None
self._task_gather = None
self._esphome_api = None
self._esphome_name = None
self._watchdog = CircuitSetup._DEFAULT_WATCHDOG
async def start(self) -> bool:
"""Initialize the CS/ESPHome API."""
def _start_influxdb(config) -> bool:
success = False
if 'influxdb2' in config.keys():
self._influxdb_client = InfluxDB(config)
success = self._influxdb_client.start()
if not success:
self._influxdb_client = None
else:
_LOGGER.error(
"'influxdb2' options missing from YAML configuration file")
return success
_LOGGER.info(
f"CS/ESPHome energy collection utility {version.get_version()}, PID is {os.getpid()}"
)
config = self._config
if 'settings' in config.keys():
if 'watchdog' in config.settings.keys():
self._watchdog = config.settings.get(
'watchdog', CircuitSetup._DEFAULT_WATCHDOG)
if not _start_influxdb(config=config):
return False
self._esphome_api = ESPHomeApi(config=config)
if not await self._esphome_api.start():
return False
self._esphome_name = self._esphome_api.name()
self._task_manager = TaskManager(config=config,
influxdb_client=self._influxdb_client)
if not await self._task_manager.start(
by_location=self._esphome_api.sensors_by_location(),
by_integration=self._esphome_api.sensors_by_integration()):
return False
return True
async def run(self):
try:
queue = asyncio.Queue()
self._task_gather = asyncio.gather(
self._task_manager.run(),
self.task_deletions(),
self.task_esphome_sensor_post(queue),
self.task_esphome_sensor_gather(queue),
)
await self._task_gather
except FailedInitialization as e:
_LOGGER.error(f"run(): {e}")
except WatchdogTimer as e:
_LOGGER.debug(f"run(): {e}")
raise
except Exception as e:
_LOGGER.error(f"Unexpected exception in run(): {e}")
async def stop(self):
"""Shutdown."""
if self._task_manager:
await self._task_manager.stop()
self._task_manager = None
if self._esphome_api:
await self._esphome_api.disconnect()
self._esphome_api = None
if self._influxdb_client:
self._influxdb_client.stop()
self._influxdb_client = None
if self._task_gather:
self._task_gather.cancel()
self._task_gather = None
await asyncio.sleep(1.0)
async def task_deletions(self) -> None:
"""Task to remove older database entries."""
delete_api = self._influxdb_client.delete_api()
bucket = self._influxdb_client.bucket()
org = self._influxdb_client.org()
pruning_tasks = []
config = self._config
if 'influxdb2' in config.keys():
if 'pruning' in config.influxdb2.keys():
for pruning_task in config.influxdb2.pruning:
for task in pruning_task.values():
name = task.get('name', None)
keep_last = task.get('keep_last', 30)
predicate = task.get('predicate', None)
if name and predicate:
new_task = {
'name': name,
'predicate': predicate,
'keep_last': keep_last
}
pruning_tasks.append(new_task)
_LOGGER.debug(
f"Added database pruning task: {new_task}")
while True:
right_now = datetime.datetime.now()
midnight = datetime.datetime.combine(
right_now + datetime.timedelta(days=1), datetime.time(1, 30))
await asyncio.sleep((midnight - right_now).total_seconds())
try:
start = datetime.datetime(1970, 1, 1).isoformat() + 'Z'
for task in pruning_tasks:
stop = datetime.datetime.combine(
datetime.datetime.now() -
datetime.timedelta(days=keep_last), datetime.time(
0, 0)).isoformat() + 'Z'
predicate = task.get('predicate')
keep_last = task.get('keep_last')
delete_api.delete(start,
stop,
predicate,
bucket=bucket,
org=org)
_LOGGER.info(
f"Pruned database '{bucket}': {predicate}, kept last {keep_last} days"
)
except Exception as e:
_LOGGER.debug(f"Unexpected exception in task_deletions(): {e}")
async def task_esphome_sensor_post(self, queue):
"""Process the subscribed data."""
batch_ts = 0
batch_sensors = []
try:
watchdog = self._watchdog
while True:
try:
packet = queue.get_nowait()
except asyncio.QueueEmpty:
_PACKET_DELAY = 0.1
watchdog -= _PACKET_DELAY
if watchdog < 0.0:
raise WatchdogTimer(
f"Lost connection to {self._esphome_name}")
await asyncio.sleep(_PACKET_DELAY)
continue
sensor = packet.get('sensor', None)
state = packet.get('state', None)
queue.task_done()
watchdog = self._watchdog
if sensor and state and self._influxdb_client:
ts = packet.get('ts', None)
if batch_ts != ts:
try:
self._influxdb_client.write_batch_sensors(
batch_sensors=batch_sensors,
timestamp=batch_ts)
batch_ts = ts
batch_sensors = [packet]
except InfluxDBFormatError as e:
_LOGGER.warning(f"{e}")
else:
batch_sensors.append(packet)
except WatchdogTimer:
raise
except Exception as e:
_LOGGER.error(f"task_esphome_sensor_post(): {e}")
async def task_esphome_sensor_gather(self, queue):
"""Post the subscribed data."""
def sensor_callback(state):
CircuitSetup._WATCHDOG += 1
if type(state) == SensorState:
ts = (int(time.time()) // 10) * 10
sensor = sensors_by_key.get(state.key, None)
if sensor:
queue.put_nowait({
'sensor': sensor,
'state': state.state,
'ts': ts
})
# if sensor.get('location') == 'basement':
# _LOGGER.debug(f": device='{sensor.get('device')}' name='{sensor.get('sensor_name')}' state='{state.state}' ts='{ts}'")
try:
sensors_by_key = self._esphome_api.sensors_by_key()
await self._esphome_api.subscribe_states(sensor_callback)
except Exception as e:
_LOGGER.error(f"task_esphome_sensor_gather(): {e}")
from registry import Register
from discovery import ConsumerManager
from tasks import TaskManager
_consumer_manager = ConsumerManager()
_task_manager = TaskManager()
_register = Register()
def register(calendar_id, transformers):
_register.register(calendar_id, transformers)
_consumer_manager.autodiscover()
_task_manager.setup_tasks(_register.consumers)
from tasks import TaskManager tm = TaskManager() tm.do_recrawl()