Covalent integration

[jmmshn]

I am interested in integrating covalent as a runner similar to fireworks and jobflow-remote.
I think they have some features that can make debugging complex workflows a bit more manageable. Also, I have encountered problems with serializing large workflow on certain implementations of MongoDB due to (possible) node value size limits.
I have a basic implementation here:

from __future__ import annotations

from typing import TYPE_CHECKING
import covalent as ct

if TYPE_CHECKING:
    from collections.abc import Sequence

    import jobflow
    from jobflow.core.job import Job
    from covalent._workflow.electron import Electron

@ct.lattice
@ct.electron
def job_to_electron(
    job: jobflow.Job,
    store: jobflow.JobStore = None,
    **kwargs,
) -> Electron:
    """Run a job wrapped in an electron.
    
    The outputs of parents are passed as keyword arguments to the function.
    This allows Covalent to automatically determine the job dependencies.

    Parameters
    ----------
    job
        A job to run.
    store
        A job store. Alternatively, if set to None, :obj:`JobflowSettings.JOB_STORE`
        will be used. Note, this could be different on the computer that submits the
        workflow and the computer which runs the workflow. The value of ``JOB_STORE`` on
        the computer that runs the workflow will be used.
    **kwargs
        {parent_job_uuid: parent_electron_output} dictionary.
    """
    from jobflow.core.reference import OnMissing
    from jobflow import SETTINGS, initialize_logger
    initialize_logger()
    with store as connected_store:
        response = job.run(store=connected_store)
    return response

@ct.lattice(executor="local")
def flow_to_lattice(
    flow: jobflow.Flow | jobflow.Job | list[jobflow.Job],
    store: jobflow.JobStore = None,
    allow_external_references: bool = False,
):
    """Convert a :obj:`Flow` or a :obj:`Job` to a Covalent :obj:`Lattice` object.
    
    Parameters
    ----------
    flow
        A flow or job.
    store
        A job store. Alternatively, if set to None, :obj:`JobflowSettings.JOB_STORE`
        will be used. Note, this could be different on the computer that submits the
        workflow and the computer which runs the workflow. The value of ``JOB_STORE`` on
        the computer that runs the workflow will be used.
    allow_external_references
        If False all the references to other outputs should be from other Jobs
        of the Flow.
    """
    from jobflow.core.flow import get_flow

    flow = get_flow(flow, allow_external_references=allow_external_references)
    job_uuid_to_electron_output = dict()
    
    for job_, parents_ in flow.iterflow():
        parent_outputs = {p_: job_uuid_to_electron_output[p_] for p_ in parents_}
        new_electron = job_to_electron(
            job_, store=store, **parent_outputs
        )
        job_uuid_to_electron_output[job_.uuid] = new_electron
        job_uuid_to_electron_output[job_.uuid].function.__name__ = job_.name
    return flow.output

It's still very rough as dynamic workflows still need to be implemented. I'm interested in some feedback
Looks like there were some thoughts on this last year:
AgnostiqHQ/covalent#1642

[davidwaroquiers]

Hi @jmmshn

Thanks a lot for opening this discussion. I agree with you that we should have something easy to interact with, especially for debugging/developping workflows. Could you maybe elaborate on the features in Covalent that helps in this ? As jobflow-remote is a "dedicated" jobflow manager (dedicated in the sense that it was developed explicitly FOR jobflow), we can see if/how we can add some of these features.

[jmmshn]

Hi @davidwaroquiers
A challenging point when working with jobflow is when you have WF that fails it's hard to enter into the job that failed and examine interactively interact with the code.
As an example, if I used Fireworks it becomes I can deserialize the FW.json and run that interactively to see why a Python function broke.
Covalent spits the entire Python function call into a directory where you have all the inputs and the function call in python format which can help with debugging.

Also, for more complex workflow it really helps to have a dashboard. These are the main features that I'm interested in.
So fireworks is OK most of the time, but there is always the issue with air gaps and firewalls that makes some situations more difficult to deal with.

Hope these points help. I'm following the jobflow-remote package, so if there is some kind of dashboard implemented I am very interested in using it and adding the debugging features as I encounter the need. But in the meantime, I think having some kind of basic adapter for Covalent will still be good for jobflow as a whole.

[gpetretto]

Hi @jmmshn,
jobflow-remote actually has a file that could be considered more or less an equivalent of the FW.json.
If you look in the running folder of a job you will find the jfremote_in.json file. That is all you need to run the job on its own. You can just copy that to a folder and run jf execute run in that folder. Otherwise you can use the run_remote_job function (see https://github.com/Matgenix/jobflow-remote/blob/4870f02eeb02ed71c167dc046dced0f306d7006b/src/jobflow_remote/jobs/run.py#L31), that is what is executed by the previous CLI command. Or look into that and just run the part that you need. It boils down to deserializing the jfremote_in.json and call the run method of the job contained in it. All the references are already resolved, so there is no need to access the DB when you do that.

Adding a dashboard to jobflow remote is definitely something we plan to do. However, since it will require some efforts it will still take some time for us to implement it.