449 lines
15 KiB
Python
Executable File
449 lines
15 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import zmq
|
|
import time
|
|
import signal
|
|
from uuid import uuid4
|
|
from collections import namedtuple
|
|
from collections import deque
|
|
from multiprocessing import Process, TimeoutError
|
|
from datetime import datetime
|
|
|
|
# strat 1: script to copy files
|
|
# strat 2: build pip packages around these
|
|
# could be its own pip package, which we'd need to build and release
|
|
from cereal import log as capnp_log
|
|
from cereal.services import service_list
|
|
from cereal.messaging import pub_sock, MultiplePublishersError
|
|
from common import realtime
|
|
|
|
from tools.lib.kbhit import KBHit
|
|
from tools.lib.logreader import MultiLogIterator
|
|
from tools.lib.route import Route
|
|
from tools.lib.route_framereader import RouteFrameReader
|
|
|
|
# Commands.
|
|
SetRoute = namedtuple("SetRoute", ("name", "start_time", "data_dir"))
|
|
SeekAbsoluteTime = namedtuple("SeekAbsoluteTime", ("secs",))
|
|
SeekRelativeTime = namedtuple("SeekRelativeTime", ("secs",))
|
|
TogglePause = namedtuple("TogglePause", ())
|
|
StopAndQuit = namedtuple("StopAndQuit", ())
|
|
|
|
|
|
class UnloggerWorker(object):
|
|
def __init__(self):
|
|
self._frame_reader = None
|
|
self._cookie = None
|
|
self._readahead = deque()
|
|
|
|
def run(self, commands_address, data_address, pub_types):
|
|
zmq.Context._instance = None
|
|
commands_socket = zmq.Context.instance().socket(zmq.PULL)
|
|
commands_socket.connect(commands_address)
|
|
|
|
data_socket = zmq.Context.instance().socket(zmq.PUSH)
|
|
data_socket.connect(data_address)
|
|
|
|
poller = zmq.Poller()
|
|
poller.register(commands_socket, zmq.POLLIN)
|
|
|
|
# We can't publish frames without encodeIdx, so add when it's missing.
|
|
if "frame" in pub_types:
|
|
pub_types["encodeIdx"] = None
|
|
|
|
# gc.set_debug(gc.DEBUG_LEAK | gc.DEBUG_OBJECTS | gc.DEBUG_STATS | gc.DEBUG_SAVEALL |
|
|
# gc.DEBUG_UNCOLLECTABLE)
|
|
|
|
# TODO: WARNING pycapnp leaks memory all over the place after unlogger runs for a while, gc
|
|
# pauses become huge because there are so many tracked objects solution will be to switch to new
|
|
# cython capnp
|
|
try:
|
|
route = None
|
|
while True:
|
|
while poller.poll(0.) or route is None:
|
|
cookie, cmd = commands_socket.recv_pyobj()
|
|
route = self._process_commands(cmd, route, pub_types)
|
|
|
|
# **** get message ****
|
|
self._read_logs(cookie, pub_types)
|
|
self._send_logs(data_socket)
|
|
finally:
|
|
if self._frame_reader is not None:
|
|
self._frame_reader.close()
|
|
data_socket.close()
|
|
commands_socket.close()
|
|
|
|
def _read_logs(self, cookie, pub_types):
|
|
fullHEVC = capnp_log.EncodeIndex.Type.fullHEVC
|
|
lr = self._lr
|
|
while len(self._readahead) < 1000:
|
|
route_time = lr.tell()
|
|
msg = next(lr)
|
|
typ = msg.which()
|
|
if typ not in pub_types:
|
|
continue
|
|
|
|
# **** special case certain message types ****
|
|
if typ == "encodeIdx" and msg.encodeIdx.type == fullHEVC:
|
|
# this assumes the encodeIdx always comes before the frame
|
|
self._frame_id_lookup[
|
|
msg.encodeIdx.frameId] = msg.encodeIdx.segmentNum, msg.encodeIdx.segmentId
|
|
#print "encode", msg.encodeIdx.frameId, len(self._readahead), route_time
|
|
self._readahead.appendleft((typ, msg, route_time, cookie))
|
|
|
|
def _send_logs(self, data_socket):
|
|
while len(self._readahead) > 500:
|
|
typ, msg, route_time, cookie = self._readahead.pop()
|
|
smsg = msg.as_builder()
|
|
|
|
if typ == "frame":
|
|
frame_id = msg.frame.frameId
|
|
|
|
# Frame exists, make sure we have a framereader.
|
|
# load the frame readers as needed
|
|
s1 = time.time()
|
|
img = self._frame_reader.get(frame_id, pix_fmt="rgb24")
|
|
fr_time = time.time() - s1
|
|
if fr_time > 0.05:
|
|
print("FRAME(%d) LAG -- %.2f ms" % (frame_id, fr_time*1000.0))
|
|
|
|
if img is not None:
|
|
img = img[:, :, ::-1] # Convert RGB to BGR, which is what the camera outputs
|
|
img = img.flatten()
|
|
smsg.frame.image = img.tobytes()
|
|
|
|
data_socket.send_pyobj((cookie, typ, msg.logMonoTime, route_time), flags=zmq.SNDMORE)
|
|
data_socket.send(smsg.to_bytes(), copy=False)
|
|
|
|
def _process_commands(self, cmd, route, pub_types):
|
|
seek_to = None
|
|
if route is None or (isinstance(cmd, SetRoute) and route.name != cmd.name):
|
|
seek_to = cmd.start_time
|
|
route = Route(cmd.name, cmd.data_dir)
|
|
self._lr = MultiLogIterator(route.log_paths(), wraparound=True)
|
|
if self._frame_reader is not None:
|
|
self._frame_reader.close()
|
|
if "frame" in pub_types or "encodeIdx" in pub_types:
|
|
# reset frames for a route
|
|
self._frame_id_lookup = {}
|
|
self._frame_reader = RouteFrameReader(
|
|
route.camera_paths(), None, self._frame_id_lookup, readahead=True)
|
|
|
|
# always reset this on a seek
|
|
if isinstance(cmd, SeekRelativeTime):
|
|
seek_to = self._lr.tell() + cmd.secs
|
|
elif isinstance(cmd, SeekAbsoluteTime):
|
|
seek_to = cmd.secs
|
|
elif isinstance(cmd, StopAndQuit):
|
|
exit()
|
|
|
|
if seek_to is not None:
|
|
print("seeking", seek_to)
|
|
if not self._lr.seek(seek_to):
|
|
print("Can't seek: time out of bounds")
|
|
else:
|
|
next(self._lr) # ignore one
|
|
return route
|
|
|
|
def _get_address_send_func(address):
|
|
sock = pub_sock(address)
|
|
return sock.send
|
|
|
|
|
|
def unlogger_thread(command_address, forward_commands_address, data_address, run_realtime,
|
|
address_mapping, publish_time_length, bind_early, no_loop):
|
|
# Clear context to avoid problems with multiprocessing.
|
|
zmq.Context._instance = None
|
|
context = zmq.Context.instance()
|
|
|
|
command_sock = context.socket(zmq.PULL)
|
|
command_sock.bind(command_address)
|
|
|
|
forward_commands_socket = context.socket(zmq.PUSH)
|
|
forward_commands_socket.bind(forward_commands_address)
|
|
|
|
data_socket = context.socket(zmq.PULL)
|
|
data_socket.bind(data_address)
|
|
|
|
# Set readahead to a reasonable number.
|
|
data_socket.setsockopt(zmq.RCVHWM, 10000)
|
|
|
|
poller = zmq.Poller()
|
|
poller.register(command_sock, zmq.POLLIN)
|
|
poller.register(data_socket, zmq.POLLIN)
|
|
|
|
if bind_early:
|
|
send_funcs = {
|
|
typ: _get_address_send_func(address)
|
|
for typ, address in address_mapping.items()
|
|
}
|
|
|
|
# Give subscribers a chance to connect.
|
|
time.sleep(0.1)
|
|
else:
|
|
send_funcs = {}
|
|
|
|
start_time = float("inf")
|
|
printed_at = 0
|
|
generation = 0
|
|
paused = False
|
|
reset_time = True
|
|
prev_msg_time = None
|
|
while True:
|
|
evts = dict(poller.poll())
|
|
if command_sock in evts:
|
|
cmd = command_sock.recv_pyobj()
|
|
if isinstance(cmd, TogglePause):
|
|
paused = not paused
|
|
if paused:
|
|
poller.modify(data_socket, 0)
|
|
else:
|
|
poller.modify(data_socket, zmq.POLLIN)
|
|
else:
|
|
# Forward the command the the log data thread.
|
|
# TODO: Remove everything on data_socket.
|
|
generation += 1
|
|
forward_commands_socket.send_pyobj((generation, cmd))
|
|
if isinstance(cmd, StopAndQuit):
|
|
return
|
|
|
|
reset_time = True
|
|
elif data_socket in evts:
|
|
msg_generation, typ, msg_time, route_time = data_socket.recv_pyobj(flags=zmq.RCVMORE)
|
|
msg_bytes = data_socket.recv()
|
|
if msg_generation < generation:
|
|
# Skip packets.
|
|
continue
|
|
|
|
if no_loop and prev_msg_time is not None and prev_msg_time > msg_time + 1e9:
|
|
generation += 1
|
|
forward_commands_socket.send_pyobj((generation, StopAndQuit()))
|
|
return
|
|
prev_msg_time = msg_time
|
|
|
|
msg_time_seconds = msg_time * 1e-9
|
|
if reset_time:
|
|
msg_start_time = msg_time_seconds
|
|
real_start_time = realtime.sec_since_boot()
|
|
start_time = min(start_time, msg_start_time)
|
|
reset_time = False
|
|
|
|
if publish_time_length and msg_time_seconds - start_time > publish_time_length:
|
|
generation += 1
|
|
forward_commands_socket.send_pyobj((generation, StopAndQuit()))
|
|
return
|
|
|
|
# Print time.
|
|
if abs(printed_at - route_time) > 5.:
|
|
print("at", route_time)
|
|
printed_at = route_time
|
|
|
|
if typ not in send_funcs:
|
|
if typ in address_mapping:
|
|
# Remove so we don't keep printing warnings.
|
|
address = address_mapping.pop(typ)
|
|
try:
|
|
print("binding", typ)
|
|
send_funcs[typ] = _get_address_send_func(address)
|
|
except Exception as e:
|
|
print("couldn't replay {}: {}".format(typ, e))
|
|
continue
|
|
else:
|
|
# Skip messages that we are not registered to publish.
|
|
continue
|
|
|
|
# Sleep as needed for real time playback.
|
|
if run_realtime:
|
|
msg_time_offset = msg_time_seconds - msg_start_time
|
|
real_time_offset = realtime.sec_since_boot() - real_start_time
|
|
lag = msg_time_offset - real_time_offset
|
|
if lag > 0 and lag < 30: # a large jump is OK, likely due to an out of order segment
|
|
if lag > 1:
|
|
print("sleeping for", lag)
|
|
time.sleep(lag)
|
|
elif lag < -1:
|
|
# Relax the real time schedule when we slip far behind.
|
|
reset_time = True
|
|
|
|
# Send message.
|
|
try:
|
|
send_funcs[typ](msg_bytes)
|
|
except MultiplePublishersError:
|
|
del send_funcs[typ]
|
|
|
|
def timestamp_to_s(tss):
|
|
return time.mktime(datetime.strptime(tss, '%Y-%m-%d--%H-%M-%S').timetuple())
|
|
|
|
def absolute_time_str(s, start_time):
|
|
try:
|
|
# first try if it's a float
|
|
return float(s)
|
|
except ValueError:
|
|
# now see if it's a timestamp
|
|
return timestamp_to_s(s) - start_time
|
|
|
|
def _get_address_mapping(args):
|
|
if args.min is not None:
|
|
services_to_mock = [
|
|
'thermal', 'can', 'health', 'sensorEvents', 'gpsNMEA', 'frame', 'encodeIdx',
|
|
'model', 'features', 'liveLocation', 'gpsLocation'
|
|
]
|
|
elif args.enabled is not None:
|
|
services_to_mock = args.enabled
|
|
else:
|
|
services_to_mock = service_list.keys()
|
|
|
|
address_mapping = {service_name: service_name for service_name in services_to_mock}
|
|
address_mapping.update(dict(args.address_mapping))
|
|
|
|
for k in args.disabled:
|
|
address_mapping.pop(k, None)
|
|
|
|
non_services = set(address_mapping) - set(service_list)
|
|
if non_services:
|
|
print("WARNING: Unknown services {}".format(list(non_services)))
|
|
|
|
return address_mapping
|
|
|
|
def keyboard_controller_thread(q, route_start_time):
|
|
print("keyboard waiting for input")
|
|
kb = KBHit()
|
|
while 1:
|
|
c = kb.getch()
|
|
if c=='m': # Move forward by 1m
|
|
q.send_pyobj(SeekRelativeTime(60))
|
|
elif c=='M': # Move backward by 1m
|
|
q.send_pyobj(SeekRelativeTime(-60))
|
|
elif c=='s': # Move forward by 10s
|
|
q.send_pyobj(SeekRelativeTime(10))
|
|
elif c=='S': # Move backward by 10s
|
|
q.send_pyobj(SeekRelativeTime(-10))
|
|
elif c=='G': # Move backward by 10s
|
|
q.send_pyobj(SeekAbsoluteTime(0.))
|
|
elif c=="\x20": # Space bar.
|
|
q.send_pyobj(TogglePause())
|
|
elif c=="\n":
|
|
try:
|
|
seek_time_input = input('time: ')
|
|
seek_time = absolute_time_str(seek_time_input, route_start_time)
|
|
|
|
q.send_pyobj(SeekAbsoluteTime(seek_time))
|
|
except Exception as e:
|
|
print("Time not understood: {}".format(e))
|
|
|
|
def get_arg_parser():
|
|
parser = argparse.ArgumentParser(
|
|
description="Mock openpilot components by publishing logged messages.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
|
|
parser.add_argument("route_name", type=(lambda x: x.replace("#", "|")), nargs="?",
|
|
help="The route whose messages will be published.")
|
|
parser.add_argument("data_dir", nargs='?', default=os.getenv('UNLOGGER_DATA_DIR'),
|
|
help="Path to directory in which log and camera files are located.")
|
|
|
|
parser.add_argument("--no-loop", action="store_true", help="Stop at the end of the replay.")
|
|
|
|
def key_value_pair(x):
|
|
return x.split("=")
|
|
|
|
parser.add_argument("address_mapping", nargs="*", type=key_value_pair,
|
|
help="Pairs <service>=<zmq_addr> to publish <service> on <zmq_addr>.")
|
|
|
|
def comma_list(x):
|
|
return x.split(",")
|
|
|
|
to_mock_group = parser.add_mutually_exclusive_group()
|
|
to_mock_group.add_argument("--min", action="store_true", default=os.getenv("MIN"))
|
|
to_mock_group.add_argument("--enabled", default=os.getenv("ENABLED"), type=comma_list)
|
|
|
|
parser.add_argument("--disabled", type=comma_list, default=os.getenv("DISABLED") or ())
|
|
|
|
parser.add_argument(
|
|
"--tl", dest="publish_time_length", type=float, default=None,
|
|
help="Length of interval in event time for which messages should be published.")
|
|
|
|
parser.add_argument(
|
|
"--no-realtime", dest="realtime", action="store_false", default=True,
|
|
help="Publish messages as quickly as possible instead of realtime.")
|
|
|
|
parser.add_argument(
|
|
"--no-interactive", dest="interactive", action="store_false", default=True,
|
|
help="Disable interactivity.")
|
|
|
|
parser.add_argument(
|
|
"--bind-early", action="store_true", default=False,
|
|
help="Bind early to avoid dropping messages.")
|
|
|
|
return parser
|
|
|
|
def main(argv):
|
|
args = get_arg_parser().parse_args(sys.argv[1:])
|
|
|
|
command_address = "ipc:///tmp/{}".format(uuid4())
|
|
forward_commands_address = "ipc:///tmp/{}".format(uuid4())
|
|
data_address = "ipc:///tmp/{}".format(uuid4())
|
|
|
|
address_mapping = _get_address_mapping(args)
|
|
|
|
command_sock = zmq.Context.instance().socket(zmq.PUSH)
|
|
command_sock.connect(command_address)
|
|
|
|
if args.route_name is not None:
|
|
route_name_split = args.route_name.split("|")
|
|
if len(route_name_split) > 1:
|
|
route_start_time = timestamp_to_s(route_name_split[1])
|
|
else:
|
|
route_start_time = 0
|
|
command_sock.send_pyobj(
|
|
SetRoute(args.route_name, 0, args.data_dir))
|
|
else:
|
|
print("waiting for external command...")
|
|
route_start_time = 0
|
|
|
|
subprocesses = {}
|
|
try:
|
|
subprocesses["data"] = Process(
|
|
target=UnloggerWorker().run,
|
|
args=(forward_commands_address, data_address, address_mapping.copy()))
|
|
|
|
subprocesses["control"] = Process(
|
|
target=unlogger_thread,
|
|
args=(command_address, forward_commands_address, data_address, args.realtime,
|
|
_get_address_mapping(args), args.publish_time_length, args.bind_early, args.no_loop))
|
|
|
|
for p in subprocesses.values():
|
|
p.daemon = True
|
|
|
|
subprocesses["data"].start()
|
|
subprocesses["control"].start()
|
|
|
|
# Exit if any of the children die.
|
|
def exit_if_children_dead(*_):
|
|
for name, p in subprocesses.items():
|
|
if not p.is_alive():
|
|
[p.terminate() for p in subprocesses.values()]
|
|
exit()
|
|
signal.signal(signal.SIGCHLD, signal.SIG_IGN)
|
|
signal.signal(signal.SIGCHLD, exit_if_children_dead)
|
|
|
|
if args.interactive:
|
|
keyboard_controller_thread(command_sock, route_start_time)
|
|
else:
|
|
# Wait forever for children.
|
|
while True:
|
|
time.sleep(10000.)
|
|
finally:
|
|
for p in subprocesses.values():
|
|
if p.is_alive():
|
|
try:
|
|
p.join(3.)
|
|
except TimeoutError:
|
|
p.terminate()
|
|
continue
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv[1:]))
|