schrodinger.application.matsci.restart module

Infrastructure for restarting JobDJ workflows from lightweight checkpoints.

The minimal restart-aware Job and JobDJ classes are kept separate from the matsci checkpoint and artifact policy. Job objects are deliberately not serialized; clients rebuild and configure the executable graph on every invocation before applying the saved job state.

Copyright Schrodinger, LLC. All rights reserved.

class schrodinger.application.matsci.restart.RestartStateMixin(*args, restart_data=None, **kwargs)

Bases: object

Hold arbitrary in-memory data that can be restored by a later job.

__init__(*args, restart_data=None, **kwargs)

Initialize the restart data.

Parameters:

restart_data (object) – Optional pickleable data transferred from a previous job or driver invocation.

getRestartData()

Return data to transfer through the restart checkpoint.

Return type:

object or None

Returns:

Optional pickleable restart data.

loadRestartData(restart_data)

Load data transferred from a previous job or driver invocation.

Parameters:

restart_data (object) – Optional pickleable restart data.

class schrodinger.application.matsci.restart.RestartableJob(*args, **kwargs)

Bases: RobustSubmissionJob

Minimal restart-aware RobustSubmissionJob.

This class contains execution behavior and generic graph/state checkpoint records that could eventually move to job infrastructure. Matsci restart data and artifact policy belong in MatsciRestartableJob.

RESTART_STATE_SKIPPED = 'skipped'

JOB_STATE_KEY = 'state'

PARENT_JOBNAMES_KEY = 'parent_jobnames'

DEPENDENT_JOBNAMES_KEY = 'dependent_jobnames'

__init__(*args, **kwargs)

Initialize a restartable job.

getRestartState()

Return the state to persist for this job.

Return type:

JobState or str

Returns:

Current JobDJ state, or “skipped” when restart handling intentionally bypassed this job.

setRestartState(restart_state)

Restore a terminal checkpoint state.

Parameters:

restart_state (JobState or str) – queue.JobState.DONE or “skipped”.

resetForRun()

Clear restored state and make the job eligible for scheduling.

getInitialJobInfo()

Return graph relationships captured before JobDJ starts.

Return type:

dict

Returns:

Immediate parent and dependent job names.

getRuntimeJobInfo()

Return state that can change while JobDJ is running.

Return type:

dict

Returns:

Current checkpoint state.

restoreJobInfo(job_info)

Restore reusable terminal state on a rebuilt job.

Parameters:

job_info (dict) – Saved job information.

Return type:

bool

Returns:

Whether the job was restored as complete or skipped.

requestStop()

Mark every downstream restartable job as skipped.

Scheduling continues normally; each affected job’s own run/finalize become no-ops once the queue reaches it.

run(*args, **kwargs)

Skip submission when this job is already restored complete or skipped.

See jobutils.RobustSubmissionJob.run for standard arguments.

finalize()

Skip finalizers for restored or deliberately skipped jobs.

class schrodinger.application.matsci.restart.MatsciRestartableJob(*args, input_files=None, required_output_files=None, **kwargs)

Bases: RestartStateMixin, RestartableJob

Restartable job with matsci artifact checkpoint policy.

REQUIRED_OUTPUT_FILES_KEY = 'required_output_files'

RESTART_DATA_KEY = 'restart_data'

__init__(*args, input_files=None, required_output_files=None, **kwargs)

Initialize matsci restart artifacts.

Parameters:

• input_files (list(str)) – Declared launch inputs.

• required_output_files (list(str)) – Outputs required to restore a completed job.

getInitialJobInfo()

Return immutable job configuration captured before submission.

Launch inputs belong to the executable graph reconstructed by the driver and are not checkpoint state.

Return type:

dict

Returns:

Job outputs and graph relationships.

getRuntimeJobInfo()

Return state that can change while the JobDJ is running.

Return type:

dict

Returns:

Current checkpoint state.

restoreJobInfo(job_info, validate_outputs=True)

Validate and restore completion state from a saved job record.

Parameters:

• job_info (dict) – Saved job information.

• validate_outputs (bool) – Whether to confirm required outputs still exist on disk. Callers may pass False for jobs whose artifacts are not required to still be present, such as non-front jobs whose outputs were already cleaned up.

Return type:

bool

Returns:

Whether the job was restored as complete or skipped.

validateRequiredOutputs()

Validate outputs required after job completion.

Return type:

bool

Returns:

Whether every required output exists.

class schrodinger.application.matsci.restart.RestartableJobDJ(restart_filename, *args, **kwargs)

