Stack Exchange
MPI in python - classification
I'm having a trouble with the following code:
Testcode.py
from mpi4py import MPI
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier

def classify_data(classifier, data):

    X_train, y_train, X_test = data
    classifier.fit(X_train, y_train)
    y_pred = classifier.predict(X_test)
    return y_pred


comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

iris = load_iris()
X, y = iris.data, iris.target

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

comm.bcast((X_train, y_train), root=0)
    
classifiers = {
    0: RandomForestClassifier(n_estimators=100),
    1: DecisionTreeClassifier()
}
   
classifier = classifiers.get(rank, None)

if classifier is not None:    
    y_pred = classify_data(classifier, (X_train, y_train, X_test))
    comm.Barrier()
    if rank == 0:
        all_predictions = comm.gather(y_pred, root=0)

I'm executing the code from console by using:
mpiexec -n 2 python -m mpi4py TestCode.py

I see that if I try to execute the code it has a trouble to calculate for
rank == 0

Why is that happening?

Could you reduce your example further to obtain a MWE?

@DimaChubarov I reduced it as much as I could

Please elaborate on "has trouble to calculate". does the program work if you run a single MPI task?

Yes, when I use single MPI task it works, if I use it on two (for two methods) it freezes after calculating one method.

gather() is a collective operation and must hence be invoked by all the ranks of the communicator.

Thanks, indeed it worked, I fixed the code already and provided the solution I implemented.

Thanks @Giles-gouaillardet for advice above. Indeed, the entire problem was fixed by calling the function below for all tasks.
all_predictions = comm.gather(y_pred, root=0)

I found one more "minor" problem, with the function:
comm.bcast((X_train, y_train), root=0)

This line wasn't even necessary in my implementation because the code already assigned values for each task. Therefore for it to work I was assigning values as follows:
if rank == 0:
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
else:
    X_train, X_test, y_train, y_test = None, None, None, None

X_train = comm.bcast(X_train, root=0)
...

