<commit_before>from functools import partial
from openforcefield.typing.chemistry import *
from openforcefield.utils import get_data_file_path
from unittest import TestCase
import pytest
from openforcefield.utils.toolkits import OPENEYE_AVAILABLE

# TODO: Bring these back online once OpenEye dependence is resolved
@pytest.mark.skip
class TestChemicalEnvironments(TestCase):
    def test_createEnvironments(self):
        """
        Test all types of ChemicalEnvironment objects with defined atoms and bonds
        Each will be tetrahedral carbons connected by ring single bonds
        """
        carbon = [['#6'], ['X4']]
        singleBond = [['-'], ['@']]
        atom = AtomChemicalEnvironment('[#6X4:1]','CT')
        bond = BondChemicalEnvironment('[#6X4:1]-[#6X4:2]', 'CT-CT')
        angle = AngleChemicalEnvironment('[#6X4:1]-[#6X4:2]-[#6X4:3]', 'CT-CT-CT')
        torsion = TorsionChemicalEnvironment('[#6X4:1]-[#6X4:2]-[#6X4:3]-[#6X4:4]', 'CT-CT-CT-CT')
        improper = ImproperChemicalEnvironment('[#6X4:1]-[#6X4:2](-[#6X4:3])-[#6X4:4]', 'CT-CT(-CT)-CT')

    # TODO: Can we remove explicit OE dependence from this?
    @pytest.mark.skip
    def test_complicatedTorsion(self):
        """
        Test ChemicalEnvironment objects with complicated torsion
        test methods that add atoms, remove atoms
        add ORtypes and ANDtypes to existing atoms

        This is the SMIRK for the final torsion
        "[*:1] - [#6:2](=[#8,#7;H0]) - [#6:3](-[#7X3,#8X2;+0]-[#1])(-[#1]) - [*:4]"
        """
        #from openeye.oechem import *
        torsion_smirks = "[*:1]-[#6:2]-[#6:3]-[*:4]"
        torsion = TorsionChemicalEnvironment(torsion_smirks)
        # save atoms (use selectAtom)
        atom1 = torsion.selectAtom(1)
        atom2 = torsion.selectAtom(2)
        atom3 = torsion.selectAtom(3)

        # Add atoms with names so I can try to remove them
        atom2alpha = torsion.addAtom(atom2, [('=',[])], None, [('#8',[]),('#7',[])], ['H0'])
        atom3alpha1 = torsion.addAtom(atom3)
        atom3beta1 = torsion.addAtom(atom3alpha1, [('-',[])], None, [('#1',[])])
        atom3alpha2 = torsion.addAtom(atom3, [('-',[])], None, [('#1',[])])

        # Get bond for atom3 and alpha and add ANDtype
        bond = torsion.getBond(atom3, atom3alpha1)
        self.assertIsNotNone(bond, "could not find bond between atom3 and it's alpha atom")
        bond.addORtype('-', [])

        # Add ORtypes and ANDtypes to atom3 alpha atom
        atom3alpha1.addORtype('#7', ['X3'])
        atom3alpha1.addORtype('#8', ['X2'])
        atom3alpha1.addANDtype('+0')

        # Call getAtoms and getBonds just to make sure they work
        torsion.getAtoms()
        torsion.getBonds()

        # get smarts and smirks for the large torsion
        smarts = torsion.asAtomtypeSMARTS()
        smirks = torsion.asSMIRKS()
        qmol = OEQMol()
        if not OEParseSmarts(qmol, smirks):
            raise Exception(f"could not parse created SMIRKS {smirks}")

        # Try removing atoms
        # if it was labeled:
        removed = torsion.removeAtom(atom1)
        self.assertFalse(removed, f"removed labeled atom ({atom1.asSMIRKS()}) in {torsion.asSMIRKS()}")
        removed = torsion.removeAtom(atom3alpha1)
        self.assertFalse(removed, f"removed bridging atom ({atom3alpha1.asSMIRKS()}) in {torsion.asSMIRKS()}")
        removed = torsion.removeAtom(atom3beta1)
        self.assertTrue(removed, f"failed to remove atom ({atom3beta1.asSMIRKS()}) in {torsion.asSMIRKS()} which is removeable")

    def test_parseSMIRKS(self):
        """
        Test creating environments with SMIRKS
        """
        replacements = [ ('ewg1', '[#7!-1,#8!-1,#16!-1,#9,#17,#35,#53]'),
                ('ewg2', '[#7!-1,#8,#16]') ]

        smirksList = [ ["[#6](-[#1])-[#8]", None, ChemicalEnvironment],
                ["[#6&X4&H0:1](-[#1])-[#6&X4]", 'Atom', AtomChemicalEnvironment],
                [ "[#6&X4&H0:1](-[#1])-[#6&X4:2]", 'Bond', BondChemicalEnvironment],
                [ "[*:1]-[*:2](-[#6&X4])-[*:3]", 'Angle', AngleChemicalEnvironment],
                [ "[#6&X4&H0:1](-[#1])-[#6&X4:2]-[#6&X4&H0:3](-[#1])-[#6&X4:4]", 'ProperTorsion', TorsionChemicalEnvironment],
                [ "[#1:1]-[#6&X4:2](-[#8:3])-[#1:4]", 'ImproperTorsion', ImproperChemicalEnvironment],
                [ "[#1:1]-[#6&X4:2](-[#8:3])-[*:4](-[#6&H1])-[#8:5]", None, ChemicalEnvironment],
                [ "[#6$(*~[#6]=[#8])$(*-,=[$ewg2,#7])]", None, ChemicalEnvironment],
                [ "CCC", None, ChemicalEnvironment],
                [ "[#6:1]1(-;!@[#1,#6])=;@[#6]-;@[#6]1", 'Atom', ChemicalEnvironment],
                [ "C(O-[#7,#8])CC=[*]", None, ChemicalEnvironment],
                [ "[#6$([#6X4](~[$ewg1])(~[#8]~[#1])):1]-[#6X2H2;+0:2]-,=,:;!@;!#[$ewg2:3]-[#4:4]", 'ProperTorsion', TorsionChemicalEnvironment],
                [ "[#6$([#6X4](~[$ewg1])(~[#8]~[#1])):1]1=CCCC1", 'Atom', AtomChemicalEnvironment],
                [ "[*:1]-[#7X3:2](-[#6a$(*1ccc(-[#8-1X1])cc1):3])-[*:4]", 'ImproperTorsion', ImproperChemicalEnvironment],
                [ "[#6X4:1]1~[*:2]~[*$(*~[#1]):3]1", 'Angle', AngleChemicalEnvironment],
                [ "[$([#7]1~[#6]-CC1)]", None, ChemicalEnvironment],
                [ "[$(c1ccccc1)]", None, ChemicalEnvironment],
                ]

        for toolkit in ['openeye', 'rdkit']:
            for [smirks, checkType, chemEnv] in smirksList:
                env = chemEnv(smirks = smirks, replacements = replacements, toolkit=toolkit)
                Type = env.getType()
                self.assertEqual(Type, checkType, f"SMIRKS ({smirks}) clasified as {Type} instead of {checkType} using {toolkit} toolkit")

    def test_environment_functions(self):
        """
        Test all minor functions in environments
        """
        # example 3 indexed, 1 unindexed, 1 alpha, 0 beta
        angle_smirks = "[#6X3;R1:1]=,:;@[#6X3;R1;a:2](-,:;@[#7])-;!@[#8X2H1;!R:3]"
        angle = AngleChemicalEnvironment(angle_smirks)

        # check selectAtom and selectBond
        testSelect = [(4,True), ('Beta', True), (2, False), ('Indexed', False),
                ('Unindexed', False), ('Alpha', False), (None, False)]
        for (descriptor, isNone) in testSelect:
            atom = angle.selectAtom(descriptor)
            bond = angle.selectBond(descriptor)
            if isNone:
                self.assertIsNone(atom, f"Found atom with descriptor {descriptor} in angle {angle_smirks}")
                self.assertIsNone(bond, f"Found bond with descriptor {descriptor} in angle {angle_smirks}")

            else:
                self.assertIsNotNone(atom, f"Could not find a {descriptor} atom in angle {angle_smirks}")
                self.assertIsNotNone(bond, f"Could not find a {descriptor} bond in angle {angle_smirks}")

        # Check function getComponentList which uses all functions of the form
        # get_Atoms and get_Bonds
        compOptions = {'atom': [ ('all', 4), ('Indexed', 3), ('Unindexed', 1), ('Alpha', 1), ('Beta', 0)], 'bond': [ ('all', 3), ('Indexed', 2), ('Unindexed', 1), ('Alpha', 1), ('Beta', 0)]}

        for component, desList in compOptions.items():
            for (descriptor, amount) in desList:
                compList = angle.getComponentList(component, descriptor)
                ncomps = len(compList)
                assert ncomps == amount, f"Found {ncomps} {descriptor} {component}s instead of {amount} in angle {angle_smirks}"

        # Check is__ descriptors
        atom2 = angle.selectAtom(2)
        bond1 = angle.selectBond(1)
        alpha_atom = angle.selectAtom('Alpha')
        beta_atom = angle.addAtom(alpha_atom)
        alpha_bond = angle.getBond(atom2, alpha_atom)
        beta_bond = angle.getBond(alpha_atom, beta_atom)

        # list of lists:
        # [ [[components], [(method, expected)]], [...]]
        checkIsMethods = [
                [[atom2,bond1], [(angle.isAlpha, False), (angle.isBeta, False),
                    (angle.isIndexed, True), (angle.isUnindexed, False)]],
                [[alpha_atom, alpha_bond], [(angle.isAlpha, True), (angle.isIndexed, False),
                    (angle.isUnindexed, True)]],
                [[beta_atom, beta_bond], [(angle.isBeta, True)]]]

        for compSet, methodList in checkIsMethods:
            for comp in compSet:
                for (method, expected) in methodList:
                    # same message
                    classify = method(comp)
                    msg = (f"{method} wrongly returned {classify} for {comp.asSMIRKS()} in angle {angle.asSMIRKS()}")
                    if expected:
                        self.assertTrue(classify, msg)
                    else:
                        self.assertFalse(classify, msg)

        # Check getBond when atoms aren't bonded
        atom1 = angle.selectAtom(1)
        beta_to_atom1 = angle.getBond(beta_atom, atom1)
        self.assertIsNone(beta_to_atom1, f"Incorrect bond was returned connecting the beta atom ({beta_atom.asSMIRKS()}) and atom 1 ({atom1.asSMIRKS()}) in angle {angle_smirks}")

        # Check valence: should be 3 for atom2
        val = angle.getValence(atom2)
        self.assertEqual(val, 3, f"Atom 2 in angle {angle_smirks} should have a valence of 3, but a valence of {val} was found")

        # Check bond order
        # For bond1 =,:;@ it should be 1.5 because order returns lowest possible
        order = bond1.getOrder()
        self.assertEqual(order, 1.5, f"Bond 1 ({bond1.asSMIRKS()}) should have a minimum order of 1.5, but {order:.1f} was assigned")

        # For atom
        order = angle.getBondOrder(atom2)
        self.assertEqual(order, 3.5, f"Atom 2 in angle ({angle_smirks}) should have a minimum bond order of 3.5, but {order:.1f} was assigned")

    def test_creating_wrong_environments(self):
        """
        Test exceptions for making environments with the wrong smirks
        """

        wrongDict = {'[*]': [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]": [BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]~[*:2]": [AtomChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:3]~[*:2]~[*:1]":[AtomChemicalEnvironment, BondChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]~[*:2]~[*:3]~[*:4]": [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]~[*:2](~[*:3])~[*:4]": [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment],
                "[*:1]~[*:2]~[*:3]~[*:4]~[*:5]": [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment] }

        for smirks, Environments in wrongDict.items():
            for Environment in Environments:
                msg = (f"SMIRKS ({smirks}) is in appropriate input for {str(Environment)}, but no error was raised")
                with self.assertRaises(SMIRKSMismatchError, msg = msg):
                    env = Environment(smirks)

    def test_wrong_smirks_error(self):
        """
        Check that an imparseable SMIRKS raises errors
        """
        smirks = "[*;X:1]"
        msg = f"SMIRKS ({smirks}) should not be parseable, but an environment was successfully created"
        with self.assertRaises(SMIRKSParsingError, msg = msg):
            env = ChemicalEnvironment(smirks)

    def test_embedded_atoms_smirks(self):
        """
        Check embedded atom parsing works
        """
        smirks = "[#1$(*-[#6](-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]):1]~[$([#1]~[#6])]"
        env = ChemicalEnvironment(smirks)
