performance changes

app.py

@@ -96,6 +96,37 @@ class IKServer:
 
         self._register_routes()
         logger.info("Server ready.")
+        # NEW warmup thread to JIT both solvers early
+        if getattr(self.args, 'warmup', True):
+            threading.Thread(target=self._warmup_all, daemon=True).start()
+
+    # ---- Warmup helpers (NEW) ----
+    def _warmup_agent(self):
+        try:
+            mand = [DistanceObjTraj(target_points=np.array([0.0,0.2,0.35]), bone_name=self.current_end_effector, use_head=True, weight=1.0)]
+            opt = [BoneZeroRotationObj(weight=0.01)]
+            _ = self.solver.solve(initial_rotations=self.initial_rotations, learning_rate=self.args.learning_rate,
+                                   mandatory_objective_functions=tuple(mand), optional_objective_functions=tuple(opt),
+                                   ik_points=1, verbose=False)
+            logger.info("Agent warmup complete")
+        except Exception as e:
+            logger.warning(f"Agent warmup failed: {e}")
+    def _warmup_urdf(self):
+        try:
+            mand = [DistanceObjTraj(target_points=np.array([0.3,0.3,0.35]), bone_name=self.urdf_current_end_effector, use_head=True, weight=1.0)]
+            opt = [BoneZeroRotationObj(weight=0.01)]
+            _ = self.urdf_solver.solve(initial_rotations=self.urdf_initial_rotations, learning_rate=self.args.learning_rate,
+                                        mandatory_objective_functions=tuple(mand), optional_objective_functions=tuple(opt),
+                                        ik_points=1, verbose=False)
+            logger.info("URDF warmup complete")
+        except Exception as e:
+            logger.warning(f"URDF warmup failed: {e}")
+    def _warmup_all(self):
+        t0 = time.time()
+        logger.info("Starting background warmup...")
+        self._warmup_agent()
+        self._warmup_urdf()
+        logger.info(f"Warmup finished in {time.time()-t0:.2f}s")
 
     # ---------- Cache ----------
     def _evict_lru(self, is_urdf=False):
@@ -347,7 +378,7 @@ class IKServer:
             frames.append({"vertices": verts.tolist(), "faces": faces.tolist()})
         return frames
 
-    def solve_agent(self, payload):
+    def solve_agent(self, payload, last_only=False):
         mand, opt, sub = self._build_agent_objectives(payload)
         start = time.time()
         best_angles, obj_val, steps = self.solver.solve(
@@ -360,18 +391,20 @@ class IKServer:
         )
         self.best_angles = best_angles[-1].copy()
         self.initial_rotations = self.best_angles.copy()
-        self.animation_frames_agent = self._frames_from_angles(best_angles, False)
-        print(f"Solved in {time.time()-start:.2f}s over {steps} steps. Obj: {obj_val:.6f}")
-        self.agent_solve_counter += 1  # added
+        if last_only:
+            self.animation_frames_agent = self._frames_from_angles([best_angles[-1]], False)
+        else:
+            self.animation_frames_agent = self._frames_from_angles(best_angles, False)
+        self.agent_solve_counter += 1
         return {
             "solve_time": time.time()-start,
             "iterations": steps,
             "objective": obj_val,
             "frames": len(self.animation_frames_agent),
-            "solve_id": self.agent_solve_counter,  # added
+            "solve_id": self.agent_solve_counter,
         }
 
-    def solve_urdf(self, payload):
+    def solve_urdf(self, payload, last_only=False):
         mand, opt, sub = self._build_urdf_objectives(payload)
         start = time.time()
         best_angles, obj_val, steps = self.urdf_solver.solve(
@@ -384,17 +417,17 @@ class IKServer:
         )
         self.urdf_best_angles = best_angles[-1].copy()
         self.urdf_initial_rotations = self.urdf_best_angles.copy()
-        self.animation_frames_urdf = self._frames_from_angles(best_angles, True)
-        print(
-            f"Solved in {time.time() - start:.2f}s over {steps} steps. Obj: {obj_val:.6f}"
-        )
-        self.urdf_solve_counter += 1  # added
+        if last_only:
+            self.animation_frames_urdf = self._frames_from_angles([best_angles[-1]], True)
+        else:
+            self.animation_frames_urdf = self._frames_from_angles(best_angles, True)
+        self.urdf_solve_counter += 1
         return {
             "solve_time": time.time()-start,
             "iterations": steps,
             "objective": obj_val,
             "frames": len(self.animation_frames_urdf),
-            "solve_id": self.urdf_solve_counter,  # added
+            "solve_id": self.urdf_solve_counter,
         }
 
     # ---------- Housekeeping ----------
@@ -467,6 +500,8 @@ class IKServer:
             model = payload.get("model","agent").lower()
             return_mode = payload.get("frames_mode", "auto")
             num_steps = int(payload.get("num_steps", self.args.num_steps))  # NEW
+            subpoints = int(payload.get("subpoints",1))
+            last_only = (return_mode != 'all' and subpoints == 1)
             self._cleanup()
             with self.solve_lock:
                 try:
@@ -474,21 +509,25 @@ class IKServer:
                         self.configure_urdf(payload.get("controlled_bones", []),
                                             payload.get("end_effector", self.urdf_current_end_effector),
                                             num_steps=num_steps)
-                        result = self.solve_urdf(payload); result["model"]="pepper"; result["num_steps"] = num_steps
-                        if return_mode == "all" or (return_mode == "auto" and payload.get("subpoints",1) > 1):
-                            frames = list(self.animation_frames_urdf)  # ensure fresh list
+                        result = self.solve_urdf(payload, last_only=last_only); result["model"]="pepper"; result["num_steps"] = num_steps
+                        if last_only:
+                            frames = list(self.animation_frames_urdf)  # only one frame
+                        elif return_mode == "all" or (return_mode == "auto" and subpoints > 1):
+                            frames = list(self.animation_frames_urdf)
                         else:
-                            frames = [self.animation_frames_urdf[-1].copy()] if hasattr(self.animation_frames_urdf[-1],'copy') else [dict(self.animation_frames_urdf[-1])]
+                            frames = [self.animation_frames_urdf[-1]]
                         result["frames_data"] = frames
                     else:
                         self.configure_agent(payload.get("controlled_bones", []),
                                              payload.get("end_effector", self.current_end_effector),
                                              num_steps=num_steps)
-                        result = self.solve_agent(payload); result["model"]="agent"; result["num_steps"] = num_steps
-                        if return_mode == "all" or (return_mode == "auto" and payload.get("subpoints",1) > 1):
+                        result = self.solve_agent(payload, last_only=last_only); result["model"]="agent"; result["num_steps"] = num_steps
+                        if last_only:
+                            frames = list(self.animation_frames_agent)
+                        elif return_mode == "all" or (return_mode == "auto" and subpoints > 1):
                             frames = list(self.animation_frames_agent)
                         else:
-                            frames = [self.animation_frames_agent[-1].copy()] if hasattr(self.animation_frames_agent[-1],'copy') else [dict(self.animation_frames_agent[-1])]
+                            frames = [self.animation_frames_agent[-1]]
                         result["frames_data"] = frames
                     return JSONResponse({"status":"ok","result":result})
                 except Exception as e:
@@ -518,6 +557,7 @@ def main():
     parser.add("--learning_rate", type=float, default=0.2)
     parser.add("--subpoints", type=int, default=1)
     parser.add("--api_port", type=int, default=17861)
+    parser.add("--warmup", action='store_true', default=True, help="Enable background JIT warmup")
     args = parser.parse_args()
 
     download_and_setup_files()