source: fedd/compose.py @ 023e79b

#!/usr/local/bin/python

import sys
import re
import os
import random
import copy

from optparse import OptionParser, OptionValueError
from federation.remote_service import service_caller
from federation.service_error import service_error
from federation import topdl

class constraint:
    """
    This is just a struct to hold constraint fields, really
    """
    def __init__(self, name=None, required=False, accepts=None, provides=None,
            match=None):
        self.name = name
        self.required = required
        self.accepts = accepts or []
        self.provides = provides or []
        self.match = match

    def __str__(self):
        return "%s:%s:%s:%s" % (self.name, self.required,
                ",".join(self.provides), ",".join(self.accepts))

        
class ComposeOptionParser(OptionParser):
    """
    This class encapsulates the options to this script in one place.  It also
    holds the callback for the multifile choice.
    """
    def __init__(self):
        OptionParser.__init__(self)
        self.add_option('--url', dest='url', default="http://localhost:13235",
                help='url of ns2 to topdl service')
        self.add_option('--certfile', dest='cert', default=None,
                help='Certificate to use as identity')
        self.add_option('--seed', dest='seed', type='int', default=None,
                help='Random number seed')
        self.add_option('--multifile', dest='files', default=[], type='string', 
                action='callback', callback=self.multi_callback, 
                help="Include file multiple times")
        self.add_option('--output', dest='outfile', default=None,
                help='Output file name')
        self.add_option('--format', dest='format', default="xml",
                help='Output file format')
        self.add_option('--add_testbeds', dest='add_testbeds', default=False,
                action='store_true',
                help='add testbed attributes to each component')
        self.add_option('--lax', dest='lax', default=False,
                action='store_true',
                help='allow solutions where unrequired constraints fail')

    @staticmethod
    def multi_callback(option, opt_str, value, parser):
        """
        Parse a --multifile command line option.  The parameter is of the form
        filename,count.  This splits the argument at the rightmost comma and
        inserts the filename, count tuple into the "files" option.  It handles
        a couple error cases, too.
        """
        idx = value.rfind(',')
        if idx != -1:
            try:
                parser.values.files.append((value[0:idx], int(value[idx+1:])))
            except ValueError, e:
                raise OptionValueError(
                        "Can't convert %s to int in multifile (%s)" % \
                                (value[idx+1:], value))
        else:
            raise OptionValueError(
                    "Bad format (need a comma) for multifile: %s" % value)

def warn(msg):
    """
    Exactly what you think.  Print a message to stderr
    """
    print >>sys.stderr, msg


def make_new_name(names, prefix="name"):
    """
    Generate an identifier not present in names by appending an integer to
    prefix.  The new name is inserted into names and returned.
    """
    i = 0
    n = "%s%d" % (prefix,i)
    while n in names:
        i += 1
        n = "%s%d" % (prefix,i)
    names.add(n)
    return n

def add_interfaces(top, mark):
    """
    Add unconnected interfaces to the nodes whose names are keys in the mark
    dict.  As the interface is added change the name field of the contstraint
    in mark into a tuple containing the node name, interface name, and topology
    in which they reside.
    """
    for e in top.elements:
        if e.name in mark:
            con = mark[e.name]
            ii = make_new_name(set([x.name for x in e.interface]), "inf")
            e.interface.append(
                    topdl.Interface(substrate=[], name=ii, element=e))
            con.name = (e.name, ii, top)

