feat: add custom reducers to estimators [DET-3098]

[rb-determined-ai]

Description

Introduce custom reducers for estimator trial. From the docstring in the PR:

During distributed evaluation, many types of metrics calculated via ``tf.metrics`` or
``tf.keras.metrics`` cannot be aggregated properly from the per-slot final metrics
calculated by each separate Estimator replica. One example is ``tf.metrics.auc``, where
the ROC AUC calculated over predictions and labels from a full dataset cannot be derived
from a set of ROC AUC metrics evaluated over the shards of a dataset. However, with
``make_metric``, the ROC AUC could be calculated in distributed training by calling
``sklearn.metrics.roc_auc_score`` in a custom ``reducer`` function.

Test Plan

Lots of manual testing, in addition to adding a new unit test and a new parallel test.

[rb-determined-ai]

@aaron276h This is now ready for a "for-realsies" review. The incremental update since the last time you looked is:

• added unit and parallel e2e test
• fixed a bug where the tf.control_dependencies() was somehow triggering the call to _DistributedMetric.result() twice, resulting in building two allgather ops for every metric in the graph. I fixed this by pre-building every result op and returning the cached op when result() is called.
• the _allgather_ops list in the context has to be reset after every .evaluate() call, because every .evaluate() is going to build a clean graph, and you don't want to call tf.control_dependencies() on ops from the old graph.

[aaron276h]

non-blocking: can we rename this from Debug to something else

[rb-determined-ai]

sure, done

[aaron276h]

I think worth calling out that this is for validation metrics

[rb-determined-ai]

I made the following edits to the docstrings:

 Reducing Metrics
 ~~~~~~~~~~~~~~~~

-Determined supports proper reduction of arbitrary metrics during distributed
-training by allowing users to define custom reducers for their metrics. Custom
-reducers can be either a function or an implementation of the
+Determined supports proper reduction of arbitrary validation metrics during
+distributed training by allowing users to define custom reducers for their
+metrics. Custom reducers can be either a function or an implementation of the

     def make_metric(..) -> tf.keras.metrics.Metric:
         """
-        Return an estimator-compatible metric which will be calculated properly, even during
-        distributed training.
+        Return an estimator-compatible validation metric which will be calculated properly, even
+        during distributed evaluation.

 class MetricReducer:
     """
-    Efficiently aggregating metrics across a multi-slot distributed evaluation is done in two steps:
+    Efficiently aggregating validation metrics across a multi-slot distributed evaluation is done
+    in two steps:

[aaron276h]

question: why make this a decorator and not just a function call?

[rb-determined-ai]

To keep the MetricReducer a pure python API.

A natural tensorflowy way to write this would be to set the granularity of py_func such that the final metric reduction is done in two steps: 1. allgather the final outputs of accumulate(), 2. apply the user's cross_slot_reduce to the algathered stuff.

That would be natural because the only parts of the graph which have to be serialized are as small as possible; only step 1, which is the network communication. Also, the allgather op would just do allgather and you could easily have a function to build a generic allgather that other ops would connect to. The drawback is that you would have to convert the output of the accumulate() function to tensorflow types since those outputs would have to pass through the graph. The output of the final allgather call would also have to have a declared dtype, since that's a requirement of py_func, which adds another layer of configurability we would need in the interface.

What I did was set the granularity of the py_func such both of the above two steps were accomplished within a single py_func. The input to cross_slot_reduce() is then much, much easier to reason about for the user, since the user will get their exact outputs rather than some tensorflow-casted outputs. The cost is that the entire cross_slot_reduce becomes part of the serialized section of graph operations.

Then with the granularity of py_func I chose, I think a decorator is the best way to represent how to parameterize the _build_allgather_op, but it's definitely a little confusing.

[rb-determined-ai]

Talked offline, the question was about calling _build_allgather_op directly rather than applying it as a decorator. I don't care either way, so I took the direct call approach.