jobs

Tools for setting up the JobDB.

TODO: Add some quick tools for inspecting the jobdb here.

fail(job, errcode, msg, logger)

Mark and job as failed.

Parameters:

Name	Type	Description	Default
job	Job	The job to fail.	required
errcode	ErrCode	The error code that thi job failed with.	required
msg	str	The detailed error message.	required
logger	LoggerLike	The logger to log the error to.	required

Source code in lat_beams/utils/jobs.py

def fail(job: jobdb.Job, errcode: ErrCode, msg: str, logger: Optional[LoggerLike]):
    """
    Mark and job as failed.

    Parameters
    ----------
    job : jobdb.Job
        The job to fail.
    errcode : ErrCode
        The error code that thi job failed with.
    msg : str
        The detailed error message.
    logger : LoggerLike
        The logger to log the error to.
    """
    if logger is not None:
        logger.error("%s (Err %d: %s)", msg, errcode.value, errcode.name)
    set_tag(job, "message", msg)
    if "errcode" in job.tags:
        set_tag(job, "errcode", errcode.value)
    job.jstate = cast(sqy.Column[str], jobdb.JState.failed)

make_jobdb(comm, data_dir, append='')

Create or load a JobDB at {data_dir}/jobdb{append}.db.

Parameters:

Name	Type	Description	Default
comm	Optional[Comm]	The communicator if we are running with MPI or None if not. If provided then rank 0 will load the database first so that it may create it if it doesn't exist.	required
data_dir	str	The directory to load the db from.	required
append	str	String appended to db name. See docstring for details.	''

Returns:

Name	Type	Description
jobdb	JobManager	The loaded database. For better MPI support this has a timeout of 10 and uses NullPool.

Source code in lat_beams/utils/jobs.py

def make_jobdb(
    comm: Optional["Comm"], data_dir: str, append: str = ""
) -> jobdb.JobManager:
    """
    Create or load a `JobDB` at `{data_dir}/jobdb{append}.db`.

    Parameters
    ----------
    comm : Optional[Comm]
        The communicator if we are running with MPI or None if not.
        If provided then rank 0 will load the database first so that it may
        create it if it doesn't exist.
    data_dir : str
        The directory to load the db from.
    append : str
        String appended to db name.
        See docstring for details.

    Returns
    -------
    jobdb : jobdb.JobManager
        The loaded database.
        For better MPI support this has a timeout of 10 and uses NullPool.
    """
    path = os.path.join(data_dir, f"jobdb{append}.db")
    myrank = 0
    if comm is not None:
        myrank = comm.Get_rank()
    # Let rank 0 make jobdb first to avoid race conditions
    if myrank == 0:
        engine = sqy.create_engine(
            f"sqlite:///{path}",
            connect_args={"timeout": 10},
            poolclass=NullPool,
        )
        jdb = jobdb.JobManager(engine=engine)
        jdb.clear_locks(jobs="all")
        if comm is None:
            return jdb
    if comm is not None:
        comm.barrier()
    if myrank != 0:
        engine = sqy.create_engine(
            f"sqlite:///{path}",
            connect_args={"timeout": 10},
            poolclass=NullPool,
        )
        jdb = jobdb.JobManager(engine=engine)
    return jdb

setup_jobs(comm, data_dir, jclass, get_jobdict, get_jobit, get_jobstr, get_tags, source_list, overwrite, retry_failed, job_memory, job_memory_buffer, replot, logger, append='')

Discover, create, and select jobs for execution across MPI ranks.

Existing jobs are loaded from the job database and matched against a collection of candidate jobs generated by get_jobit. Missing jobs are created, eligible jobs are reopened when requested, and a consolidated list of jobs to process is returned.

Jobs may be filtered by lock status, source tag, recent visit time, and job state. Database writes are serialized across MPI ranks to avoid contention.

Jobs are created and database updates are committed serially across MPI ranks to avoid database locking contention.

Parameters:

