Python firework.Workflow类代码示例

 def test_remove_root_fws(self):
     fw4 = Firework(Task1(), parents=[self.fw2, self.fw3])
     fws = [self.fw1, self.fw2, self.fw3, fw4]
     wflow = Workflow(fws)
     root_ids = wflow.root_fw_ids
     children = []
     for i in root_ids:
         children.extend(wflow.links[i])
     wflow.remove_fws(wflow.root_fw_ids)
     self.assertEqual(sorted(wflow.root_fw_ids), sorted(children))

 def test_remove_leaf_fws(self):
     fw4 = Firework(Task1(), parents=[self.fw2, self.fw3])
     fws = [self.fw1, self.fw2, self.fw3, fw4]
     wflow = Workflow(fws)
     leaf_ids = wflow.leaf_fw_ids
     parents = []
     for i in leaf_ids:
         parents.extend(wflow.links.parent_links[i])
     wflow.remove_fws(wflow.leaf_fw_ids)
     self.assertEqual(wflow.leaf_fw_ids, parents)

    def add_wf(self, wf, reassign_all=True):
        """

        :param wf: a Workflow object.
        """
        m_timer.start("add_wf")
        if isinstance(wf, FireWork):
            wf = Workflow.from_FireWork(wf)

        # sets the root FWs as READY
        # prefer to wf.refresh() for speed reasons w/many root FWs
        for fw_id in wf.root_fw_ids:
            wf.id_fw[fw_id].state = 'READY'

        # insert the FireWorks and get back mapping of old to new ids
        old_new = self._upsert_fws(list(wf.id_fw.values()), reassign_all=reassign_all)

        # update the Workflow with the new ids
        wf._reassign_ids(old_new)

        # insert the WFLinks
        self.workflows.insert(wf.to_db_dict())

        m_timer.stop("add_wf")
        self.m_logger.info('Added a workflow. id_map: {}'.format(old_new))
        return old_new

    def test_copy(self):
        """Test that we can produce a copy of a Workflow but that the copy
        has unique fw_ids.

        """
        fws = []
        for i in range(5):
            fw = Firework([PyTask(func="print", args=[i])], fw_id=i,
                          name=i)
            fws.append(fw)

        wf = Workflow(fws, links_dict={0: [1, 2, 3], 1: [4], 2: [4]})

        wf_copy = Workflow.from_wflow(wf)
        
        # now we compare to the original to make sure dependencies are same.
        # have to do gymnastics because ids will NOT be the same
        # but names are retained
        for fw in wf_copy.fws:
            children = wf_copy.links.get(fw.fw_id, list())
            orig_id = fw.name

            orig_children = wf.links.get(orig_id, list())

            for child_id, orig_child_id in zip(children, orig_children):
                self.assertEqual(orig_child_id, wf_copy.id_fw[child_id].name)

def add_wf(args):
    lp = get_lp(args)
    if args.dir:
        files = []
        for f in args.wf_file:
            files.extend([os.path.join(f, i) for i in os.listdir(f)])
    else:
        files = args.wf_file
    for f in files:
        fwf = Workflow.from_file(f)
        lp.add_wf(fwf)

def task_dict_to_wf(task_dict, launchpad):
    fw_id = launchpad.get_new_fw_id()
    l_id = launchpad.get_new_launch_id()

    spec = {'task_type': task_dict['task_type'], 'run_tags': task_dict['run_tags'],
            'vaspinputset_name': None, 'vasp': None, 'mpsnl': task_dict['snl'],
            'snlgroup_id': task_dict['snlgroup_id']}
    tasks = [DummyLegacyTask()]

    launch_dir = task_dict['dir_name_full']

    stored_data = {'error_list': []}
    update_spec = {'prev_vasp_dir': task_dict['dir_name'],
                   'prev_task_type': spec['task_type'],
                   'mpsnl': spec['mpsnl'], 'snlgroup_id': spec['snlgroup_id'],
                   'run_tags': spec['run_tags']}

    fwaction = FWAction(stored_data=stored_data, update_spec=update_spec)

    if task_dict['completed_at']:
        complete_date = datetime.datetime.strptime(task_dict['completed_at'], "%Y-%m-%d %H:%M:%S")
        state_history = [{"created_on": complete_date, 'state': 'COMPLETED'}]
    else:
        state_history = []

    launches = [Launch('COMPLETED', launch_dir, fworker=None, host=None, ip=None, action=fwaction,
                       state_history=state_history, launch_id=l_id, fw_id=fw_id)]

    f = Composition(task_dict['pretty_formula']).alphabetical_formula


    fw = Firework(tasks, spec, name=get_slug(f + '--' + spec['task_type']), launches=launches, state='COMPLETED', created_on=None,
                 fw_id=fw_id)

    wf_meta = get_meta_from_structure(Structure.from_dict(task_dict['snl']))
    wf_meta['run_version'] = 'preproduction (0)'

    wf = Workflow.from_FireWork(fw, name=f, metadata=wf_meta)

    launchpad.add_wf(wf, reassign_all=False)
    launchpad._upsert_launch(launches[0])

    print 'ADDED', fw_id
    # return fw_id
    return fw_id

def add_wf_dir(args):
    lp = get_lp(args)
    for filename in os.listdir(args.wf_dir):
        fwf = Workflow.from_file(filename)
        lp.add_wf(fwf)

def task_dict_to_wf(task_dict, launchpad):
    fw_id = launchpad.get_new_fw_id()
    l_id = launchpad.get_new_launch_id()

    spec = {
        "task_type": task_dict["task_type"],
        "run_tags": task_dict["run_tags"],
        "vaspinputset_name": None,
        "vasp": None,
        "mpsnl": task_dict["snl"],
        "snlgroup_id": task_dict["snlgroup_id"],
    }
    tasks = [DummyLegacyTask()]

    launch_dir = task_dict["dir_name_full"]

    stored_data = {"error_list": []}
    update_spec = {
        "prev_vasp_dir": task_dict["dir_name"],
        "prev_task_type": spec["task_type"],
        "mpsnl": spec["mpsnl"],
        "snlgroup_id": spec["snlgroup_id"],
        "run_tags": spec["run_tags"],
    }

    fwaction = FWAction(stored_data=stored_data, update_spec=update_spec)

    if task_dict["completed_at"]:
        complete_date = datetime.datetime.strptime(task_dict["completed_at"], "%Y-%m-%d %H:%M:%S")
        state_history = [{"created_on": complete_date, "state": "COMPLETED"}]
    else:
        state_history = []

    launches = [
        Launch(
            "COMPLETED",
            launch_dir,
            fworker=None,
            host=None,
            ip=None,
            action=fwaction,
            state_history=state_history,
            launch_id=l_id,
            fw_id=fw_id,
        )
    ]

    f = Composition(task_dict["pretty_formula"]).alphabetical_formula

    fw = Firework(
        tasks,
        spec,
        name=get_slug(f + "--" + spec["task_type"]),
        launches=launches,
        state="COMPLETED",
        created_on=None,
        fw_id=fw_id,
    )

    wf_meta = get_meta_from_structure(Structure.from_dict(task_dict["snl"]))
    wf_meta["run_version"] = "preproduction (0)"

    wf = Workflow.from_FireWork(fw, name=f, metadata=wf_meta)

    launchpad.add_wf(wf, reassign_all=False)
    launchpad._upsert_launch(launches[0])

    print "ADDED", fw_id
    # return fw_id
    return fw_id