Expand instructions for deploying a new agent server instance (#2561)

lib/livebook/config.ex

@@ -44,15 +44,20 @@ defmodule Livebook.Config do
   @doc """
   Returns docker images to be used when generating sample Dockerfiles.
   """
-  @spec docker_images() :: list(%{tag: String.t(), name: String.t(), env: keyword()})
+  @spec docker_images() ::
+          list(%{
+            tag: String.t(),
+            name: String.t(),
+            env: list({String.t(), String.t()})
+          })
   def docker_images() do
     version = app_version()
     base = if version =~ "dev", do: "latest", else: version
 
     [
       %{tag: base, name: "Livebook", env: []},
-      %{tag: "#{base}-cuda11.8", name: "Livebook + CUDA 11.8", env: [XLA_TARGET: "cuda118"]},
+      %{tag: "#{base}-cuda11.8", name: "Livebook + CUDA 11.8", env: [{"XLA_TARGET", "cuda118"}]},
-      %{tag: "#{base}-cuda12.1", name: "Livebook + CUDA 12.1", env: [XLA_TARGET: "cuda120"]}
+      %{tag: "#{base}-cuda12.1", name: "Livebook + CUDA 12.1", env: [{"XLA_TARGET", "cuda120"}]}
     ]
   end
 

lib/livebook/hubs/dockerfile.ex

@@ -29,6 +29,19 @@ defmodule Livebook.Hubs.Dockerfile do
     }
   end
 
+  @doc """
+  Builds Dockerfile configuration with defaults from deployment group.
+  """
+  @spec from_deployment_group(Livebook.Teams.DeploymentGroup.t()) :: config()
+  def from_deployment_group(deployment_group) do
+    %{
+      config_new()
+      | clustering: deployment_group.clustering,
+        zta_provider: deployment_group.zta_provider,
+        zta_key: deployment_group.zta_key
+    }
+  end
+
   @doc """
   Builds a changeset for app Dockerfile configuration.
   """
@@ -53,7 +66,7 @@ defmodule Livebook.Hubs.Dockerfile do
   @doc """
   Builds Dockerfile definition for app deployment.
   """
