Support persisting Fly runtime config in a workspace secret (#2714)

lib/livebook_web/live/hub/secret_list_component.ex

@@ -75,10 +75,16 @@ defmodule LivebookWeb.Hub.SecretListComponent do
       end
     end
 
+    assigns = %{name: attrs["name"]}
+
+    description = ~H"""
+    Are you sure you want to delete this secret - <span class="font-semibold"><%= @name %></span>?
+    """
+
     {:noreply,
      confirm(socket, on_confirm,
-       title: "Delete secret - #{attrs["name"]}",
-       description: "Are you sure you want to delete this secret?",
+       title: "Delete secret",
+       description: description,
        confirm_text: "Delete",
        confirm_icon: "delete-bin-6-line"
      )}

lib/livebook_web/live/session_live/fly_runtime_component.ex

@@ -5,6 +5,8 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
 
   alias Livebook.{Session, Runtime}
 
+  @config_secret_prefix "FLY_RUNTIME_"
+
   @impl true
   def mount(socket) do
     unless Livebook.Config.runtime_enabled?(Livebook.Runtime.Fly) do
@@ -21,31 +23,38 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
        app_check: %{status: :initial, error: nil},
        volumes: nil,
        region: nil,
-       specs_changeset: specs_changeset(%{}),
+       specs_changeset: specs_changeset(),
        volume_id: nil,
-       volume_action: nil
+       volume_action: nil,
+       save_config: nil
      )}
   end
 
   @impl true
   def update(assigns, socket) do
-    socket =
-      case assigns.runtime do
-        %Runtime.Fly{config: config} when not is_map_key(socket.assigns, :runtime) ->
-          assign(socket,
-            token: config.token,
-            app_name: config.app_name,
-            specs_changeset: specs_changeset(config)
-          )
-          |> load_org_and_regions()
-          |> load_app()
-
-        _ ->
-          socket
-      end
-
     socket = assign(socket, assigns)
 
+    socket =
+      cond do
+        is_map_key(socket.assigns, :config_defaults) ->
+          socket
+
+        is_struct(assigns.runtime, Runtime.Fly) ->
+          %{config: config} = assigns.runtime
+
+          config_defaults =
+            Map.new(config, fn {key, value} ->
+              {Atom.to_string(key), value}
+            end)
+
+          socket
+          |> assign(config_defaults: config_defaults)
+          |> load_config_defaults()
+
+        true ->
+          assign(socket, config_defaults: nil)
+      end
+
     {:ok, socket}
   end
 
@@ -58,72 +67,83 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
         The machine is automatically destroyed, once you disconnect the runtime.
       </p>
 
-      <form class="mt-4 flex flex-col gap-4" phx-change="set_token" phx-nosubmit phx-target={@myself}>
-        <.password_field name="token" value={@token} label="Token" />
-        <.message_box :if={@token == nil} kind={:info}>
-          Go to <a
-            class="text-blue-600 hover:text-blue-700"
-            href="https://fly.io/dashboard"
-            phx-no-format
-          >Fly dashboard</a>, click "Tokens" in the left sidebar and create a new
-          token for your organization of choice. This functionality is restricted
-          to organization admins. Alternatively, you can create an app in the
-          organization by running <code>fly app create</code>
-          and generate a deploy token in
-          the app dashboard.
-        </.message_box>
-        <.loader :if={@token_check.status == :inflight} />
-        <.message_box
-          :if={error = @token_check.error}
-          kind={:error}
-          message={"Error: " <> error.message}
-        />
-      </form>
+      <.save_config_form :if={@save_config} save_config={@save_config} hub={@hub} myself={@myself} />
 
-      <.app_config
-        :if={@token_check.status == :ok}
-        org_name={@org.name}
-        regions={@regions}
-        app_name={@app_name}
-        app_check={@app_check}
-        volumes={@volumes}
-        region={@region}
-        myself={@myself}
-      />
+      <div :if={@save_config == nil}>
+        <.config_actions hub_secrets={@hub_secrets} myself={@myself} />
 
