少しタイムアウトがどこで(AWS IoT or Lambda)されているかわからず、、一旦IoT側と考えて回答いたします。
現在使用されているAWS IoTのドキュメントを眺めてたのですが永続的セッション(リンク先のコード見る限りMQTTを使ってメッセージのやり取りをしています。)には有効時間がございます。こちらを引き上げることで接続時間を伸ばせるかとはおもいますがpython側で制御すると良いと思います。
https://docs.aws.amazon.com/ja_jp/iot/latest/developerguide/mqtt.html
デフォルトの永続的セッションの有効時間は、標準の制限引き上げプロセスを使用して、1 時間を延長できます。セッションの有効期限を長くすると、オフラインデバイス用にさらに多くのメッセージが保存され、それらのメッセージは標準のメッセージレートでアカウントに課金されるため、メッセージの料金が増加する可能性があります。セッションの有効期限は概算であり、セッションはアカウントの制限よりも最大 30 分間持続する可能性があります。ただし、セッションはアカウントの制限よりも短くなることはありません。セッション制限の詳細については、」
ざっくり制御例
python
1try:
2 while True:
3 time.sleep(5)
4 # ここに接続がされているか確認するメソッドと、接続が切れた場合に再接続する処理を用意
5except KeyboardInterrupt:
6 pig.stop()