<commit_before>from util import ticks_ms

story = None

def loadStory(storyIdString):
    storyModulePath = "stories." + storyIdString
    print("Loading story from " + storyModulePath)
    loaded = __import__(storyModulePath)
    module = getattr(loaded, storyIdString)
    return module.story

def getStoryContext():
    global story
    assert story != None, "Code needs to execute in a `with story:` block"
    return story

class Uid:
    def __init__(self, idString):
        self.idString = idString

    """Override the default Equals behavior"""
    def __eq__(self, other):
        if isinstance(other, self.__class__):
            return self.idString == other.idString
        return False

    def __neq__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(type(self).__name__) ^ hash(self.idString)

# valid types which can be passed as the 'uid' value when creating an item
uidInitTypes = [type(None), Uid, str]

'''
Base class which treats all named arguments as attributes and
all positional arguments as dicts containing named attributes
'''
class Item(object):
    required=[]
    defaults={}

    # populate attributes from positional dicts and keyword args
    def __init__(self, *args, **kwargs):
        # assume positional arguments are dicts
        for data in args: # exclude self
            assert type(data)==dict, "non-dict passed to Item#__init__()"
            for key in data:
                setattr(self, key, data[key])
        # assume keyword arguments specify attributes
        for key in kwargs:
            suffix = "uid"
            # check that uids are typed as Uid
            if key.lower().endswith(suffix):
                uidType = type(kwargs[key])
                assert uidType == Uid or uidType == type(None), "Attribute '{0}' with name ending '1' is actually of type {2}".format(key, suffix, str(uidType))
            setattr(self, key, kwargs[key])

        # raise error if any 'required' attributes are missing
        missing = list()
        for name in type(self).required:
            if not(hasattr(self, name)):
                missing.append(name)
        if len(missing) > 0:
            raise AssertionError(type(self).__name__ + " missing required attributes " + str(missing) )

        for key in type(self).defaults:
            if not hasattr(self, key):
                setattr(self, key, type(self).defaults[key])

# Class which simplifies handling a list of idStrings or Uids and exposing them as attributes on an object, named by idString
class UidRegistry(Item):
        def __init__(self,  *a):
            uidDict = dict()
            for uid in a:
                if type(uid)==str:
                    uid = Uid(uid)
                if type(uid)!=Uid:
                    import pdb; pdb.set_trace()
                assert type(uid)==Uid
                uidDict[uid.idString]=uid
            super().__init__(**uidDict)

'''
    A common superclass for items with ids. Accepts a string id argument
    and silently replaces it with a typed Uid object which is required
    by other Uid oriented references, this supports implicit validation
    of data structures (e.g. encouraging you to refer to myPassage.uid in
    preference to the possibly "theuid")
'''
class UidItem(Item):
    required = Item.required + ["uid"]
    def __init__(self, *a, **k):
        # intercept string uids, turn into Uid objects
        if "uid" in k:
            uid = k["uid"]
            uidType = type(uid)
            assert uidType in uidInitTypes, "'uid' is of '{}'. Uids must be one of ()".format(str(uidType),str(uidInitTypes))
            if type(uid) == str: # cast strings to Uid
                k["uid"]=Uid(uid)
        super().__init__(*a,**k)

class Container(UidItem):

    def _get_table(self, cls):
        if not(hasattr(self, 'registry')):
            self.registry = dict()
        if not(cls.__name__ in self.registry):
             self.registry[cls.__name__]=dict()
        return self.registry[cls.__name__]

    '''register any subclass of item to be looked up later by its id'''
    def _register(self, cls, item):
        table = self._get_table(cls)
        if not(item.uid.idString in table):
            table[item.uid.idString]=item
        else:
            raise AssertionError("Lookup table for " + cls.__name__ + " already contains item with uid" + item.uid.idString)

    ''''''
    def _lookup(self, cls, uid):
        if not(type(uid) == Uid):
            raise AssertionError("Attempted lookup with type" + type(uid) + "; should be Uid")
        else:
            if hasattr(self, "registry"):
                if cls.__name__ in self.registry:
                    table = self.registry[cls.__name__]
                    if uid.idString in table:
                        return table[uid.idString]
            raise AssertionError("'" + uid.idString + "' not in " + cls.__name__ + " lookup table")

class AnonymousContainer(Container):
    required = [n for n in Container.required if n!="uid"]

class Story(Container):
    required = Container.required + ["startPassageUid"]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __enter__(self):
        global story
        story = self

    def __exit__(self, type, value, traceback):
        global story
        story = None

    def registerPassage(self, passage):
        return self._register(Passage, passage)

    def lookupPassage(self, passageUid):
        return self._lookup(Passage, passageUid)

    def registerBox(self, box):
        return self._register(Box, box)

    def lookupBox(self, boxUid):
        return self._lookup(Box, boxUid)

    def createBlankCard(self, cardUid):
        return Card(
            uid=cardUid,
            storyUid = self.uid,
            passageUid = self.startPassageUid,
            sack = dict(
                money=100
            )
        )

class Card(UidItem):
    required = UidItem.required + ["storyUid", "passageUid", "sack"]

class Box(UidItem):
    required = UidItem.required + ["label"]
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        getStoryContext().registerBox(self)

class Passage(UidItem):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        global story
        story.registerPassage(self)

    def activate(self,  engine):
        raise AssertionError("Not yet implemented")

    def getRenderedDict(self, engine):
        d = dict()

        # populate using all primitive attributes of Passage
        for key,value in self.__dict__.items():
            valType = type(value)
            if valType==str or valType==int or valType==float or valType==bool:
                d[key]=str(value)

        # populate using all attributes of sack
        for key,value in engine.card.sack.items():
            assert key not in d, "Error processing sack value. RenderedDict already contains '" + key + "'"
            d[key]=value

        return d

    def getRenderedText(self, engine):
        d = self.getRenderedDict(engine)
        nextT = self.getRenderedTemplate(engine)
        prevT = None
        try:
            while nextT != prevT: # repeats until formatted text is the same as unformatted
                prevT = nextT
                nextT = prevT.format(**d)
            return nextT
        except KeyError as k:
            import pdb; pdb.set_trace()
            raise k

    def getRenderedTemplate(self):
        raise AssertionError("Not yet implemented")

    def handleTap(self, engine, card):
        raise AssertionError("Not yet implemented")

class BoxPassage(Passage):
    required = Passage.required + ["rightBoxUid", "rightBoxText"]
    defaults = dict(Passage.defaults,
        wrongBoxText="Please go to {rightBoxLabel} to continue your adventure"
    )

    def getRenderedDict(self, engine):
        d = super().getRenderedDict(engine)
        d.update(
            rightBoxLabel = engine.story.lookupBox(self.rightBoxUid).label
        )
        return d

    def getRenderedTemplate(self, engine):
        box = engine.box
        if box.uid == self.rightBoxUid:
            return "{rightBoxText}"
        else:
            return "{wrongBoxText}"

class PagePassage(BoxPassage):
    required = BoxPassage.required + ["nextPassageUid"]

    def getNextPassage(self, engine):
        return engine.story.lookupPassage(self.nextPassageUid)

    def getNextBox(self, engine):
        nextPassage = self.getNextPassage(engine)
        return engine.story.lookupBox(nextPassage.rightBoxUid)

    def getRenderedDict(self, engine):
        d = super().getRenderedDict(engine)
        d.update(
            nextBoxLabel=self.getNextBox(engine).label
        )
        return d

    def getRenderedTemplate(self, engine):
        template = super().getRenderedTemplate(engine)
        nextPassage = self.getNextPassage(engine)
        if self.rightBoxUid == nextPassage.rightBoxUid:
            return template + "\n...tap to continue"
        else:
            return template + "\n...now go to {nextBoxLabel}"

    def progress_player(self, engine):
        engine.card.passageUid = self.nextPassageUid

    def handleTap(self, engine):
        self.progress_player(engine)
        engine.renderText(self.getRenderedText(engine))

class ConfirmationPassage(PagePassage):
    defaults = dict(PagePassage.defaults,
        tapTime=4000
    )

    def handleTap(self, engine):
        now = ticks_ms()
        # progress only if tap follows quickly
        if now - self.last_tap_ms < self.tapTime:
            self.progress_player(engine.card)
        engine.renderText(self.getRenderedText(engine))
        self.last_tap_ms = now

'''
# superceded by 'script' attribute which is evaluated in the context of a tap (e.g. "sack.money = sack.money + 5")
# See also https://twine2.neocities.org/1.html for Twine reference features around variables and execution
class ExecutePassage(PagePassage):
    pass
'''

'''
class WaitPassage(PagePassage):
    pass
'''

class ChoicePassage(PagePassage):
    pass

class ConditionalPassage(Passage):
        def checkCondition(self,  engine):
            raise AssertionError("Not yet implemented")

# Condition render
# and text rendering from sack
# Condition routing
# Test Verification of story
# Rendering as graph
<commit_msg>Added post-story-completion validation logic. Changed Passage to Node. Extended use of storyContext. Moved templating logic into Engine. Added Fork nodes, BoolMethod, BoolExpression, ChoicePassage. Drafted linearSequence constructor. Implemented SackChangePassage<commit_after>from util import ticks_ms
from boilerplate import Resolver

story = None

def loadStory(storyIdString):
    storyModulePath = "stories." + storyIdString
    print("Loading story from " + storyModulePath)
    loaded = __import__(storyModulePath)
    module = getattr(loaded, storyIdString)
    return module.story

def getStoryContext():
    global story
    assert story != None, "Code needs to execute in a `with story:` block"
    return story

class Uid:
    def __init__(self, idString):
        self.idString = idString

    """Override the default Equals behavior"""
    def __eq__(self, other):
        if isinstance(other, self.__class__):
            return self.idString == other.idString
        return False

    def __neq__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(type(self).__name__) ^ hash(self.idString)