-      <div :if={@token_check.status == :ok and @app_check.status == :ok}>
-        <.specs_config specs_changeset={@specs_changeset} myself={@myself} />
+        <form
+          class="mt-1 flex flex-col gap-4"
+          phx-change="set_token"
+          phx-nosubmit
+          phx-target={@myself}
+        >
+          <.password_field name="token" value={@token} label="Token" />
+          <.message_box :if={@token == nil} kind={:info}>
+            Go to <a
+              class="text-blue-600 hover:text-blue-700"
+              href="https://fly.io/dashboard"
+              phx-no-format
+            >Fly dashboard</a>, click "Tokens" in the left sidebar and create a new
+            token for your organization of choice. This functionality is restricted
+            to organization admins. Alternatively, you can create an app in the
+            organization by running <code>fly app create</code>
+            and generate a deploy token in
+            the app dashboard.
+          </.message_box>
+          <.loader :if={@token_check.status == :inflight} />
+          <.message_box
+            :if={error = @token_check.error}
+            kind={:error}
+            message={"Error: " <> error.message}
+          />
+        </form>
 
-        <.storage_config
+        <.app_config
+          :if={@token_check.status == :ok}
+          org_name={@org.name}
+          regions={@regions}
+          app_name={@app_name}
+          app_check={@app_check}
           volumes={@volumes}
-          volume_id={@volume_id}
           region={@region}
-          volume_action={@volume_action}
           myself={@myself}
         />
 
-        <div class="mt-8">
-          <.button
-            phx-click="init"
-            phx-target={@myself}
-            disabled={
-              @runtime_status == :connecting or not @specs_changeset.valid? or
-                volume_errors(@volume_id, @volumes, @region) != []
-            }
-          >
-            <%= label(@app_name, @runtime, @runtime_status) %>
-          </.button>
-          <div
-            :if={reconnecting?(@app_name, @runtime) && @runtime_connect_info}
-            class="mt-4 scroll-mb-8"
-            phx-mounted={JS.dispatch("lb:scroll_into_view", detail: %{behavior: "instant"})}
-          >
-            <.message_box kind={:info}>
-              <div class="flex items-center gap-2">
-                <.spinner />
-                <span>Step: <%= @runtime_connect_info %></span>
-              </div>
-            </.message_box>
+        <div :if={@token_check.status == :ok and @app_check.status == :ok}>
+          <.specs_config specs_changeset={@specs_changeset} myself={@myself} />
+
+          <.storage_config
+            volumes={@volumes}
+            volume_id={@volume_id}
+            region={@region}
+            volume_action={@volume_action}
+            myself={@myself}
+          />
+
+          <div class="mt-8">
+            <.button
+              phx-click="init"
+              phx-target={@myself}
+              disabled={
+                @runtime_status == :connecting or not @specs_changeset.valid? or
+                  volume_errors(@volume_id, @volumes, @region) != []
+              }
+            >
+              <%= label(@app_name, @runtime, @runtime_status) %>
+            </.button>
+            <div
+              :if={reconnecting?(@app_name, @runtime) && @runtime_connect_info}
+              class="mt-4 scroll-mb-8"
+              phx-mounted={JS.dispatch("lb:scroll_into_view", detail: %{behavior: "instant"})}
+            >
+              <.message_box kind={:info}>
+                <div class="flex items-center gap-2">
+                  <.spinner />
+                  <span>Step: <%= @runtime_connect_info %></span>
+                </div>
+              </.message_box>
+            </div>
           </div>
         </div>
       </div>