-  @spec build_dockerfile(
+  @spec airgapped_dockerfile(
           config(),
           Hubs.Provider.t(),
           list(Livebook.Secrets.Secret.t()),
@@ -62,7 +75,15 @@ defmodule Livebook.Hubs.Dockerfile do
           list(Livebook.Notebook.file_entry()),
           Livebook.Session.Data.secrets()
         ) :: String.t()
-  def build_dockerfile(config, hub, hub_secrets, hub_file_systems, file, file_entries, secrets) do
+  def airgapped_dockerfile(
+        config,
+        hub,
+        hub_secrets,
+        hub_file_systems,
+        file,
+        file_entries,
+        secrets
+      ) do
     base_image = Enum.find(Livebook.Config.docker_images(), &(&1.tag == config.docker_tag))
 
     image = """
@@ -121,8 +142,13 @@ defmodule Livebook.Hubs.Dockerfile do
     RUN /app/bin/warmup_apps
     """
 
-    random_secret_key_base = Livebook.Utils.random_secret_key_base()
+    secret_key =
-    random_cookie = Livebook.Utils.random_cookie()
+      case hub_type do
+        "team" -> hub.teams_key
+        "personal" -> hub.secret_key
+      end
+
+    {secret_key_base, cookie} = deterministic_skb_and_cookie(secret_key)
 
     startup =
       if config.clustering == :fly_io do
@@ -130,8 +156,8 @@ defmodule Livebook.Hubs.Dockerfile do
         # --- Clustering ---
 
         # Set the same Livebook secrets across all nodes
-        ENV LIVEBOOK_SECRET_KEY_BASE "#{random_secret_key_base}"
+        ENV LIVEBOOK_SECRET_KEY_BASE "#{secret_key_base}"
-        ENV LIVEBOOK_COOKIE "#{random_cookie}"
+        ENV LIVEBOOK_COOKIE "#{cookie}"
         ENV LIVEBOOK_CLUSTER "fly"
         """
       end
@@ -149,33 +175,16 @@ defmodule Livebook.Hubs.Dockerfile do
     |> Enum.join("\n")
   end
 
-  @doc """
+  defp deterministic_skb_and_cookie(secret_key) do
-  Builds Dockerfile definition for Livebook Agent app deployment.
+    hash = :crypto.hash(:sha256, secret_key)
-  """
-  @spec build_agent_dockerfile(config(), Hubs.Provider.t()) :: String.t()
-  def build_agent_dockerfile(config, hub) do
-    base_image = Enum.find(Livebook.Config.docker_images(), &(&1.tag == config.docker_tag))
 
-    image = """
+    <<left::48-binary, right::39-binary>> =
-    FROM ghcr.io/livebook-dev/livebook:#{base_image.tag}
+      Plug.Crypto.KeyGenerator.generate(hash, "dockerfile",
-    """
+        cache: Plug.Crypto.Keys,
+        length: 48 + 39
+      )
 
-    image_envs = format_envs(base_image.env)
+    {Base.url_encode64(left, padding: false), "c_" <> Base.url_encode64(right, padding: false)}
-
-    hub_config = """
-    # Teams Hub configuration for Livebook Agent deployment
-    ENV LIVEBOOK_AGENT_NAME ""
-    ENV LIVEBOOK_TEAMS_KEY "#{hub.teams_key}"
-    ENV LIVEBOOK_TEAMS_AUTH "online:#{hub.hub_name}:#{hub.org_id}:#{hub.org_key_id}:${LIVEBOOK_AGENT_KEY}"
-    """
-
-    [
-      image,
-      image_envs,
-      hub_config
-    ]
-    |> Enum.reject(&is_nil/1)
-    |> Enum.join("\n")
   end
 
   defp format_hub_config("team", config, hub, hub_file_systems, used_secrets) do
@@ -289,12 +298,54 @@ defmodule Livebook.Hubs.Dockerfile do
     config.zta_provider != nil and config.zta_key != nil
   end
 
+  @doc """
+  Returns information for deploying Livebook Agent using Docker.
+  """
+  @spec online_docker_info(config(), Hubs.Provider.t(), Livebook.Teams.AgentKey.t()) :: %{
+          image: String.t(),
+          env: list({String.t(), String.t()})
+        }
+  def online_docker_info(config, %Hubs.Team{} = hub, agent_key) do
+    base_image = Enum.find(Livebook.Config.docker_images(), &(&1.tag == config.docker_tag))
+
+    image = "ghcr.io/livebook-dev/livebook:#{base_image.tag}"
+
+    env = [
+      {"LIVEBOOK_AGENT_NAME", "default"},
+      {"LIVEBOOK_TEAMS_KEY", "#{hub.teams_key}"},
+      {"LIVEBOOK_TEAMS_AUTH",
+       "online:#{hub.hub_name}:#{hub.org_id}:#{hub.org_key_id}:#{agent_key.key}"}
+    ]
+
+    hub_env =
+      if zta_configured?(config) do
+        [{"LIVEBOOK_IDENTITY_PROVIDER", "#{config.zta_provider}:#{config.zta_key}"}]
+      else
+        []
+      end
+
+    clustering_env =
+      if config.clustering == :fly_io do
+        {secret_key_base, cookie} = deterministic_skb_and_cookie(hub.teams_key)
+
+        [
+          {"LIVEBOOK_CLUSTER", "fly"},
+          {"LIVEBOOK_SECRET_KEY_BASE", secret_key_base},
+          {"LIVEBOOK_COOKIE", cookie}
+        ]
+      else
+        []
+      end
+
+    %{image: image, env: base_image.env ++ env ++ hub_env ++ clustering_env}
+  end
+
   @doc """
   Returns a list of Dockerfile-related warnings.
 
   The returned messages may include HTML.
   """
-  @spec warnings(
+  @spec airgapped_warnings(
           config(),
           Hubs.Provider.t(),
           list(Livebook.Secrets.Secret.t()),
@@ -303,14 +354,22 @@ defmodule Livebook.Hubs.Dockerfile do
           list(Livebook.Notebook.file_entry()),
           Livebook.Session.Data.secrets()
         ) :: list(String.t())
-  def warnings(config, hub, hub_secrets, hub_file_systems, app_settings, file_entries, secrets) do
+  def airgapped_warnings(
+        config,
+        hub,
+        hub_secrets,
+        hub_file_systems,
+        app_settings,
+        file_entries,
+        secrets
+      ) do
     common_warnings =
       [
         if Livebook.Session.Data.session_secrets(secrets, hub.id) != [] do
           "The notebook uses session secrets, but those are not available to deployed apps." <>
             " Convert them to Hub secrets instead."
         end
-      ]
+      ] ++ config_warnings(config)
 
     hub_warnings =
       case Hubs.Provider.type(hub) do
@@ -354,4 +413,23 @@ defmodule Livebook.Hubs.Dockerfile do
 
     Enum.reject(common_warnings ++ hub_warnings, &is_nil/1)
   end
+
+  defp config_warnings(config) do
+    [
+      if config.clustering == nil do
+        "The deployment is not configured for clustering. Make sure to run only one instance" <>
+          " of Livebook, or configure clustering."
+      end
+    ]
+  end
+
+  @doc """
+  Returns warnings specific to agent Docker deployment.
+  """
+  @spec online_warnings(config()) :: list(String.t())
+  def online_warnings(config) do
+    warnings = config_warnings(config)
+
+    Enum.reject(warnings, &is_nil/1)
+  end
 end

lib/livebook/hubs/team_client.ex

@@ -405,7 +405,7 @@ defmodule Livebook.Hubs.TeamClient do
       | name: deployment_group_updated.name,
         secrets: secrets,
         agent_keys: agent_keys,
-        clustering: nullify(deployment_group_updated.clustering),
+        clustering: atomize(deployment_group_updated.clustering),
         zta_provider: atomize(deployment_group_updated.zta_provider),
         zta_key: nullify(deployment_group_updated.zta_key)
     }

lib/livebook_web/components/app_components.ex

@@ -91,7 +91,7 @@ defmodule LivebookWeb.AppComponents do
 
   def deployment_group_form_content(assigns) do
     ~H"""
-    <div class="grid grid-cols-1 md:grid-cols-2">
+    <div class="grid grid-cols-1 md:grid-cols-2 gap-3">
       <.select_field
         label="Clustering"
         help={

lib/livebook_web/components/core_components.ex

@@ -686,7 +686,7 @@ defmodule LivebookWeb.CoreComponents do
     """
   end
 
-  @doc ~S"""
+  @doc """
   Renders a table with generic styling.
 
   ## Examples
@@ -758,7 +758,7 @@ defmodule LivebookWeb.CoreComponents do
     """
   end
 
-  @doc ~S"""
+  @doc """
   Renders a button.
 
   ## Examples
@@ -836,7 +836,7 @@ defmodule LivebookWeb.CoreComponents do
     ]
   end
 
-  @doc ~S"""
+  @doc """
   Renders an icon button.
 
   ## Examples
@@ -886,6 +886,57 @@ defmodule LivebookWeb.CoreComponents do
     ]
   end
 
+  @doc """
+  Renders stateful tabs with content.
+
+  ## Examples
+
+      <.tabs id="animals" default="cat">
+        <:tab id="cat" label="Cat">
+          This is a cat.
+        </:tab>
+        <:tab id="dog" label="Dog">
+          This is a dog.
+        </:tab>
+      </.tabs>
+
+  """
+
+  attr :id, :string, required: true
+  attr :default, :string, required: true
+
+  slot :tab do
+    attr :id, :string, required: true
+    attr :label, :string, required: true
+  end
+
+  def tabs(assigns) do
+    ~H"""
+    <div id={@id} class="flex flex-col gap-4">
+      <div class="tabs">
+        <button
+          :for={tab <- @tab}
+          class={["tab", @default == tab.id && "active"]}
+          phx-click={
+            JS.remove_class("active", to: "##{@id} .tab.active")
+            |> JS.add_class("active")
+            |> JS.add_class("hidden", to: "##{@id} [data-tab]")
+            |> JS.remove_class("hidden", to: "##{@id} [data-tab='#{tab.id}']")
+          }
+        >
+          <span class="font-medium">
+            <%= tab.label %>
+          </span>
+        </button>
+      </div>
+
+      <div :for={tab <- @tab} data-tab={tab.id} class={@default == tab.id || "hidden"}>
+        <%= render_slot(tab) %>
+      </div>
+    </div>
+    """
+  end
+
   # JS commands
 
   @doc """

lib/livebook_web/live/hub/teams/deployment_group_component.ex

@@ -1,7 +1,6 @@
 defmodule LivebookWeb.Hub.Teams.DeploymentGroupComponent do
   use LivebookWeb, :live_component
 
-  alias Livebook.Hubs
   alias LivebookWeb.NotFoundError
 
   @impl true
@@ -61,7 +60,7 @@ defmodule LivebookWeb.Hub.Teams.DeploymentGroupComponent do
               patch={~p"/hub/#{@hub.id}/groups/#{@deployment_group.id}/agents/new"}
               class="pl-2 text-blue-600"
             >
-              + Add new
+              + Deploy
             </.link>
           </.labeled_text>
           <.labeled_text class="grow mt-6 lg:border-l lg:pl-4" label="Apps deployed">
@@ -148,80 +147,12 @@ defmodule LivebookWeb.Hub.Teams.DeploymentGroupComponent do
           width={:big}
           patch={~p"/hub/#{@hub.id}"}
         >
-          <div class="p-6 max-w-4xl flex flex-col space-y-3">
+          <.live_component
-            <h3 class="text-2xl font-semibold text-gray-800">
+            module={LivebookWeb.Hub.Teams.DeploymentGroupInstanceComponent}
-              Deployment group instance setup
+            id="deployment-group-agent-instance"
-            </h3>
+            hub={@hub}
-
+            deployment_group={@deployment_group}
-            <p class="text-gray-700">
+          />
-              A deployment group instance is an instance of Livebook where you can
-              deploy Livebook apps via Livebook Teams.
-            </p>
-
-            <.table id="hub-agent-keys-table" rows={@deployment_group.agent_keys}>
-              <:col :let={agent_key} label="ID"><%= agent_key.id %></:col>
-              <:col :let={agent_key} label="Key">
-                <div class="flex flex-nowrap gap-2">
-                  <div class="grow">
-                    <.password_field
-                      id={"agent-key-#{agent_key.id}"}
-                      name="agent_key"
-                      value={agent_key.key}
-                      readonly
-                    />
-                  </div>
-
-                  <span
-                    data-tooltip="Copied to clipboard"
-                    aria-label="copy to clipboard"
-                    phx-click={
-                      JS.dispatch("lb:clipcopy", to: "#agent-key-#{agent_key.id}")
-                      |> JS.transition("tooltip top", time: 2000)
-                    }
-                  >
-                    <.button color="gray" small type="button">
-                      <.remix_icon icon="clipboard-line" class="text-xl leading-none py-1" />
-                    </.button>
-                  </span>
-                </div>
-              </:col>
-            </.table>
-
-            <p class="text-gray-700">
-              Use the Dockerfile below to set up an instance in your own infrastructure.
-              Once the instance is running, it will connect to Livebook Teams and become
-              available for app deployments.
-            </p>
-
-            <div class="flex flex-col gap-4">
-              <div>
-                <div class="flex items-end mb-1 gap-1">
-                  <span class="text-sm text-gray-700 font-semibold">Dockerfile</span>
-                  <div class="grow" />
-                  <.button
-                    color="gray"
-                    small
-                    data-tooltip="Copied to clipboard"
-                    type="button"
-                    aria-label="copy to clipboard"
-                    phx-click={
-                      JS.dispatch("lb:clipcopy", to: "#agent-dockerfile-source")
-                      |> JS.add_class("", transition: {"tooltip top", "", ""}, time: 2000)
-                    }
-                  >
-                    <.remix_icon icon="clipboard-line" />
-                    <span>Copy source</span>
-                  </.button>
-                </div>
-
-                <.code_preview
-                  source_id="agent-dockerfile-source"
-                  source={Hubs.Dockerfile.build_agent_dockerfile(Hubs.Dockerfile.config_new(), @hub)}
-                  language="dockerfile"
-                />
-              </div>
-            </div>
-          </div>
         </.modal>
 
         <.modal

lib/livebook_web/live/hub/teams/deployment_group_instance_component.ex

@@ -0,0 +1,220 @@
+defmodule LivebookWeb.Hub.Teams.DeploymentGroupInstanceComponent do
+  use LivebookWeb, :live_component
+
+  alias Livebook.Hubs
+
+  @impl true
+  def mount(socket) do
+    {:ok, assign(socket, messages: [])}
+  end
+
+  @impl true
+  def update(assigns, socket) do
+    socket = assign(socket, assigns)
+
+    socket =
+      if socket.assigns[:agent_key_id] do
+        socket
+      else
+        agent_key_id =
+          case assigns.deployment_group.agent_keys do
+            [%{id: id} | _] -> id
+            _ -> nil
+          end
+
+        assign(socket, :agent_key_id, agent_key_id)
+      end
+
+    socket =
+      assign_new(socket, :changeset, fn ->
+        Hubs.Dockerfile.config_changeset(base_config(socket))
+      end)
+
+    {:ok, update_instructions(socket)}
+  end
+
+  defp base_config(socket) do
+    Hubs.Dockerfile.from_deployment_group(socket.assigns.deployment_group)
+  end
+
+  @impl true
+  def render(assigns) do
+    ~H"""
+    <div class="p-6 max-w-4xl flex flex-col gap-3">
+      <h3 class="text-2xl font-semibold text-gray-800">
+        App server setup
+      </h3>
+
+      <p class="text-gray-700">
+        App server is an instance of Livebook where you can deploy Livebook
+        apps via Livebook Teams. Each app server belongs to a specific
+        deployment group.
+      </p>
+
+      <p class="mb-5 text-gray-700">
+        Use the instructions below to set up an instance in your own infrastructure.
+        Once the instance is running, it will connect to Livebook Teams and become
+        available for app deployments.
+      </p>
+
+      <div :if={@messages != []} class="flex flex-col gap-2">
+        <.message_box :for={{kind, message} <- @messages} kind={kind}>
+          <%= raw(message) %>
+        </.message_box>
+      </div>
+
+      <.form
+        :let={f}
+        for={%{"id" => @agent_key_id}}
+        as={:agent_key}
+        phx-change="select_agent_key"
+        phx-target={@myself}
+      >
+        <div class="grid grid-cols-1 md:grid-cols-2">
+          <.select_field
+            label="Server key"
+            field={f[:id]}
+            options={
+              for key <- @deployment_group.agent_keys do
+                {value_preview(key.key), key.id}
+              end
+            }
+          />
+        </div>
+      </.form>
+
+      <.form :let={f} for={@changeset} as={:data} phx-change="validate" phx-target={@myself}>
+        <.radio_field
+          label="Base image"
+          field={f[:docker_tag]}
+          options={LivebookWeb.AppComponents.docker_tag_options()}
+        />
+      </.form>
+
+      <%= if @agent_key_id do %>
+        <div class="mt-5">
+          <.tabs id="deployment-instruction" default={default_tab(@deployment_group)}>
+            <:tab id="docker" label="Docker">
+              <div class="flex flex-col gap-3">
+                <p class="text-gray-700">
+                  Deploy an app server to any Docker-based infrastructure. You may want
+                  to set the environment variables as secrets, if applicable. Below is
+                  an example calling Docker CLI directly, adapt it as necessary.
+                </p>
+                <div>
+                  <div class="flex items-end mb-1 gap-1">
+                    <span class="text-sm text-gray-700 font-semibold">CLI</span>
+                  </div>
+
+                  <.code_preview
+                    source_id="agent-dockerfile-source"
+                    source={@instructions.docker_instructions}
+                    language="shell"
+                  />
+                </div>
+              </div>
+            </:tab>
+            <:tab id="fly_io" label="Fly.io">
+              <div class="flex flex-col gap-3">
+                <p class="text-gray-700">
+                  Deploy an app server to Fly.io with a few simple commands.
+                </p>
+                <div>
+                  <div class="flex items-end mb-1 gap-1">
+                    <span class="text-sm text-gray-700 font-semibold">CLI</span>
+                  </div>
+
+                  <.code_preview
+                    source_id="agent-dockerfile-source"
+                    source={@instructions.fly_instructions}
+                    language="shell"
+                  />
+                </div>
+              </div>
+            </:tab>
+          </.tabs>
+        </div>
+      <% end %>
+    </div>
+    """
+  end
+
+  defp value_preview(string) do
+    preview_length = 10
+    length = String.length(string)
+    String.slice(string, 0, preview_length) <> String.duplicate("•", length - preview_length)
+  end
+
+  defp default_tab(%{clustering: :fly_io}), do: "fly_io"
+  defp default_tab(_deloyment_group), do: "docker"
+
+  @impl true
+  def handle_event("select_agent_key", %{"agent_key" => %{"id" => id}}, socket) do
+    id = if(id != "", do: id)
+    {:noreply, assign(socket, agent_key_id: id) |> update_instructions()}
+  end
+
+  def handle_event("validate", %{"data" => data}, socket) do
+    changeset =
+      socket
+      |> base_config()
+      |> Hubs.Dockerfile.config_changeset(data)
+      |> Map.replace!(:action, :validate)
+
+    {:noreply, assign(socket, changeset: changeset) |> update_instructions()}
+  end
+
+  defp update_instructions(socket) do
+    config = Ecto.Changeset.apply_changes(socket.assigns.changeset)
+    warnings = Hubs.Dockerfile.online_warnings(config)
+    messages = Enum.map(warnings, &{:warning, &1})
+    assign(socket, instructions: instructions(socket), messages: messages)
+  end
+
+  defp instructions(%{assigns: %{agent_key_id: nil}}), do: nil
+
+  defp instructions(socket) do
+    hub = socket.assigns.hub
+
+    agent_key =
+      Enum.find(
+        socket.assigns.deployment_group.agent_keys,
+        &(&1.id == socket.assigns.agent_key_id)
+      )
+
+    config = Ecto.Changeset.apply_changes(socket.assigns.changeset)
+
+    %{image: image, env: env} =
+      Livebook.Hubs.Dockerfile.online_docker_info(config, hub, agent_key)
+
+    %{
+      docker_instructions: docker_instructions(image, env),
+      fly_instructions: fly_instructions(image, env)
+    }
+  end
+
+  defp docker_instructions(image, env) do
+    envs = Enum.map_join(env, "\n", fn {key, value} -> ~s/  -e #{key}="#{value}" \\/ end)
+
+    """
+    docker run -p 8080:8080 -p 8081:8081 --pull always \\
+    #{envs}
+      #{image}
+    """
+  end
+
+  defp fly_instructions(image, env) do
+    envs = Enum.map_join(env, " \\\n", fn {key, value} -> ~s/  #{key}="#{value}"/ end)
+
+    """
+    APP_NAME="my_name"
+
+    flyctl apps create $APP_NAME
+
+    flyctl secrets set --app $APP_NAME \\
+    #{envs}
+
+    flyctl deploy --app $APP_NAME --image #{image}
+    """
+  end
+end

lib/livebook_web/live/session_live/app_docker_component.ex

@@ -30,7 +30,6 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
         hub_file_systems: Hubs.get_file_systems(assigns.hub, hub_only: true),
         deployment_groups: deployment_groups,
         deployment_group: deployment_group,
-        deployment_group_form: %{"deployment_group_id" => assigns.deployment_group_id},
         deployment_group_id: assigns.deployment_group_id
       )
       |> assign_new(:messages, fn -> [] end)
@@ -50,12 +49,7 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
 
   defp base_config(socket) do
     if deployment_group = socket.assigns.deployment_group do
-      %{
+      Hubs.Dockerfile.from_deployment_group(deployment_group)
-        Hubs.Dockerfile.config_new()
-        | clustering: deployment_group.clustering,
-          zta_provider: deployment_group.zta_provider,
-          zta_key: deployment_group.zta_key
-      }
     else
       Hubs.Dockerfile.config_new()
     end
@@ -74,7 +68,6 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
         hub={@hub}
         deployment_group={@deployment_group}
         deployment_groups={@deployment_groups}
-        deployment_group_form={@deployment_group_form}
         deployment_group_id={@deployment_group_id}
         changeset={@changeset}
         session={@session}
@@ -138,18 +131,18 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
         <%= if @deployment_groups do %>
           <%= if @deployment_groups != [] do %>
             <.form
-              for={@deployment_group_form}
+              :let={f}
+              for={%{"id" => @deployment_group_id}}
+              as={:deployment_group}
               phx-change="select_deployment_group"
               phx-target={@myself}
               id="select_deployment_group_form"
             >
               <.select_field
                 help={deployment_group_help()}
-                field={@deployment_group_form[:deployment_group_id]}
+                field={f[:id]}
                 options={deployment_group_options(@deployment_groups)}
                 label="Deployment Group"
-                name="deployment_group_id"
-                value={@deployment_group_id}
               />
             </.form>
           <% else %>
@@ -311,7 +304,7 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
     end
   end
 
-  def handle_event("select_deployment_group", %{"deployment_group_id" => id}, socket) do
+  def handle_event("select_deployment_group", %{"deployment_group" => %{"id" => id}}, socket) do
     id = if(id != "", do: id)
     Livebook.Session.set_notebook_deployment_group(socket.assigns.session.pid, id)
 
@@ -348,7 +341,7 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
       end
 
     dockerfile =
-      Hubs.Dockerfile.build_dockerfile(
+      Hubs.Dockerfile.airgapped_dockerfile(
         config,
         hub,
         hub_secrets,
@@ -359,7 +352,7 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
       )
 
     warnings =
-      Hubs.Dockerfile.warnings(
+      Hubs.Dockerfile.airgapped_warnings(
         config,
         hub,
         hub_secrets,

test/livebook/hubs/dockerfile_test.exs

@@ -11,13 +11,13 @@ defmodule Livebook.Hubs.DockerfileTest do
                 do: "latest",
                 else: Livebook.Config.app_version()
 
-  describe "build_dockerfile/7" do
+  describe "airgapped_dockerfile/7" do
     test "deploying a single notebook in personal hub" do
       config = dockerfile_config()
       hub = personal_hub()
       file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, [], %{})
 
       assert dockerfile == """
              FROM ghcr.io/livebook-dev/livebook:#{@docker_tag}
@@ -43,7 +43,8 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub_secrets = [secret, unused_secret]
       secrets = %{"TEST" => secret, "SESSION" => session_secret}
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, hub_secrets, [], file, [], secrets)
+      dockerfile =
+        Dockerfile.airgapped_dockerfile(config, hub, hub_secrets, [], file, [], secrets)
 
       assert dockerfile =~
                """
@@ -59,7 +60,7 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub = personal_hub()
       file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, [], %{})
 
       assert dockerfile =~ """
              # Notebooks and files