# valid types which can be passed as the 'uid' value when creating an item
uidInitTypes = [type(None), Uid, str]

'''
Base class which treats all named arguments as attributes and
all positional arguments as dicts containing named attributes
'''
class Item(object):
    required=[]
    defaults={}

    # populate attributes from positional dicts and keyword args
    def __init__(self, *args, **kwargs):
        # assume positional arguments are dicts
        for data in args: # exclude self
            assert type(data)==dict, "non-dict passed to Item#__init__()"
            for key in data:
                setattr(self, key, data[key])
        # assume keyword arguments specify attributes
        for key in kwargs:
            suffix = "uid"
            # check that uids are typed as Uid
            if key.lower().endswith(suffix):
                uidType = type(kwargs[key])
                assert uidType == Uid or uidType == type(None), "Attribute '{0}' with name ending '1' is actually of type {2}".format(key, suffix, str(uidType))
            setattr(self, key, kwargs[key])

        # raise error if any 'required' attributes are missing
        missing = list()
        for name in type(self).required:
            if not(hasattr(self, name)):
                missing.append(name)
        if len(missing) > 0:
            raise AssertionError(type(self).__name__ + " missing required attributes " + str(missing) )

        for key in type(self).defaults:
            if not hasattr(self, key):
                setattr(self, key, type(self).defaults[key])

# Class which simplifies handling a list of idStrings or Uids and exposing them as attributes on an object, named by idString
class UidRegistry(Item):
        def __init__(self,  *a):
            uidDict = dict()
            for uid in a:
                if type(uid)==str:
                    uid = Uid(uid)
                if type(uid)!=Uid:
                    import pdb; pdb.set_trace()
                assert type(uid)==Uid
                uidDict[uid.idString]=uid
            super().__init__(**uidDict)

'''
    A common superclass for items with ids. Accepts a string id argument
    and silently replaces it with a typed Uid object which is required
    by other Uid oriented references, this supports implicit validation
    of data structures (e.g. encouraging you to refer to myNode.uid in
    preference to the possibly "theuid")
'''
class UidItem(Item):
    required = Item.required + ["uid"]
    def __init__(self, *a, **k):
        # intercept string uids, turn into Uid objects
        if "uid" in k:
            uid = k["uid"]
            uidType = type(uid)
            assert uidType in uidInitTypes, "'uid' is of '{}'. Uids must be one of ()".format(str(uidType),str(uidInitTypes))
            if type(uid) == str: # cast strings to Uid
                k["uid"]=Uid(uid)
        super().__init__(*a,**k)

class Container(UidItem):

    def _get_table(self, cls):
        if not(hasattr(self, 'registry')):
            self.registry = dict()
        if not(cls.__name__ in self.registry):
             self.registry[cls.__name__]=dict()
        return self.registry[cls.__name__]

    '''register any subclass of item to be looked up later by its id'''
    def _register(self, cls, item):
        table = self._get_table(cls)
        if not(item.uid.idString in table):
            table[item.uid.idString]=item
        else:
            raise AssertionError("Lookup table for " + cls.__name__ + " already contains item with uid" + item.uid.idString)

    ''''''
    def _lookup(self, cls, uid):
        if not(type(uid) == Uid):
            raise AssertionError("Attempted lookup with type" + type(uid) + "; should be Uid")
        else:
            if hasattr(self, "registry"):
                if cls.__name__ in self.registry:
                    table = self.registry[cls.__name__]
                    if uid.idString in table:
                        return table[uid.idString]
            raise AssertionError("'" + uid.idString + "' not in " + cls.__name__ + " lookup table")

class AnonymousContainer(Container):
    required = [n for n in Container.required if n!="uid"]

class Story(Container):
    required = Container.required + ["startNodeUid", "startSack"]
    defaults = dict(Container.defaults,
        startSack = dict()
    )

    def __enter__(self):
        global story
        story = self

    def __exit__(self, type, value, traceback):
        self.validate()
        global story
        story = None

    def validate(self):
        assert type(self.startSack) == dict, "'sack' must be of type dict"
        # delegate validation to nodes
        for nodeUid, node in self._getTable(Node).items():
            node.validate(self)

    def registerNode(self, node):
        return self._register(Node, node)

    def lookupNode(self, nodeUid):
        return self._lookup(Node, nodeUid)

    def registerBox(self, box):
        return self._register(Box, box)

    def lookupBox(self, boxUid):
        return self._lookup(Box, boxUid)

    def createBlankCard(self, cardUid):
        return Card(
            uid=cardUid,
            storyUid = self.uid,
            nodeUid = self.startNodeUid,
            sack = dict(self.startSack)
        )

class Card(UidItem):
    required = UidItem.required + ["storyUid", "nodeUid", "sack"]

class Box(UidItem):
    required = UidItem.required + ["label"]
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        getStoryContext().registerBox(self)

class Node(UidItem):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        getStoryContext().registerNode(self)

    def validate(self, story):
        pass

    def activate(self,  engine):
        pass

    def deactivate(self, engine):
        pass

    def handleTap(self, engine, card):
        raise AssertionError("Not yet implemented")

class BoolMethodFork(ThroughNode):
    required = Node.required + ["trueNodeUid", "falseNodeUid"]

    def activate(self, engine):
        if self.evaluate(engine):
            engine.setNodeUid(self.trueNodeUid)
        else:
            engine.setNodeUid(self.falseNodeUid)

    def evaluate(self, engine):
        raise AssertionError("Not yet implemented")

class BoolExpressionFork(ThroughNode):
    required = BooleanMethodRouter.required + ["expression"]

    def evaluate(self, engine):
        return engine.evaluateExpression(self.expression)

class TemplatedNode(Node):
    required = Node.required + ["template"]

    def getTemplateString():
        return self.template

    def getTemplateResolver():
        return TemplateResolver(**self.__dict__)

    def displayTemplate(self, engine):
        engine.render(self.getTemplateString(), self.getTemplateResolver())

class VisitBoxPassage(TemplatedNode):
    required = TemplatedNode.required + ["visitBoxUid", "visitBoxText"]
    defaults = dict(TemplatedNode.defaults,
        template="{% if box.uid == node.visitBoxUid %}{{node.visitBoxText}}{% else %}{{node.otherBoxText}}{% endif %}"
        otherBoxText="Please go to {node.visitBoxLabel} to continue your adventure"
    )

    def validate(self, story):
        self.visitBoxLabel = story.lookupBox(self.visitBoxUid).label

class LinearPassage(VisitBoxPassage):
    required = VisitBoxPassage.required + ["nextNodeUid"]

    def getTemplateString(self, engine):
        nextNode = engine.story.lookupNode(self.getNextNodeUid())
        if isinstance(nextNode, VisitBoxPassage):
            self.nextBoxLabel = engine.story.lookupBox(nextNode.visitBoxUid).label
            if nextNode.visitBoxUid != self.visitBoxUid:
                return super().getTemplateString(engine) + "\n find {{node.nextBoxLabel}} and tap in to continue"
        return super().getTemplateString(engine) = "\n...tap to continue"

    def gotoNextNode(self, engine):
        engine.setNodeUid(self.nextNodeUid)

    def handleTap(self, engine):
        self.gotoNextNode(engine)
        self.displayTemplate(engine)

class ConfirmationPassage(LinearPassage):
    defaults = dict(StraightPassage.defaults,
        timeout=4000
    )

    def handleTap(self, engine):
        now = ticks_ms()
        # progress only if tap follows quickly
        if now - self.last_tap_ms < self.timeout:
            self.progressPlayer(engine.card)
        engine.renderText(self.getRenderedText(engine))
        self.last_tap_ms = now

