Base image with CUDA

lib/livebook/config.ex

@@ -44,6 +44,17 @@ defmodule Livebook.Config do
 
   @identity_provider_read_only Enum.filter(@identity_providers, & &1.read_only)
 
+  def docker_tags do
+    version = app_version()
+    base = if version =~ "dev", do: "latest", else: version
+
+    [
+      {base, []},
+      {"#{base}-cuda11.8", [XLA_TARGET: "cuda118"]},
+      {"#{base}-cuda12.1", [XLA_TARGET: "cuda120"]}
+    ]
+  end
+
   @doc """
   Returns the longname if the distribution mode is configured to use long names.
   """

lib/livebook_web/live/hub/edit/team_component.ex

@@ -30,6 +30,8 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
           raise(NotFoundError, "could not find file system matching #{inspect(file_system_id)}")
       end
 
+    [{default_base_image, _} | _] = Livebook.Config.docker_tags()
+
     {:ok,
      socket
      |> assign(
@@ -43,7 +45,8 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
        hub_metadata: Provider.to_metadata(assigns.hub),
        is_default: is_default?,
        zta: %{"provider" => "", "key" => ""},
-       zta_metadata: nil
+       zta_metadata: nil,
+       base_image: default_base_image
      )
      |> assign_dockerfile()
      |> assign_form(changeset)}
@@ -206,6 +209,17 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
                 the steps below. First, configure your deployment:
               </p>
 
+              <div class="grid grid-cols-1">
+                <form phx-change="base_image" phx-target={@myself} phx-nosubmit>
+                  <.radio_field
+                    name="base_image"
+                    label="Base image"
+                    value={@base_image}
+                    options={for {k, _v} <- Livebook.Config.docker_tags(), do: {k, k}}
+                  />
+                </form>
+              </div>
+
               <.form :let={f} class="py-2" for={@zta} phx-change="change_zta" phx-target={@myself}>
                 <div class="grid grid-cols-1 md:grid-cols-2 gap-3">
                   <.select_field
@@ -526,6 +540,10 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     {:noreply, push_navigate(socket, to: ~p"/hub/#{socket.assigns.hub.id}")}
   end
 
+  def handle_event("base_image", %{"base_image" => base_image}, socket) do
+    {:noreply, assign(socket, base_image: base_image) |> assign_dockerfile()}
+  end
+
   defp is_default?(hub) do
     Hubs.get_default_hub().id == hub.id
   end
@@ -535,14 +553,22 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
   end
 
   defp assign_dockerfile(socket) do
-    version = to_string(Application.spec(:livebook, :vsn))
-    version = if version =~ "dev", do: "edge", else: version
+    {base_image, xla_target} =
+      List.keyfind(Livebook.Config.docker_tags(), socket.assigns.base_image, 0)
+
+    image = """
+    FROM ghcr.io/livebook-dev/livebook:#{base_image}
+    """
 
-    base =
+    xla_target = xla_target(xla_target)
+
+    base_args = """
+    ARG APPS_PATH=/path/to/my/notebooks
+    ARG TEAMS_KEY="#{socket.assigns.hub.teams_key}"
+    """
+
+    base_env =
       """
-      FROM ghcr.io/livebook-dev/livebook:#{version}
-      ARG APPS_PATH=/path/to/my/notebooks
-      ARG TEAMS_KEY="#{socket.assigns.hub.teams_key}"
 
       ENV LIVEBOOK_TEAMS_KEY ${TEAMS_KEY}
       ENV LIVEBOOK_TEAMS_NAME "#{socket.assigns.hub.hub_name}"
@@ -565,7 +591,7 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     zta = zta_env(socket.assigns.zta)
 
     dockerfile =
-      [base, secrets, file_systems, zta, apps]
+      [image, base_args, xla_target, base_env, secrets, file_systems, zta, apps]
       |> Enum.reject(&is_nil/1)
       |> Enum.join()
 
@@ -630,4 +656,12 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     ENV LIVEBOOK_TEAMS_FS "#{encrypt_file_systems_to_dockerfile(socket)}"
     """
   end
+
+  defp xla_target([]), do: nil
+
+  defp xla_target(XLA_TARGET: xla_target) do
+    """
+    ENV XLA_TARGET #{xla_target}
+    """
+  end
 end