More updates to script

Author: holtskinner

.cloud-build/executor/test_notebooks.sh

@@ -1,104 +1,109 @@
 #!/bin/bash
-set -e
 
 alias gcurl='curl -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json"'
 
-readarray -t notebooks < .cloud-build/Notebooks.txt
-NOTEBOOK_RUNTIME_TEMPLATE=$(<NOTEBOOK_RUNTIME_TEMPLATE)
-OUTPUT_URI=$(<OUTPUT_URI)
-SA=$(<SA)
-PROJECT_ID=$(<PROJECT_ID)
-REGION=$(<REGION)
-PUBSUB_TOPIC=$(<PS_TOPIC)
-
-DATE=$(date +%Y-%m-%d)
-TIME=$(date +%H-%M-%S)
-TIMESTAMP=$(date "+%B %d %Y %H:%M:%S")
-
-declare -A operation_map
-declare -a pending completed_success completed_failure
-
-launch_notebook() {
-  local path="$1"
-  local name="${path##*/}"
-  local display_name="${name%.ipynb}-$DATE-$TIME"
-
-  echo "Launching: $path"
-  local operation_id=$(gcloud colab executions create \
-    --display-name="$display_name" \
+TARGET=$(cat .cloud-build/Notebooks.txt)
+
+current_date=$(date +%Y-%m-%d)
+current_time=$(date +%H-%M-%S)
+current_time_readable=$(date "+%B %d %Y %H:%M:%S")
+
+NOTEBOOK_RUNTIME_TEMPLATE=$(cat NOTEBOOK_RUNTIME_TEMPLATE)
+OUTPUT_URI=$(cat OUTPUT_URI)
+SA=$(cat SA)
+PROJECT_ID=$(cat PROJECT_ID)
+REGION=$(cat REGION)
+PUBSUB_TOPIC=$(cat PS_TOPIC)
+
+failed_count=0
+successful_count=0
+declare -a failed_notebooks
+declare -a successful_notebooks
+
+MAX_PARALLEL_JOBS=5
+joblist=()
+
+# Function to run a single notebook
+run_notebook() {
+  local x="$1"
+  local current_date="$2"
+  local current_time="$3"
+
+  DISPLAY_NAME="${x##generative-ai/}"
+  DISPLAY_NAME="${DISPLAY_NAME%.ipynb}-$current_date-$current_time"
+  echo "Starting execution for ${x}"
+
+  OPERATION_ID=$(gcloud colab executions create \
+    --display-name="$DISPLAY_NAME" \
     --notebook-runtime-template="$NOTEBOOK_RUNTIME_TEMPLATE" \
-    --direct-content="$path" \
+    --direct-content="$x" \
     --gcs-output-uri="$OUTPUT_URI" \
     --project="$PROJECT_ID" \
     --region="$REGION" \
     --service-account="$SA" \
     --execution-timeout="1h30m" \
     --format="value(name)")
 
-  local id=$(basename "$operation_id")
-  operation_map["$path"]="$id"
-  pending+=("$path")
-}
-
-monitor_executions() {
-  while [[ ${#pending[@]} -gt 0 ]]; do
-    echo "Waiting for ${#pending[@]} notebooks..."
-
-    local still_pending=()
-
-    for path in "${pending[@]}"; do
-      local id="${operation_map["$path"]}"
-
-      local status=$(gcloud colab executions describe "$id" --region="$REGION" --format="value(jobState)" 2>/dev/null || echo "JOB_STATE_FAILED")
-
-      case "$status" in
-        JOB_STATE_SUCCEEDED)
-          echo "Success: $path"
-          completed_success+=("$path")
-          ;;
-        JOB_STATE_FAILED | *_CANCELLED | *_UNSPECIFIED)
-          echo "Failure: $path ($status)"
-          completed_failure+=("$path")
-          ;;
-        *)
-          echo "Still running: $path ($status)"
-          still_pending+=("$path")
-          ;;
-      esac
-    done
-
-    pending=("${still_pending[@]}")
-    [[ ${#pending[@]} -gt 0 ]] && sleep 60
-  done
-}
+  TRUNCATED_OPERATION_ID=$(echo "$OPERATION_ID" | cut -c 67-85)
 
-publish_results() {
-  local total=${#notebooks[@]}
-  local failed=${#completed_failure[@]}
-  local passed=${#completed_success[@]}
+  if ! EXECUTION_DETAILS=$(gcloud colab executions describe "$TRUNCATED_OPERATION_ID" --region="$REGION"); then
+    echo "Error describing execution for ${x}" >&2
+    echo "fail:$x"
+    return
+  fi
 
-  printf "%s\n" "${completed_failure[@]}" > /workspace/Failure.txt
+  JOB_STATE=$(echo "$EXECUTION_DETAILS" | grep "jobState:" | awk '{print $2}')
+  if [[ "$JOB_STATE" == "JOB_STATE_SUCCEEDED" ]]; then
+    echo "success:$x"
+  else
+    echo "fail:$x"
+  fi
+}
 
-  local fail_list=$(IFS=, ; echo "${completed_failure[*]}")
-  local pass_list=$(IFS=, ; echo "${completed_success[*]}")
+# Parallel runner
+for x in $TARGET; do
+  run_notebook "$x" "$current_date" "$current_time" > "result_$total_count.txt" 2>&1 &
 
-  local message="{\"total_count\":$total,\"failed_count\":$failed,\"failed_notebooks\":\"$fail_list\",\"successful_notebooks\":\"$pass_list\",\"successful_count\":$passed,\"execution_date\":\"$TIMESTAMP\"}"
+  joblist+=($!)
+  total_count=$((total_count + 1))
 
-  echo "$(date) - INFO - Publishing results..."
-  if ! gcloud pubsub topics publish "$PUBSUB_TOPIC" --message="$message" --project="$PROJECT_ID"; then
-    echo "$(date) - ERROR - Failed to publish to Pub/Sub topic $PUBSUB_TOPIC"
+  # Control concurrency
+  if [[ ${#joblist[@]} -ge $MAX_PARALLEL_JOBS ]]; then
+    wait -n
+    joblist=($(jobs -p))  # prune finished jobs
   fi
-}
+done
 
-echo "--- Launching notebooks ---"
-for nb in "${notebooks[@]}"; do
-  [[ -n "$nb" ]] && launch_notebook "$nb"
+# Wait for all remaining jobs
+wait
+
+# Collect results
+for result_file in result_*.txt; do
+  if grep -q "^success:" "$result_file"; then
+    notebook=$(grep "^success:" "$result_file" | cut -d':' -f2-)
+    successful_notebooks+=("$notebook")
+    successful_count=$((successful_count + 1))
+  elif grep -q "^fail:" "$result_file"; then
+    notebook=$(grep "^fail:" "$result_file" | cut -d':' -f2-)
+    failed_notebooks+=("$notebook")
+    failed_count=$((failed_count + 1))
+    echo "- $notebook" | tee -a /workspace/Failure.txt
+  fi
+  rm "$result_file"
 done
 
-echo "--- Monitoring executions ---"
-monitor_executions
+# Summary
+echo "Total successful notebook executions: $successful_count"
+echo "Total failed notebook executions: $failed_count"
+
+# Publish result to Pub/Sub
+failed_notebooks_str=$(IFS=', '; echo "${failed_notebooks[*]}")
+successful_notebooks_str=$(IFS=', '; echo "${successful_notebooks[*]}")
+
+message_data="{\"total_count\":$((total_count)),\"failed_count\":$((failed_count)),\"failed_notebooks\":\"${failed_notebooks_str}\",\"successful_notebooks\":\"${successful_notebooks_str}\",\"successful_count\":$((successful_count)),\"execution_date\":\"${current_time_readable}\"}"
 
-echo "--- Publishing summary ---"
-publish_results
+echo "$(date) - INFO - Publishing to Pub/Sub topic: $PUBSUB_TOPIC"
+gcloud pubsub topics publish "$PUBSUB_TOPIC" --message="$message_data" --project="$PROJECT_ID" || \
+  echo "$(date) - ERROR - Failed to publish to Pub/Sub"
 
-echo "Done. Success: ${#completed_success[@]}, Failures: ${#completed_failure[@]}"
+echo "All notebook executions completed."