From fb23199295ce26ae07289df672b81d5d7df99a50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonatan=20K=C5=82osko?= Date: Fri, 7 Mar 2025 11:51:09 +0100 Subject: [PATCH] Dispatch Python output rendering to kino (#2954) --- lib/livebook/runtime/definitions.ex | 4 ++ lib/livebook/runtime/evaluator.ex | 48 +++++++++++++++++++-- lib/livebook/runtime/evaluator/formatter.ex | 26 ++++++++++- lib/livebook/session.ex | 7 ++- 4 files changed, 78 insertions(+), 7 deletions(-) diff --git a/lib/livebook/runtime/definitions.ex b/lib/livebook/runtime/definitions.ex index 02802d24c..018bbf551 100644 --- a/lib/livebook/runtime/definitions.ex +++ b/lib/livebook/runtime/definitions.ex @@ -509,5 +509,9 @@ defmodule Livebook.Runtime.Definitions do %{dep: {:pythonx, "~> 0.4.2"}, config: []} end + def kino_pythonx_dependency() do + %{dep: {:kino_pythonx, github: "livebook-dev/kino_pythonx"}, config: []} + end + def pythonx_requirement(), do: "~> 0.4.0" end diff --git a/lib/livebook/runtime/evaluator.ex b/lib/livebook/runtime/evaluator.ex index c7ae1e366..69d503221 100644 --- a/lib/livebook/runtime/evaluator.ex +++ b/lib/livebook/runtime/evaluator.ex @@ -498,6 +498,8 @@ defmodule Livebook.Runtime.Evaluator do state = put_context(state, ref, new_context) output = Evaluator.Formatter.format_result(language, result) + after_evaluation(language) + metadata = %{ errored: error_result?(result), interrupted: interrupt_result?(result), @@ -934,6 +936,7 @@ defmodule Livebook.Runtime.Evaluator do end @compile {:no_warn_undefined, {Pythonx, :eval, 2}} + @compile {:no_warn_undefined, {Pythonx, :uv_init, 1}} @compile {:no_warn_undefined, {Pythonx, :decode, 1}} defp eval_python(code, binding, env) do @@ -1027,14 +1030,30 @@ defmodule Livebook.Runtime.Evaluator do defp eval_pyproject_toml(code, binding, env) do with :ok <- ensure_pythonx() do - quoted = {{:., [], [{:__aliases__, [alias: false], [:Pythonx]}, :uv_init]}, [], [code]} - {result, _diagnostics} = Code.with_diagnostics([log: true], fn -> try do - {value, binding, env} = - Code.eval_quoted_with_env(quoted, binding, env, prune_binding: true) + Pythonx.uv_init(code) + # The default matplotlib backend relies on OS-specific GUI + # and crashes when embedding Python. For this reason, we + # configure a non-interactive backend that only allows + # exporting figures as images. In general we want to avoid + # special casing like this, but given how common matplotlib + # is, it does make sense to streamline the experience. + # We set the backend using env var, instead of calling + # plt.backend(...), because importing the module for the + # first time is slow, so we prefer to avoid that as part + # of setup. + Pythonx.eval( + """ + import os + os.environ["MPLBACKEND"] = "Agg" + """, + %{} + ) + + value = :ok result = {:ok, value, binding, env} code_markers = [] {result, code_markers} @@ -1081,6 +1100,27 @@ defmodule Livebook.Runtime.Evaluator do defp pythonx_version(), do: List.to_string(Application.spec(:pythonx)[:vsn]) + defp after_evaluation(:python) do + if ensure_pythonx() == :ok do + # With matplotlib the charts are built imperatively, by modifying + # a global figure state. We clear the global state after the + # evaluation, otherwise re-evaluating cells draws on top of the + # previous figure. We do this only if matplotlib is imported. + Pythonx.eval( + """ + import sys + + if "matplotlib" in sys.modules: + import matplotlib.pyplot as plt + plt.close("all") + """, + %{} + ) + end + end + + defp after_evaluation(_language), do: :ok + defp identifier_dependencies(context, tracer_info, prev_context) do identifiers_used = MapSet.new() identifiers_defined = %{} diff --git a/lib/livebook/runtime/evaluator/formatter.ex b/lib/livebook/runtime/evaluator/formatter.ex index 4a359654f..b3ff2401f 100644 --- a/lib/livebook/runtime/evaluator/formatter.ex +++ b/lib/livebook/runtime/evaluator/formatter.ex @@ -2,6 +2,7 @@ defmodule Livebook.Runtime.Evaluator.Formatter do require Logger @compile {:no_warn_undefined, {Kino.Render, :to_livebook, 1}} + @compile {:no_warn_undefined, {Kino.Render, :impl_for, 1}} @compile {:no_warn_undefined, {Pythonx, :eval, 2}} @compile {:no_warn_undefined, {Pythonx, :decode, 1}} @@ -57,8 +58,22 @@ defmodule Livebook.Runtime.Evaluator.Formatter do end def format_result(:python, {:ok, value}) do - repr_string = Pythonx.eval("repr(value)", %{"value" => value}) |> elem(0) |> Pythonx.decode() - %{type: :terminal_text, text: repr_string, chunk: false} + if Code.ensure_loaded?(Kino.Render) do + try do + if Kino.Render.impl_for(value) == Kino.Render.Any do + to_repr_output(value) + else + Kino.Render.to_livebook(value) + end + catch + kind, error -> + formatted = format_error(kind, error, __STACKTRACE__) + Logger.error(formatted) + to_repr_output(value) + end + else + to_repr_output(value) + end end def format_result(:python, {:error, _kind, error, _stacktrace}) @@ -108,6 +123,13 @@ defmodule Livebook.Runtime.Evaluator.Formatter do end end + defp to_repr_output(value) do + repr_string = + Pythonx.eval("repr(value)", %{"value" => value}) |> elem(0) |> Pythonx.decode() + + %{type: :terminal_text, text: repr_string, chunk: false} + end + defp to_inspect_output(value, opts \\ []) do try do inspected = inspect(value, inspect_opts(opts)) diff --git a/lib/livebook/session.ex b/lib/livebook/session.ex index e84fa3b2a..c99b70fd0 100644 --- a/lib/livebook/session.ex +++ b/lib/livebook/session.ex @@ -1243,7 +1243,12 @@ defmodule Livebook.Session do end def handle_cast({:enable_language, client_pid, language}, state) do - case do_add_dependencies(state, [Livebook.Runtime.Definitions.pythonx_dependency()]) do + dependencies = [ + Livebook.Runtime.Definitions.pythonx_dependency(), + Livebook.Runtime.Definitions.kino_pythonx_dependency() + ] + + case do_add_dependencies(state, dependencies) do {:ok, state} -> client_id = client_id(state, client_pid)