Serially setup apps before deploying from directory (#2115)

Co-authored-by: José Valim <jose.valim@dashbit.co>

README.md

@@ -172,9 +172,10 @@ The following environment variables can be used to configure Livebook on boot:
     Livebook instance within the cloud provider platform.
 
   * LIVEBOOK_APPS_PATH - the directory with app notebooks. When set, the apps
-    are deployed on Livebook startup with the persisted settings.
+    are deployed on Livebook startup with the persisted settings. Password-protected
-    Password-protected notebooks will receive a random password,
+    notebooks will receive a random password, unless LIVEBOOK_APPS_PATH_PASSWORD
-    unless LIVEBOOK_APPS_PATH_PASSWORD is set.
+    is set. When deploying using Livebook's Docker image, consider using
+    `LIVEBOOK_APPS_PATH_WARMUP`.
 
   * LIVEBOOK_APPS_PATH_HUB_ID - deploy only the notebooks in
     LIVEBOOK_APPS_PATH that belong to the given Hub ID
@@ -182,6 +183,12 @@ The following environment variables can be used to configure Livebook on boot:
   * LIVEBOOK_APPS_PATH_PASSWORD - the password to use for all protected apps
     deployed from LIVEBOOK_APPS_PATH.
 
+  * LIVEBOOK_APPS_PATH_WARMUP - sets the warmup mode for apps deployed from
+    LIVEBOOK_APPS_PATH. Must be either "auto" (apps are warmed up on Livebook
+    startup, right before app deployment) or "manual" (apps are warmed up when
+    building the Docker image; to do so add "RUN /app/bin/warmup_apps.sh" to
+    your image). Defaults to "auto".
+
   * LIVEBOOK_BASE_URL_PATH - sets the base url path the web application is
     served on. Useful when deploying behind a reverse proxy.
 

lib/livebook.ex

@@ -168,6 +168,10 @@ defmodule Livebook do
       config :livebook, :apps_path_password, apps_path_password
     end
 
+    if apps_path_warmup = Livebook.Config.apps_path_warmup!("LIVEBOOK_APPS_PATH_WARMUP") do
+      config :livebook, :apps_path_warmup, apps_path_warmup
+    end
+
     if force_ssl_host = Livebook.Config.force_ssl_host!("LIVEBOOK_FORCE_SSL_HOST") do
       config :livebook, :force_ssl_host, force_ssl_host
     end

lib/livebook/application.ex

@@ -43,6 +43,9 @@ defmodule Livebook.Application do
         # Start the supervisor dynamically managing connections
         {DynamicSupervisor, name: Livebook.HubsSupervisor, strategy: :one_for_one}
       ] ++
+        if serverless?() do
+          []
+        else
           iframe_server_specs() ++
             identity_provider() ++
             [
@@ -50,6 +53,7 @@ defmodule Livebook.Application do
               # We skip the access url as we do our own logging below
               {LivebookWeb.Endpoint, log_access_url: false}
             ] ++ app_specs()
+        end
 
     opts = [strategy: :one_for_one, name: Livebook.Supervisor]
 
@@ -61,7 +65,11 @@ defmodule Livebook.Application do
         clear_env_vars()
         display_startup_info()
         Livebook.Hubs.connect_hubs()
+
+        unless serverless?() do
           deploy_apps()
+        end
+
         result
 
       {:error, error} ->
@@ -179,7 +187,8 @@ defmodule Livebook.Application do
   end
 
   defp display_startup_info() do
-    if Phoenix.Endpoint.server?(:livebook, LivebookWeb.Endpoint) do
+    if Process.whereis(LivebookWeb.Endpoint) &&
+         Phoenix.Endpoint.server?(:livebook, LivebookWeb.Endpoint) do
       IO.puts("[Livebook] Application running at #{LivebookWeb.Endpoint.access_url()}")
     end
   end
@@ -239,7 +248,12 @@ defmodule Livebook.Application do
 
   defp deploy_apps() do
     if apps_path = Livebook.Config.apps_path() do
-      Livebook.Apps.deploy_apps_in_dir(apps_path, password: Livebook.Config.apps_path_password())
+      warmup = Livebook.Config.apps_path_warmup() == :auto
+
+      Livebook.Apps.deploy_apps_in_dir(apps_path,
+        password: Livebook.Config.apps_path_password(),
+        warmup: warmup
+      )
     end
   end
 
@@ -294,4 +308,8 @@ defmodule Livebook.Application do
     {module, key} = Livebook.Config.identity_provider()
     [{module, name: LivebookWeb.ZTA, identity: [key: key]}]
   end
+
+  defp serverless?() do
+    Application.get_env(:livebook, :serverless, false)
+  end
 end

lib/livebook/apps.ex