<commit_msg>refactor smirks parsing test, add ring cases<commit_after>from functools import partial
from openforcefield.typing.chemistry import *
from openforcefield.utils import get_data_file_path
#from unittest import TestCase
import pytest
from openforcefield.utils.toolkits import OPENEYE_AVAILABLE

if OPENEYE_AVAILABLE:
    toolkits = ['rdkit', 'openeye']
else:
    toolkits = ['rdkit']
# TODO: Bring these back online once OpenEye dependence is resolved
#@pytest.mark.skip
class TestChemicalEnvironments():#TestCase):
    def test_createEnvironments(self):
        """
        Test all types of ChemicalEnvironment objects with defined atoms and bonds
        Each will be tetrahedral carbons connected by ring single bonds
        """
        carbon = [['#6'], ['X4']]
        singleBond = [['-'], ['@']]
        atom = AtomChemicalEnvironment('[#6X4:1]','CT')
        bond = BondChemicalEnvironment('[#6X4:1]-[#6X4:2]', 'CT-CT')
        angle = AngleChemicalEnvironment('[#6X4:1]-[#6X4:2]-[#6X4:3]', 'CT-CT-CT')
        torsion = TorsionChemicalEnvironment('[#6X4:1]-[#6X4:2]-[#6X4:3]-[#6X4:4]', 'CT-CT-CT-CT')
        improper = ImproperChemicalEnvironment('[#6X4:1]-[#6X4:2](-[#6X4:3])-[#6X4:4]', 'CT-CT(-CT)-CT')

    # TODO: Can we remove explicit OE dependence from this?
    @pytest.mark.skip
    def test_complicatedTorsion(self):
        """
        Test ChemicalEnvironment objects with complicated torsion
        test methods that add atoms, remove atoms
        add ORtypes and ANDtypes to existing atoms

        This is the SMIRK for the final torsion
        "[*:1] - [#6:2](=[#8,#7;H0]) - [#6:3](-[#7X3,#8X2;+0]-[#1])(-[#1]) - [*:4]"
        """
        #from openeye.oechem import *
        torsion_smirks = "[*:1]-[#6:2]-[#6:3]-[*:4]"
        torsion = TorsionChemicalEnvironment(torsion_smirks)
        # save atoms (use selectAtom)
        atom1 = torsion.selectAtom(1)
        atom2 = torsion.selectAtom(2)
        atom3 = torsion.selectAtom(3)

        # Add atoms with names so I can try to remove them
        atom2alpha = torsion.addAtom(atom2, [('=',[])], None, [('#8',[]),('#7',[])], ['H0'])
        atom3alpha1 = torsion.addAtom(atom3)
        atom3beta1 = torsion.addAtom(atom3alpha1, [('-',[])], None, [('#1',[])])
        atom3alpha2 = torsion.addAtom(atom3, [('-',[])], None, [('#1',[])])

        # Get bond for atom3 and alpha and add ANDtype
        bond = torsion.getBond(atom3, atom3alpha1)
        self.assertIsNotNone(bond, "could not find bond between atom3 and it's alpha atom")
        bond.addORtype('-', [])

        # Add ORtypes and ANDtypes to atom3 alpha atom
        atom3alpha1.addORtype('#7', ['X3'])
        atom3alpha1.addORtype('#8', ['X2'])
        atom3alpha1.addANDtype('+0')

        # Call getAtoms and getBonds just to make sure they work
        torsion.getAtoms()
        torsion.getBonds()

        # get smarts and smirks for the large torsion
        smarts = torsion.asAtomtypeSMARTS()
        smirks = torsion.asSMIRKS()
        qmol = OEQMol()
        if not OEParseSmarts(qmol, smirks):
            raise Exception(f"could not parse created SMIRKS {smirks}")

        # Try removing atoms
        # if it was labeled:
        removed = torsion.removeAtom(atom1)
        self.assertFalse(removed, f"removed labeled atom ({atom1.asSMIRKS()}) in {torsion.asSMIRKS()}")
        removed = torsion.removeAtom(atom3alpha1)
        self.assertFalse(removed, f"removed bridging atom ({atom3alpha1.asSMIRKS()}) in {torsion.asSMIRKS()}")
        removed = torsion.removeAtom(atom3beta1)
        self.assertTrue(removed, f"failed to remove atom ({atom3beta1.asSMIRKS()}) in {torsion.asSMIRKS()} which is removeable")

    @pytest.mark.parametrize(['smirks', 'expected_valence', 'expected_chemenv_class'],
                             [["[#6](-[#1])-[#8]", None, ChemicalEnvironment],
                              ["[#6&X4&H0:1](-[#1])-[#6&X4]", 'Atom', AtomChemicalEnvironment],
                              ["[#6&X4&H0:1](-[#1])-[#6&X4:2]", 'Bond', BondChemicalEnvironment],
                              ["[*:1]-[*:2](-[#6&X4])-[*:3]", 'Angle', AngleChemicalEnvironment],
                              ["[#6&X4&H0:1](-[#1])-[#6&X4:2]-[#6&X4&H0:3](-[#1])-[#6&X4:4]", 'ProperTorsion',
                               TorsionChemicalEnvironment],
                              ["[#1:1]-[#6&X4:2](-[#8:3])-[#1:4]", 'ImproperTorsion', ImproperChemicalEnvironment],
                              ["[#1:1]-[#6&X4:2](-[#8:3])-[*:4](-[#6&H1])-[#8:5]", None, ChemicalEnvironment],
                              ["[#6$(*~[#6]=[#8])$(*-,=[$ewg2,#7])]", None, ChemicalEnvironment],
                              ["CCC", None, ChemicalEnvironment],
                              ["[#6:1]1(-;!@[#1,#6])=;@[#6]-;@[#6]1", 'Atom', ChemicalEnvironment],
                              ["C(O-[#7,#8])CC=[*]", None, ChemicalEnvironment],
                              ["[#6$([#6X4](~[$ewg1])(~[#8]~[#1])):1]-[#6X2H2;+0:2]-,=,:;!@;!#[$ewg2:3]-[#4:4]",
                               'ProperTorsion', TorsionChemicalEnvironment],
                              ["[#6$([#6X4](~[$ewg1])(~[#8]~[#1])):1]1=CCCC1", 'Atom', AtomChemicalEnvironment],
                              ["[*:1]-[#7X3:2](-[#6a$(*1ccc(-[#8-1X1])cc1):3])-[*:4]", 'ImproperTorsion',
                               ImproperChemicalEnvironment],
                              ["[#6X4:1]1~[*:2]~[*$(*~[#1]):3]1", 'Angle', AngleChemicalEnvironment],
                              ["[$([#7]1~[#6]-CC1)]", None, ChemicalEnvironment],
                              ["[$(c1ccccc1)]", None, ChemicalEnvironment],
                              # The next two tests are for ring-closing bonds
                              ["[H][C@:4]1(C(C([C:3]([N:2]1[C:1](=O)C([H])([H])[H])([H])[H])([H])[H])([H])[H])C=O",
                               'ImproperTorsion', ChemicalEnvironment],
                              ["[P:1]=1=[P]=[P]=[P]=[P:2]=1", 'Bond', BondChemicalEnvironment],
                              ]
                             )
    @pytest.mark.parametrize('toolkit', toolkits)
    def test_parseSMIRKS(self, toolkit, smirks, expected_valence, expected_chemenv_class):
        """
        Test creating environments with SMIRKS
        """
        replacements = [ ('ewg1', '[#7!-1,#8!-1,#16!-1,#9,#17,#35,#53]'),
                ('ewg2', '[#7!-1,#8,#16]') ]

        #smirksList =
        #for toolkit in ['openeye', 'rdkit']:
        #for [smirks, checkType, chemEnv] in smirksList:
        env = expected_chemenv_class(smirks = smirks, replacements = replacements, toolkit=toolkit)
        actual_type = env.getType()
        #self.assertEqual(Type, checkType, f"SMIRKS ({smirks}) classified as {Type} instead of {checkType} using {toolkit} toolkit")
        assert actual_type == expected_valence, \
               f"SMIRKS ({smirks}) classified as {actual_type} instead of {expected_valence} using {toolkit} toolkit"

    def test_environment_functions(self):
        """
        Test all minor functions in environments
        """
        # example 3 indexed, 1 unindexed, 1 alpha, 0 beta
        angle_smirks = "[#6X3;R1:1]=,:;@[#6X3;R1;a:2](-,:;@[#7])-;!@[#8X2H1;!R:3]"
        angle = AngleChemicalEnvironment(angle_smirks)

        # check selectAtom and selectBond
        testSelect = [(4,True), ('Beta', True), (2, False), ('Indexed', False),
                ('Unindexed', False), ('Alpha', False), (None, False)]
        for (descriptor, isNone) in testSelect:
            atom = angle.selectAtom(descriptor)
            bond = angle.selectBond(descriptor)
            if isNone:
                self.assertIsNone(atom, f"Found atom with descriptor {descriptor} in angle {angle_smirks}")
                self.assertIsNone(bond, f"Found bond with descriptor {descriptor} in angle {angle_smirks}")

            else:
                self.assertIsNotNone(atom, f"Could not find a {descriptor} atom in angle {angle_smirks}")
                self.assertIsNotNone(bond, f"Could not find a {descriptor} bond in angle {angle_smirks}")

        # Check function getComponentList which uses all functions of the form
        # get_Atoms and get_Bonds
        compOptions = {'atom': [ ('all', 4), ('Indexed', 3), ('Unindexed', 1), ('Alpha', 1), ('Beta', 0)], 'bond': [ ('all', 3), ('Indexed', 2), ('Unindexed', 1), ('Alpha', 1), ('Beta', 0)]}

        for component, desList in compOptions.items():
            for (descriptor, amount) in desList:
                compList = angle.getComponentList(component, descriptor)
                ncomps = len(compList)
                assert ncomps == amount, f"Found {ncomps} {descriptor} {component}s instead of {amount} in angle {angle_smirks}"

        # Check is__ descriptors
        atom2 = angle.selectAtom(2)
        bond1 = angle.selectBond(1)
        alpha_atom = angle.selectAtom('Alpha')
        beta_atom = angle.addAtom(alpha_atom)
        alpha_bond = angle.getBond(atom2, alpha_atom)
        beta_bond = angle.getBond(alpha_atom, beta_atom)

        # list of lists:
        # [ [[components], [(method, expected)]], [...]]
        checkIsMethods = [
                [[atom2,bond1], [(angle.isAlpha, False), (angle.isBeta, False),
                    (angle.isIndexed, True), (angle.isUnindexed, False)]],
                [[alpha_atom, alpha_bond], [(angle.isAlpha, True), (angle.isIndexed, False),
                    (angle.isUnindexed, True)]],
                [[beta_atom, beta_bond], [(angle.isBeta, True)]]]

        for compSet, methodList in checkIsMethods:
            for comp in compSet:
                for (method, expected) in methodList:
                    # same message
                    classify = method(comp)
                    msg = (f"{method} wrongly returned {classify} for {comp.asSMIRKS()} in angle {angle.asSMIRKS()}")
                    if expected:
                        self.assertTrue(classify, msg)
                    else:
                        self.assertFalse(classify, msg)

        # Check getBond when atoms aren't bonded
        atom1 = angle.selectAtom(1)
        beta_to_atom1 = angle.getBond(beta_atom, atom1)
        self.assertIsNone(beta_to_atom1, f"Incorrect bond was returned connecting the beta atom ({beta_atom.asSMIRKS()}) and atom 1 ({atom1.asSMIRKS()}) in angle {angle_smirks}")

        # Check valence: should be 3 for atom2
        val = angle.getValence(atom2)
        self.assertEqual(val, 3, f"Atom 2 in angle {angle_smirks} should have a valence of 3, but a valence of {val} was found")

        # Check bond order
        # For bond1 =,:;@ it should be 1.5 because order returns lowest possible
        order = bond1.getOrder()
        self.assertEqual(order, 1.5, f"Bond 1 ({bond1.asSMIRKS()}) should have a minimum order of 1.5, but {order:.1f} was assigned")

        # For atom
        order = angle.getBondOrder(atom2)
        self.assertEqual(order, 3.5, f"Atom 2 in angle ({angle_smirks}) should have a minimum bond order of 3.5, but {order:.1f} was assigned")

    def test_creating_wrong_environments(self):
        """
        Test exceptions for making environments with the wrong smirks
        """

        wrongDict = {'[*]': [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]": [BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]~[*:2]": [AtomChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:3]~[*:2]~[*:1]":[AtomChemicalEnvironment, BondChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]~[*:2]~[*:3]~[*:4]": [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, ImproperChemicalEnvironment],
                "[*:1]~[*:2](~[*:3])~[*:4]": [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment],
                "[*:1]~[*:2]~[*:3]~[*:4]~[*:5]": [AtomChemicalEnvironment, BondChemicalEnvironment, AngleChemicalEnvironment, TorsionChemicalEnvironment, ImproperChemicalEnvironment] }

        for smirks, Environments in wrongDict.items():
            for Environment in Environments:
                msg = (f"SMIRKS ({smirks}) is in appropriate input for {str(Environment)}, but no error was raised")
                with self.assertRaises(SMIRKSMismatchError, msg = msg):
                    env = Environment(smirks)

    def test_wrong_smirks_error(self):
        """
        Check that an imparseable SMIRKS raises errors
        """
        smirks = "[*;X:1]"
        msg = f"SMIRKS ({smirks}) should not be parseable, but an environment was successfully created"
        with self.assertRaises(SMIRKSParsingError, msg = msg):
            env = ChemicalEnvironment(smirks)

    def test_embedded_atoms_smirks(self):
        """
        Check embedded atom parsing works
        """
        smirks = "[#1$(*-[#6](-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]):1]~[$([#1]~[#6])]"
        env = ChemicalEnvironment(smirks)
