simple_pmi.c

mpi并行计算的c++代码 可用vc或gcc编译通过 可以用来搭建并行计算试验环境

/* -*- Mode: C; c-basic-offset:4 ; -*- */
/*  $Id: simple_pmi.c,v 1.84 2005/11/01 16:32:31 gropp Exp $
 *
 *  (C) 2001 by Argonne National Laboratory.
 *      See COPYRIGHT in top-level directory.
 */

/*********************** PMI implementation ********************************/

#include "pmiconf.h" 

#define PMI_VERSION    1
#define PMI_SUBVERSION 1

#include <stdio.h>
#ifdef HAVE_UNISTD_H
#include <unistd.h>
#endif
#ifdef HAVE_STDLIB_H
#include <stdlib.h>
#endif
#ifdef HAVE_STRING_H
#include <string.h>
#endif
#ifdef HAVE_STRINGS_H
#include <strings.h>
#endif
#ifdef USE_PMI_PORT
#ifndef MAXHOSTNAME
#define MAXHOSTNAME 256
#endif
#endif
/* This should be moved to pmiu for shutdown */
#if defined(HAVE_SYS_SOCKET_H)
#include <sys/socket.h>
#endif
/* mpimem includes the definitions for MPIU_Snprintfl, MPIU_Malloc, and 
   MPIU_Free */
#include "mpimem.h"

/* Temporary debug definitions */
#if 0
#define DBG_PRINTF(args) printf args ; fflush(stdout)
#define DBG_FPRINTF(args) fprintf args 
#else
#define DBG_PRINTF(args)
#define DBG_FPRINTF(args)
#endif

#include "pmi.h"
#include "simple_pmiutil.h"

/* 
   These are global variable used *ONLY* in this file, and are hence
   declared static.
 */


static int PMI_fd = -1;
static int PMI_size = 1;
static int PMI_rank = 0;

/* Set PMI_initialized to 1 for singleton init but no process manager
   to help.  Initialized to 2 for normal initialization.  Initialized
   to values higher than 2 when singleton_init by a process manager.
   All values higher than 1 invlove a PM in some way.
*/
#define SINGLETON_INIT_BUT_NO_PM 1
#define NORMAL_INIT_WITH_PM      2
#define SINGLETON_INIT_WITH_PM   3
static int PMI_initialized = 0;

/* ALL GLOBAL VARIABLES MUST BE INITIALIZED TO AVOID POLLUTING THE 
   LIBRARY WITH COMMON SYMBOLS */
static int PMI_kvsname_max = 0;
static int PMI_keylen_max = 0;
static int PMI_vallen_max = 0;

static int PMI_iter_next_idx = 0;
static int PMI_debug = 0;
static int PMI_spawned = 0;

static int PMII_getmaxes( int *kvsname_max, int *keylen_max, int *vallen_max );
static int PMII_iter( const char *kvsname, const int idx, int *nextidx, char *key, int key_len, char *val, int val_len );
static int PMII_Set_from_port( int, int );
static int PMII_Connect_to_pm( char *, int );

#ifdef USE_PMI_PORT
static int PMII_singinit(void);
static int PMI_totalview = 0;
#endif
static int accept_one_connection(int);
static char cached_singinit_key[PMIU_MAXLINE];
static char cached_singinit_val[PMIU_MAXLINE];

/******************************** Group functions *************************/