Bases: JobDJ

JobDJ with lightweight graph and state checkpoints.

JobDJ continues to own scheduling, dependency pruning, and failure handling. Checkpoints contain stable job names, immediate graph relationships, and runtime state. Artifact policy belongs in MatsciRestartableJobDJ.

RESTARTABLE_JOB_CLASS

alias of RestartableJob

VERSION_KEY = 'version'

JOBS_KEY = 'jobs'

__init__(restart_filename, *args, **kwargs)

Initialize a restartable queue.

Parameters:

restart_filename (str) – Restart checkpoint path.

getJobsByName()

Return restartable jobs indexed by unique stable name.

Return type:

dict(str, RestartableJob)

Returns:

Jobs keyed by name.

requestStop()

Mark every not-yet-completed restartable job as skipped.

checkpointJobs()

Write missing jobs and graph relationships before submission.

checkpoint(job)

Persist a job lifecycle transition.

classmethod readRestartState(filename)

Read and validate generic restart state.

Parameters:

filename (str) – Checkpoint path.

Return type:

dict

Returns:

Persisted graph and job state.

loadState(restart_state)

Apply saved terminal state to jobs in a rebuilt graph.

Jobs and dependencies must be reconstructed before this method is called. The checkpoint never serializes executable job objects.

Parameters:

restart_state (dict) – Validated checkpoint state.

run(**kwargs)

Run jobs and checkpoint lifecycle transitions.

See JobDJ.run for standard arguments.

class schrodinger.application.matsci.restart.MatsciRestartableJobDJ(input_files, options, restart_filename, required_output_files=None, append_output_files=None, all_jobs_are_front_jobs=False, *args, **kwargs)

Bases: RestartableJobDJ

RestartableJobDJ with matsci artifact and checkpoint policy.

Every status change checkpoints the most advanced job whose restart state and every prerequisite’s restart state are DONE. A restarted process rebuilds the same jobs and restores completion up to that point.

RESTARTABLE_JOB_CLASS

alias of MatsciRestartableJob

OPTIONS_KEY = 'options'

INPUT_FILES_KEY = 'input_files'

REQUIRED_OUTPUT_FILES_KEY = 'required_output_files'

FRONT_JOBS_KEY = 'front_jobs'

APPEND_OFFSETS_KEY = 'append_offsets'

LAUNCH_ENVIRONMENT_KEY = 'launch_environment'

LAUNCH_ENVIRONMENT_KEYS = ('OPLS_DIR',)

APPEND_BACKUP_SUFFIX = '.bak'

__init__(input_files, options, restart_filename, required_output_files=None, append_output_files=None, all_jobs_are_front_jobs=False, *args, **kwargs)

Initialize a restartable job queue.

Parameters:

• input_files (list(str)) – Required workflow input files.

• options (dict or SubscriptableNameSpace) – Resolved pickleable workflow options.

• restart_filename (str) – Restart checkpoint path.

• required_output_files (list(str)) – Shared workflow outputs.

• append_output_files (list(str)) – Shared append-only outputs.

• all_jobs_are_front_jobs (bool) – Whether every completed job is written as a front job instead of only completed DAG leaves.

classmethod getLaunchEnvironment()

Return the launch environment relevant to a resumed workflow.

Return type:

dict(str, str)

Returns:

Allowlisted launch environment values present in this process.

static validateInputFiles(input_files)

Validate the immutable workflow inputs required for restart.

Parameters:

input_files (list(str)) – Workflow input paths.

Raises:

ValueError – If the paths are empty, duplicated, malformed, or missing.

checkpoint(job)

Write the current checkpoint state atomically.

checkpointToJobServer()

Preserve the latest recoverable front jobs outside the compute node.

Directories are expanded because Job Server checkpointing accepts files. Original workflow inputs are already restored by Job Server.

classmethod readRestartState(filename)

Read and validate restart state.

Parameters:

filename (str) – Checkpoint path.

Return type:

dict

Returns:

Persisted workflow state.

restoreAppendOutputs(restart_state)

Roll shared append outputs back to their committed byte offsets.

The original file is preserved as a backup before being replaced, since truncation is destructive and the original contents cannot be recovered otherwise.

Parameters:

restart_state (dict) – Loaded checkpoint state.

loadState(restart_state)

Mark jobs already added to this queue as complete when the checkpoint shows they finished in a previous run.

Current workflow inputs are validated when this queue is constructed. Their paths may differ from the prior invocation when a workflow intentionally replaces an input while rebuilding its graph.

