From 95ab135a4ca056c961c1fd7cc181b7ab52a470f7 Mon Sep 17 00:00:00 2001
From: Alexey Novikov <former.token@gmail.com>
Date: Wed, 6 Nov 2024 20:24:41 +1300
Subject: [PATCH] Remove molden from status reports. Copy molden via gcloud
 CLI.

---
 QUICK/job_manager.sh | 68 ++++++++++++++++++++++++--------------------
 1 file changed, 37 insertions(+), 31 deletions(-)

diff --git a/QUICK/job_manager.sh b/QUICK/job_manager.sh
index da41b1c..ece3984 100644
--- a/QUICK/job_manager.sh
+++ b/QUICK/job_manager.sh
@@ -2,6 +2,8 @@
 DEFAULT_APP_EXECUTABLE="quick.cuda.MPI"
 DEFAULT_STATUS_CHECK_INTERVAL=60
 DEFAULT_SHARED_DIR="/mnt/shared-disk"
+DEFAULT_GCP_BUCKET="gs://mps-andromeda.appspot.com"
+DEFAULT_GCP_BUCKET_PREFIX="job-results"
 
 # Trap SIGCHLD to ensure terminated background processes are cleaned up
 trap 'wait' SIGCHLD
@@ -11,6 +13,16 @@ if [[ -z "${SHARED_DIR}" ]]; then
     echo "WARNING: SHARED_DIR not set, using default value $DEFAULT_SHARED_DIR"
 fi
 
+# Check if GCP_BUCKET or GCP_BUCKET_PREFIX are set or use default
+if [[ -z "${GCP_BUCKET}" ]]; then
+    echo "WARNING: GCP_BUCKET not set, using default value $DEFAULT_GCP_BUCKET"
+fi
+if [[ -z "${GCP_BUCKET_PREFIX}" ]]; then
+    echo "WARNING: GCP_BUCKET_PREFIX not set, using default value $DEFAULT_GCP_BUCKET_PREFIX"
+fi
+GCP_BUCKET="${GCP_BUCKET:-$DEFAULT_GCP_BUCKET}"
+GCP_BUCKET_PREFIX="${GCP_BUCKET_PREFIX:-$DEFAULT_GCP_BUCKET_PREFIX}"
+
 # Directory for shared data
 SHARED_DIR="${SHARED_DIR:-$DEFAULT_SHARED_DIR}"
 
@@ -148,45 +160,39 @@ report_status() {
         # Prepare and send report
         local response_code
         
-        if [[ -f "$molden_file" ]]; then
-            # Report with molden file
-            molden_content=$(base64 "$molden_file")
-            response_code=$(curl -s -w "%{http_code}" -o /dev/null \
-                -X POST "${STATUS_REPORT_ENDPOINT}" \
-                -H "Content-Type: application/json" \
-                -d @- << EOF
-        {
-            "filename": "$base_name",
-            "status": "$status",
-            "new_content": "$new_lines",
-            "offset": $last_offset,
-            "molden": "$molden_content"
-        }
-EOF
-        )
-        else
-            # Report without molden file
-            response_code=$(curl -s -w "%{http_code}" -o /dev/null \
-                -X POST "${STATUS_REPORT_ENDPOINT}" \
-                -H "Content-Type: application/json" \
-                -d @- << EOF
-        {
-            "filename": "$base_name",
-            "status": "$status",
-            "new_content": "$new_lines",
-            "offset": $last_offset
-        }
+        response_code=$(curl -s -w "%{http_code}" -o /dev/null \
+            -X POST "${STATUS_REPORT_ENDPOINT}" \
+            -H "Content-Type: application/json" \
+            -d @- << EOF
+{
+    "filename": "$base_name",
+    "status": "$status",
+    "new_content": "$new_lines",
+    "offset": $last_offset
+}
 EOF
         )
-        fi
 
         # Handle successful report
         if [[ "$response_code" == "204" ]]; then
+            
+            log_message "Successfully reported status $status for $file_key"
+
             # Update offset file if still running
             if [[ "$status" == "RUNNING" ]]; then
                 sed -i.bak "/^${file_key}:/d" "${OFFSET_FILE}"
                 echo "${file_key}:${current_offset}" >> "${OFFSET_FILE}"
             else
+                # Try to move the molden file to the GCP bucket if it exists, otherwise just move the input and output files
+                if [[ -f "$molden_file" ]]; then
+                    gcloud storage cp "$molden_file" "${GCP_BUCKET}/${GCP_BUCKET_PREFIX}/" && rm -f "$molden_file"
+                fi
+                # Don't proceed with the cleanup if molden file is not successfully moved
+                if [[ -f "$molden_file" ]]; then
+                    log_message "Failed to move $molden_file to GCP bucket, will try again later"
+                    continue
+                fi
+
                 # Move files to processed directory and cleanup
                 mv "$input_file" "$processed_dir/"
                 [[ -f "$output_file" ]] && mv "$output_file" "$processed_dir/"
@@ -195,9 +201,9 @@ EOF
                 rm -f "${OFFSET_FILE}.bak"
             fi
             
-            log_message "Successfully reported status $status for $file_key"
+            log_message "Successfully completed status $status operations for $file_key"
         else
-            log_message "Failed to report status for $file_key (HTTP $response_code)"
+            log_message "Failed to report status $status for $file_key (HTTP $response_code)"
         fi
     done
 }