Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[wip] Storage manager update #679

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[wip] Storage manager update #679

Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
19991a3
WIP of new storage backend in RFE protocol
richardjgowers Jan 1, 2024
6b81bb6
register files I'm expecting to be created
richardjgowers Jan 4, 2024
0571d28
more WIP of new filesystem access
richardjgowers Jan 9, 2024
ecd6716
more WIP of new filesystem access
richardjgowers Jan 10, 2024
37d448c
pass StagingPath objects back in Unit return
richardjgowers Jan 15, 2024
fb85a58
fixed StagingPath type annotations
richardjgowers Jan 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 41 additions & 27 deletions openfe/protocols/openmm_rfe/equil_rfe_methods.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
settings, ChemicalSystem, LigandAtomMapping, Component, ComponentMapping,
SmallMoleculeComponent, ProteinComponent, SolventComponent,
)
from gufe.storage import stagingregistry

from .equil_rfe_settings import (
RelativeHybridTopologyProtocolSettings, SystemSettings,
Expand Down Expand Up @@ -578,8 +579,10 @@ def __init__(self, *,
)

def run(self, *, dry=False, verbose=True,
scratch_basepath=None,
shared_basepath=None) -> dict[str, Any]:
scratch_basepath: pathlib.Path,
shared_basepath: stagingregistry.StagingRegistry,
permanent_basepath: stagingregistry.StagingRegistry,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think these are StagingRegistry ... aren't they StagingPath?

The input to this method was, e.g., root/$DAG_LABEL/shared_$UNIT_LABEL. StagingRegistry would correspond to root/. The full path comes in as a StagingPath, I believe (this is something that changed at one point, because there was no need to expose both classes to protocol developers; it also simplified serialization, IIRC.)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm working off these hints here, these three vars are just the context unpacked: https://github.com/OpenFreeEnergy/gufe/blob/staging-execute-dag/gufe/protocols/protocolunit.py#L30

) -> dict[str, Any]:
"""Run the relative free energy calculation.

Parameters
Expand All @@ -591,10 +594,12 @@ def run(self, *, dry=False, verbose=True,
verbose : bool
Verbose output of the simulation progress. Output is provided via
INFO level logging.
scratch_basepath: Pathlike, optional
Where to store temporary files, defaults to current working directory
shared_basepath : Pathlike, optional
Where to run the calculation, defaults to current working directory
scratch_basepath: StagingDirectory
Where to store temporary files
shared_basepath : StagingDirectory
Where to run the calculation
permanent_basepath : StagingDirectory
Where to store files that must persist beyond the DAG
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

update typing in docstring


Returns
-------
Expand All @@ -609,11 +614,6 @@ def run(self, *, dry=False, verbose=True,
"""
if verbose:
self.logger.info("Preparing the hybrid topology simulation")
if scratch_basepath is None:
scratch_basepath = pathlib.Path('.')
if shared_basepath is None:
# use cwd
shared_basepath = pathlib.Path('.')

# 0. General setup and settings dependency resolution step

