Python launchpad.LaunchPad类代码示例

def born():

    fw = Firework(
        [
            Cry(),
            Eat(),
            Sleep(),
            Poop(),
        ],
        spec={"name": "BabyJain"}
    )


    launchpad = LaunchPad()
    launchpad.add_wf(Workflow([fw]))

    def run_task(self, fw_spec):
        # the FW.json/yaml file is mandatory to get the fw_id
        # no need to deserialize the whole FW

        if '_add_launchpad_and_fw_id' in fw_spec:
            lp = self.launchpad
            fw_id = self.fw_id
        else:
            try:
                fw_dict = loadfn('FW.json')
            except IOError:
                try:
                    fw_dict = loadfn('FW.yaml')
                except IOError:
                    raise RuntimeError("Launchpad/fw_id not present in spec and No FW.json nor FW.yaml file present: "
                                       "impossible to determine fw_id")
            lp = LaunchPad.auto_load()
            fw_id = fw_dict['fw_id']

        wf = lp.get_wf_by_fw_id_lzyfw(fw_id)

        deleted_files = []
        # iterate over all the fws and launches
        for fw_id, fw in wf.id_fw.items():
            for l in fw.launches+fw.archived_launches:
                l_dir = l.launch_dir

                deleted_files.extend(self.delete_files(os.path.join(l_dir, TMPDIR_NAME)))
                deleted_files.extend(self.delete_files(os.path.join(l_dir, INDIR_NAME)))
                deleted_files.extend(self.delete_files(os.path.join(l_dir, OUTDIR_NAME), self.out_exts))

        logging.info("Deleted files:\n {}".format("\n".join(deleted_files)))

        return FWAction(stored_data={'deleted_files': deleted_files})

def do_launch(args):
    if not args.launchpad_file and os.path.exists(
            os.path.join(args.config_dir, 'my_launchpad.yaml')):
        args.launchpad_file = os.path.join(args.config_dir, 'my_launchpad.yaml')
    elif not args.launchpad_file:
        args.launchpad_file = LAUNCHPAD_LOC

    if not args.fworker_file and os.path.exists(
            os.path.join(args.config_dir, 'my_fworker.yaml')):
        args.fworker_file = os.path.join(args.config_dir, 'my_fworker.yaml')
    elif not args.fworker_file:
        args.fworker_file = FWORKER_LOC

    if not args.queueadapter_file and os.path.exists(
            os.path.join(args.config_dir, 'my_qadapter.yaml')):
        args.queueadapter_file = os.path.join(args.config_dir, 'my_qadapter.yaml')
    elif not args.queueadapter_file:
        args.queueadapter_file = QUEUEADAPTER_LOC

    launchpad = LaunchPad.from_file(
        args.launchpad_file) if args.launchpad_file else LaunchPad(
        strm_lvl=args.loglvl)
    fworker = FWorker.from_file(
        args.fworker_file) if args.fworker_file else FWorker()
    queueadapter = load_object_from_file(args.queueadapter_file)
    args.loglvl = 'CRITICAL' if args.silencer else args.loglvl

    if args.command == 'rapidfire':
        rapidfire(launchpad, fworker=fworker, qadapter=queueadapter, launch_dir=args.launch_dir,
                  nlaunches=args.nlaunches, njobs_queue=args.maxjobs_queue,
                  njobs_block=args.maxjobs_block, sleep_time=args.sleep,
                  reserve=args.reserve, strm_lvl=args.loglvl, timeout=args.timeout, fill_mode=args.fill_mode)
    else:
        launch_rocket_to_queue(launchpad, fworker, queueadapter,
                               args.launch_dir, args.reserve, args.loglvl, False, args.fill_mode)

    def run_task(self, fw_spec):
        # the FW.json/yaml file is mandatory to get the fw_id
        # no need to deserialize the whole FW
        try:
            fw_dict = loadfn('FW.json')
        except IOError:
            try:
                fw_dict = loadfn('FW.yaml')
            except IOError:
                raise RuntimeError("No FW.json nor FW.yaml file present: impossible to determine fw_id")

        fw_id = fw_dict['fw_id']
        lp = LaunchPad.auto_load()
        wf = lp.get_wf_by_fw_id_lzyfw(fw_id)
        wf_module = importlib.import_module(wf.metadata['workflow_module'])
        wf_class = getattr(wf_module, wf.metadata['workflow_class'])

        get_results_method = getattr(wf_class, 'get_final_structure_and_history')
        #TODO: make this more general ... just to test right now ...
        results = get_results_method(wf)

        database = MongoDatabase.from_dict(fw_spec['mongo_database'])

        database.insert_entry({'structure': results['structure'], 'history': results['history']})

        logging.info("Inserted data:\n something")


        return FWAction()

