test: triangular attention snapshots

[jandom]

Summary

Adding our 1st snapshot test. I'm not sure how this will work across different GPU types, especially with jit.

Changes

• added comments to the test to ease the reader in
• converted the test into pytest
• added a snapshot test for TriangleAttention (.npz backed)
• new test dep pytest-regressions

Question for reviewers

I think we should mutualize the test code, somewhat, otherwise the test cases have a lot of the same boilerplate

• seed torch
• eval
• call with no_grad()
• assert that outputs are not all zeros
• snapshot comparison

Related Issues

Testing

Other Notes

[jnwei]

Overall, this is awesome. I did not know about pytest-regressions, but it seems perfect for our use case.

As this PR is hopefully the first of a few PRs to introduce snapshot tests, I have a few general questions regarding testing / organization

• Which system did you create the snapshots on? I can help test the snapshots on other machines to see if they pass.

• How do you envision organizing the snapshot files? It seems that by default, a folder is created for each test. If we end up with many snapshot files, I wonder if it is better to place all of the snapshots in one directory so that we can zip the files.

[jandom]

Which system did you create the snapshots on? I can help test the snapshots on other machines to see if they pass.

These were all created on my CPU VM but passed on g5.xlarge, so either this ran CPU-only or within tolerance on the GPUs. Maybe we can include something in the fixture filename such that it's arch specific. Ideally we can automate generating for the hardware we need, it was very tedious in my past life to keep these not-stale manually.

How do you envision organizing the snapshot files? It seems that by default, a folder is created for each test. If we end up with many snapshot files, I wonder if it is better to place all of the snapshots in one directory so that we can zip the files.

Indeed, i don't have a good answer. Some things are better local, some things are better global. For tests, we chose to keep them separately but organized like the codebase. I'll explore how easy it'd be to keep it in a separate location.

The compression is a separate thing, ideally we don't expose the internals to the developer and "hide" pytest-regression behind a fixture that handles compression/location for them.

[jandom]

Now that we have the VCR snapshots, this could actually live alongside those

[jandom]

This is what allows us to place the snapshots along other test_data, sibling to the 'cassettes'

[jnwei]

Nice! Does it make sense to update the snapshot paths for the templates test you added in https://github.com/aqlaboratory/openfold-3/tree/main/openfold3/tests/test_data/cassettes/test_rscb

Can be a different PR if it is cumbersome to update here.

[jandom]

These are different snapshots – we have two types

• the VCR snapshots for web requests
• ndarrays_regression snapshots for numerics

I'd be weary of mixing those up – they mean a different thing and solve a different problem

[jnwei]

This all looks good, just one suggestion regarding a context manager for the random state.

[jnwei]

Nice! Does it make sense to update the snapshot paths for the templates test you added in https://github.com/aqlaboratory/openfold-3/tree/main/openfold3/tests/test_data/cassettes/test_rscb

Can be a different PR if it is cumbersome to update here.

[jnwei]

I think you can use a context manager to handle the rng state, something like this (I have not tested this myself)

#conftest.py

import random
import numpy as np
import torch
import pytest
from torch.random import fork_rng

@pytest.fixture()
def seeded_rng():
    """Isolate all RNG state (torch, numpy, python) for the duration of the test."""
    py_state = random.getstate()
    np_state = np.random.get_state()
    with fork_rng():          # saves/restores torch (+ CUDA) state
        torch.manual_seed(123)
        # set other random states if neeeded
        random.seed(123)
        np.random.seed(123)
        yield
    # torch state restored by freeze_rng_state on exit
    random.setstate(py_state)         # restore python state manually
    np.random.set_state(np_state)     # restore numpy state manually

Docs for fork_rng

[jandom]

Oh, this is great – in general I think something like this will be a good fit for these tests, because they're very similar to each other, and we'll duplicate a lot of code between different module tests.

[jandom]

Good news here on the cross-platform, there will definitely be some hiccups – so far i'm testing on DGX/aarch64, and all these are passing – no update needed.