<|endoftext|>


<commit_before>#!/usr/bin/env python
from Arduino import Arduino
import time

def Blink(led_pin,baud):
    """
    Blinks an LED in 1 sec intervals
    """
    board = Arduino(baud)
    while True:
        board.digitalWrite(led_pin,"LOW")
        print board.digitalRead(led_pin) #confirm LOW (0)
        time.sleep(1)
        board.digitalWrite(led_pin,"HIGH")
        print board.digitalRead(led_pin) #confirm HIGH (1)
        time.sleep(1)

def softBlink(led_pin,baud):
    """
    Fades an LED off and on, using
    Arduino's analogWrite (PWM) function
    """
    board=Arduino(baud)
    i=0
    while True:
        i+=1
        k=i%510
        if k%5==0:
            if k>255:
                k=510-k
            board.analogWrite(led_pin,k)

def adjustBrightness(pot_pin,led_pin,baud):
    """
    Adjusts brightness of an LED using a
    potentiometer
    """
    board=Arduino(baud)
    while True:
        time.sleep(0.01)
        val=board.analogRead(pot_pin)/4
        print val
        board.analogWrite(led_pin,val)

def PingSonar(pw_pin,baud):
    """
    Gets distance measurement from Ping)))
    ultrasonic rangefinder connected to pw_pin
    """
    board = Arduino(baud)
    pingPin=pw_pin
    while True:
        duration = board.pulseIn(pingPin, "HIGH")
        inches = duration/72./2.
        cent = duration/29./2.
        print inches,"inches"
        time.sleep(0.1)

def LCD(tx,baud,ssbaud,message):
    """
    Prints to two-line LCD connected to
    pin tx
    """
    board = Arduino(baud)
    board.SoftwareSerial.begin(0,tx,ssbaud)
    while True:
        board.SoftwareSerial.write(" test ")

if __name__=="__main__":
    Blink(13,9600)
    #LCD(5,9600,9600," test ")
    #adjustBrightness(5,11,9600)
    #softBlink(11,9600)<commit_msg>made the com port/device name set-able in the examples<commit_after>#!/usr/bin/env python
from Arduino import Arduino
import time

def Blink(led_pin,baud, port = ""):
    """
    Blinks an LED in 1 sec intervals
    """
    board = Arduino(baud, port=port)
    while True:
        board.digitalWrite(led_pin,"LOW")
        print board.digitalRead(led_pin) #confirm LOW (0)
        time.sleep(1)
        board.digitalWrite(led_pin,"HIGH")
        print board.digitalRead(led_pin) #confirm HIGH (1)
        time.sleep(1)

def softBlink(led_pin,baud, port = ""):
    """
    Fades an LED off and on, using
    Arduino's analogWrite (PWM) function
    """
    board=Arduino(baud, port=port)
    i=0
    while True:
        i+=1
        k=i%510
        if k%5==0:
            if k>255:
                k=510-k
            board.analogWrite(led_pin,k)

def adjustBrightness(pot_pin,led_pin,baud, port=""):
    """
    Adjusts brightness of an LED using a
    potentiometer
    """
    board=Arduino(baud, port=port)
    while True:
        time.sleep(0.01)
        val=board.analogRead(pot_pin)/4
        print val
        board.analogWrite(led_pin,val)

def PingSonar(pw_pin,baud, port=""):
    """
    Gets distance measurement from Ping)))
    ultrasonic rangefinder connected to pw_pin
    """
    board = Arduino(baud, port=port)
    pingPin=pw_pin
    while True:
        duration = board.pulseIn(pingPin, "HIGH")
        inches = duration/72./2.
        cent = duration/29./2.
        print inches,"inches"
        time.sleep(0.1)

def LCD(tx,baud,ssbaud,message, port=""):
    """
    Prints to two-line LCD connected to
    pin tx
    """
    board = Arduino(baud, port=port)
    board.SoftwareSerial.begin(0,tx,ssbaud)
    while True:
        board.SoftwareSerial.write(" test ")

if __name__=="__main__":
    Blink(13,9600)
    #LCD(5,9600,9600," test ")
    #adjustBrightness(5,11,9600)
    #softBlink(11,9600)<|endoftext|>


<commit_before># Copyright 2013 Rackspace Hosting.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from troveclient.v1 import backups
from troveclient.v1 import instances

from openstack_dashboard.test.test_data import utils