@@ -75,7 +76,8 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub_secrets = [secret, unused_secret]
       secrets = %{"TEST" => secret, "SESSION" => session_secret}
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, hub_secrets, [], file, [], secrets)
+      dockerfile =
+        Dockerfile.airgapped_dockerfile(config, hub, hub_secrets, [], file, [], secrets)
 
       assert dockerfile =~
                """
@@ -92,7 +94,7 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub = team_hub()
       file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, [], %{})
 
       assert dockerfile == """
              FROM ghcr.io/livebook-dev/livebook:#{@docker_tag}
@@ -123,7 +125,8 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub_secrets = [secret]
       secrets = %{"TEST" => secret, "SESSION" => session_secret}
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, hub_secrets, [], file, [], secrets)
+      dockerfile =
+        Dockerfile.airgapped_dockerfile(config, hub, hub_secrets, [], file, [], secrets)
 
       assert dockerfile =~ "ENV LIVEBOOK_TEAMS_SECRETS"
       refute dockerfile =~ "ENV TEST"
@@ -134,7 +137,7 @@ defmodule Livebook.Hubs.DockerfileTest do
       file_system = Livebook.Factory.build(:fs_s3)
       file_systems = [file_system]
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], file_systems, file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], file_systems, file, [], %{})
 
       assert dockerfile =~ "ENV LIVEBOOK_TEAMS_FS"
     end
