From 1817e5389ae99998dfa686395f6599539d9bb022 Mon Sep 17 00:00:00 2001
From: zhangtianning <zhangtianning110@gmail.com>
Date: Wed, 25 Sep 2024 18:42:45 +0800
Subject: [PATCH] sync

---
 batch_running_task/batch_run.sh               | 23 +++++++++++---
 batch_running_task/get_data_utils.py          |  2 +-
 .../task_rec/batch_deal_with_rec.py           |  6 ++--
 batch_running_task/task_rec/run_rec.sh        | 31 +++++++++++++------
 batch_running_task/task_schedule.sh           | 12 ++++---
 5 files changed, 51 insertions(+), 23 deletions(-)

diff --git a/batch_running_task/batch_run.sh b/batch_running_task/batch_run.sh
index a7d4377..97e5647 100644
--- a/batch_running_task/batch_run.sh
+++ b/batch_running_task/batch_run.sh
@@ -1,20 +1,33 @@
 
-export LD_LIBRARY_PATH=/mnt/cache/share/gcc/gcc-7.5.0/lib64:${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-export PATH=/mnt/cache/share/gcc/gcc-7.5.0/bin:$PATH
-
 TOTALNUM=30
 CPU_NUM=$1 # Automatically get the number of CPUs
 if [ -z "$CPU_NUM" ]; then
     CPU_NUM=$TOTALNUM
 fi
+# check hostname: if it start with SH than use 
+
+if [[ $(hostname) == SH* ]]; then
+    PARA="--quotatype=spot -p AI4Chem -N1 -c8 --gres=gpu:1"
+
+    export LD_LIBRARY_PATH=/mnt/cache/share/gcc/gcc-7.5.0/lib64:${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+    export PATH=/mnt/cache/share/gcc/gcc-7.5.0/bin:$PATH
+
+else
+
+    PARA="-p vip_gpu_ailab_low -N1 -c8 --gres=gpu:1"
+fi
+SCRIPT="batch_running_task/task_rec/run_rec.sh"
+FILELIST="physics_collection/wait_for_ocr.filelist"
+
+
 START=0
 for ((CPU=0; CPU<CPU_NUM; CPU++));
 do
     
     #sbatch --quotatype=spot -p AI4Chem -N1 -c8 --gres=gpu:1  run.sh sci_index_files.addon.filelist $(($CPU+$START)) $TOTALNUM
     #sbatch --quotatype=spot -p AI4Chem -N1 -c8 --gres=gpu:1  run_mfr.sh physics_collection/sci_index_files.remain.filelist 0 1
-    sbatch --quotatype=spot -p AI4Chem -N1 -c8 --gres=gpu:1  batch_running_task/task_rec/run_rec.sh physics_collection/wait_for_ocr.filelist $(($CPU+$START)) $TOTALNUM
-    #sbatch --quotatype=spot -p AI4Chem -N1 -c8 --gres=gpu:1  batch_running_task/task_mfr/run_mfr.sh physics_collection/sci_index_files.finished.filelist $(($CPU+$START)) $TOTALNUM
+    sbatch $PARA $SCRIPT $FILELIST $(($CPU+$START)) $TOTALNUM
+    #sbatch --quotatype=spot -p AI4Chem -N1 -c8 --gres=gpu:1   physics_collection/sci_index_files.finished.filelist $(($CPU+$START)) $TOTALNUM
     #sbatch --quotatype=spot -p AI4Chem -N1 -c8 --gres=gpu:1  batch_running_task/task_layout/run_layout_for_missing_page.sh physics_collection/analysis/not_complete_pdf_page_id.pairlist.remain.filelist $(($CPU+$START)) $TOTALNUM
     ## lets sleep 20s every 10 job start
     if [ $(($CPU % 10)) -eq 9 ]; then
diff --git a/batch_running_task/get_data_utils.py b/batch_running_task/get_data_utils.py
index 8fbfa25..18ae51b 100644
--- a/batch_running_task/get_data_utils.py
+++ b/batch_running_task/get_data_utils.py
@@ -167,7 +167,7 @@ def build_client():
     return client
 
 def check_path_exists(path,client):
-    print(path)
+    #print(path)
     if "s3:" in path:
         return client.contains(path)
     elif path.startswith('http'):
diff --git a/batch_running_task/task_rec/batch_deal_with_rec.py b/batch_running_task/task_rec/batch_deal_with_rec.py
index 1cdba47..902360d 100644
--- a/batch_running_task/task_rec/batch_deal_with_rec.py
+++ b/batch_running_task/task_rec/batch_deal_with_rec.py
@@ -21,8 +21,8 @@
 
 @dataclass
 class BatchRECConfig(BatchModeConfig):
-    inner_batch_size: int = 16
-    batch_size: int = 16
+    image_batch_size: int = 256
+    pdf_batch_size: int = 32
     num_workers: int = 4
     result_save_path: str=RESULT_SAVE_PATH
     check_lock: bool = True
@@ -140,7 +140,7 @@ class BatchRECConfig(BatchModeConfig):
             try:
                 deal_with_one_dataset(inputs_path, result_path, ocrmodel,
                           #batch_size  = args.batch_size,
-                          pdf_batch_size=32, image_batch_size=256,
+                          pdf_batch_size=args.pdf_batch_size, image_batch_size=args.image_batch_size,
                           num_workers = args.num_workers,
                           partion_num = partion_num,
                           partion_idx = partion_idx,update_origin=args.update_origin)
diff --git a/batch_running_task/task_rec/run_rec.sh b/batch_running_task/task_rec/run_rec.sh
index a992fa6..190dc6f 100644
--- a/batch_running_task/task_rec/run_rec.sh
+++ b/batch_running_task/task_rec/run_rec.sh
@@ -2,20 +2,31 @@
 #SBATCH -J ParseSciHUB
 #SBATCH -o .log/%j-ParseSciHUB.out  
 #SBATCH -e .log/%j-ParseSciHUB.out  
-export LD_LIBRARY_PATH=/mnt/cache/share/gcc/gcc-7.5.0/lib64:${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-export PATH=/mnt/cache/share/gcc/gcc-7.5.0/bin:$PATH
 
-GCC_VERSION=$(gcc -v 2>&1 | grep "gcc version" | awk '{print $3}')
 
-# Required version
-REQUIRED_VERSION="7.5.0"
+
 
 # Check if the version matches
-if [ "$GCC_VERSION" != "$REQUIRED_VERSION" ]; then
-  echo "[`hostname`] GCC version is not $REQUIRED_VERSION. Exiting."
-  exit 1
+
+if [[ $(hostname) == SH* ]]; then
+    IMAGE_BATCH_SIZE=256
+    PDF_BATCH_SIZE=32
+    export LD_LIBRARY_PATH=/mnt/cache/share/gcc/gcc-7.5.0/lib64:${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+    export PATH=/mnt/cache/share/gcc/gcc-7.5.0/bin:$PATH
+    GCC_VERSION=$(gcc -v 2>&1 | grep "gcc version" | awk '{print $3}')
+    # Required version
+    REQUIRED_VERSION="7.5.0"
+    if [ "$GCC_VERSION" != "$REQUIRED_VERSION" ]; then
+      echo "[`hostname`] GCC version is not $REQUIRED_VERSION. Exiting."
+      exit 1
+    else
+      echo "[`hostname`] GCC version is $REQUIRED_VERSION."
+    fi
 else
-  echo "[`hostname`] GCC version is $REQUIRED_VERSION."
+    IMAGE_BATCH_SIZE=128
+    PDF_BATCH_SIZE=16
+
 fi
 
-python batch_running_task/task_rec/batch_deal_with_rec.py --root_path $1 --index_part $2 --num_parts $3 --num_workers 8 --update_origin --replace --shuffle #--compile 
\ No newline at end of file
+
+python batch_running_task/task_rec/batch_deal_with_rec.py --image_batch_size $IMAGE_BATCH_SIZE --pdf_batch_size $PDF_BATCH_SIZE --root_path $1 --index_part $2 --num_parts $3 --num_workers 8 --update_origin --replace --shuffle #--compile 
\ No newline at end of file
diff --git a/batch_running_task/task_schedule.sh b/batch_running_task/task_schedule.sh
index 9b38311..9ab1b76 100644
--- a/batch_running_task/task_schedule.sh
+++ b/batch_running_task/task_schedule.sh
@@ -1,12 +1,16 @@
 
 #!/bin/bash
-TASKLIMIT=70
+TASKLIMIT=30
 PENDINGLIMIT=2
 # Function to get the count of pending tasks
 user=`whoami`
-partition='AI4Chem'
-jobscript="batch_running_task/task_layout/run_layout_for_missing_page.sh"
-filelist='scihub_collection/analysis/not_complete_pdf_page_id.pairlist.filelist'
+if [[ $(hostname) == SH* ]]; then
+    partition='AI4Chem'
+else
+    partition='vip_gpu_ailab_low'
+fi
+# jobscript="batch_running_task/task_layout/run_layout_for_missing_page.sh"
+# filelist='scihub_collection/analysis/not_complete_pdf_page_id.pairlist.filelist'
 jobname='ParseSciHUB'
 get_pending_count() {
     squeue -u $user -p $partition -n $jobname | grep PD | wc -l