def testQuantum(self):
     registry = Registry.fromConfig(self.configFile)
     run = registry.makeRun(collection="test")
     storageClass = makeNewStorageClass("testQuantum")()
     # Make two predicted inputs
     datasetType1 = DatasetType(name="dst1",
                                dataUnits=("camera", ),
                                storageClass=storageClass)
     registry.registerDatasetType(datasetType1)
     ref1 = registry.addDataset(datasetType1,
                                dataId={"camera": "DummyCam"},
                                run=run)
     datasetType2 = DatasetType(name="dst2",
                                dataUnits=("camera", ),
                                storageClass=storageClass)
     registry.registerDatasetType(datasetType2)
     ref2 = registry.addDataset(datasetType2,
                                dataId={"camera": "DummyCam"},
                                run=run)
     # Create and add a Quantum
     quantum = Quantum(run=run,
                       task="some.fully.qualified.SuperTask",
                       startTime=datetime(2018, 1, 1),
                       endTime=datetime(2018, 1, 2),
                       host="localhost")
     quantum.addPredictedInput(ref1)
     quantum.addPredictedInput(ref2)
     # Quantum is not yet in Registry, so can't mark input as actual
     with self.assertRaises(KeyError):
         registry.markInputUsed(quantum, ref1)
     registry.addQuantum(quantum)
     # Now we can
     registry.markInputUsed(quantum, ref1)
     outQuantum = registry.getQuantum(quantum.id)
     self.assertEqual(outQuantum, quantum)
 def testConstructor(self):
     """Test of constructor.
     """
     # Quantum specific arguments
     run = None  # TODO add Run
     task = "some.task.object"  # TODO Add a `SuperTask`?
     # Base class arguments
     startTime = datetime(2018, 1, 1)
     endTime = datetime(2018, 1, 2)
     host = "localhost"
     quantum = Quantum(task, run, startTime, endTime, host)
     self.assertEqual(quantum.task, task)
     self.assertEqual(quantum.run, run)
     self.assertEqual(quantum.predictedInputs, dict())
     self.assertEqual(quantum.actualInputs, dict())
     self.assertIsNone(quantum.id)
     self.assertEqual(quantum.startTime, startTime)
     self.assertEqual(quantum.endTime, endTime)
     self.assertEqual(quantum.host, host)
    def _makeQuanta(self, config):
        """Create set of Quanta
        """
        run = Run(collection=1, environment=None, pipeline=None)

        dstype0 = supertask.SuperTask.makeDatasetType(config.input)
        dstype1 = supertask.SuperTask.makeDatasetType(config.output)

        quanta = []
        for visit in range(100):
            quantum = Quantum(run=run, task=None)
            quantum.addPredictedInput(self._makeDSRefVisit(dstype0, visit))
            quantum.addOutput(self._makeDSRefVisit(dstype1, visit))
            quanta.append(quantum)

        return quanta
Esempio n. 4
0
    def testAddInputsOutputs(self):
        """Test of addPredictedInput() method.
        """
        quantum = Quantum(task="some.task.object", run=None)

        # start with empty
        self.assertEqual(quantum.predictedInputs, dict())
        universe = DimensionUniverse.fromConfig()
        instrument = "DummyCam"
        datasetTypeName = "test_ds"
        storageClass = StorageClass("testref_StructuredData")
        datasetType = DatasetType(datasetTypeName,
                                  universe.extract(("instrument", "visit")),
                                  storageClass)

        # add one ref
        ref = DatasetRef(datasetType, dict(instrument=instrument, visit=42))
        quantum.addPredictedInput(ref)
        self.assertIn(datasetTypeName, quantum.predictedInputs)
        self.assertEqual(len(quantum.predictedInputs[datasetTypeName]), 1)
        # add second ref
        ref = DatasetRef(datasetType, dict(instrument=instrument, visit=43))
        quantum.addPredictedInput(ref)
        self.assertEqual(len(quantum.predictedInputs[datasetTypeName]), 2)

        # mark last ref as actually used
        self.assertEqual(quantum.actualInputs, dict())
        quantum._markInputUsed(ref)
        self.assertIn(datasetTypeName, quantum.actualInputs)
        self.assertEqual(len(quantum.actualInputs[datasetTypeName]), 1)

        # add couple of outputs too
        self.assertEqual(quantum.outputs, dict())
        ref = DatasetRef(datasetType, dict(instrument=instrument, visit=42))
        quantum.addOutput(ref)
        self.assertIn(datasetTypeName, quantum.outputs)
        self.assertEqual(len(quantum.outputs[datasetTypeName]), 1)

        ref = DatasetRef(datasetType, dict(instrument=instrument, visit=43))
        quantum.addOutput(ref)
        self.assertEqual(len(quantum.outputs[datasetTypeName]), 2)
