code: def _collect_metrics(
      task: BenchmarkTask,
      data_cache_dir: Path,
      random_seed: int,
  ) -> Optional[Mapping[Metric, float]]:
      print("Loading dataset:", task.dataset_name)
      dataset_fac = DATASET_FACTORIES.get(task.dataset_name)
      dataset = dataset_fac.create_dataset(data_cache_dir)
      print(tabulate(dataset.stats, showindex="never"))
      train_data = dataset.train
      test_data = dataset.test
  
      rng = np.random.default_rng(random_seed)
      model_fac = MODEL_FACTORIES.get(task.model_name)
      model = model_fac.create_model(train_data, rng)
  
      metrics = {}
  
      model.predict_y(test_data.X)  # Warm-up TF.
  
      print("Model before training:")
      gpflow.utilities.print_summary(model)
  
      if task.do_optimise:
          t_before = perf_counter()
          loss_fn = gpflow.models.training_loss_closure(
              model, train_data.XY, compile=task.do_compile
          )
          opt_log = gpflow.optimizers.Scipy().minimize(
              loss_fn,
              variables=model.trainable_variables,
              compile=task.do_compile,
              options={"disp": 10, "maxiter": 1_000},
          )
          t_after = perf_counter()
          n_iter = opt_log.nit
          t_train = t_after - t_before
          metrics[mt.n_training_iterations] = n_iter
          metrics[mt.training_time] = t_train
          metrics[mt.training_iteration_time] = t_train / n_iter
          print(
              f"Training took {t_after - t_before}s for {n_iter} iterations."
          )
          print("Model after training:")
          gpflow.utilities.print_summary(model)
  
      likelihood = model.likelihood
  
      if task.do_predict:
          t_before = perf_counter()
          f_m, f_v = model.predict_f(test_data.X)
          t_after = perf_counter()
          metrics[mt.prediction_time] = t_after - t_before
  
          metrics[mt.nlpd] = -np.sum(
              likelihood.predict_log_density(
                  test_data.X, f_m, f_v, test_data.Y
              )
          )
  
          y_m, _y_v = likelihood.predict_mean_and_var(
              test_data.X, f_m, f_v
          )
          error = test_data.Y - y_m
          metrics[mt.mae] = np.average(np.abs(error))
          metrics[mt.rmse] = np.average(error**2) ** 0.5
  
      if task.do_posterior:
          t_before = perf_counter()
          posterior = model.posterior()
          posterior.predict_f(np.zeros((0, test_data.D)))
          t_after = perf_counter()
          metrics[mt.posterior_build_time] = t_after - t_before
  
          t_before = perf_counter()
          f_m, f_v = posterior.predict_f(test_data.X)
          t_after = perf_counter()
          metrics[mt.posterior_prediction_time] = t_after - t_before
  
          metrics[mt.posterior_nlpd] = -np.sum(
              likelihood.predict_log_density(
                  test_data.X, f_m, f_v, test_data.Y
              )
          )
  
          y_m, _y_v = likelihood.predict_mean_and_var(
              test_data.X, f_m, f_v
          )
          error = test_data.Y - y_m
          metrics[mt.posterior_mae] = np.average(np.abs(error))
          metrics[mt.posterior_rmse] = np.average(error**2) ** 0.5
  
      metrics_values = list(metrics.values())
  
      # pylint: disable=no-member
      assert np.isfinite(
          metrics_values
      ).all(), f"{metrics} not all finite."
  
      return metrics
  
imports: import argparse
  import json
  import multiprocessing
  import traceback
  from datetime import timedelta
  from pathlib import Path
  from time import perf_counter
  from typing import Mapping, Optional
  import numpy
  import pandas
  from tabulate import tabulate
  import benchmark.benchmarks
  import benchmark.metrics
  import gpflow
  from benchmark.benchmark_api import BENCHMARK_SUITES, BenchmarkSuite, BenchmarkTask
  from benchmark.dataset_api import DATASET_FACTORIES
  from benchmark.metadata import BenchmarkMetadata
  from benchmark.metric_api import Metric
  from benchmark.model_api import MODEL_FACTORIES
  from benchmark.paths import setup_dest
  from benchmark.plot import plot
  from benchmark.sharding import ShardingStrategy
label_desc: Use image node_id 1 for calling a global function with example usage: _collect_metrics(task, data_cache_dir, random_seed) and returns: metrics