@@ -131,6 +151,99 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
     """
   end
 
+  defp save_config_form(assigns) do
+    ~H"""
+    <.form
+      :let={f}
+      for={@save_config.changeset}
+      as={:secret}
+      class="mt-4 flex flex-col"
+      phx-change="validate_save_config"
+      phx-submit="save_config"
+      phx-target={@myself}
+      autocomplete="off"
+      spellcheck="false"
+    >
+      <div class="text-lg text-gray-800 font-semibold">
+        Save config
+      </div>
+      <div class="mt-1 text-gray-700">
+        Store the config in a secret in the <.workspace hub={@hub} /> workspace to reuse it later.
+      </div>
+      <div :if={error = @save_config.error} class="mt-4">
+        <.message_box kind={:error} message={error} />
+      </div>
+      <div class="mt-4 grid grid-cols-3">
+        <.text_field field={f[:name]} label="Secret name" class="uppercase" autofocus />
+      </div>
+      <div class="mt-6 flex gap-2">
+        <.button type="submit" disabled={not @save_config.changeset.valid? or @save_config.inflight}>
+          <%= if(@save_config.inflight, do: "Saving...", else: "Save") %>
+        </.button>
+        <.button
+          color="gray"
+          outlined
+          type="button"
+          phx-click="cancel_save_config"
+          phx-target={@myself}
+        >
+          Cancel
+        </.button>
+      </div>
+    </.form>
+    """
+  end
+
+  defp workspace(assigns) do
+    ~H"""
+    <span class="font-medium">
+      <span class="text-lg"><%= @hub.hub_emoji %></span>
+      <span><%= @hub.hub_name %></span>
+    </span>
+    """
+  end
+
+  defp config_actions(assigns) do
+    ~H"""
+    <div class="mt-1 flex justify-end gap-1">
+      <.button
+        color="gray"
+        outlined
+        small
+        type="button"
+        phx-click="open_save_config"
+        phx-target={@myself}
+      >
+        Save config
+      </.button>
+      <.menu id="config-secret-menu">
+        <:toggle>
+          <.button color="gray" outlined small type="button">
+            <span>Load config</span>
+            <.remix_icon icon="arrow-down-s-line" class="text-base leading-none" />
+          </.button>
+        </:toggle>
+        <div
+          :if={config_secret_names(@hub_secrets) == []}
+          class="px-3 py-1 whitespace-nowrap text-gray-600 text-sm"
+        >
+          No configs saved yet
+        </div>
+        <.menu_item :for={name <- config_secret_names(@hub_secrets)}>
+          <button
+            class="text-gray-500 text-sm"
+            type="button"
+            role="menuitem"
+            phx-click={JS.push("load_config", value: %{name: name}, target: @myself)}
+          >
+            <%= name %>
+          </button>
+        </.menu_item>
+      </.menu>
+    </div>
+    """
+  end
+
   defp loader(assigns) do
     ~H"""
     <div class="flex items-center gap-2">
@@ -351,6 +464,12 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
             <.text_field field={f[:name]} placeholder="Name" />
             <.text_field field={f[:size_gb]} placeholder="Size (GB)" type="number" min="1" />
           </div>
+          <.button
+            type="submit"
+            disabled={not @volume_action.changeset.valid? or @volume_action.inflight}
+          >
+            <%= if(@volume_action.inflight, do: "Creating...", else: "Create") %>
+          </.button>
           <.button
             type="button"
             color="gray"
@@ -360,12 +479,6 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
           >
             Cancel
           </.button>
-          <.button
-            type="submit"
-            disabled={not @volume_action.changeset.valid? or @volume_action.inflight}
-          >
-            <%= if(@volume_action.inflight, do: "Creating...", else: "Create") %>
-          </.button>
         </.form>
         <div :if={error = @volume_action[:error]}>
           <.message_box kind={:error} message={error} />
@@ -442,38 +555,64 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
 
   def handle_event("validate_specs", %{"specs" => specs}, socket) do
     changeset =
-      socket.assigns.specs_changeset.data
-      |> specs_changeset(specs)
+      specs
+      |> specs_changeset()
       |> Map.replace!(:action, :validate)
 
     {:noreply, assign(socket, specs_changeset: changeset)}
   end
 
   def handle_event("init", %{}, socket) do
-    socket.assigns.specs_changeset
-    |> apply_action(:insert)
-    |> case do
-      {:ok, specs} ->
-        config = %{
-          token: socket.assigns.token,
-          app_name: socket.assigns.app_name,
-          region: socket.assigns.region,
-          cpu_kind: specs.cpu_kind,
-          cpus: specs.cpus,
-          memory_gb: specs.memory_gb,
-          gpu_kind: specs.gpu_kind,
-          gpus: specs.gpus,
-          volume_id: socket.assigns.volume_id,
-          docker_tag: specs.docker_tag
-        }
+    config = build_config(socket)
+    runtime = Runtime.Fly.new(config)
+    Session.set_runtime(socket.assigns.session.pid, runtime)
+    Session.connect_runtime(socket.assigns.session.pid)
+    {:noreply, socket}
+  end
 