DATABASE_DATA_ONE = {
    "status": "ACTIVE",
    "updated": "2013-08-12T22:00:09",
    "name": "Test Database",
    "links": [],
    "created": "2013-08-12T22:00:03",
    "ip": [
        "10.0.0.3"
    ],
    "volume": {
        "used": 0.13,
        "size": 1
    },
    "flavor": {
        "id": "1",
        "links": []
    },
    "id": "6ddc36d9-73db-4e23-b52e-368937d72719"
}

DATABASE_DATA_TWO = {
    "status": "ACTIVE",
    "updated": "2013-08-12T22:00:09",
    "name": "Test Database With DNS",
    "links": [],
    "created": "2013-08-12T22:00:03",
    "hostname": "trove.instance-2.com",
    "volume": {
        "used": 0.13,
        "size": 1
    },
    "flavor": {
        "id": "1",
        "links": []
    },
    "id": "4d7b3f57-44f5-41d2-8e86-36b88cad572a"
}

BACKUP_ONE = {
    "instance_id": "6ddc36d9-73db-4e23-b52e-368937d72719",
    "status": "COMPLETED",
    "updated": "2013-08-13T19:39:38",
    "locationRef": "http://swift/v1/AUTH/database_backups/0edb.tar.gz",
    "name": "backup1",
    "created": "2013-08-15T18:10:14",
    "size": 0.13,
    "id": "0edb3c14-8919-4583-9add-00df9e524081",
    "description": "Long description of backup"
}

BACKUP_TWO = {
    "instance_id": "4d7b3f57-44f5-41d2-8e86-36b88cad572a",
    "status": "COMPLETED",
    "updated": "2013-08-10T20:20:44",
    "locationRef": "http://swift/v1/AUTH/database_backups/e460.tar.gz",
    "name": "backup2",
    "created": "2013-08-10T20:20:37",
    "size": 0.13,
    "id": "e4602a3c-2bca-478f-b059-b6c215510fb4",
    "description": "Longer description of backup"
}

def data(TEST):
    database1 = instances.Instance(instances.Instances(None),
                                   DATABASE_DATA_ONE)
    database2 = instances.Instance(instances.Instances(None),
                                   DATABASE_DATA_TWO)
    bkup1 = backups.Backup(backups.Backups(None), BACKUP_ONE)
    bkup2 = backups.Backup(backups.Backups(None), BACKUP_TWO)

    TEST.databases = utils.TestDataContainer()
    TEST.database_backups = utils.TestDataContainer()
    TEST.databases.add(database1)
    TEST.databases.add(database2)
    TEST.database_backups.add(bkup1)
    TEST.database_backups.add(bkup2)
<commit_msg>Users can now be deleted within trove.<commit_after># Copyright 2013 Rackspace Hosting.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from troveclient.v1 import backups
from troveclient.v1 import databases
from troveclient.v1 import instances
from troveclient.v1 import users

from openstack_dashboard.test.test_data import utils

DATABASE_DATA_ONE = {
    "status": "ACTIVE",
    "updated": "2013-08-12T22:00:09",
    "name": "Test Database",
    "links": [],
    "created": "2013-08-12T22:00:03",
    "ip": [
        "10.0.0.3",
    ],
    "volume": {
        "used": 0.13,
        "size": 1,
    },
    "flavor": {
        "id": "1",
        "links": [],
    },
    "id": "6ddc36d9-73db-4e23-b52e-368937d72719",
}

DATABASE_DATA_TWO = {
    "status": "ACTIVE",
    "updated": "2013-08-12T22:00:09",
    "name": "Test Database With DNS",
    "links": [],
    "created": "2013-08-12T22:00:03",
    "hostname": "trove.instance-2.com",
    "volume": {
        "used": 0.13,
        "size": 1,
    },
    "flavor": {
        "id": "1",
        "links": [],
    },
    "id": "4d7b3f57-44f5-41d2-8e86-36b88cad572a",
}

BACKUP_ONE = {
    "instance_id": "6ddc36d9-73db-4e23-b52e-368937d72719",
    "status": "COMPLETED",
    "updated": "2013-08-13T19:39:38",
    "locationRef": "http://swift/v1/AUTH/database_backups/0edb.tar.gz",
    "name": "backup1",
    "created": "2013-08-15T18:10:14",
    "size": 0.13,
    "id": "0edb3c14-8919-4583-9add-00df9e524081",
    "description": "Long description of backup",
}

BACKUP_TWO = {
    "instance_id": "4d7b3f57-44f5-41d2-8e86-36b88cad572a",
    "status": "COMPLETED",
    "updated": "2013-08-10T20:20:44",
    "locationRef": "http://swift/v1/AUTH/database_backups/e460.tar.gz",
    "name": "backup2",
    "created": "2013-08-10T20:20:37",
    "size": 0.13,
    "id": "e4602a3c-2bca-478f-b059-b6c215510fb4",
    "description": "Longer description of backup",
}

USER_ONE = {
    "name": "Test_User",
    "host": "%",
    "databases": [DATABASE_DATA_ONE["name"]],
}

USER_DB_ONE = {
    "name": "db1",
}

def data(TEST):
    database1 = instances.Instance(instances.Instances(None),
                                   DATABASE_DATA_ONE)
    database2 = instances.Instance(instances.Instances(None),
                                   DATABASE_DATA_TWO)
    bkup1 = backups.Backup(backups.Backups(None), BACKUP_ONE)
    bkup2 = backups.Backup(backups.Backups(None), BACKUP_TWO)
    user1 = users.User(users.Users(None), USER_ONE)
    user_db1 = databases.Database(databases.Databases(None),
                                  USER_DB_ONE)

    TEST.databases = utils.TestDataContainer()
    TEST.database_backups = utils.TestDataContainer()
    TEST.database_users = utils.TestDataContainer()
    TEST.database_user_dbs = utils.TestDataContainer()
    TEST.databases.add(database1)
    TEST.databases.add(database2)
    TEST.database_backups.add(bkup1)
    TEST.database_backups.add(bkup2)
    TEST.database_users.add(user1)
    TEST.database_user_dbs.add(user_db1)
<|endoftext|>


<commit_before>from prometheus_client import start_http_server, Gauge
import re
import time
import os
import subprocess
import argparse

TEMP_REGEX = re.compile('t=([0-9]+)')
INTERNAL_TEMP_REGEX = re.compile('temp=([0-9]+\.[0-9])')
SENSOR_PATH = "/sys/bus/w1/devices"

def read_command_arguments():
    parser = argparse.ArgumentParser(description="Exporter of connected 1-Wire Sensors to Prometheus.io monitoring.")
    parser.add_argument("--port",
                        metavar="PORT",
                        type=int,
                        default=8001,
                        help="The port used to export the temperatures to prometheus.io. Default is 8001.")
    parser.add_argument("--internal",
                        default=False,
                        action="store_true",
                        help="Export the internal CPU Temperature of the Raspberry Pi. Requires the vcgencmd.")
    return parser.parse_args()

def find_sensors():
    ret = []
    for device in os.listdir(SENSOR_PATH):
        if "w1_bus_master" not in device:
            sensor_under_test = Sensor(device)
            test_value = sensor_under_test.read_sensor()
            if test_value is not None:
                ret.append(sensor_under_test)
    return ret

def get_callback(device):
    def read():
        return read_sensor(device)
    return read

class Sensor (object):

    def __init__(self, id):
        self.id = id
        self.last_value = 0

    def __call__(self, *args, **kwargs):
        value = self.read_sensor()
        if value is None:
            return self.last_value
        else:
            self.last_value = value
            return value

    def read_sensor(self):
        with open("%s/%s/w1_slave" % (SENSOR_PATH, self.id), "r") as sensor:
            crc = l = None
            try:
                crc = sensor.readline()
                if "YES" not in crc:
                    raise Exception("Could not read sensor. Wrong CRC.")
                l = sensor.readline()
                match = TEMP_REGEX.search(l)
                ret_val = float(match.group(1))/1000
                if ret_val == 85:
                    raise Exception("Could not read sensor. Return default value.")
                else:
                    return ret_val
            except Exception, e:
                print "Ups something went wrong.", e.message
                if crc is not None:
                    print "CRC Line: ", crc
                if l is not None:
                    print "l Line: ", l
                return None

def read_raspberry_pi_temperature():
    p = subprocess.Popen(["/usr/bin/vcgencmd", "measure_temp"], stdout=subprocess.PIPE)
    match = INTERNAL_TEMP_REGEX.search(p.stdout.read())
    return float(match.group(1))

def register_prometheus_gauges(export_internal_raspberry=False):
    g = Gauge("sensor_temperature_in_celsius", "Local room temperature around the raspberry pi", ["sensor"])
    sensors = find_sensors()
    print "Found sensors:", ", ".join(sensors)
    for sensor in sensors:
        g.labels(sensor).set_function(sensor)
    if export_internal_raspberry:
        g = Gauge("cpu_temperature_in_celsius", "CPU Temperature of the Raspberry Pi")
        g.set_function(read_raspberry_pi_temperature)
    return sensors

if __name__ == "__main__":
    args = read_command_arguments()
    start_http_server(args.port)
    register_prometheus_gauges(args.internal)
    while True:
        time.sleep(10000)
<commit_msg>Added __str__ and __repr__ to Sensor class<commit_after>from prometheus_client import start_http_server, Gauge
import re
import time
import os
import subprocess
import argparse

TEMP_REGEX = re.compile('t=([0-9]+)')
INTERNAL_TEMP_REGEX = re.compile('temp=([0-9]+\.[0-9])')
SENSOR_PATH = "/sys/bus/w1/devices"

def read_command_arguments():
    parser = argparse.ArgumentParser(description="Exporter of connected 1-Wire Sensors to Prometheus.io monitoring.")
    parser.add_argument("--port",
                        metavar="PORT",
                        type=int,
                        default=8001,
                        help="The port used to export the temperatures to prometheus.io. Default is 8001.")
    parser.add_argument("--internal",
                        default=False,
                        action="store_true",
                        help="Export the internal CPU Temperature of the Raspberry Pi. Requires the vcgencmd.")
    return parser.parse_args()

def find_sensors():
    ret = []
    for device in os.listdir(SENSOR_PATH):
        if "w1_bus_master" not in device:
            sensor_under_test = Sensor(device)
            test_value = sensor_under_test.read_sensor()
            if test_value is not None:
                ret.append(sensor_under_test)
    return ret

