The first argument to 'mpi' specifies number of workers.

The first argument passed to the 'mpi' process now specifies the number of worker processes to start on each host. The remaining arguments are passed to the controller process.

mpi.cc

@@ -86,6 +86,14 @@ int main(int argc, char **argv)
     // Init MPI
     MPI_Init(&argc, &argv);
 
+    if (argc < 2) {
+	printf("You must pass at least one argument (the number of workers to start).\n");
+	printf("The remaining arguments will be passed to the controller.\n");
+
+	MPI_Finalize();
+	exit(1);
+    }
+
     char hostname[257];
     gethostname(hostname, 256);
 
@@ -100,6 +108,8 @@ int main(int argc, char **argv)
 
     char boundportaddr[259];
 
+    unsigned int num_workers = strtol(argv[1], NULL, 10);
+
     if (rank == 0) {
 	// Start the controller
 	pipe(controllerfds);
@@ -110,7 +120,9 @@ int main(int argc, char **argv)
 	    // Child; close the read half of the pipe and all other fds
 	    close_highfds_except(controllerfds[1]);
 
-	    execv("./controller", argv);
+	    char name[] = "controller";
+	    argv[1] = name;
+	    execv("./controller", argv+1);
 	    return 1;
 	} else {
 	    // Parent; close the write half of the pipe
@@ -136,17 +148,18 @@ int main(int argc, char **argv)
     sprintf(portstr, "%hu", boundport);
     const char *boundaddr = boundportaddr + 2;
 
-    // The child will spawn two of its own children to be the
-    // workers
+    // The child will spawn its own children to be the workers
+    for (int i=0; i<num_workers; i++) {
+	if (fork_and_remember(children) == 0) {
+	    close_highfds_except(-1);
+	    char gpu_index[10];
+	    snprintf(gpu_index, 10, "%d", i);
+	    execl("./worker", "./worker", boundaddr, portstr, gpu_index, NULL);
+	    return 1;
+	}
+    }
+
     if (fork_and_remember(children) == 0) {
-	close_highfds_except(-1);
-	execl("./worker", "./worker", boundaddr, portstr, "0", NULL);
-	return 1;
-    } else if (fork_and_remember(children) == 0) {
-	close_highfds_except(-1);
-	execl("./worker", "./worker", boundaddr, portstr, "1", NULL);
-	return 1;
-    } else if (fork_and_remember(children) == 0) {
 	// And a dpnode
 	close_highfds_except(-1);
 	execl("./dpnode", "./dpnode", boundaddr, portstr, NULL);