def localize_names(top, names, marks):
    """
    Take a topology and rename any substrates or elements that share a name
    with an existing computer or substrate.  Keep the original name as a
    localized_name attribute.  In addition, remap any constraints or interfaces
    that point to the old name over to the new one.  Constraints are found in
    the marks dict, indexed by node name.  Those constraints name attributes
    have already been converted to triples (node name, interface name,
    topology) so only the node name needs to be changed.
    """
    sub_map = { }
    for s in top.substrates:
        s.set_attribute('localized_name', s.name)
        if s.name in names:
            sub_map[s.name] = n = make_new_name(names, "substrate")
            s.name = n
        else:
            names.add(s.name)

    for e in [ e for e in top.elements if isinstance(e, topdl.Computer)]:
        e.set_attribute('localized_name', e.name)
        if e.name in names:
            nn= make_new_name(names, "computer")
            if e.name in marks:
                n, i, t = marks[e.name].name
                marks[e.name].name = (nn, i, t)
            e.name = nn
        else:
            names.add(e.name)

        # Interface mapping.  The list comprehension creates a list of
        # substrate names where each element in the list is replaced by the
        # entry in sub_map indexed by it if present and left alone otherwise.
        for i in e.interface:
            i.substrate = [ sub_map.get(ii, ii) for ii in i.substrate ]

def meet_constraints(candidates, provides, accepts):
    """
    Try to meet all the constraints in candidates using the information in the
    provides and accepts dicts (which index constraints that provide or accept
    the given attribute). A constraint is met if it can be matched with another
    constraint that provides an attribute that the first constraint accepts.
    Only one match per pair is allowed, and we always prefer matching a
    required constraint to an unreqired one.  If all the candidates can be
    matches, return True and return False otherwise.
    """
    got_all = True
    for c in candidates:
        if not c.match:
            rmatch = None   # Match to a required constraint
            umatch = None   # Match to an unrequired constraint
            for a in c.accepts:
                for can in provides.get(a,[]):
                    # A constraint cannot satisfy itself, nor do we allow loops
                    # within the same topology.  This also excludes matched
                    # constraints.
                    if can.name != c.name and can.name[2] != c.name[2] and \
                            not can.match:
                        # Now check that the candidate also accepts c
                        for ca in can.accepts:
                            if ca in c.provides:
                                if can.required: rmatch = can
                                else: umatch = can
                                break
                        
                        # Within providers, prefer matches against required
                        # composition points. 
                        if rmatch:
                            break
                # Within acceptance categories, prefer matches against required
                # composition points
                if rmatch:
                    break

            # Move the better match over to the match variable
            if rmatch: match = rmatch
            elif umatch: match = umatch
            else: match = None

            # Done checking all possible matches.  Record the match or note an
            # unmatched candidate.
            if match:
                match.match = c
                c.match = match
            else:
                got_all = False
    return got_all

def randomize_constraint_order(indexes):
    """
    Randomly reorder the lists of constraints that provides and accepts hold.
    """
    if not isinstance(indexes, tuple):
        indexes = (indexes,)

    for idx in indexes:
        for k in idx.keys():
            random.shuffle(idx[k])

def remote_ns2topdl(uri, desc, cert):
    """
    Call a remote service to convert the ns2 to topdl (and in fact all the way
    to a topdl.Topology.
    """

    req = { 'description' : { 'ns2description': desc }, }

    r = service_caller('Ns2Topdl')(uri, req, cert)

    if r.has_key('Ns2TopdlResponseBody'):
        r = r['Ns2TopdlResponseBody']
        ed = r.get('experimentdescription', None)
        if 'topdldescription' in ed:
            return topdl.Topology(**ed['topdldescription'])
        else:
            return None
    else:
        return None