def get_callback(device):
    def read():
        return read_sensor(device)
    return read

class Sensor (object):

    def __init__(self, id):
        self.id = id
        self.last_value = 0

    def __call__(self, *args, **kwargs):
        value = self.read_sensor()
        if value is None:
            return self.last_value
        else:
            self.last_value = value
            return value

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return "Sensor ID: ", self.id

    def read_sensor(self):
        with open("%s/%s/w1_slave" % (SENSOR_PATH, self.id), "r") as sensor:
            crc = l = None
            try:
                crc = sensor.readline()
                if "YES" not in crc:
                    raise Exception("Could not read sensor. Wrong CRC.")
                l = sensor.readline()
                match = TEMP_REGEX.search(l)
                ret_val = float(match.group(1))/1000
                if ret_val == 85:
                    raise Exception("Could not read sensor. Return default value.")
                else:
                    return ret_val
            except Exception, e:
                print "Ups something went wrong.", e.message
                if crc is not None:
                    print "CRC Line: ", crc
                if l is not None:
                    print "l Line: ", l
                return None

def read_raspberry_pi_temperature():
    p = subprocess.Popen(["/usr/bin/vcgencmd", "measure_temp"], stdout=subprocess.PIPE)
    match = INTERNAL_TEMP_REGEX.search(p.stdout.read())
    return float(match.group(1))

def register_prometheus_gauges(export_internal_raspberry=False):
    g = Gauge("sensor_temperature_in_celsius", "Local room temperature around the raspberry pi", ["sensor"])
    sensors = find_sensors()
    print "Found sensors:", ", ".join(sensors)
    for sensor in sensors:
        g.labels(sensor).set_function(sensor)
    if export_internal_raspberry:
        g = Gauge("cpu_temperature_in_celsius", "CPU Temperature of the Raspberry Pi")
        g.set_function(read_raspberry_pi_temperature)
    return sensors

if __name__ == "__main__":
    args = read_command_arguments()
    start_http_server(args.port)
    register_prometheus_gauges(args.internal)
    while True:
        time.sleep(10000)
<|endoftext|>


<commit_before>import re
from operator import itemgetter
from collections import namedtuple, defaultdict
from .dumper_interface import Dumper as BaseDumper

class Dumper(BaseDumper):

    __ignore_list__ = ['migrations', 'sqlite_sequence']

    column_record = namedtuple('Record', ['name', 'ttype', 'nullable', 'pk',
                                          'precision', 'default', 'autoincr',
                                          'unsigned', 'origin_type'])

    mapping = {
        'INT': 'integer',
        'SMALLINT': 'small_int',
        'TINYINT': 'tiny_int',
        'MEDIUMINT': 'medium_int',
        'BIGINT': 'big_integer',
        'UNSIGNED BIGINT': 'big_integer',
        'FLOAT': 'float',
        'DECIMAL': 'decimal',
        'DOUBLE': 'double',
        'REAL': 'double',
        'BLOB': 'binary',
        'CHAR': 'char',
        'VARCHAR': 'string',
        'MEDIUMTEXT': 'medium_text',
        'TEXT': 'text',
        'LONGTEXT': 'long_text',
        'DATE': 'date',
        'DATETIME': 'timestamp',
        'INTEGER': 'integer',
        'JSON': 'json',
    }

    def grep_origin_meta(self, plain_sql):
        result = []
        s = slice(plain_sql.find('(') + 1, plain_sql.rfind(')'))
        columns = plain_sql[s].split(', ')
        meta_regex = re.compile('/\*(.*?)\*/')
        for column in columns:
            match_obj = meta_regex.search(column.strip())
            if match_obj is not None:
                result.append(match_obj.groups()[0])
            else:
                result.append(None)
        return result

    def grep_unsigned(self, plain_sql):
        result = []
        s = slice(plain_sql.find('(') + 1, plain_sql.rfind(')'))
        columns = plain_sql[s].split(', ')
        meta_regex = re.compile('/\*.*\*/.*/\*(.*)\*/')
        for column in columns:
            match_obj = meta_regex.search(column.strip())
            if match_obj is not None:
                result.append(match_obj.groups()[0])
            else:
                result.append(None)
        return result

    def list_columns(self, table_name):
        """list column in table
        rtype [namedtuple]"""
        sql = self._grammar._list_columns(table_name)
        result = self._conn.select(sql)
        plain_sql = self.list_plain_sql(table_name)
        autoincrement_columns = self.list_autoincrement_columns(plain_sql)
        origin_type = self.grep_origin_meta(plain_sql)
        unsigned_col = self.grep_unsigned(plain_sql)
        columns = []
        for r, otype, is_unsigned in zip(result, origin_type, unsigned_col):
            # add autoincrement info
            if r['name'] in autoincrement_columns:
                r['autoincr'] = True
            else:
                r['autoincr'] = False

            # handle type precision
            precision = None
            match_obj = re.match(r'\w+\((\S+)\)', r['type'])
            if match_obj is not None:
                # handle like VARCHAR(20)
                precision, *_ = match_obj.groups()
                precision = int(precision)

            columns.append(self.column_record(
                name=r['name'],
                ttype=r['type'],
                precision=precision,
                pk=r['pk'],
                unsigned=is_unsigned,
                nullable=not r['notnull'],
                default=r['dflt_value'],
                autoincr=r['autoincr'],
                origin_type=otype
            ))
        return columns

    def handle_column(self, columns):
        statements = []
        for column in columns:
            column_buffer = []
            name = column.name
            ttype = self.mapping[column.ttype.upper()]

            pk = False
            if column.autoincr:
                if ttype == 'big_integer':
                    ttype = 'big_increments'
                    pk = True
                if ttype == 'integer':
                    ttype = 'increments'
                    pk = True

            unsigned = column.unsigned
            nullable = column.nullable
            default = column.default

            # dump to orator schema syntax
            if not pk and column.origin_type is not None:
                column_buffer.append(
                    'table.' + column.origin_type % repr(name))
            else:
                column_buffer.append('table.{ttype}({name})'.format(
                    ttype=ttype, name=repr(name)))
            if not pk and unsigned:
                column_buffer.append('.unsigned()')
            if nullable:
                column_buffer.append('.nullable()')
            if default is not None:
                flag = True
                if ttype == 'timestamp' and \
                        default.startswith('CURRENT_TIMESTAMP'):
                    flag = False

                if flag:
                    default = default.strip("'")
                    if default.isdigit():
                        default = int(default)
                    elif re.match("^\d+?\.\d+?$", default) is not None:
                        default = float(default)
                    else:
                        default = "'{}'".format(default)
                    column_buffer.append('.default({})'.format(default))
            statements.append(''.join(column_buffer))
        return statements

    def list_indexes(self, table_name):
        """list index in table"""
        sql = self._grammar._list_indexes(table_name)
        indexes = defaultdict(lambda: {'columns': [], 'is_unique': False})
        result = self._conn.select(sql)
        for r in result:
            index_name = r['name']
            indexes[index_name]['is_unique'] = bool(r['unique'])
            columns = map(itemgetter('name'), self._conn.select(
                self._grammar._show_index(index_name)))
            indexes[r['name']]['columns'].extend(columns)
        return indexes

    def handle_index(self, indexes):
        statements = []
        for name, index in indexes.items():
            ttype = 'index'
            if index['is_unique']:
                ttype = 'unique'
            if name == 'PRIMARY':
                ttype = 'primary'
                name = None

            statements.append(
                'table.{}({}, name={})'.format(ttype, repr(index['columns']),
                                               repr(name)))
        return statements

    def list_plain_sql(self, table_name):
        plain_sql = self._conn.select(
            self._grammar._plain_sql(table_name))[0]['sql']
        return plain_sql

    def list_autoincrement_columns(self, plain_sql):
        # SQLite autoincrement is not recommanded to use
        # see https://sqlite.org/autoinc.html
        # Besides it have none default value
        result = []
        s = slice(plain_sql.find('(') + 1, plain_sql.rfind(')'))
        columns = plain_sql[s].split(', ')
        for column in columns:
            if 'AUTOINCREMENT' in column:
                column_name, *_ = column.partition(' ')
                result.append(column_name.strip('"'))
        return result

    def list_foreign_keys(self, table_name):
        """list foreign key from specified table"""
        sql = self._grammar._list_foreign_keys(table_name)
        result = self._conn.select(sql)
        foreign_keys = []
        for r in result:
            foreign_keys.append({
                'column': r['from'],
                'ref_key': r['to'],
                'to_table': r['table'],
                'on_update': r['on_update'],
                'on_delete': r['on_delete']
            })
        return foreign_keys

    def handle_primary_key(self, primary_key):
        return ['table.primary([{}])'.format(repr(primary_key[0].name))]

    def dump(self):
        table_names = list(self.list_tables())

        table_buffer = []
        for table in table_names:
            columns = self.list_columns(table)

            indexes = self.list_indexes(table)
            primary_key = list(filter(lambda column: column.pk, columns))
            foreign_keys = self.list_foreign_keys(table)
            statement_buffer = []

            statement_buffer.extend(self.handle_column(columns))
            statement_buffer.extend(self.handle_index(indexes))
            statement_buffer.extend(self.handle_foreign_key(foreign_keys))
            statement_buffer.extend(self.handle_primary_key(primary_key))

            table_buffer.append(self.table_tmpl.render(
                table_name=table,
                table_statement=statement_buffer
            ))

        output = self.schema_tmpl.render(
            tables_created=table_buffer,
            tables_droped=table_names
        )
        return output

    def list_tables(self):
        """list all table_names from specified database
        rtype [str]
        """
        sql = self._grammar._list_tables()
        result = self._conn.select(sql)
        return filter(
            lambda table_name: table_name not in self.__ignore_list__,
            map(itemgetter('table_name'), result))

    def handle_foreign_key(self, foreign_keys):
        statements = []
        for foreign_key in foreign_keys:
            local_key = foreign_key['column']
            ref_key = foreign_key['ref_key']
            to_table = foreign_key['to_table']
            on_update = foreign_key['on_update']
            on_delete = foreign_key['on_delete']

            statement = 'table.foreign({}).references({}).on({})'.format(
                repr(local_key), repr(ref_key), repr(to_table))

            if on_update.upper() == 'CASCADEA':
                statement += ".on_update('cascadea')"
            if on_delete.upper() == 'CASCADEA':
                statement += ".on_delete('cascadea')"

            statements.append(statement)
        return statements