Expand Down Expand Up @@ -664,11 +664,13 @@ def run(self, *, dry=False, verbose=True,
else:
ffcache = None

ffcache.register()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Problem 1 (easy): this line and the ffcache.as_path() below will fail if ffcache is None.

Problem 2 (more complicated): How is ffcache supposed to work? Isn't that supposed to be something that we can pull down from some previous unit? shared_basepath is supposed to be for a specific execution of a given unit. That directory might not exist until we're about to start executing this unit. How does a file get in there? How are we currently using that? [EDIT: I'm having a call with @IAlibay tomorrow to clarify what is going on with this]

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

reading through, I think the only thing that would be cached would be the LJ parameters for the small molecules. We can probably remove this?


system_generator = system_creation.get_system_generator(
forcefield_settings=forcefield_settings,
thermo_settings=thermo_settings,
system_settings=system_settings,
cache=ffcache,
cache=ffcache.as_path(),
has_solvent=solvent_comp is not None,
)

Expand Down Expand Up @@ -812,10 +814,18 @@ def run(self, *, dry=False, verbose=True,
)

# a. Create the multistate reporter
nc = shared_basepath / sim_settings.output_filename
# TODO: Logic about keeping/not .nc files goes here
nc = (shared_basepath / sim_settings.output_filename)
checkpoint = (shared_basepath / sim_settings.checkpoint_storage)
real_time_analysis = (shared_basepath / "real_time_analysis.yaml")
# have to flag these files as being created so that they get brought back
nc.register()
checkpoint.register()
real_time_analysis.register()

chk = sim_settings.checkpoint_storage
reporter = multistate.MultiStateReporter(
storage=nc,
storage=str(nc.as_path()),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're using this str(file.as_path()) pattern a lot. When you use .as_path(), you're explicitly bypassing the automatic registration mechanism, meaning that you must explicitly register the object.

I'm not sure if this is implemented, but you should be able to instead just use str(file), which should also register the file. Equivalently, str(pathlib.Path(file)) will definitely work, which means you don't need to manually use .register().

You should only need to use .register() on paths that are created outside your control (like the checkpoint and RTA files, which are paths generated inside OpenMMTools and not exposed to us).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah once it's MultiStateReporter wants a string (not path) and the other is going into a subprocess call, so again it's bypassing file-like objects and is a string.

analysis_particle_indices=selection_indices,
checkpoint_interval=sim_settings.checkpoint_interval.m,
checkpoint_storage=chk,
Expand Down Expand Up @@ -947,13 +957,12 @@ def run(self, *, dry=False, verbose=True,
sampling_method=sampler_settings.sampler_method.lower(),
result_units=unit.kilocalorie_per_mole,
)
analyzer.plot(filepath=shared_basepath, filename_prefix="")
analyzer.plot(filepath=permanent_basepath, filename_prefix="")
analyzer.close()

else:
# clean up the reporter file
fns = [shared_basepath / sim_settings.output_filename,
shared_basepath / sim_settings.checkpoint_storage]
fns = [nc.as_path(), checkpoint.as_path()]
for fn in fns:
os.remove(fn)
finally:
Expand All @@ -980,36 +989,40 @@ def run(self, *, dry=False, verbose=True,

if not dry: # pragma: no-cover
return {
'nc': nc,
'last_checkpoint': chk,
'nc': nc.as_path(),
'last_checkpoint': checkpoint.as_path(),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably better to return the StagingPath itself, which can then be reloaded even if you move the path. Imagine something running on the EC2, where files will get stored on S3. Knowing the path on your EC2 instance isn't particularly useful.

**analyzer.unit_results_dict
}
else:
return {'debug': {'sampler': sampler}}

@staticmethod
def analyse(where) -> dict:
def analyse(where: stagingregistry.StagingRegistry) -> dict:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is probably a StagingPath

# don't put energy analysis in here, it uses the open file reporter
# whereas structural stuff requires that the file handle is closed
ret = subprocess.run(['openfe_analysis', str(where)],
trjdir = (where / '')
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what's the purpose of doing this?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was the question I was asking on slack. I needed a .as_path() of a directory and StagingRegistry doesn't have this

output = (where / 'results.json')
ret = subprocess.run(['openfe_analysis', 'RFE_analysis',
str(trjdir.as_path()),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I understand correctly, has the requirement that the files are already in the directory trjdir. Note that this will only work for situations where this analyse method is called on the same computer where the data was generated.

To run it on a different computer would require downloading all files within the directory. The best way to do that now would be to enumerate them here -- in principle there is a way to obtain all files with a given string prefix, but all of that gets deep into the weeds.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this method is called from _execute() right after run() so it should always be the same machine

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it is a little ugly that it needs to know if the trj was stored in scratch/shared/permanent, so there's probably something there I need to fix

str(output.as_path())],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
if ret.returncode:
return {'structural_analysis_error': ret.stderr}

data = json.loads(ret.stdout)
with open(output, 'r') as f:
data = json.load(f)

savedir = pathlib.Path(where)
if d := data['protein_2D_RMSD']:
fig = plotting.plot_2D_rmsd(d)
fig.savefig(savedir / "protein_2D_RMSD.png")
fig.savefig(where / "protein_2D_RMSD.png")
plt.close(fig)
f2 = plotting.plot_ligand_COM_drift(data['time(ps)'], data['ligand_wander'])
f2.savefig(savedir / "ligand_COM_drift.png")
f2.savefig(where / "ligand_COM_drift.png")
plt.close(f2)

f3 = plotting.plot_ligand_RMSD(data['time(ps)'], data['ligand_RMSD'])
f3.savefig(savedir / "ligand_RMSD.png")
f3.savefig(where / "ligand_RMSD.png")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(not really this line, but this is the closest changed line)

You may want to include the figure filenames in the return value, too. Otherwise it will be harder to download them from the cloud.

plt.close(f3)

return {'structural_analysis': data}
Expand All @@ -1020,7 +1033,8 @@ def _execute(
log_system_probe(logging.INFO, paths=[ctx.scratch])
with without_oechem_backend():
outputs = self.run(scratch_basepath=ctx.scratch,
shared_basepath=ctx.shared)
shared_basepath=ctx.shared,
permanent_basepath=ctx.permanent)

analysis_outputs = self.analyse(ctx.shared)

Expand Down