def clear_env():
    sma = SubmissionMongoAdapter.auto_load()

    lp = LaunchPad.auto_load()

    snl = SNLMongoAdapter.auto_load()

    db_dir = os.environ['DB_LOC']
    db_path = os.path.join(db_dir, 'tasks_db.json')
    with open(db_path) as f:
        db_creds = json.load(f)

    sma._reset()
    lp.reset('', require_password=False)
    snl._reset()

    conn = MongoClient(db_creds['host'], db_creds['port'])
    db = conn[db_creds['database']]
    db.authenticate(db_creds['admin_user'], db_creds['admin_password'])
    db.tasks.remove()
    db.boltztrap.remove()
    db.counter.remove()
    db['dos_fs.chunks'].remove()
    db['dos_fs.files'].remove()
    db['band_structure_fs.files'].remove()
    db['band_structure_fs.files'].remove()

def do_launch(args):
    if not args.launchpad_file and os.path.exists(
            os.path.join(args.config_dir, 'my_launchpad.yaml')):
        args.launchpad_file = os.path.join(args.config_dir, 'my_launchpad.yaml')

    if not args.fworker_file and os.path.exists(
            os.path.join(args.config_dir, 'my_fworker.yaml')):
        args.fworker_file = os.path.join(args.config_dir, 'my_fworker.yaml')

    if not args.queueadapter_file and os.path.exists(
            os.path.join(args.config_dir, 'my_qadapter.yaml')):
        args.queueadapter_file = os.path.join(args.config_dir,
                                              'my_qadapter.yaml')

    launchpad = LaunchPad.from_file(
        args.launchpad_file) if args.launchpad_file else LaunchPad(
        strm_lvl=args.loglvl)
    fworker = FWorker.from_file(
        args.fworker_file) if args.fworker_file else FWorker()
    queueadapter = load_object_from_file(args.queueadapter_file)
    args.loglvl = 'CRITICAL' if args.silencer else args.loglvl

    if args.command == 'rapidfire':
        rapidfire(launchpad, fworker, queueadapter, args.launch_dir,
                  args.nlaunches, args.maxjobs_queue,
                  args.maxjobs_block, args.sleep, args.reserve, args.loglvl)
    else:
        launch_rocket_to_queue(launchpad, fworker, queueadapter,
                               args.launch_dir, args.reserve, args.loglvl, False)

def rlaunch():

    m_description = 'This program launches one or more Rockets. A Rocket grabs a job from the central database and ' \
                    'runs it. The "single-shot" option launches a single Rocket, ' \
                    'whereas the "rapidfire" option loops until all FireWorks are completed.'

    parser = ArgumentParser(description=m_description)
    subparsers = parser.add_subparsers(help='command', dest='command')
    single_parser = subparsers.add_parser('singleshot', help='launch a single Rocket')
    rapid_parser = subparsers.add_parser('rapidfire',
                                         help='launch multiple Rockets (loop until all FireWorks complete)')

    single_parser.add_argument('-f', '--fw_id', help='specific fw_id to run', default=None, type=int)
    single_parser.add_argument('--offline', help='run in offline mode (FW.json required)', action='store_true')

    rapid_parser.add_argument('--nlaunches', help='num_launches (int or "infinite"; default 0 is all jobs in DB)', default=0)
    rapid_parser.add_argument('--sleep', help='sleep time between loops (secs)', default=None, type=int)

    parser.add_argument('-l', '--launchpad_file', help='path to launchpad file', default=LAUNCHPAD_LOC)
    parser.add_argument('-w', '--fworker_file', help='path to fworker file', default=FWORKER_LOC)
    parser.add_argument('-c', '--config_dir', help='path to a directory containing the config file (used if -l, -w unspecified)',
                        default=CONFIG_FILE_DIR)

    parser.add_argument('--loglvl', help='level to print log messages', default='INFO')
    parser.add_argument('-s', '--silencer', help='shortcut to mute log messages', action='store_true')

    args = parser.parse_args()

    signal.signal(signal.SIGINT, handle_interrupt)  # graceful exist on ^C

    if not args.launchpad_file and os.path.exists(os.path.join(args.config_dir, 'my_launchpad.yaml')):
        args.launchpad_file = os.path.join(args.config_dir, 'my_launchpad.yaml')

    if not args.fworker_file and os.path.exists(os.path.join(args.config_dir, 'my_fworker.yaml')):
        args.fworker_file = os.path.join(args.config_dir, 'my_fworker.yaml')

    args.loglvl = 'CRITICAL' if args.silencer else args.loglvl

    if args.command == 'singleshot' and args.offline:
        launchpad = None
    else:
        launchpad = LaunchPad.from_file(args.launchpad_file) if args.launchpad_file else LaunchPad(strm_lvl=args.loglvl)

    if args.fworker_file:
        fworker = FWorker.from_file(args.fworker_file)
    else:
        fworker = FWorker()

    # prime addr lookups
    _log = get_fw_logger("rlaunch", stream_level="INFO")
    _log.info("Hostname/IP lookup (this will take a few seconds)")
    get_my_host()
    get_my_ip()

    if args.command == 'rapidfire':
        rapidfire(launchpad, fworker, None, args.nlaunches, -1, args.sleep, args.loglvl)

    else:
        launch_rocket(launchpad, fworker, args.fw_id, args.loglvl)