<commit_msg>fix big_increments bug<commit_after>import re
from operator import itemgetter
from collections import namedtuple, defaultdict
from .dumper_interface import Dumper as BaseDumper

class Dumper(BaseDumper):

    __ignore_list__ = ['migrations', 'sqlite_sequence']

    column_record = namedtuple('Record', ['name', 'ttype', 'nullable', 'pk',
                                          'precision', 'default', 'autoincr',
                                          'unsigned', 'origin_type'])

    mapping = {
        'INT': 'integer',
        'SMALLINT': 'small_int',
        'TINYINT': 'tiny_int',
        'MEDIUMINT': 'medium_int',
        'BIGINT': 'big_integer',
        'UNSIGNED BIGINT': 'big_integer',
        'FLOAT': 'float',
        'DECIMAL': 'decimal',
        'DOUBLE': 'double',
        'REAL': 'double',
        'BLOB': 'binary',
        'CHAR': 'char',
        'VARCHAR': 'string',
        'MEDIUMTEXT': 'medium_text',
        'TEXT': 'text',
        'LONGTEXT': 'long_text',
        'DATE': 'date',
        'DATETIME': 'timestamp',
        'INTEGER': 'integer',
        'JSON': 'json',
    }

    def grep_origin_meta(self, plain_sql):
        result = []
        s = slice(plain_sql.find('(') + 1, plain_sql.rfind(')'))
        columns = plain_sql[s].split(', ')
        meta_regex = re.compile('/\*(.*?)\*/')
        for column in columns:
            match_obj = meta_regex.search(column.strip())
            if match_obj is not None:
                result.append(match_obj.groups()[0])
            else:
                result.append(None)
        return result

    def grep_unsigned(self, plain_sql):
        result = []
        s = slice(plain_sql.find('(') + 1, plain_sql.rfind(')'))
        columns = plain_sql[s].split(', ')
        meta_regex = re.compile('/\*.*\*/.*/\*(.*)\*/')
        for column in columns:
            match_obj = meta_regex.search(column.strip())
            if match_obj is not None:
                result.append(match_obj.groups()[0])
            else:
                result.append(None)
        return result

    def list_columns(self, table_name):
        """list column in table
        rtype [namedtuple]"""
        sql = self._grammar._list_columns(table_name)
        result = self._conn.select(sql)
        plain_sql = self.list_plain_sql(table_name)
        autoincrement_columns = self.list_autoincrement_columns(plain_sql)
        origin_type = self.grep_origin_meta(plain_sql)
        unsigned_col = self.grep_unsigned(plain_sql)
        columns = []
        for r, otype, is_unsigned in zip(result, origin_type, unsigned_col):
            # add autoincrement info
            if r['name'] in autoincrement_columns:
                r['autoincr'] = True
            else:
                r['autoincr'] = False

            # handle type precision
            precision = None
            match_obj = re.match(r'\w+\((\S+)\)', r['type'])
            if match_obj is not None:
                # handle like VARCHAR(20)
                precision, *_ = match_obj.groups()
                precision = int(precision)

            columns.append(self.column_record(
                name=r['name'],
                ttype=r['type'],
                precision=precision,
                pk=r['pk'],
                unsigned=is_unsigned,
                nullable=not r['notnull'],
                default=r['dflt_value'],
                autoincr=r['autoincr'],
                origin_type=otype
            ))
        return columns

    def handle_column(self, columns):
        statements = []
        for column in columns:
            column_buffer = []
            name = column.name
            ttype = self.mapping[column.ttype.upper()]
            origin_type = column.origin_type

            pk = False
            if column.autoincr:
                if origin_type.startswith('big_integer'):
                    ttype = 'big_increments'
                    pk = True
                if origin_type.startswith('integer'):
                    ttype = 'increments'
                    pk = True

            unsigned = column.unsigned
            nullable = column.nullable
            default = column.default

            # dump to orator schema syntax
            if not pk and origin_type is not None:
                column_buffer.append('table.' + origin_type % repr(name))
            else:
                column_buffer.append('table.{ttype}({name})'.format(
                    ttype=ttype, name=repr(name)))
            if not pk and unsigned:
                column_buffer.append('.unsigned()')
            if nullable:
                column_buffer.append('.nullable()')
            if default is not None:
                flag = True
                if ttype == 'timestamp' and \
                        default.startswith('CURRENT_TIMESTAMP'):
                    flag = False

                if flag:
                    default = default.strip("'")
                    if default.isdigit():
                        default = int(default)
                    elif re.match("^\d+?\.\d+?$", default) is not None:
                        default = float(default)
                    else:
                        default = "'{}'".format(default)
                    column_buffer.append('.default({})'.format(default))
            statements.append(''.join(column_buffer))
        return statements

    def list_indexes(self, table_name):
        """list index in table"""
        sql = self._grammar._list_indexes(table_name)
        indexes = defaultdict(lambda: {'columns': [], 'is_unique': False})
        result = self._conn.select(sql)
        for r in result:
            index_name = r['name']
            indexes[index_name]['is_unique'] = bool(r['unique'])
            columns = map(itemgetter('name'), self._conn.select(
                self._grammar._show_index(index_name)))
            indexes[r['name']]['columns'].extend(columns)
        return indexes

    def handle_index(self, indexes):
        statements = []
        for name, index in indexes.items():
            ttype = 'index'
            if index['is_unique']:
                ttype = 'unique'
            if name == 'PRIMARY':
                ttype = 'primary'
                name = None

            statements.append(
                'table.{}({}, name={})'.format(ttype, repr(index['columns']),
                                               repr(name)))
        return statements

    def list_plain_sql(self, table_name):
        plain_sql = self._conn.select(
            self._grammar._plain_sql(table_name))[0]['sql']
        return plain_sql

    def list_autoincrement_columns(self, plain_sql):
        # SQLite autoincrement is not recommanded to use
        # see https://sqlite.org/autoinc.html
        # Besides it have none default value
        result = []
        s = slice(plain_sql.find('(') + 1, plain_sql.rfind(')'))
        columns = plain_sql[s].split(', ')
        for column in columns:
            if 'AUTOINCREMENT' in column:
                column_name, *_ = column.partition(' ')
                result.append(column_name.strip('"'))
        return result

    def list_foreign_keys(self, table_name):
        """list foreign key from specified table"""
        sql = self._grammar._list_foreign_keys(table_name)
        result = self._conn.select(sql)
        foreign_keys = []
        for r in result:
            foreign_keys.append({
                'column': r['from'],
                'ref_key': r['to'],
                'to_table': r['table'],
                'on_update': r['on_update'],
                'on_delete': r['on_delete']
            })
        return foreign_keys

    def handle_primary_key(self, primary_key):
        return ['table.primary([{}])'.format(repr(primary_key[0].name))]

    def dump(self):
        table_names = list(self.list_tables())

        table_buffer = []
        for table in table_names:
            columns = self.list_columns(table)

            indexes = self.list_indexes(table)
            primary_key = list(filter(lambda column: column.pk, columns))
            foreign_keys = self.list_foreign_keys(table)
            statement_buffer = []

            statement_buffer.extend(self.handle_column(columns))
            statement_buffer.extend(self.handle_index(indexes))
            statement_buffer.extend(self.handle_foreign_key(foreign_keys))
            statement_buffer.extend(self.handle_primary_key(primary_key))

            table_buffer.append(self.table_tmpl.render(
                table_name=table,
                table_statement=statement_buffer
            ))

        output = self.schema_tmpl.render(
            tables_created=table_buffer,
            tables_droped=table_names
        )
        return output

    def list_tables(self):
        """list all table_names from specified database
        rtype [str]
        """
        sql = self._grammar._list_tables()
        result = self._conn.select(sql)
        return filter(
            lambda table_name: table_name not in self.__ignore_list__,
            map(itemgetter('table_name'), result))

    def handle_foreign_key(self, foreign_keys):
        statements = []
        for foreign_key in foreign_keys:
            local_key = foreign_key['column']
            ref_key = foreign_key['ref_key']
            to_table = foreign_key['to_table']
            on_update = foreign_key['on_update']
            on_delete = foreign_key['on_delete']

            statement = 'table.foreign({}).references({}).on({})'.format(
                repr(local_key), repr(ref_key), repr(to_table))

            if on_update.upper() == 'CASCADEA':
                statement += ".on_update('cascadea')"
            if on_delete.upper() == 'CASCADEA':
                statement += ".on_delete('cascadea')"

            statements.append(statement)
        return statements
<|endoftext|>


<commit_before>#  Numenta Platform for Intelligent Computing (NuPIC)
#  Copyright (C) 2019, Numenta, Inc.  Unless you have an agreement
#  with Numenta, Inc., for a separate license for this software code, the
#  following terms and conditions apply:
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Affero Public License version 3 as
#  published by the Free Software Foundation.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#  See the GNU Affero Public License for more details.
#
#  You should have received a copy of the GNU Affero Public License
#  along with this program.  If not, see http://www.gnu.org/licenses.
#
#  http://numenta.org/licenses/
#
import copy
import logging
import os
import traceback
from pprint import pprint

import ray
import ray.resource_spec
from ray import ray_constants
from ray.exceptions import RayTimeoutError
from ray.tune import Trainable, tune
from ray.tune.ray_trial_executor import RESOURCE_REFRESH_PERIOD, RayTrialExecutor
from ray.tune.resources import Resources

from nupic.research.frameworks.pytorch.imagenet import ImagenetExperiment

os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
logger = logging.getLogger(__name__)
TRIAL_START_ATTEMPTS = 3
NODE_CREATION_TIMEOUT = 10 * 60

def _get_resources_per_node():
    """
    Maps node id to available resources on that node
    :return: dict with available :class:`Resources` for each node
    """
    def _is_node_key(k):
        return k.startswith(ray.resource_spec.NODE_ID_PREFIX)

    # Only consider active/alive nodes
    nodes = filter(lambda node: node["Alive"], ray.nodes())
    resources = map(lambda node: node["Resources"], nodes)

    # Group resources by node
    resources_by_node = {}
    for item in resources:
        node_id = next(filter(_is_node_key, item.keys()))

        item = item.copy()
        num_cpus = item.pop("CPU", 0)
        num_gpus = item.pop("GPU", 0)
        memory = ray_constants.from_memory_units(item.pop("memory", 0))
        object_store_memory = ray_constants.from_memory_units(
            item.pop("object_store_memory", 0))
        custom_resources = item

        resources_by_node[node_id] = Resources(
            int(num_cpus),
            int(num_gpus),
            memory=int(memory),
            object_store_memory=int(object_store_memory),
            custom_resources=custom_resources)

    return resources_by_node