def main():

    descr = 'One-line application description.'
    parser = ArgumentParser(description=descr)
    parser.add_argument('-v',
                        '--verbose',
                        dest='verbose',
                        action='count',
                        default=0,
                        help='More verbose output, can use several times.')
    parser.add_argument(
        '-p',
        '--pipeline',
        dest='pipeline',
        default=None,
        type=FileType("wb"),
        help='Name of the output file for Pipeline pickle dump.')
    parser.add_argument(
        '-g',
        '--qgraph',
        dest='qgraph',
        default=None,
        type=FileType("wb"),
        help='Name of the output file for QunatumGraph pickle dump.')
    args = parser.parse_args()

    # configure logging
    _configLogger(args.verbose)

    #
    #  Application logic goes here
    #
    if args.pipeline is None and args.qgraph is None:
        parser.error("Need one of -p or -g options")

    if args.pipeline:

        pipeline = Pipeline(
            [_makeStep1TaskDef(),
             _makeStep2TaskDef(),
             _makeStep3TaskDef()])

        # make sure it is ordered, this is joust to test some methods in supertasks
        pipeline = orderPipeline(pipeline)

        # save to file
        pickle.dump(pipeline, args.pipeline)

    if args.qgraph:

        run = Run(collection=1, environment=None, pipeline=None)

        step1 = _makeStep1TaskDef()
        step2 = _makeStep2TaskDef()
        step3 = _makeStep3TaskDef()

        dstype0 = SuperTask.makeDatasetType(step1.config.input)
        dstype1 = SuperTask.makeDatasetType(step1.config.output)
        dstype2 = SuperTask.makeDatasetType(step2.config.output)
        dstype3 = SuperTask.makeDatasetType(step3.config.output)

        # quanta for first step which is 1-to-1 tasks
        quanta = []
        for visit in range(10):
            quantum = Quantum(run=run, task=None)
            quantum.addPredictedInput(_makeDSRefVisit(dstype0, visit))
            quantum.addOutput(_makeDSRefVisit(dstype1, visit))
            quanta.append(quantum)
        step1nodes = QuantumGraphNodes(step1, quanta)

        # quanta for second step which is 1-to-1 tasks
        quanta = []
        for visit in range(10):
            quantum = Quantum(run=run, task=None)
            quantum.addPredictedInput(_makeDSRefVisit(dstype1, visit))
            quantum.addOutput(_makeDSRefVisit(dstype2, visit))
            quanta.append(quantum)
        step2nodes = QuantumGraphNodes(step2, quanta)

        # quanta for third step which is M-to-N
        patch2visits = (
            (1, 1, (0, 1, 2, 3)),
            (1, 2, (2, 3, 4, 5)),
            (1, 3, (4, 5, 6, 7)),
            (2, 0, (6, 7, 8, 9)),
        )
        quanta = []
        for tract, patch, visits in patch2visits:
            quantum = Quantum(run=run, task=None)
            for visit in visits:
                quantum.addPredictedInput(_makeDSRefVisit(dstype2, visit))
            quantum.addOutput(_makeDSRefPatch(dstype3, tract, patch))
            quanta.append(quantum)
        step3nodes = QuantumGraphNodes(step2, quanta)

        qgraph = QuantumGraph([step1nodes, step2nodes, step3nodes])
        pickle.dump(qgraph, args.qgraph)