Make uv_init return a struct with all required paths to init Pythonx

[mruoss]

This is a proposition / food for discussion at this point - coming over form the discussion here: https://elixirforum.com/t/is-there-a-way-to-access-pyproject-toml/73129.

In order to support Livebook+Pythonx+FLAME, somethink like this would be a first step. We'd still have to somehow get that struct/data over to the FLAME and initialize Pythonx in Livebook's /rel/server/overlays/bin/start_flame.exs. Since this can only be done once, we can't do it inside the FLAME.call function as FLAMEs can be reused...

WDYT?

[mruoss]

Actually, I'm not so sure about the initializing once part. Especially in FLAME it can be that you reuse the FLAME but you'd want to initialize a different pyproject.toml / venv than the last user of the FLAME. Currently, if I understand correctly, that is not possible. Initializing Python and initializing UV seem to be very entangled, right?

[jonatanklosko]

Hey, sorry for the late reply.

Another approach that could work is for Livebook to set an application env with the pyproject toml, then when Pythonx boots on another node and the env is set, it would do the initialization. We already support such application env, but it is compile time (if we don't set the env, that code path doesn't exist at all):

pythonx/lib/pythonx/application.ex

Lines 30 to 34 in 57b5398

	if pyproject_toml do
	Pythonx.Uv.fetch(pyproject_toml, true, opts)
	defp maybe_uv_init(), do: Pythonx.Uv.init(unquote(pyproject_toml), true, unquote(opts))
	else
	defp maybe_uv_init(), do: :noop

cc @josevalim

Actually, I'm not so sure about the initializing once part. Especially in FLAME it can be that you reuse the FLAME but you'd want to initialize a different pyproject.toml / venv than the last user of the FLAME. Currently, if I understand correctly, that is not possible. Initializing Python and initializing UV seem to be very entangled, right?

Initializing Pythonx itself means starting the interpreter and setting up code paths. uv is a way to get Python+dependencies. So those are distinct, but since currently we only support initializing Pythonx with uv provided Python+dependencies, they are entangled.

Within a single node, Pythonx can only be initialized once. Theoretically there is a way to uninitialize the interpreter, but in practice this doesn't work, because many libraries with native code would break (e.g. numpy).

[mruoss]

Thanks for the reply and don't worry about the delay!

Another approach that could work is for Livebook to set an application env with the pyproject toml, then when Pythonx boots on another node and the env is set, it would do the initialization.

Okay, yes, this could also be built into Pythonx of course. Either way, in the case of FLAME, passing the ENV variable would have to be a manual configuration on the FLAME pool...

Within a single node, Pythonx can only be initialized once.

Okay, if I understand correctly, one FLAME runner / BEAM node can only be used for a single python project config (i.e. venv, i.e. Python+dependencies)? I mean... this limitation is probably okay if we can make sure Pythonx is actually initialized on the FLAME.

The thing that's bothering me a bit is: ENV variables are configured on a FLAME pool. But the limitation about Pythonx actually applies to a runner, not a pool. Each fresh runner of the same pool could technically be initialized with a different pyproject.toml. But... if we have to pass pyproject.toml as ENV var, we have to configure it on the pool... Then again... maybe that's only a theoretical problem and in practice, when using it with Livebook, you'd use one FLAME pool for one pyproject.toml...

[jonatanklosko]

@mruoss to be clear, I mean pythonx application env, not an env var. The application envs are automatically copied by FLAME to the runner :)

[mruoss]

Oooh I see! Yes this sounds like a good approach! I can play around with this a bit.

[mruoss]

Hmm... maybe I'm still not getting it.

The application envs are automatically copied by FLAME to the runner :)

Aren't we talking about the following?

Mix.install([{:pythonx, "~> 0.4.2"},  {:flame, "~> 0.1.5"}, {:flame_k8s_backend, "~> 0.5"}])
# ... initialize FLAME pool ...
Application.put_env(:pythonx, :pyproject_toml, "just a test")
FLAME.call(:runner, fn -> Application.get_env(:pythonx, :pyproject_toml) end)
# returns nil, was expecting "just a test"

[josevalim]

I think we copy the .app files but not the runtime values (but I may be misremembering). So anything dynamic won't work indeed.

[jonatanklosko]

Gah, sorry, I thought the persistent config is copied too but I misremembered.

[jonatanklosko]

But I think we do need to pass the information for boot in some way. @josevalim perhaps we can do explicit and pass extra config on flame pool to set extra application env?

[josevalim]

@jonatanklosko another option is to pass Pythonx as additionals path to be copied to FLAME and, because we should cache everything, then it just works on the FLAME?

[jonatanklosko]

@josevalim that's a separate point, ideally we want to copy to effectively get cache hit, but the question is how does pythonx on the new node know that it should initialize the interpreter.

[josevalim]

It would be great if we could store it in the directory we copy but I am afraid it is not straightforward. We could also look at system env vars, but I don’t think they are copied either?

[josevalim]

Or we could copy to priv, which may have other side effects, so perhaps yeah, we need new capabilities in FLAME.

[jonatanklosko]

We could also look at system env vars, but I don’t think they are copied either?

Env vars that we set dynamically are not copied. We could pass it explicitly to the pool via :env option, though it's a bit awkward to encode all necessary information in an env var. Perhaps it would be very specific like PYTHONX_FLAME_INIT_STATE with base64 encoded value. But then a better alternative would be to allow setting app config, so at least it's more structured.

Or we could copy to priv, which may have other side effects, so perhaps yeah, we need new capabilities in FLAME.

You mean the cache is still somewhere global, but FLAME copies it to Pythonx priv on the new node? It's even more implicit, since Pythonx would need to infer if/how it should initialize based on the priv contents. We also already use priv/ if compile-time uv_init: [pyproject_toml: ...] is set, so that it works in releases, so we would need to distinguish between these too.

So currently I am leaning towards adding application env config to FLAME, and perhaps a function like Pythonx.flame_init_econfig() (name either tied or not tied to flame).

[josevalim]

Right, so any of app env or system env are fine. FLAME already supports system env but adding app env should be trivial.

[jonatanklosko]

@josevalim actually, in both cases we can make it opaque if we don't expose the key name:

[
  env: ... |> Map.merge(Pythonx.flame_env())
]

# vs

[
  config: [
    pythonx: Pythonx.flame_config()
  ]
]

Then we can have a very verbose env var name with encoded state and it's all private API. Then I am fine with using env vars.

[josevalim]

as you prefer!