Enabling of MDAnalysis.analysis.align.AverageStructure parallelization

[talagayev]

Fixes #4659 attempt

Changes made in this Pull Request:

• added backends and aggregators to AlignTraj and AverageStructure in analysis.align.
• added the client_AlignTraj and client_AverageStructure in conftest.py
• added client_AlignTraj and client_AverageStructure in run() in test_align.py

Currently for AlignTraj it only accepts serial and dask with multiprocessing leading to the pytests taking forever. An additional error that appears is the following:

OSError: File opened in mode: self.mode. Reading only allow in mode "r"

For AverageStructure the Failure that appears is the following:

AttributeError: 'numpy.ndarray' object has no attribute 'load_new'

Which leads me to believe that AverageStructure can not be parallelized, but I would need additional opinions on it and on AlignTraj :)

PR Checklist

• Tests?
• Docs?
• CHANGELOG updated?
• Issue raised/referenced?

Developers certificate of origin

• I certify that this contribution is covered by the LGPLv2.1+ license as defined in our LICENSE and adheres to the Developer Certificate of Origin.

---

📚 Documentation preview 📚: https://mdanalysis--4738.org.readthedocs.build/en/4738/

[pep8speaks]

Hello @talagayev! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

• In the file testsuite/MDAnalysisTests/analysis/test_align.py:

Line 310:80: E501 line too long (86 > 79 characters)
Line 327:80: E501 line too long (87 > 79 characters)
Line 333:80: E501 line too long (97 > 79 characters)
Line 357:80: E501 line too long (85 > 79 characters)
Line 377:80: E501 line too long (91 > 79 characters)

Comment last updated at 2025-01-11 21:40:18 UTC

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.74%. Comparing base (24548e6) to head (ee8d845).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #4738   +/-   ##
========================================
  Coverage    92.73%   92.74%           
========================================
  Files          180      180           
  Lines        22475    22491   +16     
  Branches      3190     3191    +1     
========================================
+ Hits         20842    20859   +17     
  Misses        1176     1176           
+ Partials       457      456    -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[talagayev]

@marinegor I would ping you in this PR as well.

Here basically I tried different ways to see if it is possible to parallelize the AverageStructure and AlignTraj classes.
I was able to implement the parallelization for AverageStructure with the first attempt not working due to being able to read a universe so that error appeared

AttributeError: 'numpy.ndarray' object has no attribute 'load_new'

I added _first to make the class parallelizable. This works well with the test, expect the case, when in_memory=True, then the process takes very long, which I assume is connected to memory issues, so in the current case I revert back to serial for cases when in_memory=True.

As for AlignTraj due to it transforming coordinates and writing out structures it would be necessary to rewrite more parts of the code to make it parallelizable, so I didn't find an easy solution for that, which wouldn't require bigger modifications of the class. So there the question would be if we keep it as non parallelizable or should I try to modify the code to make it parallelizable, which would require bigger modifications?

[marinegor]

@talagayev

expect the case, when in_memory=True, then the process takes very long, which I assume is connected to memory issues

I think parallelization should not actually work with in_memory cases (@yuxuanzhuang please correct me if I'm wrong, afaik you've been working on this). Hence I'd explicitly raise an exception if uses asks for parallel execution and provides in_memory as well.

So there the question would be if we keep it as non parallelizable or should I try to modify the code to make it parallelizable, which would require bigger modifications?

If you're running out of ideas, I'd suggest making this PR for AverageStructure only, and create appropriate issue for AlignTraj, describing your attempts so far.

Also, I imagine there are problems with serialization of self._writer, no? Perhaps we can chat on discord about it (I'm @marinegor there)?

[talagayev]

@talagayev

expect the case, when in_memory=True, then the process takes very long, which I assume is connected to memory issues

I think parallelization should not actually work with in_memory cases (@yuxuanzhuang please correct me if I'm wrong, afaik you've been working on this). Hence I'd explicitly raise an exception if uses asks for parallel execution and provides in_memory as well.

So there the question would be if we keep it as non parallelizable or should I try to modify the code to make it parallelizable, which would require bigger modifications?

If you're running out of ideas, I'd suggest making this PR for AverageStructure only, and create appropriate issue for AlignTraj, describing your attempts so far.

Also, I imagine there are problems with serialization of self._writer, no? Perhaps we can chat on discord about it (I'm @marinegor there)?

Yes makes sense. I think the current two ones that use in_memory and are analysis related are AverageStructure and AlignTraj.

Yes that would be good, I can then rename the PR to cover only AverageStructure for now, add the missing parts for the PR (Documentation + Changelog), create an Issue and write you on Discord, so that we can brainstorm how to adjust the code to make it parallelizable. Yes self._writer is one of the difficulties. I guess for the aligntraj you can adjust the code to give it the reference, but yes the writing during parallelization is the tricky part, maybe with tmp information that is then merged or maybe just doing the calculations and the writing is then in conclude, basically keeping that part serial and only making the calculations parallel.

[marinegor]

I won't be in favor of a warning, and would rather explicitly raise ValueError because, well, how often do you switch off / ignore warnings?)

[talagayev]

True, adjusted it to raise a ValueError for that case and adjusted the test to cover the ValueError.

[marinegor]

@talagayev ok, will be waiting for your message.
I also assigned myself a reviewer here, so just re-request review when you think you're done!

[talagayev]

@talagayev ok, will be waiting for your message. I also assigned myself a reviewer here, so just re-request review when you think you're done!

Added the Documentation, CHANGELOG and adjust to raise and ValueError. The PR would be ready to be re-reviewed :)

[marinegor]

@talagayev all looking good! I initially commented on one extra line but just realized I can remove it myself.

Also, may I ask you to create a separate issue for AlignTraj parallelization? Perhaps you could describe the issues you encountered there, and suggest the direction in which one should move to actually enable it.

[talagayev]

@talagayev all looking good! I initially commented on one extra line but just realized I can remove it myself.

Also, may I ask you to create a separate issue for AlignTraj parallelization? Perhaps you could describe the issues you encountered there, and suggest the direction in which one should move to actually enable it.

Perfect :)

Yes, I can create an Issue and write some Ideas in there.

[marinegor]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[marinegor]

@talagayev sorry for postponing that but I think these comments will make the code better :)

[marinegor]

@talagayev sorry I didn't think about it earlier, but this actually feels a bit hacky to me, and I think I know why: run() isn't supposed to do any validation, it's performed by _configure_backend() method instead (docs). In your case, I'd patch it to be something like:

def _configure_backend(
      self,
      backend: Union[str, BackendBase],
      n_workers: int,
      unsupported_backend: bool = False,
    ) -> BackendBase:
    configured_backend = super()._configure_backend(backend=backend, n_workers=n_workers, unsupported_backend=unsupported_backend)
    if not isinstance(configured_backend, MDAnalysis.analysis.backends.BackendSerial) and self._in_memory:
        raise ValueError('...')

this way you don't have to patch run() with double-nested ifs, and generally write less code.

[talagayev]

yes was a hacky approach to do the go around with the serial case. Agree your approach looks cleaner, will try to adjust it in the upcoming days :)

[talagayev]

Adjusted now to use _configure_backend as suggested :)

[marinegor]

@talagayev and regarding AlignTraj -- I'd just be bold and say that it's impossible to parallelize that with current split-and-combine technique, since that would require an ability to write an arbitrary frame with self._writer, and not only sequential writing. I don't think any writers allow that.

So even in this PR, set available_backends to only serial explicitly, and note that it's impossible to parallelize. And as for the next issue, I'd target adding an option to save aligned trajectory to self.results, and writing it in _conclude.