Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add evaluation temporary directory #2204

Merged
merged 1 commit into from
Sep 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions lib/livebook/runtime/erl_dist/runtime_server.ex
Original file line number Diff line number Diff line change
Expand Up @@ -694,6 +694,12 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
end
end

defp evaluator_tmp_dir(state) do
if tmp_dir = state.tmp_dir do
Path.join(tmp_dir, "tmp")
end
end

defp ensure_evaluator(state, container_ref) do
if Map.has_key?(state.evaluators, container_ref) do
state
Expand All @@ -705,6 +711,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
runtime_broadcast_to: state.runtime_broadcast_to,
object_tracker: state.object_tracker,
ebin_path: state.ebin_path,
tmp_dir: evaluator_tmp_dir(state),
io_proxy_registry: state.io_proxy_registry
)

Expand Down
5 changes: 5 additions & 0 deletions lib/livebook/runtime/evaluator.ex
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,9 @@ defmodule Livebook.Runtime.Evaluator do
* `:ebin_path` - a directory to write modules bytecode into. When
not specified, modules are not written to disk

* `:tmp_dir` - a temporary directory for arbitrary use during
evaluation

* `:io_proxy_registry` - the registry to register IO proxy
processes in

Expand Down Expand Up @@ -266,6 +269,7 @@ defmodule Livebook.Runtime.Evaluator do
runtime_broadcast_to = Keyword.get(opts, :runtime_broadcast_to, send_to)
object_tracker = Keyword.fetch!(opts, :object_tracker)
ebin_path = Keyword.get(opts, :ebin_path)
tmp_dir = Keyword.get(opts, :tmp_dir)
io_proxy_registry = Keyword.get(opts, :io_proxy_registry)

{:ok, io_proxy} =
Expand All @@ -275,6 +279,7 @@ defmodule Livebook.Runtime.Evaluator do
runtime_broadcast_to,
object_tracker,
ebin_path,
tmp_dir,
io_proxy_registry
)

Expand Down
62 changes: 54 additions & 8 deletions lib/livebook/runtime/evaluator/io_proxy.ex
Original file line number Diff line number Diff line change
Expand Up @@ -26,23 +26,54 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
For all supported requests a message is sent to the configured
`:send_to` process, so this device serves as a proxy.
"""
@spec start(pid(), pid(), pid(), pid(), String.t() | nil, atom() | nil) :: GenServer.on_start()
def start(evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry) do
@spec start(
pid(),
pid(),
pid(),
pid(),
String.t() | nil,
String.t() | nil,
atom() | nil
) :: GenServer.on_start()
def start(
evaluator,
send_to,
runtime_broadcast_to,
object_tracker,
ebin_path,
tmp_dir,
registry
) do
GenServer.start(
__MODULE__,
{evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry}
{evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, tmp_dir, registry}
)
end

@doc """
Linking version of `start/4`.
"""
@spec start_link(pid(), pid(), pid(), pid(), String.t() | nil, atom() | nil) ::
GenServer.on_start()
def start_link(evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry) do
@spec start_link(
pid(),
pid(),
pid(),
pid(),
String.t() | nil,
String.t() | nil,
atom() | nil
) :: GenServer.on_start()
def start_link(
evaluator,
send_to,
runtime_broadcast_to,
object_tracker,
ebin_path,
tmp_dir,
registry
) do
GenServer.start_link(
__MODULE__,
{evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry}
{evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, tmp_dir, registry}
)
end

Expand Down Expand Up @@ -73,7 +104,9 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
end

@impl true
def init({evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry}) do
def init(
{evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, tmp_dir, registry}
) do
evaluator_monitor = Process.monitor(evaluator)

if registry do
Expand All @@ -94,6 +127,7 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
runtime_broadcast_to: runtime_broadcast_to,
object_tracker: object_tracker,
ebin_path: ebin_path,
tmp_dir: tmp_dir,
tracer_info: %Evaluator.Tracer{},
modules_defined: MapSet.new()
}}
Expand Down Expand Up @@ -348,6 +382,18 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
{result, state}
end

defp io_request(:livebook_get_tmp_dir, state) do
result =
with tmp_dir when is_binary(tmp_dir) <- state.tmp_dir,
:ok <- File.mkdir_p(tmp_dir) do
{:ok, state.tmp_dir}
else
_ -> {:error, :not_available}
end

{result, state}
end

defp io_request(_, state) do
{{:error, :request}, state}
end
Expand Down