Move proof gen to after final comm between lora contributor and base model user

[ppotash]

Here's some example code to make it happen (we should also lower timeouts on both users' sode accordingly):

def handle_conn(self, conn, addr):
    try:
        data = self.recv_all(conn)
        if not data:
            return
        req = pickle.loads(data)
        rtype = req.get("request_type", "lora_forward")

        if rtype == "end_inference":
            # 1) Immediately respond with ack
            resp = {
                "response_type": "end_inference_ack",
                "message": "A acknowledges end. Will run proof generation now (offline)."
            }
            conn.sendall(pickle.dumps(resp))
            # 2) Done with B, can close the conn
            conn.close()
            # 3) Run proofs offline
            print("[A] Now running finalize_proofs_and_collect offline.")
            self.lora_server.finalize_proofs_and_collect()
            print("[A] offline proof generation done.")
            return

        # else handle other request_types as normal
        ...
    except Exception as e:
        print(f"[A-Server] error: {e}")
    finally:
        conn.close()