pdbfixer references openmm.app and causes jax error

[astomer2]

hi, everyone, I am currently using MNPP to generate some structures, which requires a good initial structure. However, when I write pdbfixer and mpnn into the same script, mpnn will report an error and jax cannot be initialized correctly (I specified the gpu to use). After checking, I found that it was because pdbfixer automatically performed some operations on the gpu, causing the program to run on the GPU I did not specify, which would prevent jax from initializing (or the index of cuda was messed up). If I disable the processing of pdbfixer, there will be no error.

Now my question is, is there any way for me to specify the gpu that openmm.app will use, or run it on the cpu?

[peastman]

Not currently, but that would make sense to add. We could let you pass a Platform to the PDBFixer constructor. I'll mark this as a feature request.

[peastman]

#324 adds a platform argument to the constructor. You can specify platform=Platform.getPlatform('CPU') to make it do calculations on the CPU.

Can you test it out and confirm that it fixes your problem?

[peastman]

@astomer2 can you confirm that the change in #324 solves your problem? Thanks!

[astomer2]

I'm sorry for the late reply. After testing the new code, I found that I need to set os.environ['CUDA_VISIBLE_DEVICES'] before the code block to specify the GPU platform. Additionally, after carefully reviewing the modified code, I noticed that pdbfixer uses mm.context, where properties can be specified as properties = {'DeviceIndex': str(specify_cuda), 'Precision': 'single'}. This code should theoretically allow specifying the GPU platform directly using a variable from the function. I tested this in OpenMM simulations, and it works, but setting properties in pdbfixer does not seem to have any effect.

[astomer2]

Additional test: when using the new version of pdbfixer to specify theGPU platform, or more precisely, using os.environ['CUDA_VISIBLE_DEVICES'] to specify the graphics card, you can no longer use properties = {'DeviceIndex': str(specify_cuda), 'Precision': 'single'} in openmm to specify the GPU platform. This will cause only the GPU you specified previously to be found in your environment, resulting in a graphics card serial number naming error. Of course, you can directly use 'DeviceIndex': '0' by default, which should not cause an error

[peastman]

I don't quite understand. If you don't specify a device, everything should work normally and use device 0. Is that not happening?

When you set CUDA_VISIBLE_DEVICES, it filters the devices at a low level so applications see fewer devices as being present. The visible devices are numbered starting from 0. For example, if you set CUDA_VISIBLE_DEVICES=2, applications will only see a single device present. They'll see it as device 0, even though it really refers to device 2 in the complete, unfiltered list.

[astomer2]

Sorry for not being clear enough. I mean that the current code allows specifying the GPU using CUDA_VISIBLE_DEVICES, which is great. Only using CUDA_VISIBLE_DEVICES can specify GPU for both pdbfixer and openmm. however, theoretically, OpenMM provides a more flexible way to specify GPUs by modifying the properties in mm.Context. For example:
properties = {'DeviceIndex': str(specific_cuda), 'Precision': 'single'} mm.Context(system, integrator, platform, properties)
I can use this approach successfully in my OpenMM simulations. However, when I tried to modify the code in PDBFixer to achieve the same functionality, it didn't work.