def connect_composition_points(top, contraints, names):
    """
    top is a topology containing copies of all the topologies represented in
    the contsraints, flattened into one name space.  This routine inserts the
    additional substrates required to interconnect the topology as described by
    the constraints.  After the links are made, unused connection points are
    purged.
    """
    done = set()
    for c in constraints:
        if c not in done and c.match:
            # c is an unprocessed matched constraint.  Create a substrate
            # and attach it to the interfaces named by c and c.match.
            sn = make_new_name(names, "sub")
            s = topdl.Substrate(name=sn)
            connected = 0
            # Walk through all the computers in the topology
            for e in [ e for e in top.elements \
                    if isinstance(e, topdl.Computer)]:
                # if the computer matches the computer and interface name in
                # either c or c.match, connect it.  Once both computers have
                # been found and connected, exit the loop walking through all
                # computers.
                for comp, inf in (c.name[0:2], c.match.name[0:2]):
                    if e.name == comp:
                        for i in e.interface:
                            if i.name == inf:
                                i.substrate.append(sn)
                                connected += 1
                                break
                        break
                # Connected both, so add the substrate to the topology
                if connected == 2: 
                    top.substrates.append(s)
                    break
            # c and c.match have been processed, so add them to the done set
            done.add(c)
            done.add(c.match)

    # All interfaces with matched constraints have been connected.  Cull any
    # interfaces unassigned to a substrate
    for e in [ e for e in top.elements if isinstance(e, topdl.Computer)]:
        if any([ not i.substrate for i in e.interface]):
            e.interface = [ i for i in e.interface if i.substrate ]

    top.incorporate_elements()

def import_ns2_constraints(contents):
    """
    Contents is a list containing the lines of an annotated ns2 file.  This
    routine extracts the constraint comment lines and convertes them into
    constraints in the namespace of the tcl experiment, as well as inserting
    them in the accepts and provides indices.

    Constraints are given in lines of the form
        name:required:provides:accepts
    where name is the name of the node in the current topology, if the second
    field is "required" this is a required constraint, a list of attributes
    that this connection point provides and a list that it accepts.  The
    provides and accepts lists are comma-separated.  The constraint is added to
    the marks dict under the name key and that dict is returned.
    """

    const_re = re.compile("\s*#\s*COMPOSITION:\s*([^:]+:[^:]+:.*)")

    marks = { }
    for l in contents:
        m = const_re.search(l)
        if m:
            exp = re.sub('\s', '', m.group(1))
            nn, r, p, a = exp.split(":")
            if nn.find('(') != -1:
                # Convert array references into topdl names
                nn = re.sub('\\((\\d)\\)', '-\\1', nn)
            p = p.split(",")
            a = a.split(",")
            c = constraint(name=nn, required=(r == 'required'),
                    provides=p, accepts=a)
            marks[nn] = c
    return marks

def add_constraint_attributes(top, marks):
    """
    Take the constraints in marks and add attributes representing them to the
    nodes in top.
    """
    for e in [e for e in top.elements if isinstance(e, topdl.Computer)]:
        if e.name in marks and not e.get_attribute('composition_point'):
            c = marks[e.name]
            e.set_attribute('composition_point', 'true')
            e.set_attribute('required', '%s' % c.required)
            e.set_attribute('provides', ",".join(c.provides))
            e.set_attribute('accepts', ",".join(c.accepts))

def remove_constraint_attributes(top, marks):
    """
    Remove constraint attributes that match the ones in the dict from the
    topology.
    """
    for e in [e for e in top.elements \
            if isinstance(e, topdl.Computer) and e.name in marks]:
        e.remove_attribute('composition_point')
        e.remove_attribute('required')
        e.remove_attribute('provides')
        e.remove_attribute('accepts')

def import_ns2_component(fn):
    """
    Pull a composition component in from an ns2 description.  The Constraints
    are parsed from the comments using import_ns2_constraints and the topology
    is created using a call to a fedd exporting the Ns2Topdl function.  A
    topdl.Topology object rempresenting the component's topology and a dict
    mapping constraints from the components names to the conttraints is
    returned.  If either the file read or the conversion fails, appropriate
    Exceptions are thrown.
    """
    f = open(fn, "r")
    contents = [ l for l in f ]

    marks = import_ns2_constraints(contents)
    top = remote_ns2topdl(opts.url, "".join(contents), cert)
    if not top:
        raise RuntimeError("Cannot create topology from: %s" % fn)
    add_constraint_attributes(top, marks)

    return (top, marks)

