Fix distributed processing and various fixes (#89)

* simplified process_distributed.sh

* removed useless parameter --mmore-folder

* fixed distributed processing

* by default, processing should still not be distributed

* Fixed documentation of distributed processing

* by default, processing should not be using the fast mode

* removed the acknowledgements line

* small logic fixes

* reformatting with black and deleted leftover TODOs

* sorting imports

* changed docs with pyment

Author: fabnemEPFL

README.md

@@ -129,7 +129,3 @@ Don't hesitate to star the project :star: if you find it interesting! (you would
 ## License
 
 This project is licensed under the Apache 2.0 License, see the [LICENSE :mortar_board:](LICENSE) file for details.
-
-## Acknowledgements
-
-This project is part of the [**OpenMeditron**](https://huggingface.co/OpenMeditron) initiative developed in [LiGHT](https://www.light-laboratory.org/) lab at EPFL/Yale/CMU Africa in collaboration with the [**SwissAI**](https://www.swiss-ai.org/) initiative. Thank you Scott Mahoney, Mary-Anne Hartley

docs/distributed_processing.md

@@ -29,7 +29,7 @@ Other important configuration options:
 - `output_folder`: Where processed results will be stored
 - `use_fast_processors`: Set to `true` for faster processing (may reduce accuracy)
 
-### 2. Install Dependencies on All Nodes
+### 2. Install Dependencies on all Nodes
 
 On each node, run:
 
@@ -51,7 +51,7 @@ pip install -e .
 #### Step 1: Start the Master Node (Rank 0)
 
 ```bash
-bash scripts/process_distributed.sh --mmore-folder /path/to/mmore --config-path /path/to/config.yaml --rank 0
+bash scripts/process_distributed.sh --config-file /path/to/config.yaml --rank 0
 ```
 
 The master node will:
@@ -64,14 +64,16 @@ The master node will:
 On each additional node, run:
 
 ```bash
-bash scripts/process_distributed.sh --mmore-folder /path/to/mmore --config-path /path/to/config.yaml --rank 1
+bash scripts/process_distributed.sh --config-path /path/to/config.yaml --rank 1
 ```
 
-Replace `rank 1` with a unique rank number for each node (1, 2, 3, etc.).
+Replace `rank 1` with a unique rank number for each node (1, 2, 3, etc.). The node should be ready in a matter of 5 seconds.
 
 #### Step 3: Begin Processing
 
-Once all nodes are running, return to the master node and type `go` when prompted to start the processing.
+Once all nodes are running, return to the master node and type `go`. The master node proceeds to crawl the input folder, split the workload among connected nodes and make them start their work.
+
+The dask server will be automatically shut down by the master node at the end of the processing. This will also shut down the dask workers on all the connected nodes.
 
 ## Monitoring Progress
 

examples/process/config.yaml

@@ -5,7 +5,7 @@ dispatcher_config:
   distributed: false
   dashboard_backend_url: null
   extract_images: true
-  #scheduler_file: scheduler-file.json #put absolute path!
+  scheduler_file: /mmore/scheduler-file.json #put absolute path!
   process_batch_sizes:
     - URLProcessor: 40
     - DOCXProcessor: 100

examples/rag/evaluation/rag_evaluator_example.py

@@ -1,9 +1,9 @@
 import argparse
-from ....src.mmore.rag.evaluator import EvalConfig, RAGEvaluator
-from ....src.mmore.rag.llm import LLMConfig, LLM
-from ....src.mmore.index.indexer import DBConfig
 
 from dotenv import load_dotenv
+
+from ....src.mmore.rag.evaluator import RAGEvaluator
+
 load_dotenv()
 
 MOCK_EVALUATOR_CONFIG = './examples/rag/evaluation/rag_eval_example_config.yaml'

scripts/process_distributed.sh

@@ -2,28 +2,22 @@
 
 # Default values
 CONFIG_PATH=""
-MMORE_FOLDER=$(pwd)
 RANK=""
 
 # Helper function to show usage
 usage() {
-  echo "Usage: $0 --mmore-folder <path> --config-path <path> --rank <value>"
+  echo "Usage: $0 --mmore-folder <path> --config-file <path> --rank <value>"
   echo ""
   echo "Required arguments:"
-  echo "  --mmore-folder       Absolute path to the mmore folder."
-  echo "  --config_path    Absolute path to the config.yaml file."
+  echo "  --config-file    Absolute path to the config.yaml file."
   echo "  --rank                                       Node rank."
   exit 1
 }
 
 # Parse command-line arguments
 while [[ $# -gt 0 ]]; do
   case $1 in
-    --mmore-folder)
-      MMORE_FOLDER="$2"
-      shift 2
-      ;;
-    --config-path)
+    --config-file)
       CONFIG_PATH="$2"
       shift 2
       ;;
@@ -40,31 +34,13 @@ done
 
 
 # Check required arguments
-if [[ -z "$MMORE_FOLDER" || -z "$CONFIG_PATH" || -z "$RANK" ]]; then
+if [[ -z "$CONFIG_PATH" || -z "$RANK" ]]; then
   echo "Error: Missing required arguments."
   usage
 fi
 
 # Update and install dependencies
 echo "Updating system and installing dependencies..."
-sudo apt-get update
-sudo apt-get install -y --no-install-recommends \
-    nano curl ffmpeg libsm6 libxext6 chromium-browser libnss3 libgconf-2-4 libxi6 libxrandr2 \
-    libxcomposite1 libxcursor1 libxdamage1 libxfixes3 libxrender1 libasound2 libatk1.0-0 \
-    libgtk-3-0 libreoffice libjpeg-dev
-
-# Install Rye
-echo "Setting up UV"
-curl -LsSf https://astral.sh/uv/install.sh | sh
-
-# Navigate to the project directory
-echo "Navigating to mmore folder: $MMORE_FOLDER"
-cd "$MMORE_FOLDER" || { echo "Directory $MMORE_FOLDER does not exist! Exiting."; exit 1; }
-export PATH="$HOME/.local/bin:$PATH"
-
-# Sync Rye to install dependencies
-echo "Syncing UV (installing dependencies)"
-pip install -e '.[process]'
 
 # Extract the distributed configuration from the YAML file
 distributed=$(grep -A3 'dispatcher_config:' "$CONFIG_PATH" | grep 'distributed:' | awk '{print $2}')
@@ -73,47 +49,47 @@ scheduler_file=$(grep 'scheduler_file:' "$CONFIG_PATH" | awk '{print $2}')
 
 # Configure environment variables
 echo "Setting up environment variables"
-export PATH="/.venv/bin:$PATH"
 export DASK_DISTRIBUTED__WORKER__DAEMON=False
 
 # Dask part of the script
-source .venv/bin/activate
 
 if [ "$distributed" = "true" ]; then
-  pip list | grep dask 
   echo "Distributed mode enabled"
   # Start the Dask scheduler if the current node is the MASTER (rank 0)
   if [ "$RANK" -eq 0 ]; then
     echo "Starting the scheduler because it is the MASTER node (rank 0)"
-    dask -h
-    dask scheduler --scheduler-file "$scheduler_file" &
+    dask scheduler --scheduler-file "$scheduler_file" &> dask_scheduler.log &
+    SCHEDULER_PID=$!
   fi
 
   # Start the Dask worker
   echo "Starting the worker of every node"
-  dask worker --scheduler-file "$scheduler_file" &
+  dask worker --scheduler-file "$scheduler_file" &> "dask_scheduler_worker_$RANK.log" &
 fi
 
 
 # Run the end-to-end test if the current node is the MASTER (rank 0)
 if [ "$RANK" -eq 0 ]; then
   echo "Running the end-to-end test in the MASTER node (rank 0)"
-  echo "Command to execute: python \"$MMORE_FOLDER/src/mmore/run_process.py\" --config_file \"$CONFIG_PATH\""
+  echo "Command to execute: python -m mmore process --config-file \"$CONFIG_PATH\""
   echo "Should maybe exit here and wait until all the workers are ready!"
   echo "Type 'go' to execute the command, or type 'exit' to stop and run it manually later."
 
   # waiting for the user to type 'go' or 'exit'
   while true; do
     read -r user_input
     if [ "$user_input" = "go" ]; then
-      python "$MMORE_FOLDER/src/mmore/run_process.py" --config_file "$CONFIG_PATH"
+      echo "Starting processing"
+      python -m mmore process --config-file "$CONFIG_PATH"
       break
     elif [ "$user_input" = "exit" ]; then
       echo "Exiting without running the command. You can run it manually later:"
-      echo "python \"$MMORE_FOLDER/src/mmore/run_process.py\" --config_file \"$CONFIG_PATH\""
+      echo "python -m mmore process --config-file \"$CONFIG_PATH\""
       exit 0
     else
       echo "Invalid input. Type 'go' to run the command or 'exit' to stop."
     fi
   done
-fi
+
+  kill -9 $SCHEDULER_PID
+fi

scripts/runai/entrypoint.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-SCRIPT="run_process.py"
+SCRIPT="process"
 while getopts e:s:c:p: flag
 do
     case "${flag}" in
@@ -12,7 +12,7 @@ done
 
 # Going to repo dir
 if [ -z "$REPO_PATH" ]; then
-    REPO_PATH="/mnt/mlo/scratch/homes/$(whoami)/mmore"
+    REPO_PATH="/mmore" # change to the actual repo path
 fi
 cd $REPO_PATH
 
@@ -21,8 +21,7 @@ set -o allexport
 source .env
 set +o allexport
 
-# TODO: Update when final version is released (Install libraries)
-pip install -e '.[rag]'
+pip install -e .
 
 echo "Start time: $(date)"
 

scripts/setup.sh

@@ -12,4 +12,5 @@ sudo apt install -y ffmpeg libsm6 libxext6 chromium-browser libnss3 libgconf-2-4
 # Install UV
 curl -LsSf https://astral.sh/uv/install.sh | sh
 echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrcuv sync
+uv venv
 source .venv/bin/activate

src/mmore/__main__.py

@@ -1,4 +1,4 @@
 from .cli import main
 
 if __name__ == "__main__":
-    main()
+    main()