Stop storing references to bound methods via tf.function
Author: Rocketknight1
Pull request: yes

This is (hopefully!) the end of a long saga this week.

@ydshieh noticed that our tests runners were going OOM, after a couple of PRs I made to dummy inputs. I thought the problem was just that the new dummy inputs were too large, but eventually we figured out that the problem was actually quite complicated!

tl;dr **A circular reference exists, which is caused by us calling tf.function() on a model method and then storing the result as a model attribute. Because this reference exists, our TF models are not cleaned up immediately when they are deleted, but only after the next Python garbage collection.**

I believe the PRs triggered the issue by eliminating unneccessary calls and making TF model building much faster. This left less time for garbage collection to happen, and as a result our test suites started a second test before the first test had been cleaned up, which caused the test runner to go OOM.

We tried resolving this problem by manually calling `gc.collect()` before each test, but this made some of the test suites much slower! Obviously the real solution had to be to resolve the circular reference that was slowing down model cleanup.

~The solution is to replace `model.eager_serving` with a method `model._get_eager_serving_fn()`. This returns a function that TensorFlow can compile, but which doesn't create a hard reference to a model method in the returned `tf.function`. I confirmed through manual inspection with `gc.get_referrers` that the reference is removed and models are cleaned up immediately once they go out of scope now.~

See the update below for a full description of the solution I finally went with!

Comments:
_The documentation is not available anymore as the PR was closed or merged._

no OOM with this PR (for the models involved), but we have some errors regarding `TypeError: Binding inputs to tf.function `eager_serving` failed due to `missing a required argument: 'inputs'`` popping up for several model/tokenizer tests.

One example is
```bash
self = <tests.models.gpt2.test_modeling_tf_gpt2.TFGPT2ModelTest testMethod=test_saved_model_creation>

    @slow
    def test_saved_model_creation(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = False
        config.output_attentions = False
    
        if hasattr(config, "use_cache"):
            config.use_cache = False
    
        model_class = self.all_model_classes[0]
    
        class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
        model = model_class(config)
    
        model(class_inputs_dict)
    
        with tempfile.TemporaryDirectory() as tmpdirname:
>           model.save_pretrained(tmpdirname, saved_model=True)

tests/test_modeling_tf_common.py:268: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/transformers/modeling_tf_utils.py:2427: in save_pretrained
    self.save(saved_model_dir, include_optimizer=False, signatures=signatures)
/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py:70: in error_handler
    raise e.with_traceback(filtered_tb) from None
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <tensorflow.python.eager.polymorphic_function.function_spec.FunctionSpec object at 0x7f6e1c4db7f0>
args = ({'attention_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='attention_mask'), 'input_ids': TensorSpec(sha...tf.int32, name='input_ids'), 'token_type_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='token_type_ids')},), kwargs = {}

    def bind_function_inputs(self, args, kwargs):
      """Bind `args` and `kwargs` into a canonicalized signature args, kwargs."""
      sanitized_kwargs = {
          function_type_lib.sanitize_arg_name(k): v for k, v in kwargs.items()
      }
      if len(kwargs) != len(sanitized_kwargs):
        raise ValueError(f"Name collision after sanitization. Please rename "
                         f"tf.function input parameters. Original: "
                         f"{sorted(kwargs.keys())}, Sanitized: "
                         f"{sorted(sanitized_kwargs.keys())}")
    
      try:
        bound_arguments = self.function_type.bind_with_defaults(
            args, sanitized_kwargs, self.default_values)
      except Exception as e:
>       raise TypeError(
            f"Binding inputs to tf.function `{self._name}` failed due to `{e}`."
            f"Received args: {args} and kwargs: {sanitized_kwargs} for signature:"
            f" {self.function_type}."
        ) from e
E       TypeError: Binding inputs to tf.function `eager_serving` failed due to `missing a required argument: 'inputs'`.Received args: ({'input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='input_ids'), 'attention_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='attention_mask'), 'token_type_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='token_type_ids')},) and kwargs: {} for signature: (self, inputs: Dict(mapping={'input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='input_ids'), 'attention_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='attention_mask'), 'token_type_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='token_type_ids')})).
```

Looks like the way I'm handling the methods fails when we try to save the model with those signatures. I'll figure it out on Monday!

This should be ready for review now! The changes are pretty small, but it took me a while to figure out the details.

It turns out anything that looks like `self.serving = tf.function(self.eager_serving)` will create a circular reference between `self.serving` and `self` and inhibit cleanup. This does not apply to methods defined at the class (rather than instance) level. Something like this is fine and does not block cleanup:
```
@tf.function(input_signature=...)
def serving(self, inputs):
    ...
```
The problem with the construction above, though, is that the `tf.function` decorator has to be called with all of its arguments at the class level, before the model has been initialized with a config. This means it can't read any shapes or details from the config, which means its signature has to be very **very** general. This is why we transitioned to `self.serving = ...` in the first place.

The solution I found is the following:

- Get rid of all helper methods like `self.eager_serving`. These were only used internally anyway, to allow us to compile multiple serving signatures.
- Decorate the base `serving` method with `tf.function` and no signature at all.
- Rely on our control of `self.save_spec` to ensure that base TF methods like `model.save()` will save with the right signature even when we aren't manually defining it (I checked this and it works!)
- When we want to manually specify signatures, we just call `self.serving.get_concrete_signature` with different signatures. No need to keep `eager_serving` around anymore!

This should totally preserve functionality and backward compatibility, while resolving the memory cleanup issue and keeping the specific save signatures. The only potentially noticeable change is that `self.serving.input_signature` is no longer defined. We read that value in a couple of tests as a shortcut to find the model input names, so I just replaced it with `self.input_signature` instead. I don't think anyone outside of Hugging Face was using it, and it certainly wasn't part of our public API, so I don't expect any issues!

Thanks to @ydshieh for his patience with the tests and to @gante for digging out the old PRs that let me finally understand why a lot of this stuff was ever here in the first place!

OK, I will run a few tests and let you know @Rocketknight1 
Thank you for trying trying!

@ydshieh actually, you're right - I thought it wasn't doing anything anymore, but it's still useful in some cases when we define a broad signature that gets inherited. Let me rework that so we keep it!

No warning sign after running tests for 4 involved models. You are very good at TF!

@ydshieh I finished rebasing and I removed your `gc.collect()` change to the doctests. Are you okay for me to merge now, or do you want to run any further tests?

Either way, I think we've finally resolved this one!

It's ok, go ahead. If doctest starts to fail, I will `call` you.

Also pinging @amyeroberts for core maintainer review

> LGTM 🔥
> 
> To be safe, can you trigger the slow CI on this branch? Most of the TF serialization tests are slow tests :D

Hello @gante !

Do you mean enable slow tests but for all the models ..? Or anything else?
Can't run slow tests on CircleCI however, so need to run on a specific VM.

Good point, actually - this PR isn't specific to any one model, so we'd need to run slow tests for all models. Since it's a long time to the next release, let's just merge this PR (after review) and see if anything fails overnight?

@Rocketknight1 @ydshieh Could we run slow tests on just a handful of models ~5 popular ones from different modalities to make sure any obvious issues have been caught? 

@amyeroberts I've been running them locally on the core models - BERT and GPT-2 look good! Are you okay if I try a few more and then merge if there are no issues?

Tested BERT, GPT-2, BART, ViT, CLIP and Wav2Vec2 without issues. Merging!