1#!/bin/bash # -*- Mode: shell-script; indent-tabs-mode:nil -*-
2# (c) 2008-2013 Nathan Hjelm <hjelmn@cs.unm.edu>
3# mpirun completion v1.1
4#
5# Bash completion script for Open MPI's mpirun
6#
7# Installation:
8#   If bash completion is not already installed, follow the instructions at:
9#      http://anonscm.debian.org/gitweb/?p=bash-completion/bash-completion.git;a=blob_plain;f=README
10#
11# Then install mpirun completion by dropping this file into either the system bash_completion.d
12# (usually found in /etc) or somewhere in your home directory (~/.bash_completion.d is a good place).
13# If installing the script in a location other than the system-wide location, you will probably need
14# to add the following to either ~/.bash_profile or ~/.bash_completion:
15#
16#        [ -r /path/to/mpirun.sh ] && . /path/to/mpirun.sh
17
18# find available MCA variables
19_get_mca_variable_names() {
20    ompi_info -a --parsable | perl -e '
21my %values, %help, @uniq;
22
23$maxlen = 0;
24
25while (<>) {
26  if (/^mca:/) {
27    chop;
28    @tmp = split(":");
29    $length = length($tmp[4]);
30    $maxlen = ($length, $maxlen)[$length < $maxlen];
31    $values{$tmp[4]} = 1;
32    if ($tmp[5] eq "help") {
33      $help{$tmp[4]} = $tmp[6];
34      $help{$tmp[4]} =~ s/:/\\:/g;
35    }
36  }
37}
38
39for $exclude (split(" ", "'"$excl"'")) {
40  delete $values{$exclude};
41}
42
43for $key (sort(keys %values)) {
44  print $key;
45  if ($help{$key} ne "") {
46    print " " x ($maxlen - length($key)) . "- $help{$key}";
47  }
48
49  print "\n";
50}' 2> /dev/null
51}
52
53# given a framework (type) name print a list of components
54_get_mca_component_names() {
55    # components are printed by default (-a is not needed)
56    ompi_info -l 9 --parsable | grep "mca:$1:" | perl -e '
57my %values;
58
59while (<>) {
60  if (/^mca:/) {
61    @tmp = split(":");
62    $values{$tmp[2]} = 1;
63  }
64}
65
66@uniq = keys %values;
67print join(" ", sort(@uniq));'
68}
69
70_get_mpirun_switches() {
71    mpirun --help 2>&1 | sed 's/^\s*//g' | egrep '^-' | cut -d' ' -f1 | tr '|\n' ' '
72}
73
74# get enumerator values for a variable
75_get_enum_values() {
76    ompi_info -a -l 9 --parsable | grep ":$1:" | perl -e '
77my @values;
78
79while (<>) {
80  chop;
81  @tmp = split(":");
82
83  if ($tmp[5] eq "enumerator") {
84    $values[++$#values] = $tmp[7];
85    $values[++$#values] = $tmp[8];
86  }
87}
88print join(" ", @values);'
89}
90
91# remove items from $1 that are also in $2. lists must be sorted
92_set_remove () {
93    comm -23 <(echo $1 | sort | tr " " "\n") <(echo $2 | sort | tr " " "\n") 2>/dev/null
94}
95
96# find mca parameters specified on the command line (prevent duplicates)
97_find_mca_parameters() {
98    for ((i = 1; i < COMP_CWORD; i++)) ; do
99        # the subcommand is the first item on the command line that isn't module or a switch
100        if test ${COMP_WORDS[$i]##*-} == "mca" -a $i -lt $((COMP_CWORD-1)) ; then
101            echo ${COMP_WORDS[$i+1]}
102        fi
103    done
104}
105
106# check for matches that have the search term somewhere in the name (but not the description)
107_fuzzy_search() {
108    for match in $2 ; do
109        if [[ ${match%%[[:space:]]*} =~ .*"$1".* ]] ; then
110            echo $match
111        fi
112    done
113}
114
115# mpirun/orterun completion
116_mpirun() {
117    local cur prv tb switches already_specified all_variables avail_variables enumerations save_IFS
118    local disable_description=yes
119
120    # Enable descriptions if we can detect the completion type (so we can strip the description off
121    # if there is only one result. This is the case for normal completion (ASCII 9) or successive
122    # completion (?: ASCII 63)
123    if test "$COMP_TYPE" = "9" -o "$COMP_TYPE" = "63" ; then
124        disable_description=no
125    fi
126
127    COMPREPLY=()
128    if test $COMP_CWORD -gt 1 ; then
129        tb=${COMP_WORDS[COMP_CWORD-2]}
130    else
131        tb=""
132    fi
133
134    if test $COMP_CWORD -gt 0 ; then
135        prv=${COMP_WORDS[COMP_CWORD-1]}
136    else
137        prv=""
138    fi
139
140    cur=${COMP_WORDS[COMP_CWORD]}
141
142    if test "${prv##*-}" = "mca" ; then
143        # Complete variable name
144
145        # Temporarily change IFS to newline
146        save_IFS=$IFS
147        IFS="
148"
149
150        # Remove mca parameters already on the command line
151        avail_variables=($(_get_mca_variable_names "$(_find_mca_parameters)"))
152
153	IFS=" "
154        # Check if we should disable descriptions
155        if test "$disable_description" = "yes" ; then
156            avail_variables=(${avail_variables[*]%%[[:space:]]*})
157        fi
158	IFS="
159"
160
161        # Return a fuzzy-search of the mca parameter names
162        COMPREPLY=($(_fuzzy_search "$cur" "${avail_variables[*]}"))
163
164        # Remove the description if this is the last item
165        if test ${#COMPREPLY[@]} -eq 1 ; then
166            completion=${COMPREPLY[0]}
167            COMPREPLY=("${COMPREPLY[0]%%[[:space:]]*}")
168        fi
169        IFS=$save_IFS
170    elif test "${tb##*-}" = "mca" ; then
171        # Complete variable value
172
173        # Check if the variable is a selection variable (no _ in the name)
174        if test "${prv#_}" = "${prv}" ; then
175            # component selection variable, find available components (removing ones already selected)
176            enumerations=($(_set_remove "$(_get_mca_component_names $prv)" "$(echo $cur | tr ',' '\n')"))
177
178            # Prepend the current selection if there is one
179            if test "${cur%,*}" = "${cur}" ; then
180                COMPREPLY=($(_fuzzy_search "${cur##*,}" "${enumerations[*]}"))
181            else
182                COMPREPLY=($(_fuzzy_search "${cur##*,}" "${enumerations[*]}" | sed "s/^/${cur%,*},/g"))
183            fi
184        else
185            enumerations=($(_get_enum_values "$prv"))
186            COMPREPLY=($(_fuzzy_search "$cur" "${enumerations[*]}"))
187        fi
188    elif test "${prv}" = "--bind-to" ; then
189        COMPREPLY=($(compgen -W "none hwthread core socket numa board" -- "$cur"))
190    elif test "${prv}" = "--map-by" -o "${prv}" = "-map-by" ; then
191        COMPREPLY=($(compgen -W "slot hwthread core socket numa board node" -- "$cur"))
192    elif test "${prv##*-}" = "hostfile" ; then
193        COMPREPLY=($(compgen -f -- "$cur"))
194    elif test "${cur:0:1}" = "-" ; then
195        switches=$(_get_mpirun_switches)
196        COMPREPLY=($(compgen -W "$switches" -- "$cur"))
197    else
198        COMPREPLY=($(compgen -f -- "$cur"))
199    fi
200
201    return 0
202}
203
204complete -o nospace -F _mpirun mpirun
205complete -o nospace -F _mpirun orterun
206complete -o nospace -F _mpirun mpiexec
207