@@ -160,32 +160,58 @@ defmodule Livebook.Apps do
 
     * `:password` - a password to set for every loaded app
 
+    * `:warmup` - when `true`, run setup cell for each of the
+      notebooks before the actual deployment. The setup cells are
+      run one by one to avoid race conditions. Defaults to `true`
+
+    * `:skip_deploy` - when `true`, the apps are not deployed.
+      This can be used to warmup apps without deployment. Defaults
+      to `false`
+
   """
   @spec deploy_apps_in_dir(String.t(), keyword()) :: :ok
   def deploy_apps_in_dir(path, opts \\ []) do
-    opts = Keyword.validate!(opts, [:password])
+    opts = Keyword.validate!(opts, [:password, warmup: true, skip_deploy: false])
 
-    pattern = Path.join([path, "**", "*.livemd"])
+    infos = import_app_notebooks(path)
-    paths = Path.wildcard(pattern)
 
-    if paths == [] do
+    if infos == [] do
       Logger.warning("No .livemd files were found for deployment at #{path}")
     end
 
-    for path <- paths do
+    for %{status: {:error, message}} = info <- infos do
-      markdown = File.read!(path)
-
-      {notebook, %{warnings: warnings, verified_hub_id: verified_hub_id}} =
-        Livebook.LiveMarkdown.notebook_from_livemd(markdown)
-
-      if warnings != [] do
-        items = Enum.map(warnings, &("- " <> &1))
-
       Logger.warning(
-          "Found warnings while importing app notebook at #{path}:\n\n" <> Enum.join(items, "\n")
+        "Skipping deployment for app at #{info.relative_path}. #{Livebook.Utils.upcase_first(message)}."
       )
     end
 
+    infos = Enum.filter(infos, &(&1.status == :ok))
+
+    for info <- infos, info.import_warnings != [] do
+      items = Enum.map(info.import_warnings, &("- " <> &1))
+
+      Logger.warning(
+        "Found warnings while importing app notebook at #{info.relative_path}:\n\n" <>
+          Enum.join(items, "\n")
+      )
+    end
+
+    if infos != [] and opts[:warmup] do
+      Logger.info("Running app warmups")
+
+      for info <- infos do
+        with {:error, message} <- run_app_setup_sync(info.notebook, info.files_source) do
+          Logger.warning(
+            "Failed to run setup for app at #{info.relative_path}. #{Livebook.Utils.upcase_first(message)}."
+          )
+        end
+      end
+    end
+
+    if infos != [] and not opts[:skip_deploy] do
+      Logger.info("Deploying apps")
+
+      for %{notebook: notebook} = info <- infos do
         notebook =
           if password = opts[:password] do
             put_in(notebook.app_settings.password, password)
@@ -193,29 +219,86 @@ defmodule Livebook.Apps do
             notebook
           end
 
-      if Livebook.Notebook.AppSettings.valid?(notebook.app_settings) do
+        warnings = Enum.map(info.import_warnings, &("Import: " <> &1))
-        warnings = Enum.map(warnings, &("Import: " <> &1))
-        apps_path_hub_id = Livebook.Config.apps_path_hub_id()
 
-        if apps_path_hub_id == nil or apps_path_hub_id == verified_hub_id do
+        {:ok, _} = deploy(notebook, warnings: warnings, files_source: info.files_source)
-          notebook_file = Livebook.FileSystem.File.local(path)
-          files_dir = Livebook.Session.files_dir_for_notebook(notebook_file)
-          deploy(notebook, warnings: warnings, files_source: {:dir, files_dir})
-        else
-          Logger.warning(
-            "Skipping app deployment at #{path}. The notebook is not verified to come from hub #{apps_path_hub_id}"
-          )
-        end
-      else
-        Logger.warning(
-          "Skipping app deployment at #{path}. The deployment settings are missing or invalid. Please configure them under the notebook deploy panel."
-        )
       end
     end
 
     :ok
   end
 
+  defp import_app_notebooks(dir) do
+    pattern = Path.join([dir, "**", "*.livemd"])
+
+    for path <- Path.wildcard(pattern) do
+      markdown = File.read!(path)
+
+      {notebook, %{warnings: warnings, verified_hub_id: verified_hub_id}} =
+        Livebook.LiveMarkdown.notebook_from_livemd(markdown)
+
+      apps_path_hub_id = Livebook.Config.apps_path_hub_id()
+
+      status =
+        cond do
+          not Livebook.Notebook.AppSettings.valid?(notebook.app_settings) ->
+            {:error,
+             "the deployment settings are missing or invalid. Please configure them under the notebook deploy panel"}
+
+          apps_path_hub_id && apps_path_hub_id != verified_hub_id ->
+            {:error, "the notebook is not verified to come from hub #{apps_path_hub_id}"}
+
+          true ->
+            :ok
+        end
+
+      notebook_file = Livebook.FileSystem.File.local(path)
+      files_dir = Livebook.Session.files_dir_for_notebook(notebook_file)
+
+      %{
+        relative_path: Path.relative_to(path, dir),
+        status: status,
+        notebook: notebook,
+        import_warnings: warnings,
+        files_source: {:dir, files_dir}
+      }
+    end
+  end
+
+  defp run_app_setup_sync(notebook, files_source) do
+    notebook = %{notebook | sections: []}
+
+    opts = [
+      notebook: notebook,
+      files_source: files_source,
+      mode: :app,
+      app_pid: self()
+    ]
+
+    case Livebook.Sessions.create_session(opts) do
+      {:ok, %{id: session_id} = session} ->
+        ref = Process.monitor(session.pid)
+
+        receive do
+          {:app_status_changed, ^session_id, status} ->
+            Process.demonitor(ref)
+            Livebook.Session.close(session.pid)
+
+            if status.execution == :executed do
+              :ok
+            else
+              {:error, "setup cell finished with failure"}
+            end
+
+          {:DOWN, ^ref, :process, _, reason} ->
+            {:error, "session terminated unexpectedly, reason: #{inspect(reason)}"}
+        end
+
+      {:error, reason} ->
+        {:error, "failed to start session, reason: #{inspect(reason)}"}
+    end
+  end
+
   @doc """
   Checks if the apps directory is configured and contains no notebooks.
   """

lib/livebook/config.ex

@@ -131,6 +131,14 @@ defmodule Livebook.Config do
     Application.get_env(:livebook, :apps_path_hub_id)
   end
 
+  @doc """
+  Returns warmup mode for apps deployed from dir.
+  """
+  @spec apps_path_warmup() :: :auto | :manual
+  def apps_path_warmup() do
+    Application.get_env(:livebook, :apps_path_warmup, :auto)
+  end
+
   @doc """
   Returns the configured port for the Livebook endpoint.
 