def get_lp_and_fw_id_from_task(task, fw_spec):
    """
    Given an instance of a running task and its spec, tries to load the LaunchPad and the current fw_id.
    It will first check for "_add_launchpad_and_fw_id", then try to load from FW.json/FW.yaml file.

    Should be used inside tasks that require to access to the LaunchPad and to the whole workflow.
    Args:
        task: An instance of a running task
        fw_spec: The spec of the task

    Returns:
        an instance of LaunchPah and the fw_id of the current task
    """
    if '_add_launchpad_and_fw_id' in fw_spec:
        lp = task.launchpad
        fw_id = task.fw_id

        # lp may be None in offline mode
        if lp is None:
            raise RuntimeError("The LaunchPad in spec is None.")
    else:
        try:
            with open('FW.json', "rt") as fh:
                fw_dict = json.load(fh, cls=MontyDecoder)
        except IOError:
            try:
                with open('FW.yaml', "rt") as fh:
                    fw_dict = yaml.load(fh)
            except IOError:
                raise RuntimeError("Launchpad/fw_id not present in spec and No FW.json nor FW.yaml file present: "
                                   "impossible to determine fw_id")

        logger.warning("LaunchPad not available from spec. Generated with auto_load.")
        lp = LaunchPad.auto_load()
        fw_id = fw_dict['fw_id']

        # since it is not given that the LaunchPad is the correct one, try to verify if the workflow
        # and the fw_id are being accessed correctly
        try:
            fw = lp.get_fw_by_id(fw_id)
        except ValueError as e:
            traceback.print_exc()
            raise RuntimeError("The firework with id {} is not present in the LaunchPad {}. The LaunchPad is "
                               "probably incorrect.". format(fw_id, lp))

        if fw.state != "RUNNING":
            raise RuntimeError("The firework with id {} from LaunchPad {} is {}. There might be an error in the "
                               "selection of the LaunchPad". format(fw_id, lp, fw.state))

        if len(fw.tasks) != len(fw_dict['spec']['_tasks']):
            raise RuntimeError("The firework with id {} from LaunchPad {} is has different number of tasks "
                               "from the current.".format(fw_id, lp))

        for db_t, dict_t in zip(fw.tasks, fw_dict['spec']['_tasks']):
            if db_t.fw_name != dict_t['_fw_name']:
                raise RuntimeError("The firework with id {} from LaunchPad {} has task that don't  match: "
                                   "{} and {}.".format(fw_id, lp, db_t.fw_name, dict_t['fw_name']))

    return lp, fw_id

    def auto_load(cls):
        sma = SubmissionMongoAdapter.auto_load()

        l_dir = FWConfig().CONFIG_FILE_DIR
        l_file = os.path.join(l_dir, 'my_launchpad.yaml')
        lp = LaunchPad.from_file(l_file)

        return SubmissionProcessor(sma, lp)

def bootstrap_app(*args, **kwargs):
    """Pass instead of `app` to a forking process.

    This is so a server process will re-initialize a MongoDB client
    connection after forking. This is useful to avoid deadlock when
    using pymongo with multiprocessing.
    """
    import fireworks.flask_site.app
    fireworks.flask_site.app.lp = LaunchPad.from_dict(
        json.loads(os.environ["FWDB_CONFIG"]))
    return app(*args, **kwargs)

 def setUp(self):
     if os.path.exists(self.scratch_dir):
         shutil.rmtree(self.scratch_dir)
     os.makedirs(self.scratch_dir)
     os.chdir(self.scratch_dir)
     try:
         self.lp = LaunchPad.from_file(os.path.join(db_dir, "my_launchpad.yaml"))
         self.lp.reset("", require_password=False)
     except:
         raise unittest.SkipTest(
             'Cannot connect to MongoDB! Is the database server running? '
             'Are the credentials correct?')