-        runtime = Runtime.Fly.new(config)
-        Session.set_runtime(socket.assigns.session.pid, runtime)
-        Session.connect_runtime(socket.assigns.session.pid)
-        {:noreply, socket}
+  def handle_event("open_save_config", %{}, socket) do
+    changeset = config_secret_changeset(socket, %{name: @config_secret_prefix})
+    save_config = %{changeset: changeset, inflight: false, error: false}
+    {:noreply, assign(socket, save_config: save_config)}
+  end
+
+  def handle_event("cancel_save_config", %{}, socket) do
+    {:noreply, assign(socket, save_config: nil)}
+  end
+
+  def handle_event("validate_save_config", %{"secret" => secret}, socket) do
+    changeset =
+      socket
+      |> config_secret_changeset(secret)
+      |> Map.replace!(:action, :validate)
+
+    {:noreply, assign_nested(socket, :save_config, changeset: changeset)}
+  end
+
+  def handle_event("save_config", %{"secret" => secret}, socket) do
+    changeset = config_secret_changeset(socket, secret)
+
+    case Ecto.Changeset.apply_action(changeset, :insert) do
+      {:ok, secret} ->
+        {:noreply, save_config_secret(socket, secret, changeset)}
 
       {:error, changeset} ->
-        {:noreply, assign(socket, specs_changeset: changeset)}
+        {:noreply, assign_nested(socket, :save_config, changeset: changeset)}
+    end
+  end
+
+  def handle_event("load_config", %{"name" => name}, socket) do
+    secret = Enum.find(socket.assigns.hub_secrets, &(&1.name == name))
+
+    case Jason.decode(secret.value) do
+      {:ok, config_defaults} ->
+        {:noreply,
+         socket
+         |> assign(config_defaults: config_defaults)
+         |> load_config_defaults()}
+
+      {:error, _} ->
+        {:noreply, socket}
     end
   end
 
@@ -482,11 +621,7 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
     socket =
       case result do
         {:ok, %{orgs: [org]} = data} ->
-          region =
-            case socket.assigns.runtime do
-              %Runtime.Fly{config: config} -> config.region
-              _ -> data.closest_region
-            end
+          region = socket.assigns.config_defaults["region"] || data.closest_region
 
           socket
           |> assign(org: org, regions: data.regions, region: region)
@@ -510,13 +645,9 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
       case result do
         {:ok, volumes} ->
           volume_id =
-            case socket.assigns.runtime do
-              %Runtime.Fly{config: %{volume_id: volume_id}} ->
-                # Ignore the volume if it no longer exists
-                if Enum.any?(volumes, &(&1.id == volume_id)), do: volume_id
-
-              _ ->
-                nil
+            if volume_id = socket.assigns.config_defaults["volume_id"] do
+              # Ignore the volume if it no longer exists
+              if Enum.any?(volumes, &(&1.id == volume_id)), do: volume_id
             end
 
           socket
@@ -575,6 +706,22 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
     {:noreply, socket}
   end
 
+  def handle_async(:save_config, {:ok, result}, socket) do
+    socket =
+      case result do
+        :ok ->
+          assign(socket, save_config: nil)
+
+        {:error, %Ecto.Changeset{} = changeset} ->
+          assign_nested(socket, :save_config, changeset: changeset, inflight: false)
+
+        {:transport_error, error} ->
+          assign_nested(socket, :save_config, error: error, inflight: false)
+      end
+
+    {:noreply, socket}
+  end
+
   defp label(app_name, runtime, runtime_status) do
     reconnecting? = reconnecting?(app_name, runtime)
 
@@ -610,18 +757,40 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
         }
   end
 