Only front jobs have their output files validated on disk. A front job with a missing output is left waiting so it reruns; earlier jobs may have had their artifacts cleaned up since they completed. Jobs must already be added (see queue.JobDJ.addJob) before this is called; this method never creates jobs itself.

Parameters:

restart_state (dict) – Validated restart checkpoint state, as returned by readRestartState.

class schrodinger.application.matsci.restart.RestartableDriver(jobname, options, restart_data=None)

Bases: RestartStateMixin

Base driver for workflows persisted by MatsciRestartableJobDJ.

CHECKPOINT_FILE_OPTION = 'checkpoint_file'

RESTART_JOBNAME_OPTION = 'restart_jobname'

RESTART_FILE_SUFFIX = '.mcpt'

RESTART_ERROR

alias of RuntimeError

ALL_JOBS_ARE_FRONT_JOBS = False

SKIP_OVERRIDE_FLAG = ('-subhost', '-TPP', '-parallel')

__init__(jobname, options, restart_data=None)

Initialize a restartable workflow driver.

Parameters:

• jobname (str) – Stable workflow job name.

• options (dict, SubscriptableNameSpace, or None) – Resolved workflow options, or None when restart is disabled.

• restart_data (object) – Optional pickleable data transferred from a previous job or driver invocation.

classmethod restoreRestartEnvironment(restart_state)

Restore launch state that is not represented by workflow options.

The current launch invocation wins when it supplied an OPLS directory; otherwise retain the directory used to create the checkpoint.

Parameters:

restart_state (dict or None) – Loaded checkpoint state.

classmethod setJobServerRestart(options)

Select the deterministic restart file after an automatic resume.

Parameters:

options (dict or SubscriptableNameSpace) – Newly parsed invocation options.

Return type:

bool

Returns:

Whether Job Server supplied a checkpoint for this invocation.

classmethod loadOptions(options)

Restore saved workflow options for a checkpoint invocation.

Only the checkpoint path and invocation-specific resource options belong to the common restart lifecycle. Workflow-specific invocation options must be applied by the workflow driver after this method returns.

Parameters:

options (dict or SubscriptableNameSpace) – Newly parsed invocation options.

Return type:

tuple(SubscriptableNameSpace, dict or None)

Returns:

Resolved options and loaded checkpoint state.

classmethod getMutableRestartFiles(restart_state, root_directory=None, skip_missing=False)

Return concrete output files required as mutable restart inputs.

Parameters:

• restart_state (dict) – Validated restart checkpoint state.

• root_directory (str) – Optional directory containing all restart artifacts. Saved paths are resolved there by basename.

• skip_missing (bool) – Whether to omit missing artifacts.

Return type:

set(str)

Returns:

Shared and front-job output files. Required directories are expanded recursively.

static configureJobServerRestart(job_builder)

Mark a driver job as resumable without subjob attachment.

The checkpoint state contains the required front-job files, so a resumed driver reruns its current stage instead of attaching to old subjobs.

Parameters:

job_builder (JobSpecificationArgsBuilder) – Job specification builder for the driver.

isRestart()

Return whether this invocation uses an existing checkpoint.

Return type:

bool

Returns:

Whether a checkpoint was supplied.

getRestartOptions()

Return options to persist for a future restart.

Return type:

dict

Returns:

Saved workflow options.

getRestartInputFiles()

Return immutable workflow inputs required for restart.

Return type:

list(str)

Returns:

Input paths.

getRestartRequiredOutputFiles()

Return shared outputs required for restart.

Return type:

list(str)

Returns:

Required output paths.

getRestartAppendOutputFiles()

Return shared append-only outputs rolled back during restart.

Return type:

list(str)

Returns:

Append-only output paths.

getRestartFilename()

Return the existing or new checkpoint filename.

Return type:

str

Returns:

Checkpoint path.

createRestartJobDJ(**kwargs)

Create an empty restart queue using the common workflow metadata.

Return type:

MatsciRestartableJobDJ

Returns:

Empty restart queue.

setUpJobGraph(job_dj)

Populate a newly created JobDJ graph for this invocation.

The graph is not part of the checkpoint. Subclasses must construct all jobs and dependencies each time this method is called.

Parameters:

job_dj (MatsciRestartableJobDJ) – JobDJ to populate.

getJobDJ(restart_state=None)

Build a fresh graph and apply optional checkpoint state.

Parameters:

restart_state (dict) – Optional checkpoint state. By default, use the state loaded during driver construction.

Return type:

MatsciRestartableJobDJ

Returns:

Configured graph with saved completion state applied.