def mlaunch():

    m_description = 'This program launches multiple Rockets simultaneously'

    parser = ArgumentParser(description=m_description)

    parser.add_argument('num_jobs', help='the number of jobs to run in parallel', type=int)
    parser.add_argument('--nlaunches', help='number of FireWorks to run in series per parallel job (int or "infinite"; default 0 is all jobs in DB)', default=0)
    parser.add_argument('--sleep', help='sleep time between loops in infinite launch mode (secs)', default=None, type=int)
    parser.add_argument('--timeout', help='timeout (secs) after which to quit (default None)', default=None, type=int)

    parser.add_argument('-l', '--launchpad_file', help='path to launchpad file',
                        default=LAUNCHPAD_LOC)
    parser.add_argument('-w', '--fworker_file', help='path to fworker file',
                        default=FWORKER_LOC)
    parser.add_argument('-c', '--config_dir', help='path to a directory containing the config file (used if -l, -w unspecified)',
                        default=CONFIG_FILE_DIR)

    parser.add_argument('--loglvl', help='level to print log messages', default='INFO')
    parser.add_argument('-s', '--silencer', help='shortcut to mute log messages', action='store_true')

    parser.add_argument('--nodefile', help='nodefile name or environment variable name containing the node file name (for populating FWData only)', default=None, type=str)
    parser.add_argument('--ppn', help='processors per node (for populating FWData only)', default=1, type=int)

    args = parser.parse_args()

    if not args.launchpad_file and args.config_dir and os.path.exists(os.path.join(args.config_dir, 'my_launchpad.yaml')):
        args.launchpad_file = os.path.join(args.config_dir, 'my_launchpad.yaml')

    if not args.fworker_file and args.config_dir and os.path.exists(os.path.join(args.config_dir, 'my_fworker.yaml')):
        args.fworker_file = os.path.join(args.config_dir, 'my_fworker.yaml')

    args.loglvl = 'CRITICAL' if args.silencer else args.loglvl

    launchpad = LaunchPad.from_file(args.launchpad_file) if args.launchpad_file else LaunchPad(strm_lvl=args.loglvl)

    if args.fworker_file:
        fworker = FWorker.from_file(args.fworker_file)
    else:
        fworker = FWorker()

    total_node_list = None
    if args.nodefile:
        if args.nodefile in os.environ:
            args.nodefile = os.environ[args.nodefile]
        with open(args.nodefile, 'r') as f:
            total_node_list = [line.strip() for line in f.readlines()]

    launch_multiprocess(launchpad, fworker, args.loglvl, args.nlaunches, args.num_jobs,
                        args.sleep, total_node_list, args.ppn, timeout=args.timeout)

def init_yaml(args):
    fields = (("host", "localhost"), ("port", 27017), ("name", "fireworks"), ("username", None), ("password", None))
    doc = {}
    print("Please supply the following configuration values")
    print("(press Enter if you want to accept the defaults)\n")
    for k, v in fields:
        val = input("Enter {} (default: {}) : ".format(k, v))
        doc[k] = val if val else v
    doc["port"] = int(doc["port"])  # enforce the port as an int
    with open(args.config_file, "w") as f:
        import yaml

        yaml.dump(LaunchPad.from_dict(doc).to_dict(), f)
        print("\nConfiguration written to {}!".format(args.config_file))

def restart_fizzled():
    module_dir = os.path.dirname(os.path.abspath(__file__))
    lp_f = os.path.join(module_dir, 'my_launchpad.yaml')
    lpdb = LaunchPad.from_file(lp_f)

    for fw in lpdb.fireworks.find({"state": "FIZZLED"}, {"fw_id": 1, "spec.task_type": 1}):
        fw_id = fw['fw_id']
        task_type = fw['spec']['task_type']
        restart_id = fw_id
        if 'VASP db insertion' in task_type:
            restart_id = fw_id - 1
        elif 'Controller' in task_type:
            restart_id = fw_id - 2

        lpdb.rerun_fw(restart_id)

def get_lp(args):
    try:
        if not args.launchpad_file and os.path.exists(os.path.join(args.config_dir, DEFAULT_LPAD_YAML)):
            args.launchpad_file = os.path.join(args.config_dir, DEFAULT_LPAD_YAML)

        if args.launchpad_file:
            return LaunchPad.from_file(args.launchpad_file)
        else:
            args.loglvl = 'CRITICAL' if args.silencer else args.loglvl
            return LaunchPad(logdir=args.logdir, strm_lvl=args.loglvl)
    except:
        traceback.print_exc()
        err_message = 'FireWorks was not able to connect to MongoDB. Is the server running? The database file specified was {}.'.format(args.launchpad_file)
        if not args.launchpad_file:
            err_message += ' Type "lpad init" if you would like to set up a file that specifies location and credentials of your Mongo database (otherwise use default localhost configuration).'
        raise ValueError(err_message)