Name	Type	Description	Default
comm	Optional[Comm]	MPI communicator used to prevent deadlock between processes.	required
data_dir	str	Directory containing the job database.	required
jclass	str	Job class name passed to JobManager.create_job when creating missing jobs.	required
get_jobdict	Callable[[JobManager], dict[str, Job]]	Function that returns a mapping from job identifier strings to existing jobs in the database.	required
get_jobit	Callable[[JobManager], Iterable[Any]]	Function that returns an iterable of candidate job descriptions.	required
get_jobstr	Callable[[Any], Optional[str]]	Function that converts a candidate job description into its unique job identifier string. Returning None causes the candidate to be skipped.	required
get_tags	Callable[[Any], dict[str, str]]	Function that generates the tag dictionary used when creating a new job from a candidate job description.	required
source_list	Sequence[str]	Allowed values of the source tag. Jobs whose source tag is not present in this sequence are ignored.	required
overwrite	bool	If True, include all matching jobs regardless of their current state and reopen non-open jobs.	required
retry_failed	bool	If True, include jobs whose state is failed.	required
job_memory	Optional[float]	Number of hours for which recently visited jobs should be skipped. If None, no visit-time filtering is performed.	required
job_memory_buffer	float	Minimum age, in minutes, before the visit-time filter is applied.	required
replot	bool	If True, include jobs whose state is done.	required
logger	LoggerLike	Logger to log to.	required
append	str	String appended to db name. See make_jobdb docstring for details.	''

Returns:

Name	Type	Description
jdb	JobManager	Job database manager instance.
jobs	list[Job]	Complete list of jobs selected for processing across all MPI ranks.

Source code in lat_beams/utils/jobs.py

