Add K8s runtime (#2756)

Co-authored-by: Jonatan Kłosko <[email protected]>

config/config.exs

@@ -40,7 +40,8 @@ config :livebook,
   teams_url: "https://teams.livebook.dev",
   github_release_info: %{repo: "livebook-dev/livebook", version: Mix.Project.config()[:version]},
   update_instructions_url: nil,
-  within_iframe: false
+  within_iframe: false,
+  k8s_kubeconfig_pipeline: Kubereq.Kubeconfig.Default
 
 config :livebook, Livebook.Apps.Manager, retry_backoff_base_ms: 5_000
 

config/test.exs

@@ -24,6 +24,12 @@ end
 
 config :livebook,
   data_path: data_path,
-  agent_name: "chonky-cat"
+  agent_name: "chonky-cat",
+  k8s_kubeconfig_pipeline:
+    {Kubereq.Kubeconfig.Stub,
+     plugs: %{
+       "default" => {Req.Test, :k8s_cluster},
+       "no-permission" => {Req.Test, :k8s_cluster}
+     }}
 
 config :livebook, Livebook.Apps.Manager, retry_backoff_base_ms: 0

lib/livebook.ex

@@ -166,7 +166,8 @@ defmodule Livebook do
            [
              Livebook.Runtime.Standalone,
              Livebook.Runtime.Attached,
-             Livebook.Runtime.Fly
+             Livebook.Runtime.Fly,
+             Livebook.Runtime.K8s
            ]
 
     if home = Livebook.Config.writable_dir!("LIVEBOOK_HOME") do

lib/livebook/k8s/auth.ex

@@ -0,0 +1,65 @@
+defmodule Livebook.K8s.Auth do
+  # Implementation of Access Review checks for the authenticated user
+  # using the `SelfSubjectAccessReview` [1] resource.
+  #
+  # [1]: https://kubernetes.io/docs/reference/kubernetes-api/authorization-resources/self-subject-access-review-v1/#SelfSubjectAccessReviewSpec
+
+  @doc """
+  Concurrently reviews access according to a list of `resource_attributes`.
+
+  Expects `req` to be prepared for `SelfSubjectAccessReview`.
+  """
+  @spec batch_check(Req.Request.t(), [keyword()]) ::
+          [:ok | {:error, %Req.Response{}} | {:error, Exception.t()}]
+  def batch_check(req, resource_attribute_list) do
+    resource_attribute_list
+    |> Enum.map(&Task.async(fn -> can_i?(req, &1) end))
+    |> Task.await_many(:infinity)
+  end
+
+  @doc """
+  Reviews access according to `resource_attributes`.
+
+  Expects `req` to be prepared for `SelfSubjectAccessReview`.
+  """
+  @spec can_i?(Req.Request.t(), keyword()) ::
+          :ok | {:error, %Req.Response{}} | {:error, Exception.t()}
+  def can_i?(req, resource_attributes) do
+    resource_attributes =
+      resource_attributes
+      |> Keyword.validate!([
+        :name,
+        :namespace,
+        :path,
+        :resource,
+        :subresource,
+        :verb,
+        :version,
+        group: ""
+      ])
+      |> Enum.into(%{})
+
+    access_review = %{
+      "apiVersion" => "authorization.k8s.io/v1",
+      "kind" => "SelfSubjectAccessReview",
+      "spec" => %{
+        "resourceAttributes" => resource_attributes
+      }
+    }
+
+    create_self_subject_access_review(req, access_review)
+  end
+
+  defp create_self_subject_access_review(req, access_review) do
+    case Kubereq.create(req, access_review) do
+      {:ok, %Req.Response{status: 201, body: %{"status" => %{"allowed" => true}}}} ->
+        :ok
+
+      {:ok, %Req.Response{} = response} ->
+        {:error, response}
+
+      {:error, error} ->
+        {:error, error}
+    end
+  end
+end

lib/livebook/k8s/pod.ex

@@ -0,0 +1,182 @@
+defmodule Livebook.K8s.Pod do
+  @main_container_name "livebook-runtime"
+  @home_pvc_volume_name "livebook-home"
+
+  @default_pod_template """
+  apiVersion: v1
+  kind: Pod
+  metadata:
+    generateName: livebook-runtime-
+  spec:
+    containers:
+      - name: livebook-runtime
+        resources:
+          limits:
+            cpu: "1"
+            memory: 1Gi
+          requests:
+            cpu: "1"
+            memory: 1Gi\
+  """
+
+  @doc """
+  Returns the default pod template.
+  """
+  @spec default_pod_template() :: String.t()
+  def default_pod_template(), do: @default_pod_template
+
+  @doc """
+  Set the namespace on the given manifest.
+  """
+  @spec set_namespace(map(), String.t()) :: map()
+  def set_namespace(manifest, namespace) do
+    put_in(manifest, ["metadata", "namespace"], namespace)
+  end
+
+  @doc """
+  Adds "volume" and "volumeMount" configurations to `manifest` in order
+  to mount `home_pvc` under /home/livebook on the pod.
+  """
+  @spec set_home_pvc(map(), String.t()) :: map()
+  def set_home_pvc(manifest, home_pvc) do
+    manifest
+    |> update_in(["spec", Access.key("volumes", [])], fn volumes ->
+      volume = %{
+        "name" => @home_pvc_volume_name,
+        "persistentVolumeClaim" => %{"claimName" => home_pvc}
+      }
+
+      [volume | volumes]
+    end)
+    |> update_in(
+      ["spec", "containers", access_main_container(), Access.key("volumeMounts", [])],
+      fn volume_mounts ->
+        [%{"name" => @home_pvc_volume_name, "mountPath" => "/home/livebook"} | volume_mounts]
+      end
+    )
+  end
+
+  @doc """
+  Adds the list of `env_vars` to the main container of the given `manifest`.
+  """
+  @spec add_env_vars(map(), list()) :: map()
+  def add_env_vars(manifest, env_vars) do
+    update_in(
+      manifest,
+      ["spec", "containers", access_main_container(), Access.key("env", [])],
+      fn existing_vars -> env_vars ++ existing_vars end
+    )
+  end
+
+  @doc """
+  Sets the tag of the main container's image.
+  """
+  @spec set_docker_tag(map(), String.t()) :: map()
+  def set_docker_tag(manifest, docker_tag) do
+    image = "ghcr.io/livebook-dev/livebook:#{docker_tag}"
+    put_in(manifest, ["spec", "containers", access_main_container(), "image"], image)
+  end
+
+  @doc """
+  Adds the `port` to the main container and adds a readiness probe.
+  """
+  @spec add_container_port(map(), non_neg_integer()) :: map()
+  def add_container_port(manifest, port) do
+    readiness_probe = %{
+      "tcpSocket" => %{"port" => port},
+      "initialDelaySeconds" => 1,
+      "periodSeconds" => 1
+    }
+
+    manifest
+    |> update_in(
+      ["spec", "containers", access_main_container(), Access.key("ports", [])],
+      &[%{"containerPort" => port} | &1]
+    )
+    |> put_in(["spec", "containers", access_main_container(), "readinessProbe"], readiness_probe)
+  end
+
+  @doc """
+  Turns the given `pod_template` into a Pod manifest.
+  """
+  @spec pod_from_template(String.t()) :: map()
+  def pod_from_template(pod_template) do
+    pod_template
+    |> YamlElixir.read_from_string!()
+    |> do_pod_from_template()
+  end
+
+  defp do_pod_from_template(pod) do
+    pod
+    |> Map.merge(%{"apiVersion" => "v1", "kind" => "Pod"})
+    |> put_in(["spec", "restartPolicy"], "Never")
+  end
+
+  @doc """
+  Validates the given Pod manifest.
+  """
+  @spec validate_pod_template(map(), String.t()) :: :ok | {:error, String.t()}
+  def validate_pod_template(pod, namespace)
+
+  def validate_pod_template(%{"apiVersion" => "v1", "kind" => "Pod"} = pod, namespace) do
+    with :ok <- validate_basics(pod),
+         :ok <- validate_main_container(pod),
+         :ok <- validate_namespace(pod, namespace) do
+      validate_container_image(pod)
+    end
+  end
+
+  def validate_pod_template(_other_input, _namespace) do
+    {:error, ~s/Make sure to define a valid resource of apiVersion "v1" and kind "Pod"./}
+  end
+
+  defp validate_basics(pod) do
+    cond do
+      not match?(%{"metadata" => %{}}, pod) ->
+        {:error, ".metadata is missing in your pod template."}
+
+      not match?(%{"spec" => %{"containers" => containers}} when is_list(containers), pod) ->
+        {:error, ".spec.containers is missing in your pod template."}
+
+      pod["metadata"]["name"] in [nil, ""] and pod["metadata"]["generateName"] in [nil, ""] ->
+        {:error,
+         "Make sure to define .metadata.name or .metadata.generateName in your pod template."}
+
+      true ->
+        :ok
+    end
+  end
+
+  defp validate_main_container(pod) do
+    if get_in(pod, ["spec", "containers", access_main_container()]) do
+      :ok
+    else
+      {:error,
+       ~s/Main container is missing. The main container should be named "#{@main_container_name}"./}
+    end
+  end
+
+  defp validate_container_image(pod) do
+    if get_in(pod, ["spec", "containers", access_main_container(), "image"]) do
+      {:error,
+       "You can't set the container image of the main container. It's going to be overridden."}
+    else
+      :ok
+    end
+  end
+
+  defp validate_namespace(pod, namespace) do
+    template_ns = get_in(pod, ["metadata", "namespace"])
+
+    if template_ns == nil or template_ns == namespace do
+      :ok
+    else
+      {:error,
+       "The field .template.metadata.namespace has to be omitted or set to the namespace you selected."}
+    end
+  end
+
+  defp access_main_container() do
+    Kubereq.Access.find(&(&1["name"] == @main_container_name))
+  end
+end

lib/livebook/k8s/pvc.ex

@@ -0,0 +1,58 @@
+defmodule Livebook.K8s.PVC do
+  use Ecto.Schema
+
+  import Ecto.Changeset
+
+  @type t :: %__MODULE__{
+          name: String.t(),
+          size_gb: integer(),
+          access_mode: String.t(),
+          storage_class: String.t()
+        }
+
+  @primary_key false
+  embedded_schema do
+    field :name, :string
+    field :size_gb, :integer
+    field :access_mode, :string, default: "ReadWriteOnce"
+    field :storage_class, :string, default: nil
+  end
+
+  @fields ~w(name size_gb access_mode storage_class)a
+  @required ~w(name size_gb access_mode)a
+
+  @doc """
+  Build a PVC changeset for the given `attrs`.
+  """
+  @spec changeset(map()) :: Ecto.Changeset.t()
+  def changeset(attrs \\ %{}) do
+    %__MODULE__{}
+    |> cast(attrs, @fields)
+    |> validate_required(@required)
+  end
+
+  @doc """
+  Build PVC manifest for the given `pvc` and `namespace` to be applied to a
+  cluster.
+  """
+  @spec manifest(pvc :: t(), namespace: String.t()) :: manifest :: map()
+  def manifest(pvc, namespace) do
+    %{
+      "apiVersion" => "v1",
+      "kind" => "PersistentVolumeClaim",
+      "metadata" => %{
+        "name" => pvc.name,
+        "namespace" => namespace
+      },
+      "spec" => %{
+        "storageClassName" => pvc.storage_class,
+        "accessModes" => [pvc.access_mode],
+        "resources" => %{
+          "requests" => %{
+            "storage" => "#{pvc.size_gb}Gi"
+          }
+        }
+      }
+    }
+  end
+end