hwrf_wcoss.py

##@namespace hwrf_wcoss
# This module contains various utility functions for running on
# WCOSS, the machine that the US National Weather Service uses for its
# operational forecasting.

import os

##@var __all__
# Symbols to export by "from hwrf_wcoss import *"
__all__=['set_vars_for_coupled_hwrf','set_vars_for_gsi',
         'set_vars_for_products']

def loudenv(logger,var,value):
    """!Set an environment variable and log a message about it
    @param logger a logging.Logger for the message
    @param var the variable to set in os.environ
    @param value the value of the variable"""
    logger.info('Set ENV %s=%s',str(var),repr(value))
    os.environ[var]=value

def fcst_wcoss1(logger,coupled):
    """!Sets task geometry for Phase 1 WCOSS
    @param logger a logging.Logger for messages
    @param coupled if True, a coupled forecast, if False, uncoupled"""
    if coupled:
        logger.info('SETTING WCOSS PHASE 1 VARIABLES FOR HWRF COUPLED FORECAST')
        # Take this, you 72 character line purists:
        loudenv(logger,'LSB_PJL_TASK_GEOMETRY','{(0,1,2,3,4,5,6,7,8,9,10,11,12)(13,14,15,16,29,30,31,32,45,46,47,48,61,62,63,64)(17,18,19,20,33,34,35,36,49,50,51,52,65,66,67,68)(21,22,23,24,37,38,39,40,53,54,55,56,69,70,71,72)(25,26,27,28,41,42,43,44,57,58,59,60,73,74,75,76)(77,78,79,80,93,94,95,96,109,110,111,112,125,126,127,128)(81,82,83,84,97,98,99,100,113,114,115,116,129,130,131,132)(85,86,87,88,101,102,103,104,117,118,119,120,133,134,135,136)(89,90,91,92,105,106,107,108,121,122,123,124,137,138,139,140)(141,142,143,144,157,158,159,160,173,174,175,176,189,190,191,192)(145,146,147,148,161,162,163,164,177,178,179,180,193,194,195,196)(149,150,151,152,165,166,167,168,181,182,183,184,197,198,199,200)(153,154,155,156,169,170,171,172,185,186,187,188,201,202,203,204)(205,206,207,208,221,222,223,224,237,238,239,240,253,254,255,256)(209,210,211,212,225,226,227,228,241,242,243,244,257,258,259,260)(213,214,215,216,229,230,231,232,245,246,247,248,261,262,263,264)(217,218,219,220,233,234,235,236,249,250,251,252,265,266,267,268)(269,270,271,272,285,286,287,288,301,302,303,304,317,318,319,320)(273,274,275,276,289,290,291,292,305,306,307,308,321,322,323,324)(277,278,279,280,293,294,295,296,309,310,311,312,325,326,327,328)(281,282,283,284,297,298,299,300,313,314,315,316,329,330,331,332)(333,334,335,336,349,350,351,352,365,366,367,368,381,382,383,384)(337,338,339,340,353,354,355,356,369,370,371,372,385,386,387,388)(341,342,343,344,357,358,359,360,373,374,375,376,389,390,391,392)(345,346,347,348,361,362,363,364,377,378,379,380,393,394,395,396)(397,398,399,400,413,414,415,416,429,430,431,432,445,446,447,448)(401,402,403,404,417,418,419,420,433,434,435,436,449,450,451,452)(405,406,407,408,421,422,423,424,437,438,439,440,453,454,455,456)(409,410,411,412,425,426,427,428,441,442,443,444,457,458,459,460)(461,462,463,464,465,466,467,468,469,470,471,472)}')
    else:
        logger.info('SETTING WCOSS PHASE 1 VARIABLES FOR HWRF ATM-ONLY FORECAST')
        # Take this, you 72 character line purists:
        loudenv(logger,'LSB_PJL_TASK_GEOMETRY','{(0,1,2,3,16,17,18,19,32,33,34,35,48,49,50,51)(4,5,6,7,20,21,22,23,36,37,38,39,52,53,54,55)(8,9,10,11,24,25,26,27,40,41,42,43,56,57,58,59)(12,13,14,15,28,29,30,31,44,45,46,47,60,61,62,63)(64,65,66,67,80,81,82,83,96,97,98,99,112,113,114,115)(68,69,70,71,84,85,86,87,100,101,102,103,116,117,118,119)(72,73,74,75,88,89,90,91,104,105,106,107,120,121,122,123)(76,77,78,79,92,93,94,95,108,109,110,111,124,125,126,127)(128,129,130,131,144,145,146,147,160,161,162,163,176,177,178,179)(132,133,134,135,148,149,150,151,164,165,166,167,180,181,182,183)(136,137,138,139,152,153,154,155,168,169,170,171,184,185,186,187)(140,141,142,143,156,157,158,159,172,173,174,175,188,189,190,191)(192,193,194,195,208,209,210,211,224,225,226,227,240,241,242,243)(196,197,198,199,212,213,214,215,228,229,230,231,244,245,246,247)(200,201,202,203,216,217,218,219,232,233,234,235,248,249,250,251)(204,205,206,207,220,221,222,223,236,237,238,239,252,253,254,255)(256,257,258,259,272,273,274,275,288,289,290,291,304,305,306,307)(260,261,262,263,276,277,278,279,292,293,294,295,308,309,310,311)(264,265,266,267,280,281,282,283,296,297,298,299,312,313,314,315)(268,269,270,271,284,285,286,287,300,301,302,303,316,317,318,319)(320,321,322,323,336,337,338,339,352,353,354,355,368,369,370,371)(324,325,326,327,340,341,342,343,356,357,358,359,372,373,374,375)(328,329,330,331,344,345,346,347,360,361,362,363,376,377,378,379)(332,333,334,335,348,349,350,351,364,365,366,367,380,381,382,383)(384,385,386,387,400,401,402,403,416,417,418,419,432,433,434,435)(388,389,390,391,404,405,406,407,420,421,422,423,436,437,438,439)(392,393,394,395,408,409,410,411,424,425,426,427,440,441,442,443)(396,397,398,399,412,413,414,415,428,429,430,431,444,445,446,447)(448,449,450,451,452,453,454,455,456,457,458,459)}')

