From fd2e8a3b490707d0ad0d6383044fca00d47c69ba Mon Sep 17 00:00:00 2001 From: Ruben van de Ven Date: Sun, 22 Oct 2023 14:25:34 +0200 Subject: [PATCH] Timeout on zmq to allow gracefull exit --- trap/prediction_server.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/trap/prediction_server.py b/trap/prediction_server.py index 9217c6c..f28e084 100644 --- a/trap/prediction_server.py +++ b/trap/prediction_server.py @@ -242,6 +242,11 @@ class PredictionServer: if self.config.predict_training_data: input_dict = eval_scene.get_clipped_input_dict(timestep, hyperparams['state']) else: + zmq_ev = self.trajectory_socket.poll(timeout=3) + if not zmq_ev: + # on no data loop so that is_running is checked + continue + data = self.trajectory_socket.recv() frame: Frame = pickle.loads(data) trajectory_data = frame.trajectories # TODO: properly refractor