def import_topdl_component(fn):
    """
    Pull a component in from a topdl description.  The topology generation is
    straightforward and the constraints are pulled from the attributes of
    Computer nodes in the experiment.  The compisition_point attribute being
    true marks a node as a composition point.  The required, provides and
    accepts attributes map directly into those constraint fields.  The dict of
    constraints and the topology are returned.
    """
    top = topdl.topology_from_xml(filename=fn, top='experiment')
    marks = { }
    for e in [ e for e in top.elements if isinstance(e, topdl.Computer)]:
        if e.get_attribute('composition_point'):
            r = e.get_attribute('required') or 'false'
            a = e.get_attribute('accepts') 
            p = e.get_attribute('provides')
            if a and p:
                c = constraint(name=e.name, required=(r == 'required'),
                    provides=p.split(','), accepts=a.split(','))
                marks[e.name] = c
    return (top, marks)

def index_constraints(constraints, provides, accepts):
    """
    Add constraints to the provides and accepts indices based on the attributes
    of the contstraints.
    """
    for c in constraints:
        for attr, dict in ((c.provides, provides), (c.accepts, accepts)):
            for a in attr:
                if a not in dict: dict[a] = [c]
                else: dict[a].append(c)

def get_suffix(fn):
    """
    We get filename suffixes a couple places.  It;s worth using the same code.
    This gets the shortest . separated suffix from a filename, or None.
    """
    idx = fn.rfind('.')
    if idx != -1: return fn[idx+1:]
    else: return None


def output_composition(top, constraints, outfile=None, format=None):
    """
    Output the composition to the file named by outfile (if any) in the format
    given by format if any.  If both are None, output to stdout in topdl
    """
    def topdl_out(f, top, constraints):
        """
        Output into topdl.  Just call the topdl output, as the constraint
        attributes should be in the topology.
        """
        print >>f, topdl.topology_to_xml(comp, top='experiment')

    def ns2_out(f, top, constraints):
        """
        Reformat the constraint data structures into ns2 constraint comments
        and output the ns2 using the topdl routine.
        """
        # Inner routines
        # Deal with the possibility that the single string name is still in the
        # constraint.
        def name_format(n):
            if isinstance(n, tuple): return n[0]
            else: return n
            
            
        # Format the required field back into a string. (ns2_out inner
        def required_format(x):
            if x: return "required"
            else: return "optional"
        
        # ns2_out main line
        for c in constraints:
            print >>f, "# COMPOSITION: %s:%s:%s:%s" % (
                    topdl.to_tcl_name(name_format(c.name)),
                    required_format(c.required), ",".join(c.provides),
                    ",".join(c.accepts))
        print >>f, topdl.topology_to_ns2(top)

    # Info to map from format to output routine.  
    exporters = {
            'xml':topdl_out, 'topdl':topdl_out,
            'tcl': ns2_out, 'ns': ns2_out,
            }

    if format:
        # Explicit format set, use it
        if format in exporters:
            exporter = exporters[format]
        else:
            raise RuntimeError("Unknown format %s" % format)
    elif outfile:
        # Determine the format from the suffix (if any)
        s = get_suffix(outfile)
        if s and s in exporters:
            exporter = exporters[s]
        else:
            raise RuntimeError("Unknown format (suffix) %s" % outfile)
    else:
        # Both outfile and format are empty
        exporter = topdl_out

    # The actual output.  Open the file, if any, and call the exporter
    if outfile: f = open(outfile, "w")
    else: f = sys.stdout
    exporter(f, top, constraints)
    if outfile: f.close()

