cellmap slurm training

Author: Donglai Wei

justfile

@@ -76,6 +76,10 @@ test-with-params dataset ckpt params *ARGS='':
 infer dataset ckpt *ARGS='':
     python scripts/main.py --config tutorials/{{dataset}}.yaml --mode infer --checkpoint {{ckpt}} {{ARGS}}
 
+# Train CellMap models (e.g., just train-cellmap cos7)
+train-cellmap dataset *ARGS='':
+    python scripts/cellmap/train_cellmap.py --config tutorials/cellmap_{{dataset}}.yaml {{ARGS}}
+
 # ============================================================================
 # Monitoring Commands
 # ============================================================================
@@ -96,8 +100,14 @@ tensorboard-all port='6006':
 tensorboard-run experiment timestamp port='6006':
     tensorboard --logdir outputs/{{experiment}}/{{timestamp}}/logs --port {{port}}
 
-# Launch any just command on SLURM (e.g., just slurm weilab 8 4 "train lucchi")
+# Launch any just command on SLURM (e.g., just slurm short 8 4 "train lucchi")
+# Optional 5th parameter: GPU type (vr80g, vr40g, vr16g for V100s)
+# Examples:
+#   just slurm short 8 4 "train lucchi"              # Any available GPU
+#   just slurm short 8 4 "train lucchi" vr80g        # A100 80GB
+#   just slurm short 8 4 "train lucchi" vr40g        # A100 40GB
 # Automatically uses srun for distributed training when num_gpu > 1
+# Time limits: short=12h, medium=2d, long=5d
 slurm partition num_cpu num_gpu cmd constraint='':
     #!/usr/bin/env bash
     # Configure for multi-GPU training with PyTorch Lightning DDP
@@ -109,6 +119,28 @@ slurm partition num_cpu num_gpu cmd constraint='':
         constraint_flag="--constraint={{constraint}}"
     fi
 
+    # Set time limit to partition maximum
+    # Query SLURM for the partition's max time limit
+    time_limit=$(sinfo -p {{partition}} -h -o "%l" | head -1)
+
+    # If sinfo fails or returns empty, use safe defaults
+    if [ -z "$time_limit" ] || [ "$time_limit" = "infinite" ]; then
+        case "{{partition}}" in
+            short|interactive)
+                time_limit="12:00:00"
+                ;;
+            medium)
+                time_limit="2-00:00:00"
+                ;;
+            long)
+                time_limit="5-00:00:00"
+                ;;
+            *)
+                time_limit="7-00:00:00"  # 7 days for private partitions
+                ;;
+        esac
+    fi
+
     sbatch --job-name="pytc_{{cmd}}" \
            --partition={{partition}} \
            --output=slurm_outputs/slurm-%j.out \
@@ -118,13 +150,54 @@ slurm partition num_cpu num_gpu cmd constraint='':
            --gpus-per-node={{num_gpu}} \
            --cpus-per-task={{num_cpu}} \
            --mem=32G \
-           --time=48:00:00 \
+           --time=$time_limit \
            $constraint_flag \
            --wrap="mkdir -p \$HOME/.just && export JUST_TEMPDIR=\$HOME/.just TMPDIR=\$HOME/.just && source /projects/weilab/weidf/lib/miniconda3/bin/activate pytc && cd $PWD && srun --ntasks={{num_gpu}} --ntasks-per-node={{num_gpu}} just {{cmd}}"
 
 # Launch parameter sweep from config (e.g., just sweep tutorials/sweep_example.yaml)
 sweep config:
     python scripts/slurm_launcher.py --config {{config}}
+
+# Launch arbitrary shell command on SLURM (e.g., just slurm-sh short 8 4 "python train.py" vr40g)
+# Unlike 'slurm', this runs the command directly without wrapping in 'just'
+slurm-sh partition num_cpu num_gpu cmd constraint='':
+    #!/usr/bin/env bash
+    constraint_flag=""
+    if [ -n "{{constraint}}" ]; then
+        constraint_flag="--constraint={{constraint}}"
+    fi
+
+    # Set time limit to partition maximum
+    time_limit=$(sinfo -p {{partition}} -h -o "%l" | head -1)
+    if [ -z "$time_limit" ] || [ "$time_limit" = "infinite" ]; then
+        case "{{partition}}" in
+            short|interactive)
+                time_limit="12:00:00"
+                ;;
+            medium)
+                time_limit="2-00:00:00"
+                ;;
+            long)
+                time_limit="5-00:00:00"
+                ;;
+            *)
+                time_limit="7-00:00:00"
+                ;;
+        esac
+    fi
+
+    sbatch --job-name="pytc_{{cmd}}" \
+           --partition={{partition}} \
+           --output=slurm_outputs/slurm-%j.out \
+           --error=slurm_outputs/slurm-%j.err \
+           --nodes=1 \
+           --ntasks={{num_gpu}} \
+           --gpus-per-node={{num_gpu}} \
+           --cpus-per-task={{num_cpu}} \
+           --mem=32G \
+           --time=$time_limit \
+           $constraint_flag \
+           --wrap="source /projects/weilab/weidf/lib/miniconda3/bin/activate pytc && cd $PWD && srun --ntasks={{num_gpu}} --ntasks-per-node={{num_gpu}} {{cmd}}"
 # ============================================================================
 # Visualization Commands
 # ============================================================================

tutorials/cellmap_cos7.yaml

@@ -27,7 +27,7 @@ system:
     num_gpus: 4
     num_cpus: 8
     num_workers: 8
-    batch_size: 8                      # Per GPU batch size
+    batch_size: 8                        # Per GPU (effective = 32)
   inference:
     num_gpus: 1
     num_cpus: 1