Add accountant registry

[evgri243]

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Docs change / refactoring / dependency upgrade

Motivation and Context / Related issue

Allows to register 3rd party accountants to be used in methods like get_noised_multiplier and generally introduces an extra extension point.

Closes #780

How Has This Been Tested (if it applies)

Used in our internal implementation with the accountant from Felipe-Gomez/riskcal.

Checklist

• The documentation is up-to-date with the changes I made.
• I have read the CONTRIBUTING document and completed the CLA (see CONTRIBUTING).
• All tests passed, and additional code has been covered with new tests.

[facebook-github-bot]

@facebook-github-bot has imported this pull request. If you are a Meta employee, you can view this in D81679360. (Because this pull request was imported automatically, there will not be any future comments.)

[iden-kalemaj]

Thank you for these changes and the unit tests. Please see a minor comment. I will approve once tests finish.

[iden-kalemaj]

It looks like the "my_accountant" example is causing an error: link. Maybe easier to just remove the example.

In addition, can you please fix the import sorting issue from isort. Please follow instructions here on how to fix depending on your local isort version.

[evgri243]

@iden-kalemaj I've fixed both the ValueError on existing account and import sorting. Let's see if the pipeline passes this time

[evgri243]

Fixed formatting

[iden-kalemaj]

@evgri243 Strangely, the unittests on Github keep timing out. I tried re-running them a few times. I also tested with a dummy commit that does not have the same issue (link). I will need to check if the issue is somehow being introduced by your unit tests. If you have any ideas here that would be helpful too.

[iden-kalemaj]

Hi @evgri243, I believe the timeout issue is due to test_force_register_existing_accountant which is affecting the global accountant registry without cleanup.

This test permanently replaces "rdp" in the global registry: register_accountant("rdp", DummyAccountant, force=True).
When test_get_noise_multiplier_overshoot runs after this, it gets DummyAccountant instead of RDPAccountant via create_accountant(mechanism="rdp"). This causes an infinite loop for get_noise_multiplier.

As a suggested fix, add cleanup to test_force_register_existing_accountant to restore the original registry state:

def test_force_register_existing_accountant(self) -> None:
    from opacus.accountants.registry import _ACCOUNTANTS
    original = _ACCOUNTANTS["rdp"]
    
    try:
        register_accountant("rdp", DummyAccountant, force=True)
        self.assertIsInstance(create_accountant("rdp"), DummyAccountant)
        self.assertEqual(create_accountant("rdp").mechanism(), "dummy")
    finally:
        _ACCOUNTANTS["rdp"] = original

[evgri243]

You may be right actually... Let's try this one. I removed any dependency on existing accountants and execute tests only on dummy accountants.

Not sure whether to keep them wrapped into try/finally or experiment with setUp and teadDown methods. I am ready to follow your styling preferences.

[iden-kalemaj]

Try/finally works fine. Thank you for the refactoring and the quick turnaround. The tests are passing so I will land this.

[iden-kalemaj]

Review automatically exported from Phabricator review in Meta.

[facebook-github-bot]

This pull request has been merged in 183b1ae.