def fcst_wcoss2(logger,coupled):
    """!Sets task geometry for Phase 2 WCOSS
    @param logger a logging.Logger for messages
    @param coupled if True, a coupled forecast, if False, uncoupled"""
    if coupled:
        logger.info('SETTING WCOSS PHASE 2 VARIABLES FOR HWRF COUPLED FORECAST')
        # Take this, you 72 character line purists:
        loudenv(logger,'LSB_PJL_TASK_GEOMETRY','{(0,1,2,3,4,5,6,7,8,9,10,11,12)(13,14,15,16,29,30,31,32,45,46,47,48,61,62,63,64,77,78,79,80,93,94,95,96)(17,18,19,20,33,34,35,36,49,50,51,52,65,66,67,68,81,82,83,84,97,98,99,100)(21,22,23,24,37,38,39,40,53,54,55,56,69,70,71,72,85,86,87,88,101,102,103,104)(25,26,27,28,41,42,43,44,57,58,59,60,73,74,75,76,89,90,91,92,105,106,107,108)(109,110,111,112,125,126,127,128,141,142,143,144,157,158,159,160,173,174,175,176,189,190,191,192)(113,114,115,116,129,130,131,132,145,146,147,148,161,162,163,164,177,178,179,180,193,194,195,196)(117,118,119,120,133,134,135,136,149,150,151,152,165,166,167,168,181,182,183,184,197,198,199,200)(121,122,123,124,137,138,139,140,153,154,155,156,169,170,171,172,185,186,187,188,201,202,203,204)(205,206,207,208,221,222,223,224,237,238,239,240,253,254,255,256,269,270,271,272,285,286,287,288)(209,210,211,212,225,226,227,228,241,242,243,244,257,258,259,260,273,274,275,276,289,290,291,292)(213,214,215,216,229,230,231,232,245,246,247,248,261,262,263,264,277,278,279,280,293,294,295,296)(217,218,219,220,233,234,235,236,249,250,251,252,265,266,267,268,281,282,283,284,297,298,299,300)(301,302,303,304,317,318,319,320,333,334,335,336,349,350,351,352,365,366,367,368,381,382,383,384)(305,306,307,308,321,322,323,324,337,338,339,340,353,354,355,356,369,370,371,372,385,386,387,388)(309,310,311,312,325,326,327,328,341,342,343,344,357,358,359,360,373,374,375,376,389,390,391,392)(313,314,315,316,329,330,331,332,345,346,347,348,361,362,363,364,377,378,379,380,393,394,395,396)(397,398,399,400,413,414,415,416,429,430,431,432,445,446,447,448,461,462,463,464,477,478,479,480)(401,402,403,404,417,418,419,420,433,434,435,436,449,450,451,452,465,466,467,468,481,482,483,484)(405,406,407,408,421,422,423,424,437,438,439,440,453,454,455,456,469,470,471,472,485,486,487,488)(409,410,411,412,425,426,427,428,441,442,443,444,457,458,459,460,473,474,475,476,489,490,491,492)(493,494,495,496,497,498,499,500,501,502,503,504)}')
    else:
        logger.info('SETTING WCOSS PHASE 2 VARIABLES FOR HWRF ATM-ONLY FORECAST')
        # Take this, you 72 character line purists:
        loudenv(logger,'LSB_PJL_TASK_GEOMETRY','{(0,1,2,3,16,17,18,19,32,33,34,35,48,49,50,51,64,65,66,67,80,81,82,83)(4,5,6,7,20,21,22,23,36,37,38,39,52,53,54,55,68,69,70,71,84,85,86,87)(8,9,10,11,24,25,26,27,40,41,42,43,56,57,58,59,72,73,74,75,88,89,90,91)(12,13,14,15,28,29,30,31,44,45,46,47,60,61,62,63,76,77,78,79,92,93,94,95)(96,97,98,99,112,113,114,115,128,129,130,131,144,145,146,147,160,161,162,163,176,177,178,179)(100,101,102,103,116,117,118,119,132,133,134,135,148,149,150,151,164,165,166,167,180,181,182,183)(104,105,106,107,120,121,122,123,136,137,138,139,152,153,154,155,168,169,170,171,184,185,186,187)(108,109,110,111,124,125,126,127,140,141,142,143,156,157,158,159,172,173,174,175,188,189,190,191)(192,193,194,195,208,209,210,211,224,225,226,227,240,241,242,243,256,257,258,259,272,273,274,275)(196,197,198,199,212,213,214,215,228,229,230,231,244,245,246,247,260,261,262,263,276,277,278,279)(200,201,202,203,216,217,218,219,232,233,234,235,248,249,250,251,264,265,266,267,280,281,282,283)(204,205,206,207,220,221,222,223,236,237,238,239,252,253,254,255,268,269,270,271,284,285,286,287)(288,289,290,291,304,305,306,307,320,321,322,323,336,337,338,339,352,353,354,355,368,369,370,371)(292,293,294,295,308,309,310,311,324,325,326,327,340,341,342,343,356,357,358,359,372,373,374,375)(296,297,298,299,312,313,314,315,328,329,330,331,344,345,346,347,360,361,362,363,376,377,378,379)(300,301,302,303,316,317,318,319,332,333,334,335,348,349,350,351,364,365,366,367,380,381,382,383)(384,385,386,387,400,401,402,403,416,417,418,419,432,433,434,435,448,449,450,451,464,465,466,467)(388,389,390,391,404,405,406,407,420,421,422,423,436,437,438,439,452,453,454,455,468,469,470,471)(392,393,394,395,408,409,410,411,424,425,426,427,440,441,442,443,456,457,458,459,472,473,474,475)(396,397,398,399,412,413,414,415,428,429,430,431,444,445,446,447,460,461,462,463,476,477,478,479)(480,481,482,483,484,485,486,487,488,489,490,491)}')