@@ -144,7 +147,7 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub = team_hub()
       file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, [], %{})
 
       assert dockerfile =~ ~S/ENV LIVEBOOK_IDENTITY_PROVIDER "cloudflare:cloudflare_key"/
     end
@@ -154,7 +157,7 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub = team_hub()
       file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, [], %{})
 
       assert dockerfile =~ """
              # Notebooks and files
@@ -167,7 +170,7 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub = personal_hub()
       file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, [], %{})
 
       assert dockerfile =~ """
              FROM ghcr.io/livebook-dev/livebook:#{@docker_tag}-cuda11.8
@@ -186,7 +189,7 @@ defmodule Livebook.Hubs.DockerfileTest do
         %{type: :attachment, name: "data.csv"}
       ]
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, file_entries, %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, file_entries, %{})
 
       assert dockerfile =~
                """
@@ -200,27 +203,77 @@ defmodule Livebook.Hubs.DockerfileTest do
       hub = personal_hub()
       file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
 
-      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+      dockerfile = Dockerfile.airgapped_dockerfile(config, hub, [], [], file, [], %{})
 
       assert dockerfile =~ ~s/ENV LIVEBOOK_CLUSTER "fly"/
     end
   end
 
+  describe "online_docker_info/3" do
+    test "includes agent authentication env vars" do
+      config = dockerfile_config()
+      hub = team_hub()
+      agent_key = Livebook.Factory.build(:agent_key)
+
+      %{env: env} = Dockerfile.online_docker_info(config, hub, agent_key)
+
+      assert env == [
+               {"LIVEBOOK_AGENT_NAME", "default"},
+               {"LIVEBOOK_TEAMS_KEY", "lb_tk_fn0pL3YLWzPoPFWuHeV3kd0o7_SFuIOoU4C_k6OWDYg"},
+               {"LIVEBOOK_TEAMS_AUTH",
+                "online:org-name-387:1:1:lb_ak_zj9tWM1rEVeweYR7DbH_2VK5_aKtWfptcL07dBncqg"}
+             ]
+    end
+
+    test "deploying with zta" do
+      config = dockerfile_config(%{zta_provider: :cloudflare, zta_key: "cloudflare_key"})
+      hub = team_hub()
+      agent_key = Livebook.Factory.build(:agent_key)
+
+      %{env: env} = Dockerfile.online_docker_info(config, hub, agent_key)
+
+      assert {"LIVEBOOK_IDENTITY_PROVIDER", "cloudflare:cloudflare_key"} in env
+    end
+
+    test "deploying with different base image" do
+      config = dockerfile_config(%{docker_tag: "#{@docker_tag}-cuda11.8"})
+      hub = team_hub()
+      agent_key = Livebook.Factory.build(:agent_key)
+
+      %{image: image, env: env} = Dockerfile.online_docker_info(config, hub, agent_key)
+
+      assert image == "ghcr.io/livebook-dev/livebook:#{@docker_tag}-cuda11.8"
+      assert {"XLA_TARGET", "cuda118"} in env
+    end
+
+    test "deploying with fly.io cluster setup" do
+      config = dockerfile_config(%{clustering: :fly_io})
+      hub = team_hub()
+      agent_key = Livebook.Factory.build(:agent_key)
+
+      %{env: env} = Dockerfile.online_docker_info(config, hub, agent_key)
+
+      assert {"LIVEBOOK_CLUSTER", "fly"} in env
+    end
+  end
+
   describe "warnings/6" do
     test "warns when session secrets are used" do
-      config = dockerfile_config()
+      config = dockerfile_config(%{clustering: :fly_io})
       hub = personal_hub()
       app_settings = Livebook.Notebook.AppSettings.new()
 
       session_secret = %Secret{name: "SESSION", value: "test", hub_id: nil}
       secrets = %{"SESSION" => session_secret}
 
-      assert [warning] = Dockerfile.warnings(config, hub, [], [], app_settings, [], secrets)
+      assert [warning] =
+               Dockerfile.airgapped_warnings(config, hub, [], [], app_settings, [], secrets)
+
       assert warning =~ "The notebook uses session secrets"
     end
 
     test "warns when hub secrets are used from personal hub" do
-      config = dockerfile_config()
+      config = dockerfile_config(%{clustering: :fly_io})
       hub = personal_hub()
       app_settings = Livebook.Notebook.AppSettings.new()
 
@@ -230,13 +283,21 @@ defmodule Livebook.Hubs.DockerfileTest do
       secrets = %{"TEST" => secret}
 
       assert [warning] =
-               Dockerfile.warnings(config, hub, hub_secrets, [], app_settings, [], secrets)
+               Dockerfile.airgapped_warnings(
+                 config,
+                 hub,
+                 hub_secrets,
+                 [],
+                 app_settings,
+                 [],
+                 secrets
+               )
 
       assert warning =~ "secrets are included in the Dockerfile"
     end
 
     test "warns when there is a reference to external file system from personal hub" do
-      config = dockerfile_config()
+      config = dockerfile_config(%{clustering: :fly_io})
       hub = personal_hub()
       app_settings = Livebook.Notebook.AppSettings.new()
 
@@ -248,46 +309,68 @@ defmodule Livebook.Hubs.DockerfileTest do
       ]
 
       assert [warning] =
-               Dockerfile.warnings(config, hub, [], file_systems, app_settings, file_entries, %{})
+               Dockerfile.airgapped_warnings(
+                 config,
+                 hub,
+                 [],
+                 file_systems,
+                 app_settings,
+                 file_entries,
+                 %{}
+               )
 
       assert warning =~
                "The S3 file storage, defined in your personal hub, will not be available in the Docker image"
     end
 
     test "warns when deploying a directory in personal hub and it has any file systems" do
-      config = dockerfile_config(%{deploy_all: true})
+      config = dockerfile_config(%{clustering: :fly_io, deploy_all: true})
       hub = personal_hub()
       app_settings = Livebook.Notebook.AppSettings.new()
 
       file_system = Livebook.Factory.build(:fs_s3)
       file_systems = [file_system]
 
-      assert [warning] = Dockerfile.warnings(config, hub, [], file_systems, app_settings, [], %{})
+      assert [warning] =
+               Dockerfile.airgapped_warnings(config, hub, [], file_systems, app_settings, [], %{})
 
       assert warning =~
                "The S3 file storage, defined in your personal hub, will not be available in the Docker image"
     end
 
     test "warns when the app has no password in personal hub" do
-      config = dockerfile_config()
+      config = dockerfile_config(%{clustering: :fly_io})
       hub = personal_hub()
       app_settings = %{Livebook.Notebook.AppSettings.new() | access_type: :public}
 
-      assert [warning] = Dockerfile.warnings(config, hub, [], [], app_settings, [], %{})
+      assert [warning] = Dockerfile.airgapped_warnings(config, hub, [], [], app_settings, [], %{})
       assert warning =~ "This app has no password configuration"
     end
 
     test "warns when the app has no password and no ZTA in teams hub" do
-      config = dockerfile_config()
+      config = dockerfile_config(%{clustering: :fly_io})
       hub = team_hub()
       app_settings = %{Livebook.Notebook.AppSettings.new() | access_type: :public}
 
-      assert [warning] = Dockerfile.warnings(config, hub, [], [], app_settings, [], %{})
+      assert [warning] = Dockerfile.airgapped_warnings(config, hub, [], [], app_settings, [], %{})
       assert warning =~ "This app has no password configuration"
 
       config = %{config | zta_provider: :cloudflare, zta_key: "key"}
 
-      assert [] = Dockerfile.warnings(config, hub, [], [], app_settings, [], %{})
+      assert [] = Dockerfile.airgapped_warnings(config, hub, [], [], app_settings, [], %{})
+    end
+
+    test "warns when no clustering is configured" do
+      config = dockerfile_config(%{})
+      hub = team_hub()
+      app_settings = Livebook.Notebook.AppSettings.new()
+
+      assert [warning] = Dockerfile.airgapped_warnings(config, hub, [], [], app_settings, [], %{})
+      assert warning =~ "The deployment is not configured for clustering"
+
+      config = %{config | clustering: :fly_io}
+
+      assert [] = Dockerfile.airgapped_warnings(config, hub, [], [], app_settings, [], %{})
     end
   end
 

test/livebook_teams/web/session_live_test.exs

@@ -430,7 +430,7 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       id = deployment_group.id
 
       view
-      |> form("#select_deployment_group_form", %{deployment_group_id: id})
+      |> form("#select_deployment_group_form", %{deployment_group: %{id: id}})
       |> render_change()
 
       assert_receive {:operation, {:set_notebook_deployment_group, _client, ^id}}