class AffinityExecutor(RayTrialExecutor):
    """
    Ray Trial Executor used to set node affinity to node trials.
    This trial executor will update the trial configuration with one extra
    parameter (__ray_node_affinity__) that can be used to launch extra workers
    on the same node.
    """

    def __init__(self, queue_trials=False, reuse_actors=False, ray_auto_init=False,
                 refresh_period=RESOURCE_REFRESH_PERIOD):
        self._resources_by_node = {}
        super().__init__(queue_trials, reuse_actors, ray_auto_init, refresh_period)

    def _commit_resources(self, resources):
        resources = copy.deepcopy(resources)
        # Make sure "custom_resources" keys match "extra_custom_resources" keys
        # because `RayTrialExecutor._commit_resources` will only commit
        # resources from "extra_custom_resources" if the same resource key
        # is present in the "custom_resources" dict
        for k in resources.extra_custom_resources:
            if k not in resources.custom_resources:
                resources.custom_resources[k] = 0.0

        super()._commit_resources(resources)

    def _update_avail_resources(self, num_retries=5):
        super()._update_avail_resources(num_retries)
        self._resources_by_node = _get_resources_per_node()

    def _get_node_resources(self, resources):
        self._update_avail_resources()

        # Check for required GPUs first
        required = resources.gpu_total()
        if required > 0:
            resource_attr = "gpu"
        else:
            # No GPU, just use CPU
            resource_attr = "cpu"
            required = resources.cpu_total()

        # Compute nodes required to fulfill trial resources request
        custom_resources = {}
        for node_id, node_resource in self._resources_by_node.items():
            # Compute resource remaining on each node
            node_capacity = node_resource.get_res_total(node_id)
            committed_capacity = self._committed_resources.get_res_total(node_id)
            remaining_capacity = node_capacity - committed_capacity
            node_procs = getattr(node_resource, resource_attr, 0)
            available = node_procs * remaining_capacity

            if available == 0:
                continue

            if required <= available:
                custom_resources[node_id] = required / node_procs
                required = 0
                break
            else:
                custom_resources[node_id] = remaining_capacity
                required -= available

        if required > 0:
            # Not enough nodes
            return None

        return custom_resources

    def start_trial(self, trial, checkpoint=None):
        # Reserve node before starting trial
        resources = trial.resources

        # Use no-op remote with same resource requirements as trial to trick
        # ray into scale up the cluster before we allocate the node resources
        # for the trial
        wait_for_resources = ray.remote(
            num_cpus=resources.cpu_total(),
            num_gpus=resources.gpu_total())(lambda: None)

        error_msg = None
        for _ in range(TRIAL_START_ATTEMPTS):
            try:
                # Wait for ray to scale up the cluster
                wait_status = wait_for_resources.remote()
                ray.get(wait_status, NODE_CREATION_TIMEOUT)

                custom_resources = self._get_node_resources(resources)
                if custom_resources is not None:
                    break
                logger.warning("Trial %s: Not enough nodes, retrying...", trial)
            except RayTimeoutError:
                logger.warning("Trial %s: Timed out, retrying...", trial)
            except Exception as ex:
                logger.exception("Trial %s: Error starting trial, aborting!", trial)
                error_msg = traceback.format_exception(ex)
                break
        else:
            logger.exception(
                "Trial %s: Aborting trial after %s start "
                "attempts!", trial, TRIAL_START_ATTEMPTS)
            self._stop_trial(trial, error=True, error_msg=error_msg)
            return

        # Update trial node affinity configuration and
        # extra_custom_resources requirement
        trial.config.update(__ray_node_affinity__=custom_resources)
        trial.resources.extra_custom_resources.update(custom_resources)
        super().start_trial(trial, checkpoint)

    def has_resources(self, resources):
        if not super().has_resources(resources):
            return False
        elif self._trial_queued:
            # Allowing trial to start even though the cluster does not have
            # enough free resources
            return True

        node_resources = self._get_node_resources(resources)
        if node_resources is not None:
            return True

        if self._queue_trials:
            self._trial_queued = True
            logger.warning(
                "Allowing trial to start even though the "
                "cluster does not have enough free resources. Trial actors "
                "may appear to hang until enough resources are added to the "
                "cluster (e.g., via autoscaling). You can disable this "
                "behavior by specifying `queue_trials=False` in "
                "ray.tune.run().")
            return True
        return False

class ImagenetTrainable(Trainable):
    """
    Trainable class used to train resnet50 on Imagenet dataset using ray
    """

    @classmethod
    def default_resource_request(cls, config):
        """
        Configure the cluster resources used by this experiment
        """
        num_gpus = config.get("num_gpus", 0)
        num_cpus = config.get("num_cpus", 1)

        if num_gpus > 0:
            # Assign extra CPUs for dataloaders
            workers = config.get("workers", 0)
            num_cpus = workers * num_gpus

        resource = Resources(cpu=0, gpu=0, extra_cpu=num_cpus, extra_gpu=num_gpus)
        return resource

    def _setup(self, config):
        num_gpus = config.get("num_gpus", 0)
        num_cpus = config.get("num_cpus", 1)

        # Determine the number of distributed processes based on the number
        # GPUs and CPUs
        if num_gpus > 0:
            resource_attr = "gpu"
            world_size = num_gpus
            # Assign one GPU per remote process
            num_gpus = 1
            # Assign extra CPUs for dataloaders
            num_cpus = config.get("workers", 0)
        else:
            resource_attr = "cpu"
            world_size = num_cpus
            # Assign one CPU per remote process
            num_cpus = 1

        # Check for node affinity
        resources_by_node = _get_resources_per_node()
        ray_node_affinity = config.get("__ray_node_affinity__", {})
        node_resource = 0.0
        committed_resource = 0.0
        resources = {}

        # Create one ray remote process for each experiment in the process group
        self.procs = []
        for _ in range(world_size):
            if ray_node_affinity and node_resource <= 0:
                node_id, node_resource = ray_node_affinity.popitem()
                proc_per_node = getattr(resources_by_node[node_id], resource_attr, 0)
                committed_resource = 1.0 / proc_per_node
                resources = {node_id: committed_resource}

            experiment = ray.remote(num_cpus=num_cpus, num_gpus=num_gpus,
                                    resources=resources)(ImagenetExperiment)
            self.procs.append(experiment.remote())
            node_resource -= committed_resource

        # Use first process as head of the group
        ip = ray.get(self.procs[0].get_node_ip.remote())
        port = config.get("dist_port", 54321)
        dist_url = "tcp://{}:{}".format(ip, port)

        # Configure each process in the group
        status = []
        for i, w in enumerate(self.procs):
            worker_config = copy.deepcopy(config)
            worker_config["distributed"] = True
            worker_config["dist_url"] = dist_url
            worker_config["world_size"] = world_size
            worker_config["rank"] = i
            status.append(w.setup_experiment.remote(worker_config))

        if config.get("checkpoint_at_init", False):
            self.save()

        # Wait for remote functions to complete
        ray.get(status)

    def _train(self):
        status = []
        for w in self.procs:
            status.append(w.run_epoch.remote(self.iteration))

        # Wait for remote functions to complete
        results = ray.get(status)

        # Return the results from the first remote function
        return copy.deepcopy(results[0])

    def _save(self, _):
        # All models are synchronized. Just save the state of first model
        return ray.get(self.procs[0].get_state.remote())

    def _restore(self, state):
        # Restore the state to every process
        state_id = ray.put(state)
        ray.get([w.set_state.remote(state_id) for w in self.procs])

    def _stop(self):
        for w in self.procs:
            w.stop_experiment.remote()
            w.__ray_terminate__.remote()

def run(config):
    # Connect to ray
    address = os.environ.get("REDIS_ADDRESS", config.get("redis_address"))
    ray.init(address=address)

    # Build kwargs for `tune.run` function using merged config and command line dict
    kwargs_names = tune.run.__code__.co_varnames[:tune.run.__code__.co_argcount]
    kwargs = dict(zip(kwargs_names, [ImagenetTrainable, *tune.run.__defaults__]))

    # Update`tune.run` kwargs with config
    kwargs.update(config)
    kwargs["config"] = config

    # Update tune stop criteria with config epochs
    stop = kwargs.get("stop", {}) or dict()
    epochs = config.get("epochs", 1)
    stop.update(training_iteration=epochs)
    kwargs["stop"] = stop

    # Make sure to only select`tune.run` function arguments
    kwargs = dict(filter(lambda x: x[0] in kwargs_names, kwargs.items()))

    # Queue trials until the cluster scales up
    kwargs.update(queue_trials=True)

    # Group trial into nodes as much as possible
    kwargs.update(trial_executor=AffinityExecutor(
        queue_trials=kwargs.get("queue_trials", True),
        reuse_actors=kwargs.get("reuse_actors", False),
        ray_auto_init=kwargs.get("ray_auto_init", True)
    ))

    pprint(kwargs)
    tune.run(**kwargs)
    ray.shutdown()
<commit_msg><commit_after>#  Numenta Platform for Intelligent Computing (NuPIC)
#  Copyright (C) 2019, Numenta, Inc.  Unless you have an agreement
#  with Numenta, Inc., for a separate license for this software code, the
#  following terms and conditions apply:
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Affero Public License version 3 as
#  published by the Free Software Foundation.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#  See the GNU Affero Public License for more details.
#
#  You should have received a copy of the GNU Affero Public License
#  along with this program.  If not, see http://www.gnu.org/licenses.
#
#  http://numenta.org/licenses/
#
import copy
import logging
import os
import traceback
from pprint import pprint

import ray
import ray.resource_spec
from ray import ray_constants
from ray.exceptions import RayTimeoutError
from ray.tune import Trainable, tune
from ray.tune.ray_trial_executor import RESOURCE_REFRESH_PERIOD, RayTrialExecutor
from ray.tune.resources import Resources