def setup_jobs(
    comm: Optional["Comm"],
    data_dir: str,
    jclass: str,
    get_jobdict: Callable[[jobdb.JobManager], dict[str, jobdb.Job]],
    get_jobit: Callable[[jobdb.JobManager], Iterable[Any]],
    get_jobstr: Callable[[Any], Optional[str]],
    get_tags: Callable[[Any], dict[str, str]],
    source_list: Sequence[str],
    overwrite: bool,
    retry_failed: bool,
    job_memory: Optional[float],
    job_memory_buffer: float,
    replot: bool,
    logger: LoggerLike,
    append: str = "",
) -> tuple[jobdb.JobManager, list[jobdb.Job]]:
    """
    Discover, create, and select jobs for execution across MPI ranks.

    Existing jobs are loaded from the job database and matched against a
    collection of candidate jobs generated by `get_jobit`. Missing jobs are
    created, eligible jobs are reopened when requested, and a consolidated
    list of jobs to process is returned.

    Jobs may be filtered by lock status, source tag, recent visit time, and
    job state. Database writes are serialized across MPI ranks to avoid
    contention.

    Jobs are created and database updates are committed serially across MPI
    ranks to avoid database locking contention.

    Parameters
    ----------
    comm : Optional[Comm]
        MPI communicator used to prevent deadlock between processes.
    data_dir : str
        Directory containing the job database.
    jclass : str
        Job class name passed to `JobManager.create_job` when creating
        missing jobs.
    get_jobdict : Callable[[jobdb.JobManager], dict[str, jobdb.Job]]
        Function that returns a mapping from job identifier strings to
        existing jobs in the database.
    get_jobit : Callable[[jobdb.JobManager], Iterable[Any]]
        Function that returns an iterable of candidate job descriptions.
    get_jobstr : Callable[[Any], Optional[str]]
        Function that converts a candidate job description into its unique
        job identifier string. Returning `None` causes the candidate to be
        skipped.
    get_tags : Callable[[Any], dict[str, str]]
        Function that generates the tag dictionary used when creating a new
        job from a candidate job description.
    source_list : Sequence[str]
        Allowed values of the `source` tag. Jobs whose source tag is not
        present in this sequence are ignored.
    overwrite : bool
        If `True`, include all matching jobs regardless of their current
        state and reopen non-open jobs.
    retry_failed : bool
        If `True`, include jobs whose state is `failed`.
    job_memory : Optional[float]
        Number of hours for which recently visited jobs should be skipped.
        If `None`, no visit-time filtering is performed.
    job_memory_buffer : float
        Minimum age, in minutes, before the visit-time filter is applied.
    replot : bool
        If `True`, include jobs whose state is `done`.
    logger : LoggerLike
        Logger to log to.
    append : str
        String appended to db name.
        See `make_jobdb` docstring for details.

    Returns
    -------
    jdb : jobdb.JobManager
        Job database manager instance.
    jobs : list[jobdb.Job]
        Complete list of jobs selected for processing across all MPI ranks.
    """
    myrank, nproc = 0, 1
    if comm is not None:
        myrank = comm.Get_rank()
        nproc = comm.Get_size()
    # Get the jobs, make them if we need to
    now = time.time()
    logger.info("Setting up jobdb")
    jdb = make_jobdb(comm, data_dir, append)
    joblist = []
    jobs_to_make = []
    jobs_to_open = []
    logger.info("Getting jobdict")
    jobdict = None
    if myrank == 0:
        jobdict = get_jobdict(jdb)
    if comm is not None:
        jobdict = comm.bcast(jobdict)
    if jobdict is None:
        raise ValueError("jobdict is None!")
    logger.info("Getting potential jobs")
    it = get_jobit(jdb)
    logger.info("Processing possible jobs")
    for info in it:
        sys.stdout.flush()
        jobstr = get_jobstr(info)
        ignore_lock = False
        if jobstr is None:
            continue
        if jobstr in jobdict:
            job = jobdict[jobstr]
        else:
            tags = get_tags(info)
            job = jdb.create_job(
                jclass=jclass, tags=tags, check_existing=False, commit=False
            )
            jobs_to_make += [job]
            ignore_lock = True
        if job.lock and not ignore_lock:
            continue
        if (
            "source" in job.tags
            and job.tags["source"] not in source_list
            and job.tags["source"] != ""
        ):
            continue
        if (
            job.visit_time is not None
            and job_memory is not None
            and now - job.visit_time < 60 * 60 * job_memory
            and now - job.visit_time > 60 * job_memory_buffer
        ):
            continue
        if (
            overwrite
            or job.jstate.name == "open"
            or (job.jstate.name == "failed" and retry_failed)
        ):
            if job.jstate.name != "open":
                job.jstate = "open"
                jobs_to_open += [job]
                # joblist += [job]
            else:
                joblist += [job]
        elif replot and job.jstate.name == "done":
            joblist += [job]
    if comm is not None:
        comm.barrier()

    # Make the missing jobs
    # Doing this serially so that we don't lock up the db
    tot_missing = 0
    if comm is not None:
        tot_missing = comm.reduce(len(jobs_to_make), root=0)
    logger.info("Adding %s new jobs", tot_missing)
    tot_opening = 0
    if comm is not None:
        tot_opening = comm.reduce(len(jobs_to_open), root=0)
    logger.info("Opening %s old jobs", tot_opening)
    t0 = time.time()
    for i in range(nproc):
        if myrank == i:
            logger.debug("\tRank %s writing", i)
            jdb.commit_jobs(jobs_to_make)
            jdb.clear_locks(jobs=joblist)
            if len(jobs_to_open) > 0:
                with jdb.session_scope() as session:
                    updated_jobs = []
                for job in jobs_to_open:
                    merged_job = session.merge(job)
                    updated_jobs.append(merged_job)

                # Single commit for all jobs in this rank
                session.commit()

                for job in updated_jobs:
                    jid = job.id
                    refreshed_job = session.get(jobdb.Job, jid)
                    session.expunge(refreshed_job)
                    joblist.append(refreshed_job)
        if comm is not None:
            comm.barrier()
    t1 = time.time()
    logger.info("Took %s seconds to add", t1 - t0)

    # Get the final job list
    if comm is not None:
        all_jobs = comm.allgather(joblist)
        all_jobs = [job for jobs in all_jobs for job in jobs]
    else:
        all_jobs = joblist
    logger.info("%s jobs to run!", len(all_jobs))

    return jdb, all_jobs