feat: Transcribe using FLAME on a Fly GPU

config/runtime.exs

@@ -72,6 +72,17 @@ if config_env() == :prod do
     ],
     secret_key_base: secret_key_base
 
+  if config_env() == :prod do
+    config :flame, :backend, FLAME.FlyBackend
+
+    config :flame, FLAME.FlyBackend,
+      token: System.fetch_env!("MEDICODE_API_TOKEN"),
+      gpu_kind: "a100-pcie-40gb",
+      cpu_kind: "performance",
+      cpus: 8,
+      memory_mb: 20_480
+  end
+
   # ## SSL Support
   #
   # To get SSL working, you will need to add the `https` key

lib/medicode/application.ex

@@ -9,12 +9,15 @@ defmodule Medicode.Application do
   def start(_type, _args) do
     Logger.add_backend(Sentry.LoggerBackend)
 
+    flame_parent = FLAME.Parent.get()
+
     children =
       [
         {MedicodeWeb.MetricsStorage, MedicodeWeb.Telemetry.metrics()},
         Medicode.Repo,
         {DNSCluster, query: Application.get_env(:medicode, :dns_cluster_query) || :ignore},
         {Phoenix.PubSub, name: :medicode_pubsub},
+        {FLAME.Pool, name: Medicode.OfflineServing, min: 0, max: 1},
         Medicode.ServingSupervisor,
         {
           Medicode.TranscriptionSupervisor,
@@ -30,12 +33,22 @@ defmodule Medicode.Application do
         },
         # Start the Finch HTTP client for sending emails
         {Finch, name: Medicode.Finch}
-      ] ++ chromic_pdf_spec()
+      ] ++ chromic_pdf_spec(flame_parent)
 
     # Start a worker by calling: Medicode.Worker.start_link(arg)
     # {Medicode.Worker, arg},
     # Start to serve requests, typically the last entry
-    children = children ++ [MedicodeWeb.Endpoint]
+
+    # Since we don't need Endpoint running within FLAME nodes, we skip
+    # starting it under that condition.
+    children =
+      if flame_parent do
+        children
+      else
+        children ++ [MedicodeWeb.Endpoint]
+      end
+
+    children = Enum.filter(children, & &1)
 
     # See https://hexdocs.pm/elixir/Supervisor.html
     # for other strategies and supported options
@@ -51,8 +64,8 @@ defmodule Medicode.Application do
     :ok
   end
 
-  defp chromic_pdf_spec do
-    if Application.get_env(:medicode, :skip_chromic_pdf, false) do
+  defp chromic_pdf_spec(flame_parent) do
+    if flame_parent || Application.get_env(:medicode, :skip_chromic_pdf, false) do
       []
     else
       # NOTE: Security implications of using Chrome:

lib/medicode/transcription_server.ex

@@ -155,20 +155,23 @@ defmodule Medicode.TranscriptionServer do
   #   complete sentences based on punctuation. We may want to suggest codes for the entire audio as a single piece as
   #   well
   defp stream_transcription_and_search(audio_file_path) do
-    # audio transcription + semantic search
-    Medicode.TranscriptionServing
-    |> Nx.Serving.batched_run({:file, audio_file_path})
-    |> Enum.each(fn chunk ->
-      result = %{
-        start_mark: format_timestamp(chunk.start_timestamp_seconds),
-        end_mark: format_timestamp(chunk.end_timestamp_seconds),
-        text: chunk.text
-      }
-
-      send(self(), {:chunk, result})
+    target_pid = self()
+
+    FLAME.call(Medicode.OfflineServing, fn ->
+      Medicode.TranscriptionServing
+      |> Nx.Serving.batched_run({:file, audio_file_path})
+      |> Enum.each(fn chunk ->
+        result = %{
+          start_mark: format_timestamp(chunk.start_timestamp_seconds),
+          end_mark: format_timestamp(chunk.end_timestamp_seconds),
+          text: chunk.text
+        }
+
+        send(target_pid, {:chunk, result})
+      end)
+
+      send(target_pid, :finished)
     end)
-
-    send(self(), :finished)
   end
 
   defp format_timestamp(seconds) when is_nil(seconds), do: nil

mix.exs

@@ -78,7 +78,8 @@ defmodule Medicode.MixProject do
       # NOTE: `override: true` avoids an issue that likely originates
       # from `audio_tagger` depending on a different version of bumblebee:
       # https://elixirforum.com/t/getting-a-weird-error-message-in-mix-ive-not-see-before-because-your-app-depends-on-xxx-empty/54475/8
-      {:bumblebee, "~> 0.5.3", override: true}
+      {:bumblebee, "~> 0.5.3", override: true},
+      {:flame, "~> 0.1.12"}
     ]
   end
 

mix.lock

@@ -31,6 +31,7 @@
   "expo": {:hex, :expo, "0.5.2", "beba786aab8e3c5431813d7a44b828e7b922bfa431d6bfbada0904535342efe2", [:mix], [], "hexpm", "8c9bfa06ca017c9cb4020fabe980bc7fdb1aaec059fd004c2ab3bff03b1c599c"},
   "file_system": {:hex, :file_system, "1.0.0", "b689cc7dcee665f774de94b5a832e578bd7963c8e637ef940cd44327db7de2cd", [:mix], [], "hexpm", "6752092d66aec5a10e662aefeed8ddb9531d79db0bc145bb8c40325ca1d8536d"},
   "finch": {:hex, :finch, "0.18.0", "944ac7d34d0bd2ac8998f79f7a811b21d87d911e77a786bc5810adb75632ada4", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:mint, "~> 1.3", [hex: :mint, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 0.2.6 or ~> 1.0", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "69f5045b042e531e53edc2574f15e25e735b522c37e2ddb766e15b979e03aa65"},
+  "flame": {:hex, :flame, "0.1.12", "4d46b706d35d6eb22505d0e060fe41174052eaa38f778a4762fc74dd2c9df301", [:mix], [{:req, "~> 0.4.13", [hex: :req, repo: "hexpm", optional: false]}], "hexpm", "aa25de6614455ac01e33409c08db4560fa54dd837ad116aae16c8f7c011ccd76"},
   "floki": {:hex, :floki, "0.36.1", "712b7f2ba19a4d5a47dfe3e74d81876c95bbcbee44fe551f0af3d2a388abb3da", [:mix], [], "hexpm", "21ba57abb8204bcc70c439b423fc0dd9f0286de67dc82773a14b0200ada0995f"},
   "fss": {:hex, :fss, "0.1.1", "9db2344dbbb5d555ce442ac7c2f82dd975b605b50d169314a20f08ed21e08642", [:mix], [], "hexpm", "78ad5955c7919c3764065b21144913df7515d52e228c09427a004afe9c1a16b0"},
   "gettext": {:hex, :gettext, "0.24.0", "6f4d90ac5f3111673cbefc4ebee96fe5f37a114861ab8c7b7d5b30a1108ce6d8", [:mix], [{:expo, "~> 0.5.1", [hex: :expo, repo: "hexpm", optional: false]}], "hexpm", "bdf75cdfcbe9e4622dd18e034b227d77dd17f0f133853a1c73b97b3d6c770e8b"},