from nupic.research.frameworks.pytorch.imagenet import ImagenetExperiment

os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
logger = logging.getLogger(__name__)
TRIAL_START_ATTEMPTS = 3
NODE_CREATION_TIMEOUT = 10 * 60

def _get_resources_per_node():
    """
    Maps node id to available resources on that node
    :return: dict with available :class:`Resources` for each node
    """
    def _is_node_key(k):
        return k.startswith(ray.resource_spec.NODE_ID_PREFIX)

    # Only consider active/alive nodes
    nodes = filter(lambda node: node["Alive"], ray.nodes())
    resources = map(lambda node: node["Resources"], nodes)

    # Group resources by node
    resources_by_node = {}
    for item in resources:
        node_id = next(filter(_is_node_key, item.keys()))

        item = item.copy()
        num_cpus = item.pop("CPU", 0)
        num_gpus = item.pop("GPU", 0)
        memory = ray_constants.from_memory_units(item.pop("memory", 0))
        object_store_memory = ray_constants.from_memory_units(
            item.pop("object_store_memory", 0))
        custom_resources = item

        resources_by_node[node_id] = Resources(
            int(num_cpus),
            int(num_gpus),
            memory=int(memory),
            object_store_memory=int(object_store_memory),
            custom_resources=custom_resources)

    return resources_by_node

class AffinityExecutor(RayTrialExecutor):
    """
    Ray Trial Executor used to set node affinity to node trials.
    This trial executor will update the trial configuration with one extra
    parameter (__ray_node_affinity__) that can be used to launch extra workers
    on the same node.
    """

    def __init__(self, queue_trials=False, reuse_actors=False, ray_auto_init=False,
                 refresh_period=RESOURCE_REFRESH_PERIOD):
        self._resources_by_node = {}
        super().__init__(queue_trials, reuse_actors, ray_auto_init, refresh_period)

    def _commit_resources(self, resources):
        resources = copy.deepcopy(resources)
        # Make sure "custom_resources" keys match "extra_custom_resources" keys
        # because `RayTrialExecutor._commit_resources` will only commit
        # resources from "extra_custom_resources" if the same resource key
        # is present in the "custom_resources" dict
        for k in resources.extra_custom_resources:
            if k not in resources.custom_resources:
                resources.custom_resources[k] = 0.0

        super()._commit_resources(resources)

    def _update_avail_resources(self, num_retries=5):
        super()._update_avail_resources(num_retries)
        self._resources_by_node = _get_resources_per_node()

    def _get_node_resources(self, resources):
        self._update_avail_resources()

        # Check for required GPUs first
        required = resources.gpu_total()
        if required > 0:
            resource_attr = "gpu"
        else:
            # No GPU, just use CPU
            resource_attr = "cpu"
            required = resources.cpu_total()

        # Compute nodes required to fulfill trial resources request
        custom_resources = {}
        for node_id, node_resource in self._resources_by_node.items():
            # Compute resource remaining on each node
            node_capacity = node_resource.get_res_total(node_id)
            committed_capacity = self._committed_resources.get_res_total(node_id)
            remaining_capacity = node_capacity - committed_capacity
            node_procs = getattr(node_resource, resource_attr, 0)
            available = node_procs * remaining_capacity

            if available == 0:
                continue

            if required <= available:
                custom_resources[node_id] = required / node_procs
                required = 0
                break
            else:
                custom_resources[node_id] = remaining_capacity
                required -= available

        if required > 0:
            # Not enough nodes
            return None

        return custom_resources

    def start_trial(self, trial, checkpoint=None):
        # Reserve node before starting trial
        resources = trial.resources

        # Use no-op remote with same resource requirements as trial to trick
        # ray into scale up the cluster before we allocate the node resources
        # for the trial
        wait_for_resources = ray.remote(
            num_cpus=resources.cpu_total(),
            num_gpus=resources.gpu_total())(lambda: None)

        error_msg = None
        for _ in range(TRIAL_START_ATTEMPTS):
            try:
                # Wait for ray to scale up the cluster
                wait_status = wait_for_resources.remote()
                ray.get(wait_status, NODE_CREATION_TIMEOUT)

                custom_resources = self._get_node_resources(resources)
                if custom_resources is not None:
                    break
                logger.warning("Trial %s: Not enough nodes, retrying...", trial)
            except RayTimeoutError:
                logger.warning("Trial %s: Timed out, retrying...", trial)
            except Exception as ex:
                logger.exception("Trial %s: Error starting trial, aborting!", trial)
                error_msg = traceback.format_exception(ex)
                break
        else:
            logger.exception(
                "Trial %s: Aborting trial after %s start "
                "attempts!", trial, TRIAL_START_ATTEMPTS)
            self._stop_trial(trial, error=True, error_msg=error_msg)
            return

        # Update trial node affinity configuration and
        # extra_custom_resources requirement
        trial.config.update(__ray_node_affinity__=custom_resources)
        trial.resources.extra_custom_resources.update(custom_resources)
        super().start_trial(trial, checkpoint)

    def has_resources(self, resources):
        if not super().has_resources(resources):
            return False
        elif self._trial_queued:
            # Allowing trial to start even though the cluster does not have
            # enough free resources
            return True

        node_resources = self._get_node_resources(resources)
        if node_resources is not None:
            return True

        if self._queue_trials:
            self._trial_queued = True
            logger.warning(
                "Allowing trial to start even though the "
                "cluster does not have enough free resources. Trial actors "
                "may appear to hang until enough resources are added to the "
                "cluster (e.g., via autoscaling). You can disable this "
                "behavior by specifying `queue_trials=False` in "
                "ray.tune.run().")
            return True
        return False

class ImagenetTrainable(Trainable):
    """
    Trainable class used to train resnet50 on Imagenet dataset using ray
    """

    @classmethod
    def default_resource_request(cls, config):
        """
        Configure the cluster resources used by this experiment
        """
        num_gpus = config.get("num_gpus", 0)
        num_cpus = config.get("num_cpus", 1)

        if num_gpus > 0:
            # Assign extra CPUs for dataloaders
            workers = config.get("workers", 0)
            num_cpus = workers * num_gpus

        resource = Resources(cpu=0, gpu=0, extra_cpu=num_cpus, extra_gpu=num_gpus)
        return resource

    def _setup(self, config):
        num_gpus = config.get("num_gpus", 0)
        num_cpus = config.get("num_cpus", 1)

        # Determine the number of distributed processes based on the number
        # GPUs and CPUs
        if num_gpus > 0:
            resource_attr = "gpu"
            world_size = num_gpus
            # Assign one GPU per remote process
            num_gpus = 1
            # Assign extra CPUs for dataloaders
            num_cpus = config.get("workers", 0)
        else:
            resource_attr = "cpu"
            world_size = num_cpus
            # Assign one CPU per remote process
            num_cpus = 1

        # Check for node affinity
        resources_by_node = _get_resources_per_node()
        ray_node_affinity = config.get("__ray_node_affinity__", {})
        node_resource = 0.0
        committed_resource = 0.0
        resources = {}

        # Create one ray remote process for each experiment in the process group
        self.procs = []
        for _ in range(world_size):
            if ray_node_affinity and node_resource <= 0:
                node_id, node_resource = ray_node_affinity.popitem()
                proc_per_node = getattr(resources_by_node[node_id], resource_attr, 0)
                committed_resource = 1.0 / proc_per_node
                resources = {node_id: committed_resource}

            experiment = ray.remote(num_cpus=num_cpus, num_gpus=num_gpus,
                                    resources=resources)(ImagenetExperiment)
            self.procs.append(experiment.remote())
            node_resource -= committed_resource

        # Use first process as head of the group
        ip = ray.get(self.procs[0].get_node_ip.remote())
        port = config.get("dist_port", 54321)
        dist_url = "tcp://{}:{}".format(ip, port)

        # Configure each process in the group
        status = []
        for i, w in enumerate(self.procs):
            worker_config = copy.deepcopy(config)
            worker_config["distributed"] = True
            worker_config["dist_url"] = dist_url
            worker_config["world_size"] = world_size
            worker_config["rank"] = i
            status.append(w.setup_experiment.remote(worker_config))

        # Wait for remote functions to complete
        ray.get(status)

        # Save initialized model.
        if config.get("checkpoint_at_init", False):
            self.save()

    def _train(self):
        status = []
        for w in self.procs:
            status.append(w.run_epoch.remote(self.iteration))

        # Wait for remote functions to complete
        results = ray.get(status)

        # Return the results from the first remote function
        return copy.deepcopy(results[0])

    def _save(self, _):
        # All models are synchronized. Just save the state of first model
        return ray.get(self.procs[0].get_state.remote())

    def _restore(self, state):
        # Restore the state to every process
        state_id = ray.put(state)
        ray.get([w.set_state.remote(state_id) for w in self.procs])

    def _stop(self):
        for w in self.procs:
            w.stop_experiment.remote()
            w.__ray_terminate__.remote()

def run(config):
    # Connect to ray
    address = os.environ.get("REDIS_ADDRESS", config.get("redis_address"))
    ray.init(address=address)

    # Build kwargs for `tune.run` function using merged config and command line dict
    kwargs_names = tune.run.__code__.co_varnames[:tune.run.__code__.co_argcount]
    kwargs = dict(zip(kwargs_names, [ImagenetTrainable, *tune.run.__defaults__]))

    # Update`tune.run` kwargs with config
    kwargs.update(config)
    kwargs["config"] = config

    # Update tune stop criteria with config epochs
    stop = kwargs.get("stop", {}) or dict()
    epochs = config.get("epochs", 1)
    stop.update(training_iteration=epochs)
    kwargs["stop"] = stop

    # Make sure to only select`tune.run` function arguments
    kwargs = dict(filter(lambda x: x[0] in kwargs_names, kwargs.items()))

    # Queue trials until the cluster scales up
    kwargs.update(queue_trials=True)

    # Group trial into nodes as much as possible
    kwargs.update(trial_executor=AffinityExecutor(
        queue_trials=kwargs.get("queue_trials", True),
        reuse_actors=kwargs.get("reuse_actors", False),
        ray_auto_init=kwargs.get("ray_auto_init", True)
    ))

    pprint(kwargs)
    tune.run(**kwargs)
    ray.shutdown()
<|endoftext|>