-  defp specs_changeset(config, attrs \\ %{}) do
-    defaults = %{
+  defp config_secret_names(hub_secrets) do
+    names =
+      for %{name: name} <- hub_secrets,
+          String.starts_with?(name, @config_secret_prefix),
+          do: name
+
+    Enum.sort(names)
+  end
+
+  defp load_config_defaults(socket) do
+    config_defaults = socket.assigns.config_defaults
+
+    socket
+    |> assign(
+      token: config_defaults["token"],
+      app_name: config_defaults["app_name"],
+      specs_changeset: specs_changeset(config_defaults)
+    )
+    |> load_org_and_regions()
+    |> load_app()
+  end
+
+  defp specs_changeset(attrs \\ %{}) do
+    docker_tags = Enum.map(Livebook.Config.docker_images(), & &1.tag)
+
+    data = %{
       cpu_kind: "shared",
       cpus: 1,
       memory_gb: 1,
       gpu_kind: nil,
       gpus: nil,
-      docker_tag: Livebook.Config.docker_images() |> hd() |> Map.fetch!(:tag)
+      docker_tag: hd(docker_tags)
     }
 
-    data = for {key, default} <- defaults, into: %{}, do: {key, Map.get(config, key, default)}
-
     types = %{
       cpu_kind: :string,
       cpus: :integer,
@@ -634,6 +803,7 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
     changeset =
       cast({data, types}, attrs, Map.keys(types))
       |> validate_required([:cpu_kind, :cpus, :memory_gb, :docker_tag])
+      |> validate_inclusion(:docker_tag, docker_tags)
 
     if get_field(changeset, :gpu_kind) do
       changeset
@@ -656,6 +826,18 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
     |> validate_required([:name, :size_gb])
   end
 
+  defp config_secret_changeset(socket, attrs) do
+    hub = socket.assigns.hub
+    value = socket |> build_config() |> Jason.encode!()
+    secret = %Livebook.Secrets.Secret{hub_id: hub.id, name: nil, value: value}
+
+    secret
+    |> Livebook.Secrets.change_secret(attrs)
+    |> validate_format(:name, ~r/^#{@config_secret_prefix}\w+$/,
+      message: "must be in the format #{@config_secret_prefix}*"
+    )
+  end
+
   defp volume_errors(nil, _volumes, _region), do: []
 
   defp volume_errors(volume_id, volumes, region) do
@@ -737,12 +919,52 @@ defmodule LivebookWeb.SessionLive.FlyRuntimeComponent do
     |> assign_nested(:volume_action, inflight: true)
   end
 
+  defp save_config_secret(socket, secret, changeset) do
+    hub = socket.assigns.hub
+    exists? = Enum.any?(socket.assigns.hub_secrets, &(&1.name == secret.name))
+
+    socket
+    |> start_async(:save_config, fn ->
+      result =
+        if exists? do
+          Livebook.Hubs.update_secret(hub, secret)
+        else
+          Livebook.Hubs.create_secret(hub, secret)
+        end
+
+      with {:error, errors} <- result do
+        {:error,
+         changeset
+         |> Livebook.Utils.put_changeset_errors(errors)
+         |> Map.replace!(:action, :validate)}
+      end
+    end)
+    |> assign_nested(:save_config, inflight: true)
+  end
+
   defp assign_nested(socket, key, keyword) do
     update(socket, key, fn map ->
       Enum.reduce(keyword, map, fn {key, value}, map -> Map.replace!(map, key, value) end)
     end)
   end
 
+  defp build_config(socket) do
+    specs = apply_changes(socket.assigns.specs_changeset)
+
+    %{
+      token: socket.assigns.token,
+      app_name: socket.assigns.app_name,
+      region: socket.assigns.region,
+      cpu_kind: specs.cpu_kind,
+      cpus: specs.cpus,
+      memory_gb: specs.memory_gb,
+      gpu_kind: specs.gpu_kind,
+      gpus: specs.gpus,
+      volume_id: socket.assigns.volume_id,
+      docker_tag: specs.docker_tag
+    }
+  end
+
   defp nullify(""), do: nil
   defp nullify(value), do: value
 end

lib/livebook_web/live/session_live/render.ex

@@ -63,6 +63,8 @@ defmodule LivebookWeb.SessionLive.Render do
         runtime={@data_view.runtime}
         runtime_status={@data_view.runtime_status}
         runtime_connect_info={@data_view.runtime_connect_info}
+        hub={@data_view.hub}
+        hub_secrets={@data_view.hub_secrets}
       />
     </.modal>
 

lib/livebook_web/live/session_live/runtime_component.ex

@@ -82,6 +82,8 @@ defmodule LivebookWeb.SessionLive.RuntimeComponent do
             runtime={@runtime}
             runtime_status={@runtime_status}
             runtime_connect_info={@runtime_connect_info}
+            hub={@hub}
+            hub_secrets={@hub_secrets}
           />
         </div>
       </div>

test/livebook_web/live/session_live_test.exs

@@ -1166,6 +1166,109 @@ defmodule LivebookWeb.SessionLiveTest do
              |> element(~s/select[name="specs[cpu_kind]"] option[value="performance"][selected]/)
              |> has_element?()
     end
+
+    test "saving and loading config from secret", %{conn: conn, session: session} do
+      runtime =
+        Runtime.Fly.new(%{
+          token: "my-token",
+          app_name: "my-app",
+          region: "ams",
+          cpu_kind: "performance",
+          cpus: 1,
+          memory_gb: 1,
+          gpu_kind: nil,
+          gpus: nil,
+          volume_id: "vol_1",
+          docker_tag: "edge"
+        })
+
+      Session.set_runtime(session.pid, runtime)
+
+      Livebook.FlyAPI.stub(fn
+        conn when conn.method == "POST" ->
+          Req.Test.json(conn, %{
+            "data" => %{
+              "organizations" => %{
+                "nodes" => [
+                  %{
+                    "id" => "1",
+                    "name" => "Grumpy Cat",
+                    "rawSlug" => "grumpy-cat",
+                    "slug" => "personal"
+                  }
+                ]
+              },
+              "platform" => %{
+                "regions" => [
+                  %{"code" => "ams", "name" => "Amsterdam, Netherlands"},
+                  %{"code" => "fra", "name" => "Frankfurt, Germany"}
+                ],
+                "requestRegion" => "fra"
+              }
+            }
+          })
+
+        conn
+        when conn.method == "GET" and
+               conn.path_info == ["v1", "apps", "my-app", "volumes"] ->
+          Req.Test.json(conn, [
+            %{
+              "id" => "vol_1",
+              "name" => "new_volume",
+              "region" => "ams",
+              "size_gb" => 1,
+              "state" => "created"
+            }
+          ])
+      end)
+
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}/settings/runtime")
+
+      # The form is already filled with the runtime configuration, we
+      # just save it in a secret
+      view
+      |> element("button", "Save config")
+      |> render_click()
+
+      secret_name = "FLY_RUNTIME_#{System.unique_integer([:positive])}"
+
+      view
+      |> element(~s/form[phx-submit="save_config"]/)
+      |> render_submit(%{secret: %{name: secret_name}})
+
+      assert render_async(view) =~ "Load config"
+
+      # Set a different runtime, so there are no defaults
+      Session.set_runtime(session.pid, Runtime.Standalone.new())
+
+      # Open new runtime configuratino
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}/settings/runtime")
+
+      view
+      |> element("#runtime-settings-modal button", "Fly.io machine")
+      |> render_click()
+
+      refute render(view) =~ "CPU kind"
+
+      # Load the configuration from secret
+      view
+      |> element("#config-secret-menu-content button", secret_name)
+      |> render_click()
+
+      assert render_async(view) =~ "Grumpy Cat"
+
+      assert view
+             |> element(~s/select[name="region"] option[value="ams"][selected]/)
+             |> has_element?()
+
+      assert view
+             |> element(~s/select[name="volume_id"] option[value="vol_1"][selected]/)
+             |> has_element?()
+
+      assert view
+             |> element(~s/select[name="specs[cpu_kind]"] option[value="performance"][selected]/)
+             |> has_element?()
+    end
   end
 
   describe "persistence settings" do