@@ -509,6 +517,31 @@ defmodule Livebook.Config do
     end
   end
 
+  @doc """
+  Parses and validates apps warmup mode from env.
+  """
+  def apps_path_warmup!(env) do
+    if warmup = System.get_env(env) do
+      apps_path_warmup!(env, warmup)
+    end
+  end
+
+  @doc """
+  Parses and validates apps warmup mode within context.
+  """
+  def apps_path_warmup!(context, warmup) do
+    case warmup do
+      "auto" ->
+        :auto
+
+      "manual" ->
+        :manual
+
+      other ->
+        abort!(~s{expected #{context} to be either "auto" or "manual", got: #{inspect(other)}})
+    end
+  end
+
   @doc """
   Parses and validates allowed URI schemes from env.
   """

lib/livebook/release.ex

@@ -0,0 +1,31 @@
+defmodule Livebook.Release do
+  @moduledoc false
+
+  @doc """
+  Runs the setup for all apps deployed from directory on startup.
+  """
+  def warmup_apps() do
+    start_app()
+
+    if apps_path = Livebook.Config.apps_path() do
+      case Livebook.Config.apps_path_warmup() do
+        :manual ->
+          :ok
+
+        other ->
+          Livebook.Config.abort!(
+            "expected apps warmup mode to be :manual, got: #{inspect(other)}." <>
+              " Make sure to set LIVEBOOK_APPS_PATH_WARMUP=manual"
+          )
+      end
+
+      Livebook.Apps.deploy_apps_in_dir(apps_path, warmup: true, skip_deploy: true)
+    end
+  end
+
+  defp start_app() do
+    Application.load(:livebook)
+    Application.put_env(:livebook, :serverless, true)
+    {:ok, _} = Application.ensure_all_started(:livebook)
+  end
+end

rel/server/overlays/bin/warmup_apps.bat

@@ -0,0 +1 @@
+call "%~dp0\livebook" eval Livebook.Release.warmup_apps

rel/server/overlays/bin/warmup_apps.sh

@@ -0,0 +1,4 @@
+#!/bin/sh
+
+cd -P -- "$(dirname -- "$0")"
+exec ./livebook eval Livebook.Release.warmup_apps

test/livebook/apps_test.exs

@@ -111,7 +111,7 @@ defmodule Livebook.AppsTest do
       assert capture_log(fn ->
                Livebook.Apps.deploy_apps_in_dir(tmp_dir)
              end) =~
-               "Skipping app deployment at #{app_path}. The deployment settings are missing or invalid. Please configure them under the notebook deploy panel."
+               "Skipping deployment for app at app.livemd. The deployment settings are missing or invalid. Please configure them under the notebook deploy panel."
     end
 
     @tag :tmp_dir