Shortfin LLM Debug Ergonomics: one flag to dump them all

[renxida]

Thanks @stbaione for originally writing a lot of the debug code in here. I organized his debug logging & dumping into a separate file & made the output easier to deal with.

[ScottTodd]

Keep the full # Invoke. Logits are of shape [bs, bsl, d]. comment?

[ScottTodd]

Can you move this to an init() function so it doesn't always run at import time? This debug file could provide a class instead of a loose function that uses global variables. Then there would be no need to duplicate this os.getenv code across both files or add the Importing this should do nothing comment

# debug_service.py

class ShortfinLLMDebugService:
  def __init__(self):
    dump_id = 0
    dump_dir = ""
    dump_dir.mkdir(...)

  async def pre_invocation_debug_dump(...):
    is_decode = ...

# service.py

class GenerateService:
    def __init__(self, ...):
      ...

      SHORTFIN_DEBUG_LLM_SERVICE = os.getenv(
          "SHORTFIN_DEBUG_LLM_SERVICE", "False"
      ).lower() in ("true", "yes", "1", "y")
      if SHORTFIN_DEBUG_LLM_SERVICE:
        from .debug_service import ShortfinLLMDebugService
        llm_debug_service = ShortfinLLMDebugService

...

    async def run(self):
      ...
      if self.llm_debug_service:
        llm_debug_service.pre_invocation_debug_dump(...)