because your driver is not able to talk with your nodes for this you can add configuration where you can increase databricks heartbeat interval and you can also add rpc max size this will also help.
you can explore cluster configuration from here- https://spark.apache.org/docs/latest/configuration.html
Thanks
Aviral Bhardwaj
AviralBhardwaj