def import_components(files):
    """
    Pull in the components.  The input is a sequence of tuples where each tuple
    includes the name of the file to pull the component from and the number of
    copies to make of it.  The routine to read a file is picked by its suffix.
    On completion, a tuple containing the number of components successfully
    read, a list of the topologies, a set of names in use accross all
    topologies (for inserting later names), the constraints extracted, and
    indexes mapping the attribute provices and accepted tinto the lists of
    constraints that provide or accept them is returned.
    """
    importers = {
            'tcl': import_ns2_component, 
            'ns': import_ns2_component, 
            'xml':import_topdl_component,
            'topdl':import_topdl_component,
            }
    names = set()
    constraints = [ ]
    provides = { }
    accepts = { }
    components = 0
    topos = [ ]

    for fn, cnt in files:
        top = None
        try:
            s = get_suffix(fn)
            if s and s in importers:
                top, marks = importers[s](fn)
            else:
                warn("Unknown suffix on file %s.  Ignored" % fn)
                continue
        except service_error, e:
            warn("Remote error on %s: %s" % (fn, e))
            continue
        except EnvironmentError, e:
            warn("Error on %s: %s" % (fn, e))
            continue

        # Parsed the component and sonctraints, now work through the rest of
        # the pre-processing.  We do this once per copy of the component
        # requested, cloning topologies and constraints each time.
        for i in range(0, cnt):
            components += 1
            t = top.clone()
            m = copy.deepcopy(marks)
            index_constraints(m.values(), provides, accepts)
            add_interfaces(t, m)
            localize_names(t, names, m)
            t.incorporate_elements()
            constraints.extend(m.values())
            topos.append(t)

    return (components, topos, names, constraints, provides, accepts)

# Main line begins
parser = ComposeOptionParser()

opts, args = parser.parse_args()

if opts.cert:
    cert = opts.cert
elif os.access(os.path.expanduser("~/.ssl/emulab.pem"), os.R_OK):
    cert = os.path.expanduser("~/.ssl/emulab.pem")
else:
    cert = None

random.seed(opts.seed)

files = opts.files
files.extend([ (a, 1) for a in args])

# Pull 'em in.
components, topos, names, constraints, provides, accepts = \
        import_components(files)

# Let the user know if they messed up on specifying constraints.
if any([ not isinstance(c.name, tuple) for c in constraints]):
    warn("nodes not found for constraints on %s" % \
            ",".join([ c.name for c in constraints \
            if isinstance(c.name, basestring)]))
    constraints = [ c for c in constraints if isinstance(c.name, tuple )]

# If more than one component was given, actually do the composition, otherwise
# this is probably a format conversion. 
if components > 1:
    # Mix up the constraint indexes 
    randomize_constraint_order((provides, accepts))

    # Now the various components live in the same namespace and are marked with
    # their composition requirements. 

    if not meet_constraints([c for c in constraints if c.required], 
            provides, accepts):
        if opts.lax:
            warn("warning: could not meet all required constraints")
        else:
            sys.exit("Failed: could not meet all required constraints")

    meet_constraints([ c for c in constraints if not c.match ], 
            provides, accepts)

    # Add testbed attributes if requested
    if opts.add_testbeds:
        for i, t in enumerate(topos):
            for e in [ e for e in t.elements if isinstance(e, topdl.Computer)]:
                e.set_attribute('testbed', 'testbed%03d' % i)

    # Make a topology containing all elements and substrates from components
    # that had matches.
    comp = topdl.Topology()
    for t in set([ c.name[2] for c in constraints if c.match]):
        comp.elements.extend([e.clone() for e in t.elements])
        comp.substrates.extend([s.clone() for s in t.substrates])

    # Add substrates and connections corresponding to composition points
    connect_composition_points(comp, constraints, names)

    # Remove composition attributes for satisfied constraints
    remove_constraint_attributes(comp, dict(
        [(c.name[0], c) for c in constraints if c.match]))
elif topos == 1:
    comp = topos[0]
    if opts.add_testbeds:
        for e in [ e for e in comp.elements if isinstance(e, topdl.Computer)]:
            e.set_attribute('testbed', 'testbed001')

else:
    sys.exit("Did not read any components.")

# Put out the composition with only the unmatched constriaints
output_composition(comp, [c for c in constraints if not c.match], 
        opts.outfile, opts.format)