int PMI_Init( int *spawned )
{
    char *p;
    int notset = 1;
    char buf[PMIU_MAXLINE], cmd[PMIU_MAXLINE];

    /* setvbuf(stdout,0,_IONBF,0); */
    setbuf(stdout,NULL);
    /* PMIU_printf( 1, "PMI_INIT\n" ); */
    if ( ( p = getenv( "PMI_FD" ) ) )
	PMI_fd = atoi( p );
#ifdef USE_PMI_PORT
    else if ( ( p = getenv( "PMI_PORT" ) ) ) {
	int portnum;
	char hostname[MAXHOSTNAME];
	char *pn;
	int id = 0;
	/* Connect to the indicated port (in format hostname:portnumber) 
	   and get the fd for the socket */
	
	/* Split p into host and port */
	pn = strchr( p, ':' );

	if (PMI_debug) {
	    DBG_PRINTF( ("Connecting to %s\n", p) );
	}
	if (pn) {
	    MPIU_Strncpy( hostname, p, (pn - p) );
	    hostname[(pn-p)] = 0;
	    portnum = atoi( pn+1 );
	    /* FIXME: Check for valid integer after : */
	    /* This routine only gets the fd to use to talk to 
	       the process manager. The handshake below is used
	       to setup the initial values */
	    PMI_fd = PMII_Connect_to_pm( hostname, portnum );
	}
	/* FIXME: If PMI_PORT specified but either no valid value or
	   fd is -1, give an error return */
	if (PMI_fd < 0) return -1;

	/* We should first handshake to get size, rank, debug. */
	p = getenv( "PMI_ID" );
	if (p) {
	    id = atoi( p );
	    /* PMII_Set_from_port sets up the values that are delivered
	       by enviroment variables when a separate port is not used */
	    PMII_Set_from_port( PMI_fd, id );
	    notset = 0;
	}
    }
#endif
    else {
	PMI_fd = -1;
    }

    if ( PMI_fd == -1 ) {
	/* Singleton init: Process not started with mpiexec, 
	   so set size to 1, rank to 0 */
	PMI_size = 1;
	PMI_rank = 0;
	*spawned = 0;
	
	PMI_initialized = SINGLETON_INIT_BUT_NO_PM;
	PMI_kvsname_max = 256;
	PMI_keylen_max  = 256;
	PMI_vallen_max  = 256;
	
	return( 0 );
    }

    /* If size, rank, and debug not set from a communication port,
       use the environment */
    if (notset) {
	if ( ( p = getenv( "PMI_SIZE" ) ) )
	    PMI_size = atoi( p );
	else 
	    PMI_size = 1;
	
	if ( ( p = getenv( "PMI_RANK" ) ) ) {
	    PMI_rank = atoi( p );
	    /* Let the util routine know the rank of this process for 
	       any messages (usually debugging or error) */
	    PMIU_Set_rank( PMI_rank );
	}
	else 
	    PMI_rank = 0;
	
	if ( ( p = getenv( "PMI_DEBUG" ) ) )
	    PMI_debug = atoi( p );
	else 
	    PMI_debug = 0;

	/* Leave unchanged otherwise, which indicates that no value
	   was set */
    }

#ifdef USE_PMI_PORT
    if ( ( p = getenv( "PMI_TOTALVIEW" ) ) )
	PMI_totalview = atoi( p );
    if ( PMI_totalview ) {
	PMIU_readline( PMI_fd, buf, PMIU_MAXLINE );
	PMIU_parse_keyvals( buf );
	PMIU_getval( "cmd", cmd, PMIU_MAXLINE );
	if ( strncmp( cmd, "tv_ready", PMIU_MAXLINE ) != 0 ) {
	    PMIU_printf( 1, "expecting cmd=tv_ready, got %s\n", buf );
	    return( PMI_FAIL );
	}
    }
#endif

    PMII_getmaxes( &PMI_kvsname_max, &PMI_keylen_max, &PMI_vallen_max );

    if ( ( p = getenv( "PMI_SPAWNED" ) ) )
	PMI_spawned = atoi( p );
    else
	PMI_spawned = 0;
    if (PMI_spawned)
	*spawned = 1;
    else
	*spawned = 0;

    if ( ! PMI_initialized )
	PMI_initialized = NORMAL_INIT_WITH_PM;

    return( 0 );
}

int PMI_Initialized( PMI_BOOL *initialized )
{
    /* Turn this into a logical value (1 or 0) .  This allows us
       to use PMI_initialized to distinguish between initialized with
       an PMI service (e.g., via mpiexec) and the singleton init, 
       which has no PMI service */
    *initialized = PMI_initialized != 0 ? PMI_TRUE : PMI_FALSE;
    return PMI_SUCCESS;
}

int PMI_Get_size( int *size )
{
    if ( PMI_initialized )
	*size = PMI_size;
    else
	*size = 1;
    return( 0 );
}

int PMI_Get_rank( int *rank )
{
    if ( PMI_initialized )
	*rank = PMI_rank;
    else
	*rank = 0;
    return( 0 );
}

int PMI_Get_universe_size( int *size)
{
    char buf[PMIU_MAXLINE], cmd[PMIU_MAXLINE], size_c[PMIU_MAXLINE];
    int rc;

#ifdef USE_PMI_PORT
    if (PMI_initialized < 2)
    {
	rc = PMII_singinit();
	if (rc < 0)
	    return(-1);
	PMI_initialized = SINGLETON_INIT_WITH_PM;    /* do this right away */
	PMI_size = 1;
	PMI_rank = 0;
	PMI_debug = 0;
	PMI_spawned = 0;
	PMII_getmaxes( &PMI_kvsname_max, &PMI_keylen_max, &PMI_vallen_max );
	PMI_KVS_Put( "singinit_kvs_0", cached_singinit_key, cached_singinit_val );
    }
#endif

    if ( PMI_initialized > 1)  /* Ignore SINGLETON_INIT_BUT_NO_PM */
    {
	PMIU_writeline( PMI_fd, "cmd=get_universe_size\n" );
	PMIU_readline( PMI_fd, buf, PMIU_MAXLINE );
	PMIU_parse_keyvals( buf );
	PMIU_getval( "cmd", cmd, PMIU_MAXLINE );
	if ( strncmp( cmd, "universe_size", PMIU_MAXLINE ) != 0 ) {
	    PMIU_printf( 1, "expecting cmd=universe_size, got %s\n", buf );
	    return( PMI_FAIL );
	}
	else {
	    PMIU_getval( "size", size_c, PMIU_MAXLINE );
	    *size = atoi(size_c);
	    return( PMI_SUCCESS );
	}
    }
    else
	*size = 1;
    return( PMI_SUCCESS );
}

