code: def test_one_var_training(self, test_checkpointing, tmp_path):
      from tests.experiment.fixtures import tf_keras_one_var_model
  
      checkpoint_dir = str(tmp_path.joinpath("checkpoint"))
      latest_checkpoint = None
      steps_completed = 0
  
      # In the test_checkpointing case, we will call make_workloads() twice but batches and w
      # will persist across both calls.
      batches = enumerate([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]])
      w = 0.0
  
      trial_class = tf_keras_one_var_model.OneVarTrial
  
      def make_workloads() -> workload.Stream:
          nonlocal w
          interceptor = workload.WorkloadResponseInterceptor()
  
          for idx, batch in batches:
              yield from interceptor.send(workload.train_workload(1))
              metrics = interceptor.metrics_result()
  
              # Calculate what the loss should be.
              loss = trial_class.calc_loss(w, batch)
  
              epsilon = 0.0001
              assert (
                  abs(metrics["metrics"]["avg_metrics"]["loss"] - loss)
                  < epsilon
              )
  
              # Update what the weight should be.
              w = w - hparams[
                  "learning_rate"
              ] * trial_class.calc_gradient(w, batch)
  
              if test_checkpointing and idx == 3:
                  # Checkpoint and let the next TrialController finish the work.
                  interceptor = workload.WorkloadResponseInterceptor()
                  yield from interceptor.send(
                      workload.checkpoint_workload()
                  )
                  nonlocal latest_checkpoint, steps_completed
                  latest_checkpoint = interceptor.metrics_result()[
                      "uuid"
                  ]
                  # steps_completed is unused, but can't be 0.
                  steps_completed = 1
                  break
  
      hparams = {
          "learning_rate": 0.001,
          "global_batch_size": 3,
          "dataset_range": 10,
      }
      exp_config = utils.make_default_exp_config(
          hparams,
          scheduling_unit=100,
          searcher_metric=trial_class._searcher_metric,
      )
      exp_config["records_per_epoch"] = 100
      controller = (
          utils.make_trial_controller_from_trial_implementation(
              trial_class,
              hparams,
              make_workloads(),
              exp_config=exp_config,
              trial_seed=self.trial_seed,
              checkpoint_dir=checkpoint_dir,
          )
      )
      controller.run()
  
      # In the checkpointing case, we need to create another controller to finish training.
      if test_checkpointing:
          controller = (
              utils.make_trial_controller_from_trial_implementation(
                  trial_class,
                  hparams,
                  make_workloads(),
                  exp_config=exp_config,
                  trial_seed=self.trial_seed,
                  checkpoint_dir=checkpoint_dir,
                  latest_checkpoint=latest_checkpoint,
                  steps_completed=steps_completed,
              )
          )
          controller.run()
  
imports: import os
  from pathlib import Path
  from typing import Any, Callable, Dict, List, Optional
  import pytest
  import tensorflow
  from packaging import version
  import determined
  from determined import keras, workload
  from tests.experiment import utils
  from tests.experiment.fixtures import tf_keras_xor_model
label_desc: Use image node_id 5 for calling the TestKerasTrial obj's underlying member method code with example usage: obj.test_one_var_training(test_checkpointing, tmp_path) without return types