diff --git a/llvm/include/llvm/Analysis/Utils/TrainingLogger.h b/llvm/include/llvm/Analysis/Utils/TrainingLogger.h
--- a/llvm/include/llvm/Analysis/Utils/TrainingLogger.h
+++ b/llvm/include/llvm/Analysis/Utils/TrainingLogger.h
@@ -61,6 +61,7 @@
 #include "llvm/Support/JSON.h"
 
 #include <memory>
+#include <optional>
 #include <vector>
 
 namespace llvm {
@@ -88,19 +89,18 @@
 /// Alternatively, don't call logReward at the end of each event, just
 /// log{Float|Int32|Int64}FinalReward at the end.
 class Logger final {
+  std::unique_ptr<raw_ostream> OS;
   const std::vector<TensorSpec> FeatureSpecs;
   const TensorSpec RewardSpec;
   const bool IncludeReward;
-  std::vector<std::unique_ptr<char[]>> FeatureStorage;
-  std::vector<std::unique_ptr<char[]>> RewardStorage;
-  raw_ostream &dumpHeader(raw_ostream &OS) const;
-  raw_ostream &startContext(raw_ostream &OS, StringRef Name) const;
-  raw_ostream &startObservation(raw_ostream &OS, size_t Nr) const;
-  raw_ostream &writeOutcome(raw_ostream &OS, size_t CurrentObservationID) const;
-  char *addNewTensor(size_t FeatureID);
-  size_t getNrRecords() const;
+  StringMap<size_t> ObservationIDs;
+  std::string CurrentContext;
 
-  void logRewardImpl(const char *Value, size_t Size);
+  void writeHeader();
+  void writeTensor(const TensorSpec &Spec, const char *RawData) {
+    OS->write(RawData, Spec.getTotalTensorBufferSize());
+  }
+  void logRewardImpl(const char *RawData);
 
 public:
   /// Construct a Logger. If IncludeReward is false, then logReward or
@@ -109,44 +109,27 @@
   /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have
   /// corresponding indices) with any MLModelRunner implementations
   /// corresponding to the model being trained/logged.
-  Logger(const std::vector<TensorSpec> &FeatureSpecs,
-         const TensorSpec &RewardSpec, bool IncludeReward)
-      : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
-        IncludeReward(IncludeReward) {}
-
-  template <typename T> void logReward(T Value) {
-    logRewardImpl(reinterpret_cast<const char *>(&Value), sizeof(T));
-  }
-  void logFloatReward(float Value);
-  void logInt32Reward(int32_t Value);
-  void logInt64Reward(int64_t Value);
-
-  void logFloatFinalReward(float Value);
-  void logInt32FinalReward(int32_t Value);
-  void logInt64FinalReward(int64_t Value);
+  Logger(std::unique_ptr<raw_ostream> OS,
+         const std::vector<TensorSpec> &FeatureSpecs,
+         const TensorSpec &RewardSpec, bool IncludeReward);
 
-  void logFloatValue(size_t FeatureID, const float *Value);
-  void logInt32Value(size_t FeatureID, const int32_t *Value);
-  void logInt64Value(size_t FeatureID, const int64_t *Value);
+  void switchContext(StringRef Name);
+  void startObservation();
+  void endObservation();
 
-  void logSpecifiedTensorValue(size_t FeatureID, const char *RawData);
+  const std::string &currentContext() const { return CurrentContext; }
 
-  // Warning! For int32_t, the return is set up for int64_t, so the caller needs
-  // to piecemeal cast their int32_t values.
-  // FIXME: let's drop int32_t support. While it's supported by evaluator, it's
-  // not supported by the tensorflow::SequenceExample proto. For small values,
-  // we can consider using bytes.
-  char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);
+  bool hasObservationInProgress() const {
+    return ObservationIDs.find(CurrentContext) != ObservationIDs.end();
+  }
 
-  // Flush the content of the log to the stream, clearing the stored data in the
-  // process.
-  raw_ostream &flush(raw_ostream &OS, bool WithHeader = true,
-                     StringRef Context = "default") const;
+  template <typename T> void logReward(T Value) {
+    logRewardImpl(reinterpret_cast<const char *>(&Value));
+  }
 
-  // Flush a set of logs that are produced from the same module, e.g.
-  // per-function regalloc traces.
-  static void flushLogs(raw_ostream &OS,
-                        const StringMap<std::unique_ptr<Logger>> &Loggers);
+  void logTensorValue(size_t FeatureID, const char *RawData) {
+    writeTensor(FeatureSpecs[FeatureID], RawData);
+  }
 };
 
 } // namespace llvm
diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
--- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
@@ -102,9 +102,6 @@
   void logInlineEvent(const InlineEvent &Event,
                       const MLModelRunner &ModelRunner);
 
-  /// Print the stored tensors.
-  void print();
-
 private:
   StringRef LogFileName;
   const ModelUnderTrainingRunner *const MUTR;
@@ -150,7 +147,6 @@
 
   size_t getTotalSizeEstimate();
 
-  virtual ~DevelopmentModeMLInlineAdvisor();
   void updateNativeSizeEstimate(int64_t Change) {
     *CurrentNativeSize += Change;
   }
@@ -288,45 +284,48 @@
 
   DecisionPos = FT.size();
   FT.push_back(TensorSpec::createSpec<int64_t>(DecisionName, {1}));
+  std::error_code EC;
+  auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
+  if (EC)
+    dbgs() << (EC.message() + ":" + TrainingLog);
 
   L = std::make_unique<Logger>(
-      FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
+      std::move(OS), FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
       InlineSizeEstimatorAnalysis::isEvaluatorRequested());
+  L->switchContext("");
 }
 
 /// Log one inlining event.
 void TrainingLogger::logInlineEvent(const InlineEvent &Event,
                                     const MLModelRunner &ModelRunner) {
+  L->startObservation();
   size_t CurrentFeature = 0;
-  for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) {
-    int64_t F = *ModelRunner.getTensor<int64_t>(CurrentFeature);
-    L->logInt64Value(CurrentFeature, &F);
-  }
+  for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature)
+    L->logTensorValue(CurrentFeature,
+                      reinterpret_cast<const char *>(
+                          ModelRunner.getTensorUntyped(CurrentFeature)));
 
   if (MUTR)
     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I) {
       const char *RawData =
           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I));
-      L->logSpecifiedTensorValue(CurrentFeature, RawData);
+      L->logTensorValue(CurrentFeature, RawData);
       ++CurrentFeature;
     }
 
   assert(CurrentFeature == DefaultDecisionPos);
-  L->logInt64Value(DefaultDecisionPos, &Event.DefaultDecision);
-  L->logInt64Value(DecisionPos, &Event.AdvisedDecision);
+  L->logTensorValue(DefaultDecisionPos,
+                    reinterpret_cast<const char *>(&Event.DefaultDecision));
+  L->logTensorValue(DecisionPos,
+                    reinterpret_cast<const char *>(&Event.AdvisedDecision));
+  L->endObservation();
   if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
-    L->logInt64Reward(Event.Reward);
+    L->logReward(Event.Reward);
 
   // For debugging / later use
   Effects.push_back(Event.Effect);
 }
 
-void TrainingLogger::print() {
-  std::error_code EC;
-  raw_fd_ostream OutFile(LogFileName, EC);
-  L->flush(OutFile);
-}
-
 DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
     Module &M, ModuleAnalysisManager &MAM,
     std::unique_ptr<MLModelRunner> ModelRunner,
@@ -342,11 +341,6 @@
   assert(IsDoingInference || isLogging());
 }
 
-DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() {
-  if (isLogging())
-    Logger->print();
-}
-
 std::optional<size_t>
 DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
   if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
diff --git a/llvm/lib/Analysis/TrainingLogger.cpp b/llvm/lib/Analysis/TrainingLogger.cpp
--- a/llvm/lib/Analysis/TrainingLogger.cpp
+++ b/llvm/lib/Analysis/TrainingLogger.cpp
@@ -32,8 +32,8 @@
     UseSimpleLogger("tfutils-use-simplelogger", cl::init(true), cl::Hidden,
                     cl::desc("Output simple (non-protobuf) log."));
 
