mpirun.cpp

刚才是说明 现在是安装程序在 LINUX环境下进行编程的MPICH安装文件

// Function name	: main
// Description	    : 
// Return type		: void 
// Argument         : int argc
// Argument         : char *argv[]
int  main(int argc, char *argv[])
{
    int i;
    int iproc = 0;
    char pszEnv[MAX_CMD_LENGTH] = "";
    HANDLE *pThread = NULL;
    int nShmLow, nShmHigh;
    DWORD dwThreadID;
    bool bLogon = false;
    char pBuffer[MAX_CMD_LENGTH];
    char phrase[MPD_PASSPHRASE_MAX_LENGTH + 1];// = MPD_DEFAULT_PASSPHRASE;
    bool bLogonDots = true;
    HANDLE hStdout;
    char cMapDrive, pszMapShare[MAX_PATH];
    int nArgsToStrip;
    bool bRunLocal;
    char pszMachineFileName[MAX_PATH] = "";
    bool bUseMachineFile;
    bool bDoSMP;
    bool bPhraseNeeded;
    DWORD dwType;

    if (argc < 2)
    {
	PrintOptions();
	return 0;
    }

    SetConsoleCtrlHandler(CtrlHandlerRoutine, TRUE);

    DWORD pmi_host_length = MAX_HOST_LENGTH;
    GetComputerName(pmi_host, &pmi_host_length);

    // Set defaults
    g_bDoMultiColorOutput = true;
    bRunLocal = false;
    g_bNoMPI = false;
    bLogon = false;
    bLogonDots = true;
    GetCurrentDirectory(MAX_PATH, g_pszDir);
    bUseMachineFile = false;
    bDoSMP = true;
    phrase[0] = '\0';
    bPhraseNeeded = true;
    g_nHosts = 0;
    g_pHosts = NULL;

    // Parse mpirun options
    while (argv[1] && (argv[1][0] == '-' || argv[1][0] == '/'))
    {
	nArgsToStrip = 1;
	if (stricmp(&argv[1][1], "np") == 0)
	{
	    if (argc < 3)
	    {
		printf("Error: no number specified after -np option.\n");
		return 0;
	    }
	    g_nHosts = atoi(argv[2]);
	    if (g_nHosts < 1)
	    {
		printf("Error: must specify a number greater than 0 after the -np option\n");
		return 0;
	    }
	    nArgsToStrip = 2;
	}
	else if (stricmp(&argv[1][1], "localonly") == 0)
	{
	    bRunLocal = true;
	    if (argc > 2)
	    {
		if (isnumber(argv[2]))
		{
		    g_nHosts = atoi(argv[2]);
		    if (g_nHosts < 1)
		    {
			printf("Error: If you specify a number after -localonly option,\n        it must be greater than 0.\n");
			return 0;
		    }
		    nArgsToStrip = 2;
		}
	    }
	}
	else if (stricmp(&argv[1][1], "machinefile") == 0)
	{
	    if (argc < 3)
	    {
		printf("Error: no filename specified after -machinefile option.\n");
		return 0;
	    }
	    strcpy(pszMachineFileName, argv[2]);
	    bUseMachineFile = true;
	    nArgsToStrip = 2;
	}
	else if (stricmp(&argv[1][1], "map") == 0)
	{
	    if (argc < 3)
	    {
		printf("Error: no drive specified after -map option.\n");
		return 0;
	    }
	    if ((strlen(argv[2]) > 2) && argv[2][1] == ':')
	    {
		MapDriveNode *pNode = new MapDriveNode;
		pNode->cDrive = argv[2][0];
		strcpy(pNode->pszShare, &argv[2][2]);
		pNode->pNext = g_pDriveMapList;
		g_pDriveMapList = pNode;
	    }
	    nArgsToStrip = 2;
	}
	else if (stricmp(&argv[1][1], "dir") == 0)
	{
	    if (argc < 3)
	    {
		printf("Error: no directory after -dir option\n");
		return 0;
	    }
	    strcpy(g_pszDir, argv[2]);
	    nArgsToStrip = 2;
	}
	else if (stricmp(&argv[1][1], "env") == 0)
	{
	    if (argc < 3)
	    {
		printf("Error: no environment variables after -env option\n");
		return 0;
	    }
	    strncpy(g_pszEnv, argv[2], MAX_CMD_LENGTH);
	    g_pszEnv[MAX_CMD_LENGTH-1] = '\0';
	    if (strlen(argv[2]) >= MAX_CMD_LENGTH)
	    {
		printf("Warning: environment variables truncated.\n");
	    }
	    nArgsToStrip = 2;
	}
	else if (stricmp(&argv[1][1], "logon") == 0)
	{
	    bLogon = true;
	}
	else if (stricmp(&argv[1][1], "hosts") == 0)
	{
	    if (g_nHosts != 0)
	    {
		printf("Error: only one option is allowed to determine the number of processes.\n");
		printf("       -hosts cannot be used with -np or -localonly\n");
		return 0;
	    }
	    if (argc > 2)
	    {
		if (isnumber(argv[2]))
		{
		    g_nHosts = atoi(argv[2]);
		    if (g_nHosts < 1)
		    {
			printf("Error: You must specify a number greater than 0 after -hosts.\n");
			return 0;
		    }
		    nArgsToStrip = 2 + g_nHosts;
		    int index = 3;
		    for (i=0; i<g_nHosts; i++)
		    {
			if (index >= argc)
			{
			    printf("Error: missing host name after -hosts option.\n");
			    return 0;
			}
			HostNode *pNode = new HostNode;
			pNode->next = NULL;
			pNode->nSMPProcs = 1;
			pNode->exe[0] = '\0';
			strcpy(pNode->host, argv[index]);
			index++;
			if (argc > index)
			{
			    if (isnumber(argv[index]))
			    {
				pNode->nSMPProcs = atoi(argv[index]);
				index++;
				nArgsToStrip++;
			    }
			}
			if (g_pHosts == NULL)
			{
			    g_pHosts = pNode;
			}
			else
			{
			    HostNode *pIter = g_pHosts;
			    while (pIter->next)
				pIter = pIter->next;
			    pIter->next = pNode;
			}
		    }
		}
		else
		{
		    printf("Error: You must specify the number of hosts after the -hosts option.\n");
		    return 0;
		}
	    }
	    else
	    {
		printf("Error: not enough arguments.\n");
		return 0;
	    }
	}
	else if (stricmp(&argv[1][1], "tcp") == 0)
	{
	    bDoSMP = false;
	}
	else if (stricmp(&argv[1][1], "getphrase") == 0)
	{
	    GetMPDPassPhrase(phrase);
	    bPhraseNeeded = false;
	}
	else if (stricmp(&argv[1][1], "nocolor") == 0)
	{
	    g_bDoMultiColorOutput = false;
	}
	else if (stricmp(&argv[1][1], "nompi") == 0)
	{
	    g_bNoMPI = true;
	}
	else if (stricmp(&argv[1][1], "nodots") == 0)
	{
	    bLogonDots = false;
	}
	else if (stricmp(&argv[1][1], "nomapping") == 0)
	{
	    g_bNoDriveMapping = true;
	}
	else if (stricmp(&argv[1][1], "nopopup_debug") == 0)
	{
	    SetErrorMode(SEM_FAILCRITICALERRORS | SEM_NOGPFAULTERRORBOX);
	}
	else if (stricmp(&argv[1][1], "help") == 0 || argv[1][1] == '?')
	{
	    PrintOptions();
	    return 0;
	}
	else if (stricmp(&argv[1][1], "help2") == 0)
	{
	    PrintExtraOptions();
	    return 0;
	}
	else
	{
	    printf("Unknown option: %s\n", argv[1]);
	}
	StripArgs(argc, argv, nArgsToStrip);
    }

    if (argc < 2)
    {
	printf("Error: no executable or configuration file specified\n");
	return 0;
    }

    // The next argument is the executable or a configuration file
    strncpy(g_pszExe, argv[1], MAX_CMD_LENGTH);
    g_pszExe[MAX_CMD_LENGTH-1] = '\0';

    // All the rest of the arguments are passed to the application
    g_pszArgs[0] = '\0';
    for (i = 2; i<argc; i++)
    {
	strncat(g_pszArgs, argv[i], MAX_CMD_LENGTH - 1 - strlen(g_pszArgs));
	if (i < argc-1)
	{
	    strncat(g_pszArgs, " ", MAX_CMD_LENGTH - 1 - strlen(g_pszArgs));
	}
    }

    if (g_nHosts == 0)
    {
	// If -np or -localonly options have not been specified, check if the first
	// parameter is an executable or a configuration file
	if (GetBinaryType(g_pszExe, &dwType) || (ParseConfigFile(g_pszExe) == PARSE_ERR_NO_FILE))
	{
	    g_nHosts = 1;
	    bRunLocal = true;
	}
    }

    // Fix up the executable name
    char pszTempExe[MAX_CMD_LENGTH], *namepart;
    if (g_pszExe[0] == '\\' && g_pszExe[1] == '\\')
    {
	strncpy(pszTempExe, g_pszExe, MAX_CMD_LENGTH);
	pszTempExe[MAX_CMD_LENGTH-1] = '\0';
    }
    else
	GetFullPathName(g_pszExe, MAX_PATH, pszTempExe, &namepart);
    // Quote the executable in case there are spaces in the path
    sprintf(g_pszExe, "\"%s\"", pszTempExe);

    easy_socket_init();

    if (!bRunLocal && g_pHosts == NULL)
    {
	// Save the original file name in case we end up running locally
	strncpy(pszTempExe, g_pszExe, MAX_CMD_LENGTH);
	pszTempExe[MAX_CMD_LENGTH-1] = '\0';

	// Convert the executable to its unc equivalent. This negates
	// the need to map network drives on remote machines just to locate
	// the executable.
	ExeToUnc(g_pszExe);

	// If we are not running locally and the hosts haven't been setup with a configuration file,
	// create the host list now
	if (bUseMachineFile)
	{
	    if (!GetHostsFromFile(pszMachineFileName))
	    {
		printf("Error parsing the machine file '%s'\n", pszMachineFileName);
		return 0;
	    }
	}
	else if (!GetAvailableHosts())
	{
	    strncpy(g_pszExe, pszTempExe, MAX_CMD_LENGTH);
	    g_pszExe[MAX_CMD_LENGTH-1] = '\0';
	    bRunLocal = true;
	}
    }

    // Setup multi-color output
    if (g_bDoMultiColorOutput)
    {
	char pszTemp[10];
	DWORD len = 10;
	if (ReadMPDRegistry("color", pszTemp, &len))
	{
	    g_bDoMultiColorOutput = (stricmp(pszTemp, "yes") == 0);
	}
    }
    if (g_bDoMultiColorOutput)
    {
	CONSOLE_SCREEN_BUFFER_INFO info;
	// Save the state of the console so it can be restored
	hStdout = GetStdHandle(STD_OUTPUT_HANDLE);
	GetConsoleScreenBufferInfo(hStdout, &info);
	g_ConsoleAttribute = info.wAttributes;
    }
    
    // Check if the directory needs to be mapped on the remote machines
    if (NeedToMap(g_pszDir, &cMapDrive, pszMapShare))
    {
	MapDriveNode *pNode = new MapDriveNode;
	pNode->cDrive = cMapDrive;
	strcpy(pNode->pszShare, pszMapShare);
	pNode->pNext = g_pDriveMapList;
	g_pDriveMapList = pNode;
    }

    // If -getphrase was not specified, get the mpd passphrase from
    // the registry or use the default
    if (bPhraseNeeded)
    {
	if (!ReadMPDRegistry("phrase", phrase, NULL))
	{
	    strcpy(phrase, MPD_DEFAULT_PASSPHRASE);
	}
    }

    char pmi_port_str[100];
    DWORD port_str_length = 100;
    if (ReadMPDRegistry("port", pmi_port_str, &port_str_length))
    {
	pmi_port = atoi(pmi_port_str);
	if (pmi_port < 1)
	    pmi_port = MPD_DEFAULT_PORT;
    }
    CreatePMIDatabase(pmi_host, pmi_port, phrase, pmi_kvsname);

    if (bRunLocal)
    {
	RunLocal(bDoSMP);
	DestroyPMIDatabase(pmi_host, pmi_port, phrase, pmi_kvsname);
	easy_socket_finalize();
	return 0;
    }

    //dbg_printf("retrieving account information\n");
    if (bLogon)
	GetAccountAndPassword();
    else
    {
	char pszTemp[10] = "no";
	ReadMPDRegistry("SingleUser", pszTemp, NULL);
	if (stricmp(pszTemp, "yes"))
	{
	    if (bLogonDots)
	    {
		DWORD dwThreadId;
		HANDLE hEvent, hDotThread;
		hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
		hDotThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)PrintDots, hEvent, 0, &dwThreadId);
		if (!ReadPasswordFromRegistry(g_pszAccount, g_pszPassword))
		{
		    SetEvent(hEvent);
		    GetAccountAndPassword();
		}
		else
		    SetEvent(hEvent);
		CloseHandle(hDotThread);
	    }
	    else
	    {
		if (!ReadPasswordFromRegistry(g_pszAccount, g_pszPassword))
		{
		    GetAccountAndPassword();
		}
	    }
	    bLogon = true;
	}
    }
    
    // Figure out how many processes to launch
    int nProc = 0;
    HostNode *n = g_pHosts;
    if (g_pHosts == NULL)
	nProc = g_nHosts;
    while (n)
    {
	nProc += n->nSMPProcs;
	n = n->next;
    }
    g_nNproc = nProc;
    
    // Set the environment variables common to all processes
    if (g_bNoMPI)
	pszEnv[0] = '\0';
    else
    {
	sprintf(pszEnv, "PMI_SIZE=%d|PMI_MPD=%s:%d|PMI_KVS=%s", nProc, pmi_host, pmi_port, pmi_kvsname);
    }
    
    // Allocate an array to hold handles to the LaunchProcess threads, sockets, ids, ranks, and forward host structures
    pThread = new HANDLE[nProc];
    g_pProcessSocket = new SOCKET[nProc];
    for (i=0; i<nProc; i++)
	g_pProcessSocket[i] = INVALID_SOCKET;
    g_pProcessLaunchId = new int[nProc];
    g_pLaunchIdToRank = new int [nProc];
    g_nNumProcessSockets = 0;
    g_pForwardHost = new ForwardHostStruct[nProc];
    for (i=0; i<nProc; i++)
	g_pForwardHost[i].nPort = 0;
    
    // Start the IO redirection thread
    HANDLE hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    g_hRedirectIOListenThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)RedirectIOThread, hEvent, 0, &dwThreadID);
    if (g_hRedirectIOListenThread)
    {
	if (WaitForSingleObject(hEvent, 10000) != WAIT_OBJECT_0)
	{
	    printf("RedirectIOThread failed to initialize\n");
	    DestroyPMIDatabase(pmi_host, pmi_port, phrase, pmi_kvsname);
	    return 0;
	}
    }
    else
    {
	printf("Unable to create RedirectIOThread, error %d\n", GetLastError());
	DestroyPMIDatabase(pmi_host, pmi_port, phrase, pmi_kvsname);
	return 0;
    }
    CloseHandle(hEvent);

    strncpy(g_pForwardHost[0].pszHost, g_pszIOHost, MAX_HOST_LENGTH);
    g_pForwardHost[0].pszHost[MAX_HOST_LENGTH-1] = '\0';
    g_pForwardHost[0].nPort = g_nIOPort;
    //printf("io redirection: %s:%d\n", g_pForwardHost[0].pszHost, g_pForwardHost[0].nPort);fflush(stdout);

    // Launch the threads to launch the processes
    iproc = 0;
    while (g_pHosts)
    {
	nShmLow = iproc;
	nShmHigh = iproc + g_pHosts->nSMPProcs - 1;
	for (int i = 0; i<g_pHosts->nSMPProcs; i++)
	{
	    MPIRunLaunchProcessArg *arg = new MPIRunLaunchProcessArg;
	    sprintf(arg->pszIOHostPort, "%s:%d", g_pszIOHost, g_nIOPort);
	    strcpy(arg->pszPassPhrase, phrase);
	    arg->i = iproc;
	    arg->bLogon = bLogon;
	    if (bLogon)
	    {
		strcpy(arg->pszAccount, g_pszAccount);
		strcpy(arg->pszPassword, g_pszPassword);
	    }
	    else
	    {
		arg->pszAccount[0] = '\0';
		arg->pszPassword[0] = '\0';
	    }
	    if (strlen(g_pHosts->exe) > 0)
	    {
		strncpy(arg->pszCmdLine, g_pHosts->exe, MAX_CMD_LENGTH);
		arg->pszCmdLine[MAX_CMD_LENGTH-1] = '\0';
	    }
	    else
	    {
		strncpy(arg->pszCmdLine, g_pszExe, MAX_CMD_LENGTH);
		arg->pszCmdLine[MAX_CMD_LENGTH-1] = '\0';
	    }
	    if (strlen(g_pszArgs) > 0)
	    {
		strncat(arg->pszCmdLine, " ", MAX_CMD_LENGTH - 1 - strlen(arg->pszCmdLine));
		strncat(arg->pszCmdLine, g_pszArgs, MAX_CMD_LENGTH - 1 - strlen(arg->pszCmdLine));
	    }
	    strcpy(arg->pszDir, g_pszDir);
	    if (strlen(pszEnv) >= MAX_CMD_LENGTH)
	    {
		printf("Warning: environment variables truncated.\n");
		fflush(stdout);
	    }
	    strncpy(arg->pszEnv, pszEnv, MAX_CMD_LENGTH);
	    arg->pszEnv[MAX_CMD_LENGTH-1] = '\0';
	    strncpy(arg->pszHost, g_pHosts->host, MAX_HOST_LENGTH);
	    arg->pszHost[MAX_HOST_LENGTH-1] = '\0';

	    if (g_bNoMPI)
	    {