def set_vars_for_products(logger):
    """!Sets variables to speed up the products job on WCOSS
    @param logger a logging.Logger for log messages"""
    loudenv(logger,'MP_USE_BULK_XFER','yes')
    loudenv(logger,'MP_EUIDEVICE','sn_single')
    loudenv(logger,'MP_EUILIB','us')
    loudenv(logger,'OMP_NUM_THREADS','1')
    loudenv(logger,'MKL_NUM_THREADS','1')
    loudenv(logger,'FORT_BUFFERED','TRUE')

    loudenv(logger,'MP_POLLING_INTERVAL',str(3*3600000))
    loudenv(logger,'MP_RETRANSMIT_INTERVAL',str(3*3600000))

def set_vars_for_coupled_hwrf(logger,wave,ocean,wrf_ranks):
    """!Sets variables that seem to speed up the H15B (2015
    pre-baseline) version of HWRF on Phase I WCOSS.
    @param logger a logging.Logger for log messages
    @param wave "WW3" or the empty string
    @param ocean "POM", "HYCOM" or the empty string
    @param wrf_ranks Number of WRF ranks.  If this number is small,
      then the function assumes a custom task count is in use, and will
      not set task geometry.  This is a kludge used to prevent 2km task
      geometry settings for the 3km configuration. """

    loudenv(logger,'MP_USE_BULK_XFER','yes')
    loudenv(logger,'MP_EUIDEVICE','sn_full')
    loudenv(logger,'MP_EUILIB','us')
    loudenv(logger,'MPICH_ALLTOALL_THROTTLE','0')
    loudenv(logger,'MP_EAGER_LIMIT','100K')
    loudenv(logger,'OMP_NUM_THREADS','1')
    loudenv(logger,'MKL_NUM_THREADS','1')
    loudenv(logger,'FORT_BUFFERED','TRUE')
    which_wcoss=os.environ.get('WHICH_WCOSS','2')

    if not wave and not ocean:
        coupled=False
    elif not wave and ocean=='HYCOM':
        coupled=True
    else:
        logger.warning('No known LSB_PJL_TASK_GEOMETRY for this '
                       'configuration.  Will not set that variable.')
        return

    if wrf_ranks<400:
        logger.warning('Not setting LSB_PJL_TASK_GEOMETRY due to small '
                       'wrf rank count %d'%wrf_ranks)
    elif which_wcoss=='2':
        fcst_wcoss2(logger,coupled)
    elif which_wcoss=='1':
        fcst_wcoss1(logger,coupled)
    else:
        raise AssertionError('OOPES!  FORGOT TO SET WHICH_WCOSS (%s)'
                            %(which_wcoss,))