int PMI_Get_appnum( int *appnum )
{
    char buf[PMIU_MAXLINE], cmd[PMIU_MAXLINE], appnum_c[PMIU_MAXLINE];

    if ( PMI_initialized > 1)  /* Ignore SINGLETON_INIT_BUT_NO_PM */
    {
	PMIU_writeline( PMI_fd, "cmd=get_appnum\n" );
	PMIU_readline( PMI_fd, buf, PMIU_MAXLINE );
	PMIU_parse_keyvals( buf );
	PMIU_getval( "cmd", cmd, PMIU_MAXLINE );
	if ( strncmp( cmd, "appnum", PMIU_MAXLINE ) != 0 ) {
	    PMIU_printf( 1, "expecting cmd=appnum, got %s\n", buf );
	    return( PMI_FAIL );
	}
	else {
	    PMIU_getval( "appnum", appnum_c, PMIU_MAXLINE );
	    *appnum = atoi(appnum_c);
	    return( PMI_SUCCESS );
	}
    }
    else
	*appnum = -1;

    return( PMI_SUCCESS );
}

int PMI_Barrier( )
{
    char buf[PMIU_MAXLINE], cmd[PMIU_MAXLINE];

    if ( PMI_initialized > 1)  /* Ignore SINGLETON_INIT_BUT_NO_PM */
    {
	PMIU_writeline( PMI_fd, "cmd=barrier_in\n");
	PMIU_readline( PMI_fd, buf, PMIU_MAXLINE );
	PMIU_parse_keyvals( buf );
	PMIU_getval( "cmd", cmd, PMIU_MAXLINE );
	if ( strncmp( cmd, "barrier_out", PMIU_MAXLINE ) != 0 ) {
	    PMIU_printf( 1, "expecting cmd=barrier_out, got %s\n", buf );
	    return( -1 );
	}
	else
	    return( 0 );
    }
    else
	return( 0 );
}

int PMI_Get_clique_size( int *size )
{
    *size = 1;
    return PMI_SUCCESS;
}

int PMI_Get_clique_ranks( int ranks[], int length )
{
    if ( length < 1 )
	return PMI_ERR_INVALID_ARG;
    else
	return PMI_Get_rank( &ranks[0] );
}

/* Inform the process manager that we're in finalize */
int PMI_Finalize( )
{
    char buf[PMIU_MAXLINE], cmd[PMIU_MAXLINE];

    if ( PMI_initialized > 1)  /* Ignore SINGLETON_INIT_BUT_NO_PM */
    {
	PMIU_writeline( PMI_fd, "cmd=finalize\n" );
	PMIU_readline( PMI_fd, buf, PMIU_MAXLINE );
	PMIU_parse_keyvals( buf );
	PMIU_getval( "cmd", cmd, PMIU_MAXLINE );
	if ( strncmp( cmd, "finalize_ack", PMIU_MAXLINE ) != 0 ) {
	    PMIU_printf( 1, "expecting cmd=finalize_ack, got %s\n", buf );
	    return( -1 );
	}
	shutdown( PMI_fd, SHUT_RDWR );
	close( PMI_fd );
    }
    return( 0 );
}

int PMI_Abort(int exit_code, const char error_msg[])
{
    PMIU_printf(1, "aborting job:\n%s\n", error_msg);
    exit(exit_code);
    return -1;
}

/**************************************** Keymap functions *************************/

int PMI_KVS_Get_my_name( char kvsname[], int length )
{
    char buf[PMIU_MAXLINE], cmd[PMIU_MAXLINE];

    if (PMI_initialized == SINGLETON_INIT_BUT_NO_PM) {
	/* Return a dummy name */
	MPIU_Strncpy( kvsname, "singinit_kvs_0", PMIU_MAXLINE );
	return 0;
    }
    PMIU_writeline( PMI_fd, "cmd=get_my_kvsname\n" );
    PMIU_readline( PMI_fd, buf, PMIU_MAXLINE );
    PMIU_parse_keyvals( buf );
    PMIU_getval( "cmd", cmd, PMIU_MAXLINE );
    if ( strncmp( cmd, "my_kvsname", PMIU_MAXLINE ) != 0 ) {
	PMIU_printf( 1, "got unexpected response to get_my_kvsname :%s:\n", buf );
	return( -1 );
    }
    else {
	PMIU_getval( "kvsname", kvsname, PMI_kvsname_max );
	return( 0 );
    }
}

int PMI_KVS_Get_name_length_max( int *maxlen )
{
    if (maxlen == NULL)
	return PMI_ERR_INVALID_ARG;
    *maxlen = PMI_kvsname_max;
    return PMI_SUCCESS;