From 5a4e9293975cd5bcc2b5359925567f82bd12b993 Mon Sep 17 00:00:00 2001
From: "houhan@gmail.com" <han.hou@alleninstitute.org>
Date: Thu, 8 Aug 2024 22:25:51 +0000
Subject: [PATCH] feat: add pipeline-rerun message

---
 code/Home.py        | 10 ++++++++--
 code/util/aws_s3.py | 11 +++++++++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/code/Home.py b/code/Home.py
index 0603d44..e195a6b 100644
--- a/code/Home.py
+++ b/code/Home.py
@@ -293,6 +293,9 @@ def init():
             
     df = load_data(['sessions'], data_source='bonsai')
     
+    if not len(df):
+        return False
+    
     # --- Perform any data source-dependent preprocessing here ---
     # Because sync_URL_to_session_state() needs df to be loaded (for dynamic column filtering),
     # 'if_load_bpod_sessions' has not been synced from URL to session state yet.
@@ -447,6 +450,8 @@ def _get_data_source(rig):
        
     # Establish communication between pygwalker and streamlit
     init_streamlit_comm()
+    
+    return True
 
 
 def app():
@@ -717,6 +722,7 @@ def app():
 
 
 if 'df' not in st.session_state or 'sessions_bonsai' not in st.session_state.df.keys(): 
-    init()
+    ok = init()
 
-app()
+if ok:
+    app()
diff --git a/code/util/aws_s3.py b/code/util/aws_s3.py
index 9d38fd8..1d97e66 100644
--- a/code/util/aws_s3.py
+++ b/code/util/aws_s3.py
@@ -24,8 +24,15 @@ def load_data(tables=['sessions'], data_source = 'bonsai'):
     df = {}
     for table in tables:
         file_name = s3_processed_nwb_folder[data_source] + f'df_{table}.pkl'
-        with fs.open(file_name) as f:
-            df[table + '_bonsai'] = pd.read_pickle(f)
+        try:
+            with fs.open(file_name) as f:
+                df[table + '_bonsai'] = pd.read_pickle(f)
+        except FileNotFoundError as e:
+            st.markdown(f'''### df_{table}.pkl is missing on S3. \n'''
+                        f'''## It is very likely that Han is [rerunning the whole pipeline](https://github.com/AllenNeuralDynamics/aind-foraging-behavior-bonsai-trigger-pipeline?tab=readme-ov-file#notes-on-manually-re-process-all-nwbs-and-overwrite-s3-database-and-thus-the-streamlit-app). Please come back after an hour.''')
+            st.markdown('')
+            st.image("https://private-user-images.githubusercontent.com/24734299/356400867-bea2c7a9-c561-4c5f-afa5-92d63c040be6.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjMxNTYxNDYsIm5iZiI6MTcyMzE1NTg0NiwicGF0aCI6Ii8yNDczNDI5OS8zNTY0MDA4NjctYmVhMmM3YTktYzU2MS00YzVmLWFmYTUtOTJkNjNjMDQwYmU2LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA4MDglMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwODA4VDIyMjQwNlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTM5NDI4YjYzNTNiM2FiMDZlMzMyYmE2NjY5ZGMzNmM4ZTA4MWI4YjUxZGJlMWNlMTM0NGYwMDQ2MjRhNDcwODImWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.Sb8JeqY1yFnO6DAZLFe4eDd6AYEICaCtFar_CdTT_-Y",
+                     width=500)
     return df