Clarify ownership of runner components (#10338)

Summary:

The ownership of these components need some clarification.

* `Module` should be shared by `TextDecoderRunner` and potentially `TextPrefiller` (or `ImagePrefiller` in multimodal runner).
* `TextDecoderRunner` should be shared by the `TextPrefiller` and `TextTokenGenerator`.
* `Tokenizer` should be owned by the `Runner` as well as `TextTokenGenerator`.

Reviewed By: swolchok, kirklandsign

Differential Revision: D73399600

Author: larryliu0820

examples/models/llama/runner/runner.cpp

@@ -99,6 +99,7 @@ Error Runner::load() {
         "Failed to load %s as a Tiktoken artifact, trying BPE tokenizer",
         tokenizer_path_.c_str());
     tokenizer_.reset();
+    // @lint-ignore CLANGTIDY facebook-hte-Deprecated
     tokenizer_ = std::make_unique<::tokenizers::Llama2cTokenizer>();
     err = tokenizer_->load(tokenizer_path_);
     ET_CHECK_TK_OK_OR_RETURN_ERROR(

examples/models/llava/runner/llava_image_prefiller.h

@@ -18,7 +18,7 @@ namespace example {
 class ET_EXPERIMENTAL LlavaImagePrefiller
     : public ::executorch::extension::llm::ImagePrefiller {
  public:
-  LlavaImagePrefiller(::executorch::extension::Module* module)
+  explicit LlavaImagePrefiller(::executorch::extension::Module* module)
       : ImagePrefiller(module){};
   /**
    * Prefill an LLM Module with the given image input.

extension/llm/runner/text_decoder_runner.h

@@ -14,7 +14,6 @@
 #include <executorch/extension/module/module.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/platform/compiler.h>
-#include <functional>
 
 namespace executorch {
 namespace extension {
@@ -94,7 +93,13 @@ class ET_EXPERIMENTAL TextDecoderRunner {
   }
 
  protected:
-  // TODO: use shared_ptr for module
+  /**
+   * Note: TextDecoderRunner does not own the Module instance. It is expected
+   * that the outer class (likely Runner) manages the lifecycle of the Module.
+   * This means that the responsibility for creating, maintaining, and
+   * destroying the Module lies outside of TextDecoderRunner. Ensure that the
+   * Module remains valid for the duration of TextDecoderRunner's usage.
+   */
   Module* module_;
   bool use_kv_cache_;
   bool should_stop_{false};

extension/llm/runner/text_prefiller.h

@@ -24,6 +24,8 @@ class ET_EXPERIMENTAL TextPrefiller {
       bool use_kv_cache_,
       bool enable_parallel_prefill,
       int64_t max_seq_len = 128);
+
+  virtual ~TextPrefiller() = default;
   /**
    * Prefill an LLM Module with the given text input.
    * @param prompt_tokens The text prompt tokens to the LLM Module. Encoded by
@@ -32,7 +34,7 @@ class ET_EXPERIMENTAL TextPrefiller {
    * Module.
    * @return The next token of the LLM Module after prefill.
    */
-  ::executorch::runtime::Result<uint64_t> prefill(
+  virtual ::executorch::runtime::Result<uint64_t> prefill(
       std::vector<uint64_t>& prompt_tokens,
       int64_t& start_pos);
 
@@ -48,6 +50,12 @@ class ET_EXPERIMENTAL TextPrefiller {
       int64_t& start_pos);
 
  private:
+  /**
+   * Note: TextPrefiller does not own the TextDecoderRunner instance.
+   * The responsibility of managing the lifecycle of TextDecoderRunner
+   * lies with the outer class or entity (likely Runner) that creates
+   * and passes the TextDecoderRunner instance to TextPrefiller.
+   */
   TextDecoderRunner* text_decoder_runner_;
   bool use_kv_cache_;
   bool enable_parallel_prefill_;

extension/llm/runner/text_token_generator.h

@@ -32,6 +32,8 @@ class ET_EXPERIMENTAL TextTokenGenerator {
         use_kv_cache_(use_kv_cache),
         stats_(stats) {}
 
+  virtual ~TextTokenGenerator() = default;
+
   /**
    * Token generation loop.
    * @param tokens prompt tokens as well as the first token generated by
@@ -136,6 +138,12 @@ class ET_EXPERIMENTAL TextTokenGenerator {
   }
 
  private:
+  /**
+   * Note: TextTokenGenerator does not own the tokenizer_ and
+   * text_decoder_runner_. The lifecycle of these objects should be managed
+   * externally, likely in the Runner. This class assumes that the provided
+   * pointers remain valid for the duration of its use.
+   */
   ::tokenizers::Tokenizer* tokenizer_;
   TextDecoderRunner* text_decoder_runner_;
   std::unique_ptr<std::unordered_set<uint64_t>> eos_ids_;