-raw_ostream &Logger::dumpHeader(raw_ostream &OS) const {
-  json::OStream JOS(OS);
+void Logger::writeHeader() {
+  json::OStream JOS(*OS);
   JOS.object([&]() {
     JOS.attributeArray("features", [&]() {
       for (const auto &TS : FeatureSpecs)
@@ -45,140 +45,44 @@
       JOS.attributeEnd();
     }
   });
-  OS << "\n";
-  return OS;
+  *OS << "\n";
 }
 
-raw_ostream &Logger::startContext(raw_ostream &OS, StringRef Name) const {
-  json::OStream JOS(OS);
+void Logger::switchContext(StringRef Name) {
+  CurrentContext = Name.str();
+  json::OStream JOS(*OS);
   JOS.object([&]() { JOS.attribute("context", Name); });
-  OS << "\n";
-  return OS;
+  *OS << "\n";
 }
 
-raw_ostream &Logger::startObservation(raw_ostream &OS, size_t Nr) const {
-  json::OStream JOS(OS);
-  JOS.object([&]() { JOS.attribute("observation", static_cast<int64_t>(Nr)); });
-  OS << "\n";
-  return OS;
-}
-
-raw_ostream &Logger::writeOutcome(raw_ostream &OS,
-                                  size_t CurrentObservationID) const {
-  if (IncludeReward) {
-    OS << "\n";
-    json::OStream JOS(OS);
-    JOS.object([&]() {
-      JOS.attribute("outcome", static_cast<int64_t>(CurrentObservationID));
-    });
-    OS << "\n";
-    OS.write(RewardStorage[CurrentObservationID].get(),
-             RewardSpec.getTotalTensorBufferSize());
-  }
-  OS << "\n";
-  return OS;
-}
-
-char *Logger::addNewTensor(size_t FeatureID) {
-  return FeatureStorage
-      .emplace_back(
-          new char[FeatureSpecs[FeatureID].getTotalTensorBufferSize()])
-      .get();
-}
-
-size_t Logger::getNrRecords() const {
-  assert(FeatureStorage.size() % FeatureSpecs.size() == 0);
-  return FeatureStorage.size() / FeatureSpecs.size();
-}
-
-void Logger::logRewardImpl(const char *Value, size_t Size) {
-  std::memcpy(RewardStorage.emplace_back(new char[Size]).get(), Value, Size);
-}
-
-raw_ostream &Logger::flush(raw_ostream &OS, bool WithHeader,
-                           StringRef Context) const {
-  if (WithHeader)
-    dumpHeader(OS);
-  startContext(OS, Context);
-  size_t CurrentObservationID = 0;
-  for (size_t I = 0; I < FeatureStorage.size(); ++I) {
-    size_t TensorID = I % FeatureSpecs.size();
-    if (TensorID == 0) {
-      CurrentObservationID = I / FeatureSpecs.size();
-      startObservation(OS, CurrentObservationID);
-    }
-    OS.write(FeatureStorage[I].get(),
-             FeatureSpecs[TensorID].getTotalTensorBufferSize());
-    if (TensorID == FeatureSpecs.size() - 1) {
-      writeOutcome(OS, CurrentObservationID);
-    }
-  }
-  return OS;
-}
-
-#define LOG_REWARD(NAME, TYPE)                                                 \
-  void Logger::log##NAME##Reward(TYPE Value) {                                 \
-    assert(IncludeReward);                                                     \
-    (void)IncludeReward;                                                       \
-    logReward(Value);                                                          \
-  }
-
-LOG_REWARD(Float, float)
-LOG_REWARD(Int32, int32_t)
-LOG_REWARD(Int64, int64_t)
-#undef LOG_REWARD
-
-#define LOG_FINAL_REWARD(NAME, TYPE)                                           \
-  void Logger::log##NAME##FinalReward(TYPE Value) {                            \
-    assert(RewardSpec.isElementType<TYPE>());                                  \
-    for (size_t I = 1; I < getNrRecords(); ++I)                                \
-      log##NAME##Reward(0);                                                    \
-    log##NAME##Reward(Value);                                                  \
-  }
-
-LOG_FINAL_REWARD(Float, float)
-LOG_FINAL_REWARD(Int32, int32_t)
-LOG_FINAL_REWARD(Int64, int64_t)
-#undef LOG_FINAL_REWARD
-
-void Logger::logFloatValue(size_t FeatureID, const float *Value) {
-  assert(FeatureSpecs[FeatureID].isElementType<float>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) {
-  assert(FeatureSpecs[FeatureID].isElementType<int64_t>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
+void Logger::startObservation() {
+  auto I = ObservationIDs.insert({CurrentContext, 0});
+  size_t NewObservationID = I.second ? 0 : ++I.first->second;
+  json::OStream JOS(*OS);
+  JOS.object([&]() {
+    JOS.attribute("observation", static_cast<int64_t>(NewObservationID));
+  });
+  *OS << "\n";
 }
 
-void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) {
-  assert(FeatureSpecs[FeatureID].isElementType<int32_t>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
+void Logger::endObservation() { *OS << "\n"; }
 
-void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
-  const auto &Spec = FeatureSpecs[FeatureID];
-  char *Buff = addEntryAndGetFloatOrInt64Buffer(FeatureID);
-  if (Spec.isElementType<int32_t>())
-    for (size_t I = 0; I < Spec.getElementCount(); ++I)
-      (reinterpret_cast<int64_t *>(Buff))[I] =
-          static_cast<int64_t>((reinterpret_cast<const int32_t *>(RawData))[I]);
-  else if (Spec.isElementType<int64_t>() || Spec.isElementType<float>())
-    std::memcpy(Buff, RawData,
-                Spec.getElementCount() * Spec.getElementByteSize());
-  else
-    llvm_unreachable("Unsupported tensor type");
-}
-
-char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) {
-  return reinterpret_cast<char *>(addNewTensor(FeatureID));
+void Logger::logRewardImpl(const char *RawData) {
+  assert(IncludeReward);
+  json::OStream JOS(*OS);
+  JOS.object([&]() {
+    JOS.attribute("outcome", static_cast<int64_t>(
+                                 ObservationIDs.find(CurrentContext)->second));
+  });
+  *OS << "\n";
+  writeTensor(RewardSpec, RawData);
+  *OS << "\n";
 }
 
-void Logger::flushLogs(raw_ostream &OS,
-                       const StringMap<std::unique_ptr<Logger>> &Loggers) {
-  bool IsFirst = true;
-  for (const auto &NamedLogger : Loggers) {
-    NamedLogger.second->flush(OS, IsFirst, NamedLogger.first());
-    IsFirst = false;
-  }
+Logger::Logger(std::unique_ptr<raw_ostream> OS,
+               const std::vector<TensorSpec> &FeatureSpecs,
+               const TensorSpec &RewardSpec, bool IncludeReward)
+    : OS(std::move(OS)), FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
+      IncludeReward(IncludeReward) {
+  writeHeader();
 }
diff --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
--- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
@@ -456,19 +456,20 @@
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
-  /// get the logger for the given function, or nullptr if we didn't collect
-  /// one. This is used to inject the score by the RegAllocScoring pass.
-  Logger *getLogger(const MachineFunction &MF) const {
-    auto I = LogMap.find(MF.getName());
-    if (I == LogMap.end())
-      return nullptr;
-    return I->second.get();
-  }
-
   void logRewardIfNeeded(const MachineFunction &MF,
                          llvm::function_ref<float()> GetReward) override {
-    if (auto *Log = this->getLogger(MF))
-      Log->logFloatFinalReward(GetReward());
+    if (!Log)
+      return;
+    // The function pass manager would run all the function passes for a
+    // function, so we assume the last context belongs to this function. If
+    // this invariant ever changes, we can implement at that time switching
+    // contexts. At this point, it'd be an error
+    if (Log->currentContext() != MF.getName()) {
+      MF.getFunction().getContext().emitError(
+          "The training log context shouldn't have had changed.");
+    }
+    if (Log->hasObservationInProgress())
+      Log->logReward<float>(GetReward());
   }
 
 private:
@@ -481,8 +482,22 @@
     RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
-  // Save all the logs (when requested).
-  bool doFinalization(Module &M) override {
+  bool doInitialization(Module &M) override {
+    LLVMContext &Ctx = M.getContext();
+    if (ModelUnderTraining.empty() && TrainingLog.empty()) {
+      Ctx.emitError("Regalloc development mode should be requested with at "
+                    "least logging enabled and/or a training model");
+      return false;
+    }
+    if (ModelUnderTraining.empty())
+      Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
+    else
+      Runner = ModelUnderTrainingRunner::createAndEnsureValid(
+          Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
+    if (!Runner) {
+      Ctx.emitError("Regalloc: could not set up the model runner");
+      return false;
+    }
     if (TrainingLog.empty())
       return false;
     std::error_code EC;
@@ -491,52 +506,32 @@
       M.getContext().emitError(EC.message() + ":" + TrainingLog);
       return false;
     }
-    Logger::flushLogs(*OS, LogMap);
+    std::vector<TensorSpec> LFS = InputFeatures;
+    if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
+      append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
+    // We always log the output; in particular, if we're not evaluating, we
+    // don't have an output spec json file. That's why we handle the
+    // 'normal' output separately.
+    LFS.push_back(Output);
+
+    Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
+                                   /*IncludeReward*/ true);
     return false;
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    LLVMContext &Ctx = MF.getFunction().getContext();
-    if (ModelUnderTraining.empty() && TrainingLog.empty()) {
-      Ctx.emitError("Regalloc development mode should be requested with at "
-                    "least logging enabled and/or a training model");
+    if (!Runner)
       return nullptr;
-    }
-    if (!Runner) {
-      if (ModelUnderTraining.empty())
-        Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
-      else
-        Runner = ModelUnderTrainingRunner::createAndEnsureValid(
-            Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
-      if (!Runner) {
-        Ctx.emitError("Regalloc: could not set up the model runner");
-        return nullptr;
-      }
-    }
-
-    Logger *Log = nullptr;
-    if (!TrainingLog.empty()) {
-      std::vector<TensorSpec> LFS = InputFeatures;
-      if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
-        append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
-      // We always log the output; in particular, if we're not evaluating, we
-      // don't have an output spec json file. That's why we handle the
-      // 'normal' output separately.
-      LFS.push_back(Output);
-      auto I = LogMap.insert(std::make_pair(
-          MF.getFunction().getName(),
-          std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
-      assert(I.second);
-      Log = I.first->second.get();
-    }
+    if (Log)
+      Log->switchContext(MF.getName());
     return std::make_unique<DevelopmentModeEvictAdvisor>(
         MF, RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
-        getAnalysis<MachineLoopInfo>(), Log);
+        getAnalysis<MachineLoopInfo>(), Log.get());
   }
 
   std::unique_ptr<MLModelRunner> Runner;
-  StringMap<std::unique_ptr<Logger>> LogMap;
+  std::unique_ptr<Logger> Log;
 };
 
 #endif //#ifdef LLVM_HAVE_TFLITE
@@ -1092,23 +1087,31 @@
   }
   if (TrainingLog.empty())
     return Ret;
+  // TODO(mtrofin): when we support optional rewards, this can go away. In the
+  // meantime, we log the "pretend" reward (0) for the previous observation
+  // before starting a new one.
+  if (Log->hasObservationInProgress())
+    Log->logReward<float>(0.0);
+
+  Log->startObservation();
   size_t CurrentFeature = 0;
   size_t FeatureCount = EnableDevelopmentFeatures
                             ? FeatureIDs::FeaturesWithDevelopmentCount
                             : FeatureIDs::FeatureCount;
   for (; CurrentFeature < FeatureCount; ++CurrentFeature) {
-    Log->logSpecifiedTensorValue(
-        CurrentFeature, reinterpret_cast<const char *>(
+    Log->logTensorValue(CurrentFeature,
+                        reinterpret_cast<const char *>(
                             getRunner().getTensorUntyped(CurrentFeature)));
   }
   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner()))
     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
          ++I, ++CurrentFeature)
-      Log->logSpecifiedTensorValue(
+      Log->logTensorValue(
           CurrentFeature,
           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
   // The output is right after the features and the extra outputs
-  Log->logInt64Value(CurrentFeature, &Ret);
+  Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
+  Log->endObservation();
   return Ret;
 }
 
diff --git a/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp
--- a/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp
@@ -177,19 +177,20 @@
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
-  /// get the logger for the given function, or nullptr if we didn't collect
-  /// one. This is used to inject the score by the RegAllocScoring pass.
-  Logger *getLogger(const MachineFunction &MF) const {
-    auto I = LogMap.find(MF.getName());
-    if (I == LogMap.end())
-      return nullptr;
-    return I->second.get();
-  }
-
   void logRewardIfNeeded(const MachineFunction &MF,
                          llvm::function_ref<float()> GetReward) override {
-    if (auto *Log = this->getLogger(MF))
-      Log->logFloatFinalReward(GetReward());
+    if (!Log)
+      return;
+    // The function pass manager would run all the function passes for a
+    // function, so we assume the last context belongs to this function. If
+    // this invariant ever changes, we can implement at that time switching
+    // contexts. At this point, it'd be an error
+    if (Log->currentContext() != MF.getName()) {
+      MF.getFunction().getContext().emitError(
+          "The training log context shouldn't have had changed.");
+    }
+    if (Log->hasObservationInProgress())
+      Log->logReward<float>(GetReward());
   }
 
 private:
@@ -200,7 +201,22 @@
   }
 
   // Save all the logs (when requested).
-  bool doFinalization(Module &M) override {
+  bool doInitialization(Module &M) override {
+    LLVMContext &Ctx = M.getContext();
+    if (ModelUnderTraining.empty() && TrainingLog.empty()) {
+      Ctx.emitError("Regalloc development mode should be requested with at "
+                    "least logging enabled and/or a training model");
+      return false;
+    }
+    if (ModelUnderTraining.empty())
+      Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
+    else
+      Runner = ModelUnderTrainingRunner::createAndEnsureValid(
+          Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
+    if (!Runner) {
+      Ctx.emitError("Regalloc: could not set up the model runner");
+      return false;
+    }
     if (TrainingLog.empty())
       return false;
     std::error_code EC;
@@ -209,53 +225,33 @@
       M.getContext().emitError(EC.message() + ":" + TrainingLog);
       return false;
     }
-    Logger::flushLogs(*OS, LogMap);
+    std::vector<TensorSpec> LFS = InputFeatures;
+    if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
+      append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
+    // We always log the output; in particular, if we're not evaluating, we
+    // don't have an output spec json file. That's why we handle the
+    // 'normal' output separately.
+    LFS.push_back(Output);
+
+    Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
+                                   /*IncludeReward*/ true);
     return false;
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-
-    LLVMContext &Ctx = MF.getFunction().getContext();
-    if (ModelUnderTraining.empty() && TrainingLog.empty()) {
-      Ctx.emitError("Regalloc development mode should be requested with at "
-                    "least logging enabled and/or a training model");
+    if (!Runner)
       return nullptr;
-    }
-    if (!Runner) {
-      if (ModelUnderTraining.empty())
-        Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
-      else
-        Runner = ModelUnderTrainingRunner::createAndEnsureValid(
-            Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
-      if (!Runner) {
-        Ctx.emitError("Regalloc: could not set up the model runner");
-        return nullptr;
-      }
-    }
-
-    Logger *Log = nullptr;
-    if (!TrainingLog.empty()) {
-      std::vector<TensorSpec> LFS = InputFeatures;
-      if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
-        append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
-      // We always log the output; in particular, if we're not evaluating, we
-      // don't have an output spec json file. That's why we handle the
-      // 'normal' output separately.
-      LFS.push_back(Output);
-      auto I = LogMap.insert(std::make_pair(
-          MF.getFunction().getName(),
-          std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
-      assert(I.second);
-      Log = I.first->second.get();
+    if (Log) {
+      Log->switchContext(MF.getName());
     }
 
     return std::make_unique<DevelopmentModePriorityAdvisor>(
-        MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log);
+        MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log.get());
   }
 
   std::unique_ptr<MLModelRunner> Runner;
-  StringMap<std::unique_ptr<Logger>> LogMap;
+  std::unique_ptr<Logger> Log;
 };
 #endif //#ifdef LLVM_HAVE_TFLITE
 
@@ -307,23 +303,31 @@
   if (TrainingLog.empty())
     return Prio;
 
+  // TODO(mtrofin): when we support optional rewards, this can go away. In the
+  // meantime, we log the "pretend" reward (0) for the previous observation
+  // before starting a new one.
+  if (Log->hasObservationInProgress())
+    Log->logReward<float>(0.0);
+
+  Log->startObservation();
   size_t CurrentFeature = 0;
   for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
-    Log->logSpecifiedTensorValue(
-        CurrentFeature, reinterpret_cast<const char *>(
+    Log->logTensorValue(CurrentFeature,
+                        reinterpret_cast<const char *>(
                             getRunner().getTensorUntyped(CurrentFeature)));
   }
 
   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
          ++I, ++CurrentFeature)
-      Log->logSpecifiedTensorValue(
+      Log->logTensorValue(
           CurrentFeature,
           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
   }
 
   float Ret = static_cast<float>(Prio);
-  Log->logFloatValue(CurrentFeature, &Ret);
+  Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
+  Log->endObservation();
 
   return static_cast<unsigned>(Prio);
 }
diff --git a/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll b/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll
--- a/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll
+++ b/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll
@@ -36,13 +36,13 @@
 ; NOML: reward: 36.64
 
 
-; CHECK-TWO-FCTS: context: SyFgetsCopy
+; CHECK-TWO-FCTS: context: SyFgets
 ; CHECK-TWO-FCTS-NEXT: observation: 0
 ; CHECK-TWO-FCTS-NEXT: mask: 0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
 ; CHECK-TWO-FCTS: index_to_evict: 12
 ; CHECK-TWO-FCTS: observation: 16
 ; CHECK-TWO-FCTS: reward: 36.64
-; CHECK-TWO-FCTS: context: SyFgets
+; CHECK-TWO-FCTS: context: SyFgetsCopy
 ; CHECK-TWO-FCTS-NEXT: observation: 0
 ; CHECK-TWO-FCTS-NEXT: mask: 0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
 ; CHECK-TWO-FCTS: index_to_evict: 12