def set_vars_for_ensda_hwrf(logger):
    """!Sets variables that seem to speed up the H15B (2015
    pre-baseline) version of HWRF on Phase 2 WCOSS when running ensda
    jobs.

    @param logger a logging.Logger for log messages"""

    loudenv(logger,'MP_USE_BULK_XFER','yes')
    loudenv(logger,'MP_EUIDEVICE','sn_full')
    loudenv(logger,'MP_EUILIB','us')
    loudenv(logger,'MPICH_ALLTOALL_THROTTLE','0')
    loudenv(logger,'MP_EAGER_LIMIT','100K')
    loudenv(logger,'OMP_NUM_THREADS','1')
    loudenv(logger,'MKL_NUM_THREADS','1')
    loudenv(logger,'FORT_BUFFERED','TRUE')

def set_vars_for_gsi(logger):
    """!Set variables for running GSI quickly on WCOSS.

    This routine is copied verbatim from the old ksh-based
    hwrf_nmm_gsi_d*.sh scripts and modified line-by-line for Python.
    It speeds up the 2014 HWRF GSI jobs by more than a factor of
    three.
    @param logger a logging.Logger for log messages"""

    logger.info('SETTING WCOSS VARIABLES FOR GSI')

    ################# set environment variables ##################
    # Set environment variables for NCEP IBM
    loudenv(logger,'MEMORY_AFFINITY','MCM')
    loudenv(logger,'MP_SHARED_MEMORY','yes')

    # Set environment variables for threads
    loudenv(logger,'AIXTHREAD_SCOPE','S')
    loudenv(logger,'XLSMPOPTS',"parthds=1:stack=128000000")

    # Recommended MPI environment variable setttings from IBM
    # (Appendix E, HPC Clusters Using InfiniBand on IBM Power Systems
    # Servers).  
    loudenv(logger,'LAPI_DEBUG_ENABLE_AFFINITY','YES')
    loudenv(logger,'MP_FIFO_MTU','4K')
    loudenv(logger,'MP_SYNC_QP','YES')
    loudenv(logger,'MP_SHM_ATTACH_THRESH','500000')
    loudenv(logger,'MP_EUIDEVELOP','min')
    loudenv(logger,'MP_USE_BULK_XFER','yes')
    loudenv(logger,'MP_BULK_MIN_MSG_SIZE','64k')
    loudenv(logger,'MP_RC_MAX_QP','8192')
    loudenv(logger,'LAPI_DEBUG_RC_DREG_THRESHOLD','1000000')
    loudenv(logger,'LAPI_DEBUG_QP_NOTIFICATION','no')
    loudenv(logger,'LAPI_DEBUG_RC_INIT_SETUP','yes')

    # Set environment variables for user preferences
    loudenv(logger,'XLFRTEOPTS',"nlwidth=80")
    loudenv(logger,'MP_LABELIO','yes')
    loudenv(logger,'MP_INFOLEVEL','1')

    loudenv(logger,'MP_USE_TOKEN_FLOW_CONTROL','yes')