diff --git a/requirements.txt b/requirements.txt index 23bb4f4e..de0f4bc5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ -monty==2021.8.17 +monty==2022.1.19 networkx==2.6.3 pydash==5.1.0 -maggma==0.32.1 -pydantic==1.8.2 +maggma==0.32.3 +pydantic==1.9.0 PyYAML==6.0 diff --git a/setup.py b/setup.py index 17516db4..78f45982 100644 --- a/setup.py +++ b/setup.py @@ -33,20 +33,20 @@ ], extras_require={ "docs": [ - "sphinx==4.3.0", - "furo==2021.11.15", - "m2r2==0.3.1", - "ipython==7.29.0", - "nbsphinx==0.8.7", + "sphinx==4.4.0", + "furo==2022.1.2", + "m2r2==0.3.2", + "ipython==8.0.1", + "nbsphinx==0.8.8", "nbsphinx-link==1.3.0", - "FireWorks==1.9.7", - "autodoc_pydantic==1.5.1", + "FireWorks==1.9.8", + "autodoc_pydantic==1.6.0", ], "tests": [ "pytest==6.2.5", "pytest-cov==3.0.0", - "FireWorks==1.9.7", - "matplotlib==3.4.3", + "FireWorks==1.9.8", + "matplotlib==3.5.1", "pydot==1.4.2", ], "dev": ["pre-commit>=2.12.1"], diff --git a/src/jobflow/managers/myqueue.py b/src/jobflow/managers/myqueue.py new file mode 100644 index 00000000..8448fb0c --- /dev/null +++ b/src/jobflow/managers/myqueue.py @@ -0,0 +1,110 @@ +"""Tools for running :obj:`Flow` and :obj:`Job` objects using the Myqueue package. + +Notes +----- +Myqueue heavily relies on the file system. To submit a workflow, one has to run: +mq workflow workflow.py DIRECTORY_PATTERNS +where workflow.py is a python script defining one workflow. For jobflow Flows, the +workflow.py file in myqueue_scripts has to be used. +""" + +from __future__ import annotations + +import json +import os +import typing +from datetime import datetime +from pathlib import Path +from random import randint + +from monty.json import MontyDecoder, MontyEncoder +from monty.os import cd + +if typing.TYPE_CHECKING: + from typing import List, Union + + import jobflow + +from maggma.stores import JSONStore + +from jobflow import JobStore + +__all__ = ["flow_to_myqueue", "run_myqueue_task"] + +FLOW_JSON = "flow.json" +JOB_STORE_JSON = "job_store.json" + + +def flow_to_myqueue( + flow: Union[jobflow.Flow, jobflow.Job, List[jobflow.Job]], +): + """ + Convert a jobflow Flow to myqueue. + + This is basically just dumping the jobflow Flow to a flow.json file. + The flow.json file is then read again when the user wants to submit + the workflow using myqueue. + + Parameters + ---------- + flow + A flow or job. + + """ + from jobflow.core.flow import get_flow + + flow = get_flow(flow) + with open(FLOW_JSON, "w") as f: + json.dump(flow, f, cls=MontyEncoder, indent=2) + + +def run_myqueue_task(uuid): + """ + Run a job in myqueue. + + Parameters + ---------- + uuid + Unique identifier of the job that needs to be executed. + """ + root_dir = Path.cwd() + # First get the jobflow Flow from the flow.json file + with open("flow.json", "r") as f: + flow = json.load(f, cls=MontyDecoder) + + # Get the jobflow Job corresponding to the uuid + job = _get_job(flow, uuid) + job_dir = _get_job_dir(root_dir=root_dir) + + # Initialize the store for output references + job_store_json_path = os.path.join(root_dir, JOB_STORE_JSON) + if not os.path.exists(job_store_json_path): + with open(job_store_json_path, "w") as f: + json.dump([], f) + store = JobStore(JSONStore(job_store_json_path, writable=True)) + store.connect() + + # Run the job + with cd(job_dir): + job.run(store=store) + + +def _get_job(flow, uuid): + myjob = None + for job, _ in flow.iterflow(): + if job.uuid == uuid: + if myjob is not None: + raise RuntimeError(f"Multiple jobs with uuid {uuid}") + myjob = job + return myjob + + +def _get_job_dir(root_dir): + time_now = datetime.utcnow().strftime("%Y-%m-%d-%H-%M-%S-%f") + job_dir = root_dir / f"job_{time_now}-{randint(10000, 99999)}" + job_dir.mkdir() + return job_dir + + +if __name__ == "__main__": + pass diff --git a/src/jobflow/managers/myqueue_scripts/workflow.py b/src/jobflow/managers/myqueue_scripts/workflow.py new file mode 100644 index 00000000..206c4154 --- /dev/null +++ b/src/jobflow/managers/myqueue_scripts/workflow.py @@ -0,0 +1,26 @@ +"""Template workflow.py file for running jobflow workflows in myqueue.""" + + +import json + +from monty.json import MontyDecoder +from myqueue.task import task + + +def create_tasks(): + """Create tasks for myqueue.""" + # First reconstruct the jobflow Flow object + with open("flow.json", "r") as f: + flow = json.load(f, cls=MontyDecoder) + + tasks = [] + uuid2task = {} + for job, parents in flow.iterflow(): + deps = [uuid2task[parent_uuid] for parent_uuid in parents] + t = task( + "jobflow.managers.myqueue@run_myqueue_task", args=[job.uuid], deps=deps + ) + uuid2task[job.uuid] = t + tasks.append(t) + + return tasks