class SackChangePassage(LinearPassage):
    defaults = dict(LinearPassage.defaults,
        changeCondition = None
        add = None
        remove = None
    )

    def validate(self, story):
        super().validate(story)
        assert self.add == None or type(self.add) == dict
        assert self.remove == None or type(self.remove) == dict
        assert hasattr(self, "add") or hasattr(self, "remove"), "SackChangePassage should have one of either 'add' or 'remove'"

    def handleTap(self, engine):
        if self.changeCondition == None or engine.evaluateExpression(self.changeCondition):
            sack = engine.card.sack
            # handle addition (list of entries or a number)
            if(hasattr(self, "add")):
                for key,val in self.add.items():
                    if key in sack:
                        sack[key] = sack[key] + val
                    else:
                        sack[key] = val
            # handle removal (list of entries or a number)
            if(hasattr(self, "remove)):
                for key,val in self.remove:
                    if key in sack:
                        if type(val) == list:
                            for entry in val:
                                sack[key].remove(entry)
                        else:
                            sack[key] = sack[key] - entry
        # now run the main routine
        super().handleTap(engine)

'''
# superceded by 'script' attribute which is evaluated in the context of a tap (e.g. "sack.money = sack.money + 5")
# See also https://twine2.neocities.org/1.html for Twine reference features around variables and execution
class ExecutePassage(StraightPassage):
    pass
'''

'''
class WaitPassage(StraightPassage):
    pass
'''

class ChoicePassageFork(VisitBoxPassage):
    required = VisitBoxPassage.required + ["choices"]

    def validate(self, story):
        print( "Validating Node:" + self.uid.idString)
        choices = self.choices
        assert type(choices) == dict, "'choices' is not a dict mapping box uids to string templates as expected"
        assert all([isinstance(val, str) for val in choices]), "'choices' property template values are not all of type 'str'"
        choiceNodeUids = choices.keys()
        assert all([isinstance(key, Uid) for key in choiceNodeUids]), "'choices' property uid keys are not all of type 'Uid'"
        choiceNodes = [story.lookupNode(nodeUid) for nodeUid in choiceNodeUids]
        assert all([isinstance(node, VisitBoxNode) for node in choiceNodes]), "ChoiceFork choices must point to subclasses of VisitBoxNode"
        choiceBoxUids = [node.visitBoxUid for node in choiceNodes]
        assert len(choiceBoxUids) == len(set(choiceBoxUids)), "Some VisitBoxNodes referenced by 'choices' share the same box" + str(choiceBoxUids)
        choiceBoxes = [engine.story.lookupBox(boxUid) for boxUid in choiceBoxUids]
        # record nodes and boxes
        self.choiceNodes = choiceNodes
        self.choiceBoxes = choiceBoxes

    def getTemplateString(self, engine):
        try:
            # find the matching choice, if it exists
            choiceBoxUids = [box.uid for box in self.choiceBoxes]
            chosenPos = choiceBoxUids.index(engine.box.uid)
            chosenNode = chosenNodes[choiceIndex]
            engine.setNodeUid(chosenNode.uid)
            return "You chose: " + self.choices[chosenNode.uid] + "\n... now tap to continue"
        except ValueError:
            # box not in list of choices
            # render the choice string
            choiceString = ""
            for choiceNode,choicePos in enumerate(self.choiceNodes):
                choiceString += "\n" + self.choices[choiceNode.uid] + " : " + self.choiceBoxes[choicePos].label
            if engine.box.uid == self.visitBoxUid:
                # it's the visit box, render the page
                return super().getTemplateString(engine) + choiceString
            else:
                # it's not the visit box
                return "This box is not among your choices. To continue the game..." + choiceString

'''
def constructLinearSequence(inboundUid, outboundUid, templates):
    templates = list(templates)
    templates.reverse()
    toUid = outboundUid
    for template,pos in enumerate(templates):
'''

# Condition render
# and text rendering from sack
# Condition routing
# Test Verification of story
# Rendering as graph
<|endoftext|>


<commit_before>#!/usr/bin/python

import pika
import sys

#Create a connection with the broker
connection = pika.BlockingConnection(pika.ConnectionParameters(host='localhost'))

channel = connection.channel()

#Declare a queue
channel.queue_declare(queue='myQueue')

msg = sys.argv[1] if len(sys.argv) > 1 else "Hello World"

channel.basic_publish(exchange='',
                      routing_key='myQueue',
                      body=msg)

print " [x] Sent {}".format(msg)

connection.close()
<commit_msg>Added connection to remote server<commit_after>#!/usr/bin/python

import pika
import sys

#Create a connection with the broker
host = '10.12.1.53' #'localhost'

credentials = pika.PlainCredentials('guest', 'guest')
parameters=pika.ConnectionParameters(host, 5672, '/', credentials)

connection = pika.BlockingConnection(parameters)

channel = connection.channel()

#Declare a queue
channel.queue_declare(queue='myQueue')

msg = sys.argv[1] if len(sys.argv) > 1 else "Hello World"

channel.basic_publish(exchange='',
                      routing_key='myQueue',
                      body=msg)

print " [x] Sent {}".format(msg)

connection.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python

# Copyright 2015 University of Chicago
# Licensed under the APL 2.0 license

# Process jobs that have been uploaded and create
# a pegasus workflow and submit
import argparse
import re
import sys
import os
import time
import subprocess
import fcntl

import cStringIO
import psycopg2

import Pegasus.DAX3
import fsurfer
import fsurfer.helpers
import fsurfer.log

PEGASUSRC_PATH = '/etc/fsurf/pegasusconf/pegasusrc'
VERSION = fsurfer.__version__
MAX_RUNNING_WORKFLOWS = 200

def pegasus_submit(dax, workflow_directory, output_directory):
    """
    Submit a workflow to pegasus

    :param dax:  path to xml file with DAX, used for submit
    :param workflow_directory:  directory for workflow information
    :param output_directory:  directory for workflow output
    :return:            the output from pegasus
    """
    try:
        output = subprocess.check_output(['/usr/bin/pegasus-plan',
                                          '--sites',
                                          'condorpool',
                                          '--dir',
                                          workflow_directory,
                                          '--conf',
                                          PEGASUSRC_PATH,
                                          '--output-dir',
                                          output_directory,
                                          '--dax',
                                          dax,
                                          '--submit'],
                                         stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as err:
        return err.returncode, err.output

    return 0, output

def submit_workflow(subject_files, version, subject_name, user, job_run_id,
                    multicore=True, options=None, workflow='diamond'):
    """
    Submit a workflow to OSG for processing

    :param subject_files:  lists with paths to files with MRI
                           data (mgz or subject dir)
    :param version:        FreeSurfer version to use
    :param subject_name:   name of subject being processed
    :param user:          freesurfer user that workflow is being run for
    :param job_run_id:         job run id for the workflow
    :param multicore:     boolean indicating whether to use a multicore
                          workflow or not
    :param options:       Options to pass to FreeSurfer
    :param workflow:      string indicating type of workflow to run (serial,
                          diamond, single)
    :return:              pegasus workflow id  on success, None on error
    """
    if multicore:
        cores = 8
    else:
        cores = 2
    logger = fsurfer.log.get_logger()

    logger.debug("Processing workflow using {0} as input".format(subject_files))
    dax = Pegasus.DAX3.ADAG('freesurfer')
    dax_subject_files = []
    for input_file in subject_files:
        dax_subject_file = Pegasus.DAX3.File(os.path.basename(input_file))
        dax_subject_file.addPFN(Pegasus.DAX3.PFN("file://{0}".format(input_file),
                                                 "local"))
        dax_subject_files.append(dax_subject_file)
        dax.addFile(dax_subject_file)
    workflow_directory = os.path.join(fsurfer.FREESURFER_SCRATCH, user, 'workflows')
    output_directory = os.path.join(fsurfer.FREESURFER_BASE, user, 'workflows', 'output')
    job_invoke_cmd = "/usr/bin/task_completed.py --id {0}".format(job_run_id)
    if workflow == 'serial':
        created = fsurfer.create_serial_workflow(dax,
                                                 version,
                                                 cores,
                                                 dax_subject_files,
                                                 subject_name,
                                                 invoke_cmd=job_invoke_cmd)
    elif workflow == 'diamond':
        created = fsurfer.create_diamond_workflow(dax,
                                                  version,
                                                  cores,
                                                  dax_subject_files,
                                                  subject_name,
                                                  invoke_cmd=job_invoke_cmd)
    elif workflow == 'single':
        created = fsurfer.create_single_workflow(dax,
                                                 version,
                                                 cores,
                                                 dax_subject_files,
                                                 subject_name)
    elif workflow == 'custom':
        created = fsurfer.create_custom_workflow(dax,
                                                 version,
                                                 2,  # custom workflows get 2 cores
                                                 dax_subject_files[0],
                                                 subject_name,
                                                 options)
    else:
        created = fsurfer.create_serial_workflow(dax,
                                                 version,
                                                 cores,
                                                 dax_subject_files,
                                                 subject_name)
    if created:
        curr_date = time.strftime("%Y%m%d_%H%M%S", time.gmtime(time.time()))
        dax.invoke('on_success', "/usr/bin/workflow_completed.py --success --id {0}".format(job_run_id))
        dax.invoke('on_error', "/usr/bin/workflow_completed.py --failure --id {0}".format(job_run_id))
        dax_name = "freesurfer_{0}.xml".format(curr_date)
        with open(dax_name, 'w') as f:
            dax.writeXML(f)
        exit_code, output = pegasus_submit(dax="{0}".format(dax_name),
                                           workflow_directory=workflow_directory,
                                           output_directory=output_directory)
        logger.info("Submitted workflow, got exit code {0}".format(exit_code))
        logger.info("Pegasus output: {0}".format(output))
        if exit_code != 0:
            os.unlink(dax_name)
            return None
        os.unlink(dax_name)
        capture_id = False
        for line in cStringIO.StringIO(output).readlines():
            if 'Your workflow has been started' in line:
                capture_id = True
            if capture_id and workflow_directory in line:
                id_match = re.search(r'([T\d]+-\d+)'.format(workflow_directory),
                                     line)
                if id_match is not None:
                    workflow_id = id_match.group(1)
                    return workflow_id

    return None

def exceeded_running_limit(conn):
    """
    Check to see if the number of workflows in
    RUNNING state is more than MAX_RUNNING_WORKFLOWS

    :param conn: database connection to use
    :return: True if condition holds, False otherwise
    """
    running_workflow_query = "SELECT COUNT(*) " \
                             "FROM freesurfer_interface.jobs " \
                             "WHERE state = 'RUNNING"
    try:
        fsurfer.log.initialize_logging()
        logger = fsurfer.log.get_logger()
        cursor = conn.cursor()
        cursor.execute(running_workflow_query)
        running_workflows = cursor.fetchone()[0]
        if running_workflows >= MAX_RUNNING_WORKFLOWS:
            logger.warn("Number of running workflows at or above" +
                        "MAX_RUNNING_WORKFLOWS: " +
                        "{0} vs {1}".format(running_workflows,
                                            MAX_RUNNING_WORKFLOWS))
            return True
        return False
    except psycopg2.Error as e:
        logger.exception("Got pgsql error: {0}".format(e))
        return False

def process_images():
    """
    Process uploaded images

    :return: exit code (0 for success, non-zero for failure)
    """
    fsurfer.log.initialize_logging()
    logger = fsurfer.log.get_logger()
    parser = argparse.ArgumentParser(description="Generate and submit "
                                                 "workflows to process jobs")
    # version info
    parser.add_argument('--version', action='version', version='%(prog)s ' + VERSION)
    # Arguments for action
    parser.add_argument('--dry-run', dest='dry_run',
                        action='store_true', default=False,
                        help='Mock actions instead of carrying them out')
    parser.add_argument('--debug', dest='debug',
                        action='store_true', default=False,
                        help='Output debug messages')

    args = parser.parse_args(sys.argv[1:])
    if args.debug:
        fsurfer.log.set_debugging()
    if args.dry_run:
        sys.stdout.write("Doing a dry run, no changes will be made\n")
    try:
        x = open('/tmp/fsurf_process.lock', 'w+')
        fcntl.flock(x, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        logger.warn('Lock file present, exiting')
        sys.exit(1)

    conn = fsurfer.helpers.get_db_client()
    cursor = conn.cursor()
    job_query = "SELECT id, username, num_inputs, subject, options, version " \
                "FROM freesurfer_interface.jobs " \
                "WHERE state = 'QUEUED'"
    input_file_query = "SELECT filename, path, subject_dir " \
                       "FROM freesurfer_interface.input_files " \
                       "WHERE job_id = %s AND NOT purged"
    job_update = "UPDATE freesurfer_interface.jobs " \
                 "SET state = 'RUNNING' " \
                 "WHERE id = %s;"
    job_run_update = "UPDATE freesurfer_interface.job_run " \
                     "SET pegasus_ts = %s " \
                     "WHERE id = %s;"
    job_error = "UPDATE freesurfer_interface.jobs " \
                "SET state = 'ERROR' " \
                "WHERE id = %s;"

    account_start = "INSERT INTO freesurfer_interface.job_run(job_id, " \
                    "                                         tasks) " \
                    "VALUES(%s, %s) " \
                    "RETURNING id"
    try:
        cursor.execute(job_query)
        if exceeded_running_limit(conn):
            logger.warn("Max number of running workflows reached, exiting")
            fcntl.flock(x, fcntl.LOCK_UN)
            x.close()
            os.unlink('/tmp/fsurf_process.lock')
            return 0
        for row in cursor.fetchall():
            workflow_id = row[0]
            username = row[1]
            logger.info("Processing workflow {0} for user {1}".format(workflow_id,
                                                                      username))

            workflow_directory = os.path.join('/local-scratch',
                                              'fsurf',
                                              username,
                                              'workflows')
            if not os.path.exists(workflow_directory):
                if args.dry_run:
                    sys.stdout.write("Would have created {0}".format(workflow_directory))
                else:
                    os.makedirs(workflow_directory)
            cursor2 = conn.cursor()
            cursor2.execute(input_file_query, [workflow_id])
            input_files = []
            custom_workflow = False
            if cursor2.rowcount < 1:
                logger.error("No input files, skipping workflow {0}".format(workflow_id))
                continue
            for input_info in cursor2.fetchall():
                input_file = os.path.join(fsurfer.FREESURFER_BASE,
                                          input_info[0],
                                          input_info[1])
                if not os.path.isfile(input_file):
                    logger.warn("Input file {0} missing, skipping".format(input_file))
                    continue
                elif bool(input_info[2]):
                    # input file is a subject dir
                    custom_workflow = True
                    if cursor2.rowcount != 1:
                        # can't give multiple subject dirs at one time
                        logger.error("Subject dir combined with multiple inputs, skipping!")
                        cursor3 = conn.cursor()
                        if args.dry_run:
                            sys.stdout.write("Would have changed workflow "
                                             "{0} to ERROR state\n".format(workflow_id))
                        else:
                            logger.error("Changed {0} to ERROR state".format(workflow_id))
                            cursor3.execute(job_error, [workflow_id])
                            continue
                    input_files.append(input_file)
                else:
                    input_files.append(input_file)

            pegasus_ts = None
            if custom_workflow and not args.dry_run:
                cursor.execute(account_start, [workflow_id, 1])
                job_run_id = cursor.fetchone()[0]
                pegasus_ts = submit_workflow(input_files,
                                             version=row[5],
                                             subject_name=row[3],
                                             user=username,
                                             job_run_id=job_run_id,
                                             options=row[4],
                                             workflow='custom')
            elif not args.dry_run:
                cursor.execute(account_start, [workflow_id, 4])
                job_run_id = cursor.fetchone()[0]
                pegasus_ts = submit_workflow(input_files,
                                             version=row[5],
                                             subject_name=row[3],
                                             user=username,
                                             job_run_id=job_run_id)
            if pegasus_ts:
                cursor.execute(job_run_update, [pegasus_ts,
                                                job_run_id])
            if pegasus_ts and not args.dry_run:
                cursor.execute(job_update, [workflow_id])
                conn.commit()
                logger.info("Set workflow {0} status to RUNNING".format(workflow_id))
            else:
                conn.rollback()
                logger.info("Rolled back transaction")
    except psycopg2.Error as e:
        logger.exception("Got pgsql error: {0}".format(e))
        pass
    finally:
        conn.close()

    fcntl.flock(x, fcntl.LOCK_UN)
    x.close()
    os.unlink('/tmp/fsurf_process.lock')
    return 0

if __name__ == '__main__':
    # workaround missing subprocess.check_output
    if "check_output" not in dir(subprocess):  # duck punch it in!
        def check_output(*popenargs, **kwargs):
            """
            Run command with arguments and return its output as a byte string.

            Backported from Python 2.7 as it's implemented as pure python
            on stdlib.

            """
            process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
            output, unused_err = process.communicate()
            retcode = process.poll()
            if retcode:
                cmd = kwargs.get("args")
                if cmd is None:
                    cmd = popenargs[0]
                error = subprocess.CalledProcessError(retcode, cmd)
                error.output = output
                raise error
            return output

        subprocess.check_output = check_output
    sys.exit(process_images())
<commit_msg>Fix sql syntax<commit_after>#!/usr/bin/env python

# Copyright 2015 University of Chicago
# Licensed under the APL 2.0 license

# Process jobs that have been uploaded and create
# a pegasus workflow and submit
import argparse
import re
import sys
import os
import time
import subprocess
import fcntl

import cStringIO
import psycopg2

import Pegasus.DAX3
import fsurfer
import fsurfer.helpers
import fsurfer.log

PEGASUSRC_PATH = '/etc/fsurf/pegasusconf/pegasusrc'
VERSION = fsurfer.__version__
MAX_RUNNING_WORKFLOWS = 200

def pegasus_submit(dax, workflow_directory, output_directory):
    """
    Submit a workflow to pegasus

    :param dax:  path to xml file with DAX, used for submit
    :param workflow_directory:  directory for workflow information
    :param output_directory:  directory for workflow output
    :return:            the output from pegasus
    """
    try:
        output = subprocess.check_output(['/usr/bin/pegasus-plan',
                                          '--sites',
                                          'condorpool',
                                          '--dir',
                                          workflow_directory,
                                          '--conf',
                                          PEGASUSRC_PATH,
                                          '--output-dir',
                                          output_directory,
                                          '--dax',
                                          dax,
                                          '--submit'],
                                         stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as err:
        return err.returncode, err.output

    return 0, output

def submit_workflow(subject_files, version, subject_name, user, job_run_id,
                    multicore=True, options=None, workflow='diamond'):
    """
    Submit a workflow to OSG for processing

    :param subject_files:  lists with paths to files with MRI
                           data (mgz or subject dir)
    :param version:        FreeSurfer version to use
    :param subject_name:   name of subject being processed
    :param user:          freesurfer user that workflow is being run for
    :param job_run_id:         job run id for the workflow
    :param multicore:     boolean indicating whether to use a multicore
                          workflow or not
    :param options:       Options to pass to FreeSurfer
    :param workflow:      string indicating type of workflow to run (serial,
                          diamond, single)
    :return:              pegasus workflow id  on success, None on error
    """
    if multicore:
        cores = 8
    else:
        cores = 2
    logger = fsurfer.log.get_logger()

    logger.debug("Processing workflow using {0} as input".format(subject_files))
    dax = Pegasus.DAX3.ADAG('freesurfer')
    dax_subject_files = []
    for input_file in subject_files:
        dax_subject_file = Pegasus.DAX3.File(os.path.basename(input_file))
        dax_subject_file.addPFN(Pegasus.DAX3.PFN("file://{0}".format(input_file),
                                                 "local"))
        dax_subject_files.append(dax_subject_file)
        dax.addFile(dax_subject_file)
    workflow_directory = os.path.join(fsurfer.FREESURFER_SCRATCH, user, 'workflows')
    output_directory = os.path.join(fsurfer.FREESURFER_BASE, user, 'workflows', 'output')
    job_invoke_cmd = "/usr/bin/task_completed.py --id {0}".format(job_run_id)
    if workflow == 'serial':
        created = fsurfer.create_serial_workflow(dax,
                                                 version,
                                                 cores,
                                                 dax_subject_files,
                                                 subject_name,
                                                 invoke_cmd=job_invoke_cmd)
    elif workflow == 'diamond':
        created = fsurfer.create_diamond_workflow(dax,
                                                  version,
                                                  cores,
                                                  dax_subject_files,
                                                  subject_name,
                                                  invoke_cmd=job_invoke_cmd)
    elif workflow == 'single':
        created = fsurfer.create_single_workflow(dax,
                                                 version,
                                                 cores,
                                                 dax_subject_files,
                                                 subject_name)
    elif workflow == 'custom':
        created = fsurfer.create_custom_workflow(dax,
                                                 version,
                                                 2,  # custom workflows get 2 cores
                                                 dax_subject_files[0],
                                                 subject_name,
                                                 options)
    else:
        created = fsurfer.create_serial_workflow(dax,
                                                 version,
                                                 cores,
                                                 dax_subject_files,
                                                 subject_name)
    if created:
        curr_date = time.strftime("%Y%m%d_%H%M%S", time.gmtime(time.time()))
        dax.invoke('on_success', "/usr/bin/workflow_completed.py --success --id {0}".format(job_run_id))
        dax.invoke('on_error', "/usr/bin/workflow_completed.py --failure --id {0}".format(job_run_id))
        dax_name = "freesurfer_{0}.xml".format(curr_date)
        with open(dax_name, 'w') as f:
            dax.writeXML(f)
        exit_code, output = pegasus_submit(dax="{0}".format(dax_name),
                                           workflow_directory=workflow_directory,
                                           output_directory=output_directory)
        logger.info("Submitted workflow, got exit code {0}".format(exit_code))
        logger.info("Pegasus output: {0}".format(output))
        if exit_code != 0:
            os.unlink(dax_name)
            return None
        os.unlink(dax_name)
        capture_id = False
        for line in cStringIO.StringIO(output).readlines():
            if 'Your workflow has been started' in line:
                capture_id = True
            if capture_id and workflow_directory in line:
                id_match = re.search(r'([T\d]+-\d+)'.format(workflow_directory),
                                     line)
                if id_match is not None:
                    workflow_id = id_match.group(1)
                    return workflow_id

    return None

def exceeded_running_limit(conn):
    """
    Check to see if the number of workflows in
    RUNNING state is more than MAX_RUNNING_WORKFLOWS

    :param conn: database connection to use
    :return: True if condition holds, False otherwise
    """
    running_workflow_query = "SELECT COUNT(*) " \
                             "FROM freesurfer_interface.jobs " \
                             "WHERE state = 'RUNNING'"
    try:
        fsurfer.log.initialize_logging()
        logger = fsurfer.log.get_logger()
        cursor = conn.cursor()
        cursor.execute(running_workflow_query)
        running_workflows = cursor.fetchone()[0]
        if running_workflows >= MAX_RUNNING_WORKFLOWS:
            logger.warn("Number of running workflows at or above" +
                        "MAX_RUNNING_WORKFLOWS: " +
                        "{0} vs {1}".format(running_workflows,
                                            MAX_RUNNING_WORKFLOWS))
            return True
        return False
    except psycopg2.Error as e:
        logger.exception("Got pgsql error: {0}".format(e))
        return False

def process_images():
    """
    Process uploaded images

    :return: exit code (0 for success, non-zero for failure)
    """
    fsurfer.log.initialize_logging()
    logger = fsurfer.log.get_logger()
    parser = argparse.ArgumentParser(description="Generate and submit "
                                                 "workflows to process jobs")
    # version info
    parser.add_argument('--version', action='version', version='%(prog)s ' + VERSION)
    # Arguments for action
    parser.add_argument('--dry-run', dest='dry_run',
                        action='store_true', default=False,
                        help='Mock actions instead of carrying them out')
    parser.add_argument('--debug', dest='debug',
                        action='store_true', default=False,
                        help='Output debug messages')

    args = parser.parse_args(sys.argv[1:])
    if args.debug:
        fsurfer.log.set_debugging()
    if args.dry_run:
        sys.stdout.write("Doing a dry run, no changes will be made\n")
    try:
        x = open('/tmp/fsurf_process.lock', 'w+')
        fcntl.flock(x, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        logger.warn('Lock file present, exiting')
        sys.exit(1)

    conn = fsurfer.helpers.get_db_client()
    cursor = conn.cursor()
    job_query = "SELECT id, username, num_inputs, subject, options, version " \
                "FROM freesurfer_interface.jobs " \
                "WHERE state = 'QUEUED'"
    input_file_query = "SELECT filename, path, subject_dir " \
                       "FROM freesurfer_interface.input_files " \
                       "WHERE job_id = %s AND NOT purged"
    job_update = "UPDATE freesurfer_interface.jobs " \
                 "SET state = 'RUNNING' " \
                 "WHERE id = %s;"
    job_run_update = "UPDATE freesurfer_interface.job_run " \
                     "SET pegasus_ts = %s " \
                     "WHERE id = %s;"
    job_error = "UPDATE freesurfer_interface.jobs " \
                "SET state = 'ERROR' " \
                "WHERE id = %s;"

    account_start = "INSERT INTO freesurfer_interface.job_run(job_id, " \
                    "                                         tasks) " \
                    "VALUES(%s, %s) " \
                    "RETURNING id"
    try:
        cursor.execute(job_query)
        if exceeded_running_limit(conn):
            logger.warn("Max number of running workflows reached, exiting")
            fcntl.flock(x, fcntl.LOCK_UN)
            x.close()
            os.unlink('/tmp/fsurf_process.lock')
            return 0
        for row in cursor.fetchall():
            workflow_id = row[0]
            username = row[1]
            logger.info("Processing workflow {0} for user {1}".format(workflow_id,
                                                                      username))

            workflow_directory = os.path.join('/local-scratch',
                                              'fsurf',
                                              username,
                                              'workflows')
            if not os.path.exists(workflow_directory):
                if args.dry_run:
                    sys.stdout.write("Would have created {0}".format(workflow_directory))
                else:
                    os.makedirs(workflow_directory)
            cursor2 = conn.cursor()
            cursor2.execute(input_file_query, [workflow_id])
            input_files = []
            custom_workflow = False
            if cursor2.rowcount < 1:
                logger.error("No input files, skipping workflow {0}".format(workflow_id))
                continue
            for input_info in cursor2.fetchall():
                input_file = os.path.join(fsurfer.FREESURFER_BASE,
                                          input_info[0],
                                          input_info[1])
                if not os.path.isfile(input_file):
                    logger.warn("Input file {0} missing, skipping".format(input_file))
                    continue
                elif bool(input_info[2]):
                    # input file is a subject dir
                    custom_workflow = True
                    if cursor2.rowcount != 1:
                        # can't give multiple subject dirs at one time
                        logger.error("Subject dir combined with multiple inputs, skipping!")
                        cursor3 = conn.cursor()
                        if args.dry_run:
                            sys.stdout.write("Would have changed workflow "
                                             "{0} to ERROR state\n".format(workflow_id))
                        else:
                            logger.error("Changed {0} to ERROR state".format(workflow_id))
                            cursor3.execute(job_error, [workflow_id])
                            continue
                    input_files.append(input_file)
                else:
                    input_files.append(input_file)

            pegasus_ts = None
            if custom_workflow and not args.dry_run:
                cursor.execute(account_start, [workflow_id, 1])
                job_run_id = cursor.fetchone()[0]
                pegasus_ts = submit_workflow(input_files,
                                             version=row[5],
                                             subject_name=row[3],
                                             user=username,
                                             job_run_id=job_run_id,
                                             options=row[4],
                                             workflow='custom')
            elif not args.dry_run:
                cursor.execute(account_start, [workflow_id, 4])
                job_run_id = cursor.fetchone()[0]
                pegasus_ts = submit_workflow(input_files,
                                             version=row[5],
                                             subject_name=row[3],
                                             user=username,
                                             job_run_id=job_run_id)
            if pegasus_ts:
                cursor.execute(job_run_update, [pegasus_ts,
                                                job_run_id])
            if pegasus_ts and not args.dry_run:
                cursor.execute(job_update, [workflow_id])
                conn.commit()
                logger.info("Set workflow {0} status to RUNNING".format(workflow_id))
            else:
                conn.rollback()
                logger.info("Rolled back transaction")
    except psycopg2.Error as e:
        logger.exception("Got pgsql error: {0}".format(e))
        pass
    finally:
        conn.close()

    fcntl.flock(x, fcntl.LOCK_UN)
    x.close()
    os.unlink('/tmp/fsurf_process.lock')
    return 0

if __name__ == '__main__':
    # workaround missing subprocess.check_output
    if "check_output" not in dir(subprocess):  # duck punch it in!
        def check_output(*popenargs, **kwargs):
            """
            Run command with arguments and return its output as a byte string.

            Backported from Python 2.7 as it's implemented as pure python
            on stdlib.

            """
            process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
            output, unused_err = process.communicate()
            retcode = process.poll()
            if retcode:
                cmd = kwargs.get("args")
                if cmd is None:
                    cmd = popenargs[0]
                error = subprocess.CalledProcessError(retcode, cmd)
                error.output = output
                raise error
            return output

        subprocess.check_output = check_output
    sys.exit(process_images())
<|endoftext|>


<commit_before>r"""Batch sub-classes for runs on Columbia's Yeti cluster"""

# ============================================================================
#      Copyright (C) 2013 Kyle Mandli <kyle@ices.utexas.edu>
#
#          Distributed under the terms of the MIT license
#                http://www.opensource.org/licenses/
#
# (adapted for yeti by Andrew Kaluzny <akk2141@columbia.edu>, 2015)
# ============================================================================

import os
import glob
import subprocess

import batch

class Job(batch.Job):
    r"""
    Modifications to the basic :class:`batch.Job` class for Yeti runs

    """

    def __init__(self, use_v2=True, memory=4096, time=60, nodes=1,
                       omp_num_threads=1):
        r"""
        Initialize Yeti job

        See :class:`YetiJob` for full documentation
        """

        super(Job, self).__init__()

        # Add extra job parameters
        self.omp_num_threads = omp_num_threads
        ##self.mic_omp_num_threads = 1
        ##self.mic_affinity = "none"
        self.memory = memory
        self.time = time
        self.nodes = nodes
        self.use_v2 = use_v2

        self.group = None # needs to be set at some point

class BatchController(batch.BatchController):
    r"""
    Modifications to the basic batch controller for Yeti runs

    :Ignored Attributes:

    Due to the system setup, the following controller attributes are ignored:

        *plot*, *terminal_output*, *wait*, *poll_interval*, *plotclaw_cmd*
    """

    def __init__(self, jobs=[], email=None, email_behavior="abe", output=None):
        r"""
        Initialize Yeti batch controller

        See :class:`YetiBatchController` for full documentation
        """

        super(BatchController, self).__init__(jobs)

        # Yeti specific execution controls
        self.email = email
        self.email_behavior = email_behavior
        self.output = output

    def run(self):
        r"""Run Yeti jobs from controller's *jobs* list.

        This run function is modified to run jobs through the slurm queue system
        and provides controls for running serial jobs (OpenMP only).

        Unless otherwise noted, the behavior of this function is identical to
        the base class :class:`BatchController`'s function.
        """

        # Run jobs
        paths = []
        for (i,job) in enumerate(self.jobs):
            # Create output directory
            data_dirname = ''.join((job.prefix,'_data'))
            output_dirname = ''.join((job.prefix,"_output"))
            plots_dirname = ''.join((job.prefix,"_plots"))
            run_script_name = ''.join((job.prefix,"_run.sh"))
            log_name = ''.join((job.prefix,"_log.txt"))

            if len(job.type) > 0:
                job_path = os.path.join(self.base_path,job.type,job.name)
            else:
                job_path = os.path.join(self.base_path,job.name)
            job_path = os.path.abspath(job_path)
            data_path = os.path.join(job_path,data_dirname)
            output_path = os.path.join(job_path,output_dirname)
            plots_path = os.path.join(job_path,plots_dirname)
            log_path = os.path.join(job_path,log_name)
            run_script_path = os.path.join(job_path,run_script_name)
            paths.append({'job':job_path, 'data':data_path,
                          'output':output_path, 'plots':plots_path,
                          'log':log_path})

            # Create job directory if not present
            if not os.path.exists(job_path):
                os.makedirs(job_path)

            # Clobber old data directory
            if os.path.exists(data_path):
                if not job.rundata.clawdata.restart:
                    data_files = glob.glob(os.path.join(data_path,'*.data'))
                    for data_file in data_files:
                        os.remove(data_file)
            else:
                os.mkdir(data_path)

            # Write out data
            temp_path = os.getcwd()
            os.chdir(data_path)
            job.write_data_objects()
            os.chdir(temp_path)

            # Handle restart requests
            if job.rundata.clawdata.restart:
                restart = "T"
                overwrite = "F"
            else:
                restart = "F"
                overwrite = "T"

            # Construct string commands
            run_cmd = "%s %s %s %s %s %s True\n" % (self.runclaw_cmd, job.executable, output_path,
                                                      overwrite, restart, data_path)

            # Write slurm run script
            run_script = open(run_script_path, 'w')

            hours = int(job.time / 60)
            minutes = job.time % 60

            run_script.write("#!/bin/sh\n")
            run_script.write("#PBS -N %s        # Job name\n" % job.prefix)
            run_script.write("#PBS -W group_list=%s        # Group\n" % job.group)
            run_script.write("#PBS -l mem=%smb        # Memory\n" % job.memory)
            run_script.write("#PBS -l walltime=00:%s:%s:00        # Memory\n" % (hours, minutes))
            if job.use_v2:
                run_script.write("#PBS -l nodes=%s:ppn=%s:v2         # Nodes and processers per node\n" % (job.nodes, job.omp_num_threads))
            else:
                run_script.write("#PBS -l nodes=%s:ppn=%s:v2         # Nodes and processers per node\n" % (job.nodes, job.omp_num_threads))
            run_script.write("#PBS -V    # export env. variables to the job\n")
            if self.email is not None:
                run_script.write("#PBS -M %s\n" % self.email)
                run_script.write("#PBS -m %s    # email for abort, begin, end\n")
            if self.output is not None:
                run_script.write("#PBS -o localhost:%s   # stdout\n" % self.output)
                run_script.write("#PBS -e localhost:%s   # stderr\n" % self.output)
            run_script.write("\n")
            run_script.write("# OpenMP controls\n")
            run_script.write("export OMP_NUM_THREADS=%s\n" % job.omp_num_threads)
            ## run_script.write("export MIC_ENV_PREFIX=MIC")
            ## run_script.write("export MIC_OMP_NUM_THREADS=%s\n" % job.mic_omp_num_threads)
            ## run_script.write("export MIC_KMP_AFFINITY=%s\n" % job.mic_affinity)
            run_script.write("\n")
            run_script.write("# Run command\n")
            run_script.write(run_cmd)

            run_script.close()

            # Submit job to queue
            subprocess.Popen("qsub %s > %s" % (run_script_path,log_path), shell=True).wait()

        # -- All jobs have been started --

        return paths
<commit_msg>Add queue specification option<commit_after>r"""Batch sub-classes for runs on Columbia's Yeti cluster"""

# ============================================================================
#      Copyright (C) 2013 Kyle Mandli <kyle@ices.utexas.edu>
#
#          Distributed under the terms of the MIT license
#                http://www.opensource.org/licenses/
#
# (adapted for yeti by Andrew Kaluzny <akk2141@columbia.edu>, 2015)
# ============================================================================

import os
import glob
import subprocess

import batch

class Job(batch.Job):
    r"""
    Modifications to the basic :class:`batch.Job` class for Yeti runs

    """

    def __init__(self, use_v2=True, memory=4096, time=60, nodes=1,
                       omp_num_threads=1):
        r"""
        Initialize Yeti job

        See :class:`YetiJob` for full documentation
        """

        super(Job, self).__init__()

        # Add extra job parameters
        self.omp_num_threads = omp_num_threads
        ##self.mic_omp_num_threads = 1
        ##self.mic_affinity = "none"
        self.memory = memory
        self.time = time
        self.nodes = nodes
        self.use_v2 = use_v2

        self.group = None # needs to be set at some point

class BatchController(batch.BatchController):
    r"""
    Modifications to the basic batch controller for Yeti runs

    :Ignored Attributes:

    Due to the system setup, the following controller attributes are ignored:

        *plot*, *terminal_output*, *wait*, *poll_interval*, *plotclaw_cmd*
    """

    def __init__(self, jobs=[], email=None, email_behavior="abe", output=None):
        r"""
        Initialize Yeti batch controller

        See :class:`YetiBatchController` for full documentation
        """

        super(BatchController, self).__init__(jobs)

        # Yeti specific execution controls
        self.email = email
        self.email_behavior = email_behavior
        self.output = output
        self.queue = None

    def run(self):
        r"""Run Yeti jobs from controller's *jobs* list.

        This run function is modified to run jobs through the slurm queue system
        and provides controls for running serial jobs (OpenMP only).

        Unless otherwise noted, the behavior of this function is identical to
        the base class :class:`BatchController`'s function.
        """

        # Run jobs
        paths = []
        for (i,job) in enumerate(self.jobs):
            # Create output directory
            data_dirname = ''.join((job.prefix,'_data'))
            output_dirname = ''.join((job.prefix,"_output"))
            plots_dirname = ''.join((job.prefix,"_plots"))
            run_script_name = ''.join((job.prefix,"_run.sh"))
            log_name = ''.join((job.prefix,"_log.txt"))

            if len(job.type) > 0:
                job_path = os.path.join(self.base_path,job.type,job.name)
            else:
                job_path = os.path.join(self.base_path,job.name)
            job_path = os.path.abspath(job_path)
            data_path = os.path.join(job_path,data_dirname)
            output_path = os.path.join(job_path,output_dirname)
            plots_path = os.path.join(job_path,plots_dirname)
            log_path = os.path.join(job_path,log_name)
            run_script_path = os.path.join(job_path,run_script_name)
            paths.append({'job':job_path, 'data':data_path,
                          'output':output_path, 'plots':plots_path,
                          'log':log_path})

            # Create job directory if not present
            if not os.path.exists(job_path):
                os.makedirs(job_path)

            # Clobber old data directory
            if os.path.exists(data_path):
                if not job.rundata.clawdata.restart:
                    data_files = glob.glob(os.path.join(data_path,'*.data'))
                    for data_file in data_files:
                        os.remove(data_file)
            else:
                os.mkdir(data_path)

            # Write out data
            temp_path = os.getcwd()
            os.chdir(data_path)
            job.write_data_objects()
            os.chdir(temp_path)

            # Handle restart requests
            if job.rundata.clawdata.restart:
                restart = "T"
                overwrite = "F"
            else:
                restart = "F"
                overwrite = "T"

            # Construct string commands
            run_cmd = "%s %s %s %s %s %s True\n" % (self.runclaw_cmd, job.executable, output_path,
                                                      overwrite, restart, data_path)

            # Write slurm run script
            run_script = open(run_script_path, 'w')

            hours = int(job.time / 60)
            minutes = job.time % 60

            run_script.write("#!/bin/sh\n")
            run_script.write("#PBS -N %s        # Job name\n" % job.prefix)
            run_script.write("#PBS -W group_list=%s        # Group\n" % job.group)
            run_script.write("#PBS -l mem=%smb        # Memory\n" % job.memory)
            run_script.write("#PBS -l walltime=00:%s:%s:00        # Walltime\n" % (hours, minutes))
            if job.use_v2:
                run_script.write("#PBS -l nodes=%s:ppn=%s:v2         # Nodes and processers per node\n" % (job.nodes, job.omp_num_threads))
            else:
                run_script.write("#PBS -l nodes=%s:ppn=%s            # Nodes and processers per node\n" % (job.nodes, job.omp_num_threads))
            run_script.write("#PBS -V    # export env. variables to the job\n")
            if self.queue is not None:
                run_script.write("#PBS -q %s      # Requested queue\n" % self.queue)
            if self.email is not None:
                run_script.write("#PBS -M %s\n" % self.email)
                run_script.write("#PBS -m %s # email for abort, begin, end\n" % self.email_behavior)
            if self.output is not None:
                run_script.write("#PBS -o localhost:%s   # stdout\n" % self.output)
                run_script.write("#PBS -e localhost:%s   # stderr\n" % self.output)
            run_script.write("\n")
            run_script.write("# OpenMP controls\n")
            run_script.write("export OMP_NUM_THREADS=%s\n" % job.omp_num_threads)
            ## run_script.write("export MIC_ENV_PREFIX=MIC")
            ## run_script.write("export MIC_OMP_NUM_THREADS=%s\n" % job.mic_omp_num_threads)
            ## run_script.write("export MIC_KMP_AFFINITY=%s\n" % job.mic_affinity)
            run_script.write("\n")
            run_script.write("# Location of output\n")
            run_script.write(self.output)
            run_script.write("\n")
            run_script.write("# Run command\n")
            run_script.write(run_cmd)

            run_script.close()

            # Submit job to queue
            subprocess.Popen("qsub %s > %s" % (run_script_path,log_path), shell=True).wait()

        # -- All jobs have been started --

        return paths
<|endoftext|>


<commit_before>#!/usr/bin/env python
from embeddedIpython import QIPythonWidget
import pyqtgraph as pg
from PyQt4 import QtGui, QtCore
#from pyqtgraph.Qt import QtCore, QtGui
import pyqtgraph.console
from pyqtgraph.widgets.MatplotlibWidget import MatplotlibWidget
import numpy as np
import pyqtgraph.opengl as gl
from astro import *
from sphvis import *
from optparse import OptionParser
import matplotlib

from pyqtgraph.dockarea import *

class myGLViewWidget(gl.GLViewWidget):
    def init(self, **kwargs):
        gl.GLViewWidget.__init__(self, **kwargs)
        self.reset()

    def reset(self):
        self.opts['azimuth'] = -90
        self.opts['elevation'] = 00
#        self.opts['azimuth'] = 45
#        self.opts['elevation'] = 30

        self.opts['center'] = QtGui.QVector3D(0.0, 0.0, 0.0)
        self.opts['fov'] = 60
        self.opts['distance'] = 30
        self.update()

class popupPlot(QtGui.QDialog):
    def __init__(self, parent=None):
        super(popupPlot, self).__init__()

class mainWindow(QtGui.QMainWindow):
    def __init__(self, load=None, loadsnapshot=None, options={}, **kwargs):
        super(mainWindow, self).__init__()
        self.setupUI()
        if load is not None:
            self.loaddata_gadget(load)
        if loadsnapshot is not None:
            self.loadsnapshot_gadget(loadsnapshot, basic=options.basic)

    def setupUI(self):
        area = DockArea()
        self.setCentralWidget(area)
        self.resize(1000,800)
        self.setWindowTitle('Gadget Visualization')

        self.setupMenu()

        ## Create docks, place them into the window one at a time.
        ## Note that size arguments are only a suggestion; docks will still have to
        ## fill the entire dock area and obey the limits of their internal widgets.
        d1 = Dock("Main Window", size=(800, 800))
        d2 = Dock("Console", size=(800,200))
        d3 = Dock("View", size=(200,800))
        d4 = Dock("Options", size=(200,800))
        d1.hideTitleBar()
        d2.hideTitleBar()
#        d3.hideTitleBar()

        area.addDock(d1, 'left')
        area.addDock(d3, 'right', d1)
        area.addDock(d4, 'top', d3)
        area.addDock(d2, 'bottom', d1)
        area.moveDock(d3, 'above', d4)

        ## Test ability to move docks programatically after they have been placed
        #area.moveDock(d1, 'top', d2)     ## move d4 to top edge of d2

        ## Add widgets into each dock
    #    self.sn = None
        self.filePath = '.'
        self.sn = snapshot()
        self.vis = sphvis()

        self.ViewWidget = myGLViewWidget()
        self.ViewWidget.addItem(self.vis)
        self.ViewWidget.reset()
        d1.addWidget(self.ViewWidget)

        ## setup embedded IPython terminal
        namespace={'pg': pg, 'np': np, 'vis': self.vis, 'sn': self.sn, 'vw': self.ViewWidget, 'main': self }
        #w2 = pg.console.ConsoleWidget(namespace=namespace)
        w2 = QIPythonWidget()
        w2.pushVariables(namespace)
        d2.addWidget(w2)

        ## setup control widget
        w3 = pg.LayoutWidget()
        w3.nextRow()
        d3.addWidget(w3)

        ##buttons
        buttons = [('Reset', self.ViewWidget.reset, 'r')]
        for text, action, shortcut in buttons:
            btn = QtGui.QPushButton(text)
            w3.addWidget(btn)
            w3.nextRow()
            if shortcut:
                btn.setShortcut(shortcut)
            btn.clicked.connect(action)

        ## spins
        spins = [
            ("Pointsize", pg.SpinBox(value=0.1, bounds=[0, 100]), self.sizechange),
            ("Pointsize", pg.SpinBox(value=0.1, bounds=[0, 100]), self.sizechange)
            ]
        for text, spin, slot in spins:
            label = QtGui.QLabel("%16s" % text)
            w3.addWidget(label)
            w3.addWidget(spin)
            w3.nextRow()
            spin.sigValueChanged.connect(slot)
        label =  QtGui.QLabel("")
        w3.addWidget(label)

    def sizechange(self, sb):
        self.vis.sizes[:] = sb.value()
        self.vis.update()

    def initpopup(self):
        self.popupPlot = popupPlot()
        vbox = QtGui.QVBoxLayout()
        self.popupPlot.setLayout(vbox)
        self.popupPlot.show()
        self.mw = MatplotlibWidget()
        vbox.addWidget(self.mw)
        self.subplot = self.mw.getFigure().add_subplot(111)
        self.subplot.plot(np.arange(10))
        mw.draw()

    def setupMenu(self):
        openFile = QtGui.QAction('&Open File', self)
        openFile.setShortcut('Ctrl+O')
        openFile.setStatusTip('Open new Gadget File')
        openFile.triggered.connect(self.fileSelect)

        popupAction = QtGui.QAction('&PlotPopup', self)
        popupAction.setStatusTip('Exit application')
        popupAction.triggered.connect(self.initpopup)

        exitAction = QtGui.QAction('&Exit', self)
        exitAction.setShortcut('Esc')
        exitAction.setStatusTip('Exit application')
        exitAction.triggered.connect(QtGui.qApp.quit)

        self.statusBar()

        menubar = self.menuBar()
        fileMenu = menubar.addMenu('&File')
        fileMenu.addAction(openFile)
        fileMenu.addAction(exitAction)

        self.toolbar = self.addToolBar('Exit')
        self.toolbar.addAction(openFile)
        self.toolbar.addAction(popupAction)
        self.toolbar.addAction(exitAction)

    def fileSelect(self):
        fname = QtGui.QFileDialog.getOpenFileName(self, 'Open file', self.filePath)
        if fname:
            self.filePath =  QtCore.QFileInfo(fname).path()
            self.loadsnapshot_gadget(fname)

    def loadsnapshot_gadget(self, fname, **kwargs):
        self.sn.readgadget(fname, **kwargs)
        self.loaddata_gadget(self.sn)

    def loaddata_gadget(self, *args):
        self.vis.loaddata_gadget(*args)
        self.statusBar().showMessage('Snapshot loaded')

## Start Qt event loop unless running in interactive mode or using pyside.
if __name__ == '__main__':
    import sys
    if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
        parser = OptionParser()
        parser.add_option("-b", "--basic", dest="basic", default=False,
                  help="only read basic properties")
        (options, args) = parser.parse_args()
        app = QtGui.QApplication([])
        if len(args):
            mainApp = mainWindow(loadsnapshot=args[0], options=options)
        else:
            mainApp = mainWindow(options=options)
        mainApp.show()
        app.exec_()
<commit_msg>added ability to save the rotation of the viewport in pysphviewer<commit_after>#!/usr/bin/env python
from embeddedIpython import QIPythonWidget
import pyqtgraph as pg
from PyQt4 import QtGui, QtCore
#from pyqtgraph.Qt import QtCore, QtGui
import pyqtgraph.console
from pyqtgraph.widgets.MatplotlibWidget import MatplotlibWidget
import numpy as np
import pyqtgraph.opengl as gl
from astro import *
from sphvis import *
from optparse import OptionParser
import matplotlib
import numpy.linalg
import OpenGL.GL

from pyqtgraph.dockarea import *

class myGLViewWidget(gl.GLViewWidget):
    def __init__(self, **kwargs):
        gl.GLViewWidget.__init__(self, **kwargs)
        self.reset()
        self.initview = np.array([[ 1.,  0.,  0.],
                                  [ 0.,  0.,  1.],
                                  [ 0., -1.,  0.]])

    def reset(self):
        self.opts['azimuth'] = -90
        self.opts['elevation'] = 00
#        self.opts['azimuth'] = 45
#        self.opts['elevation'] = 30

        self.opts['center'] = QtGui.QVector3D(0.0, 0.0, 0.0)
        self.opts['fov'] = 60
        self.opts['distance'] = 30
        self.update()

    def currentView(self):
        return np.dot(OpenGL.GL.glGetFloatv(OpenGL.GL.GL_MODELVIEW_MATRIX)[:3, :3], self.initview)

class popupPlot(QtGui.QDialog):
    def __init__(self, parent=None):
        super(popupPlot, self).__init__()

class mainWindow(QtGui.QMainWindow):
    def __init__(self, load=None, loadsnapshot=None, options={}, **kwargs):
        super(mainWindow, self).__init__()
        self.setupUI()
        if load is not None:
            self.loaddata_gadget(load)
        if loadsnapshot is not None:
            self.loadsnapshot_gadget(loadsnapshot, basic=options.basic)

    def setupUI(self):
        area = DockArea()
        self.setCentralWidget(area)
        self.resize(1000,800)
        self.setWindowTitle('Gadget Visualization')

        self.setupMenu()

        ## Create docks, place them into the window one at a time.
        ## Note that size arguments are only a suggestion; docks will still have to
        ## fill the entire dock area and obey the limits of their internal widgets.
        d1 = Dock("Main Window", size=(800, 800))
        d2 = Dock("Console", size=(800,200))
        d3 = Dock("View", size=(200,800))
        d4 = Dock("Options", size=(200,800))
        d1.hideTitleBar()
        d2.hideTitleBar()
#        d3.hideTitleBar()

        area.addDock(d1, 'left')
        area.addDock(d3, 'right', d1)
        area.addDock(d4, 'top', d3)
        area.addDock(d2, 'bottom', d1)
        area.moveDock(d3, 'above', d4)

        ## Test ability to move docks programatically after they have been placed
        #area.moveDock(d1, 'top', d2)     ## move d4 to top edge of d2

        ## Add widgets into each dock
    #    self.sn = None
        self.filePath = '.'
        self.sn = snapshot()
        self.vis = sphvis()

        self.ViewWidget = myGLViewWidget()
        self.ViewWidget.addItem(self.vis)
        d1.addWidget(self.ViewWidget)

        ## setup embedded IPython terminal
        namespace={'pg': pg, 'np': np, 'vis': self.vis, 'sn': self.sn, 'vw': self.ViewWidget, 'main': self }
        #w2 = pg.console.ConsoleWidget(namespace=namespace)
        w2 = QIPythonWidget()
        w2.pushVariables(namespace)
        d2.addWidget(w2)

        ## setup control widget
        w3 = pg.LayoutWidget()
        w3.nextRow()
        d3.addWidget(w3)

        ##buttons
        buttons = [('Reset', self.ViewWidget.reset, 'r'),
                   ('Save Rotation', self.saveCurrentView, None),
                   ]
        for text, action, shortcut in buttons:
            btn = QtGui.QPushButton(text)
            w3.addWidget(btn)
            w3.nextRow()
            if shortcut:
                btn.setShortcut(shortcut)
            btn.clicked.connect(action)

        ## spins
        spins = [
            ("Pointsize", pg.SpinBox(value=0.1, bounds=[0, 100]), self.sizechange),
            ("Pointsize", pg.SpinBox(value=0.1, bounds=[0, 100]), self.sizechange)
            ]
        for text, spin, slot in spins:
            label = QtGui.QLabel("%16s" % text)
            w3.addWidget(label)
            w3.addWidget(spin)
            w3.nextRow()
            spin.sigValueChanged.connect(slot)
        label =  QtGui.QLabel("")
        w3.addWidget(label)

    def sizechange(self, sb):
        self.vis.sizes[:] = sb.value()
        self.vis.update()

    def initpopup(self):
        self.popupPlot = popupPlot()
        vbox = QtGui.QVBoxLayout()
        self.popupPlot.setLayout(vbox)
        self.popupPlot.show()
        self.mw = MatplotlibWidget()
        vbox.addWidget(self.mw)
        self.subplot = self.mw.getFigure().add_subplot(111)
        self.subplot.plot(np.arange(10))
        mw.draw()

    def setupMenu(self):
        openFile = QtGui.QAction('&Open File', self)
        openFile.setShortcut('Ctrl+O')
        openFile.setStatusTip('Open new Gadget File')
        openFile.triggered.connect(self.openGadget)

        popupAction = QtGui.QAction('&PlotPopup', self)
        popupAction.setStatusTip('Exit application')
        popupAction.triggered.connect(self.initpopup)

        exitAction = QtGui.QAction('&Exit', self)
        exitAction.setShortcut('Esc')
        exitAction.setStatusTip('Exit application')
        exitAction.triggered.connect(QtGui.qApp.quit)

        self.statusBar()

        menubar = self.menuBar()
        fileMenu = menubar.addMenu('&File')
        fileMenu.addAction(openFile)
        fileMenu.addAction(exitAction)

        self.toolbar = self.addToolBar('Exit')
        self.toolbar.addAction(openFile)
        self.toolbar.addAction(popupAction)
        self.toolbar.addAction(exitAction)

    def fileSelect(self):
        fname = QtGui.QFileDialog.getOpenFileName(self, 'Open file', self.filePath)
        if fname:
            self.filePath =  QtCore.QFileInfo(fname).path()
            return fname

    def openGadget(self):
        fname = self.fileSelect()
        if fname:
            self.loadsnapshot_gadget(fname)

    def saveCurrentView(self):
        fname = self.fileSelect()
        with (open(fname, 'w')) as f:
            self.ViewWidget.currentView().transpose().astype(np.float64).tofile(f)

    def loadsnapshot_gadget(self, fname, **kwargs):
        self.sn.readgadget(fname, **kwargs)
        self.loaddata_gadget(self.sn)

    def loaddata_gadget(self, *args):
        self.vis.loaddata_gadget(*args)
        self.statusBar().showMessage('Snapshot loaded')

## Start Qt event loop unless running in interactive mode or using pyside.
if __name__ == '__main__':
    import sys
    if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
        parser = OptionParser()
        parser.add_option("-b", "--basic", dest="basic", default=False,
                  help="only read basic properties")
        (options, args) = parser.parse_args()
        app = QtGui.QApplication([])
        if len(args):
            mainApp = mainWindow(loadsnapshot=args[0], options=options)
        else:
            mainApp = mainWindow(options=options)
        mainApp.show()
        app.exec_()
<|endoftext|>


<commit_before>import os
import uuid
import shutil

from mlflow.utils.databricks_utils import is_in_databricks_runtime
from mlflow.utils._spark_utils import _get_active_spark_session

# Set spark config "spark.mlflow.nfs.rootDir" to specify a NFS (network file system) directory
# which is shared with all spark cluster nodes.
# This will help optimize routine of distributing spark driver files to remote workers.
# None represent no NFS directory available.
# Note:
#  1. If NFS directory set, you must ensure all spark cluster nodes using the same hardware and
#  installed the same OS with the same environment configured, because mlflow uses NFS directory
#  to distribute driver side virtual environment to remote workers if NFS available, heterogeneous
#  cluster nodes might cause issues under the case.
#  2. The NFS directory must be mounted before importing mlflow.
#  3. For databricks users, don't set this config, databricks product sets up internal NFS service
#  automatically.
_NFS_CACHE_ROOT_DIR = None

def get_nfs_cache_root_dir():
    if is_in_databricks_runtime():
        nfs_enabled = (
            _get_active_spark_session()
            .conf.get("spark.databricks.mlflow.nfs.enabled", "true")
            .lower()
            == "true"
        )
        if nfs_enabled:
            nfs_root_dir = "/local_disk0/.ephemeral_nfs"
            # Test whether the NFS directory is writable.
            test_path = os.path.join(nfs_root_dir, uuid.uuid4().hex)
            try:
                os.makedirs(test_path)
                return nfs_root_dir
            except Exception:
                # For databricks cluster enabled Table ACL, we have no permission to access NFS
                # directory, in this case, return None representing NFS is not available.
                return None
            finally:
                shutil.rmtree(test_path, ignore_errors=True)
        else:
            return None
    else:
        spark_session = _get_active_spark_session()
        if spark_session is not None:
            return spark_session.conf.get("spark.mlflow.nfs.rootDir", None)
<commit_msg>Fix NFS enablement check when Spark session is unavailable on Databricks workers (#7333)<commit_after>import os
import uuid
import shutil

from mlflow.utils.databricks_utils import is_in_databricks_runtime
from mlflow.utils._spark_utils import _get_active_spark_session

# Set spark config "spark.mlflow.nfs.rootDir" to specify a NFS (network file system) directory
# which is shared with all spark cluster nodes.
# This will help optimize routine of distributing spark driver files to remote workers.
# None represent no NFS directory available.
# Note:
#  1. If NFS directory set, you must ensure all spark cluster nodes using the same hardware and
#  installed the same OS with the same environment configured, because mlflow uses NFS directory
#  to distribute driver side virtual environment to remote workers if NFS available, heterogeneous
#  cluster nodes might cause issues under the case.
#  2. The NFS directory must be mounted before importing mlflow.
#  3. For databricks users, don't set this config, databricks product sets up internal NFS service
#  automatically.
_NFS_CACHE_ROOT_DIR = None

def get_nfs_cache_root_dir():
    if is_in_databricks_runtime():
        spark_sess = _get_active_spark_session()
        nfs_enabled = spark_sess and (
            spark_sess.conf.get("spark.databricks.mlflow.nfs.enabled", "true").lower() == "true"
        )
        if nfs_enabled:
            nfs_root_dir = "/local_disk0/.ephemeral_nfs"
            # Test whether the NFS directory is writable.
            test_path = os.path.join(nfs_root_dir, uuid.uuid4().hex)
            try:
                os.makedirs(test_path)
                return nfs_root_dir
            except Exception:
                # For databricks cluster enabled Table ACL, we have no permission to access NFS
                # directory, in this case, return None representing NFS is not available.
                return None
            finally:
                shutil.rmtree(test_path, ignore_errors=True)
        else:
            return None
    else:
        spark_session = _get_active_spark_session()
        if spark_session is not None:
            return spark_session.conf.get("spark.mlflow.nfs.rootDir", None)
<|endoftext|>


<commit_before>from neuroimaging.algorithms.statistics import Model

class Classifier(Model):
    def learn(self, **keywords): self.fit(**keywords)

<commit_msg>use the Model from scipy<commit_after>from scipy.sandbox.models.model import Model

class Classifier(Model):
    def learn(self, **keywords):
        self.fit(**keywords)

<|endoftext|>