From 4883bcd3a99f42da092b6172d2637646c436d2e2 Mon Sep 17 00:00:00 2001 From: Andrei Cravtov Date: Mon, 8 Jun 2026 18:08:07 +0100 Subject: [PATCH] add version from python package and expose CliArgs to python --- python/parts.nix | 3 + rust/exo_rs/exo_rs.pyi | 110 ++++++++++++++++++++++++++++++ rust/exo_rs/src/bin/stub_gen.rs | 6 +- rust/exo_rs/src/config/cli.rs | 69 +++++++++++++++---- rust/exo_rs/src/config/locator.rs | 15 ++++ rust/exo_rs/src/config/mod.rs | 16 +++++ rust/exo_rs/src/lib.rs | 52 +++++++++++++- rust/exo_rs/src/networking.rs | 2 +- src/exo/__init__.py | 3 + src/exo/main.py | 4 +- 10 files changed, 260 insertions(+), 20 deletions(-) diff --git a/python/parts.nix b/python/parts.nix index 26318c4a2..42e663866 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -211,6 +211,9 @@ let text: name: pkgs.writeShellApplication { inherit name; text = '' + unset PYTHONPATH + unset PYTHONHOME + LD_LIBRARY_PATH="${libPath}''${LD_LIBRARY_PATH:+:}''${LD_LIBRARY_PATH:-}" exec \ ${lib.optionalString cudaSupport "nixglhost "} ${text} ''; diff --git a/rust/exo_rs/exo_rs.pyi b/rust/exo_rs/exo_rs.pyi index f8105d21e..98ccbd649 100644 --- a/rust/exo_rs/exo_rs.pyi +++ b/rust/exo_rs/exo_rs.pyi @@ -2,16 +2,91 @@ # ruff: noqa: E501, F401, F403, F405 import builtins +import enum import os import pathlib import typing __all__ = [ + "CliArgs", + "ConfigArgs", + "DeprecatedArgs", "FromSwarm", + "LocatorArgs", + "LocatorConfig", "NetworkingHandle", "Pidfile", "PidfileError", + "Verbosity", ] +@typing.final +class CliArgs: + @property + def verbosity(self) -> Verbosity: ... + @property + def force_master(self) -> builtins.bool: ... + @property + def api_enabled(self) -> builtins.bool: ... + @property + def api_port(self) -> builtins.int: ... + @property + def worker_enabled(self) -> builtins.bool: ... + @property + def downloads_enabled(self) -> builtins.bool: ... + @property + def offline(self) -> builtins.bool: ... + @property + def continuous_batching_enabled(self) -> builtins.bool: ... + @property + def legacy_daemon(self) -> builtins.bool: ... + @property + def bootstrap_peers(self) -> typing.Optional[builtins.list[builtins.str]]: ... + @property + def namespace(self) -> builtins.str: ... + @property + def zenoh_port(self) -> builtins.int: ... + @property + def discovery_port(self) -> builtins.int: ... + @property + def fast_synch(self) -> typing.Optional[builtins.bool]: ... + @property + def locator(self) -> LocatorArgs: ... + @property + def config(self) -> ConfigArgs: ... + @property + def deprecated(self) -> DeprecatedArgs: ... + @staticmethod + def parse_from(argv: typing.Sequence[builtins.str]) -> CliArgs: ... + @staticmethod + def parse() -> CliArgs: ... + +@typing.final +class ConfigArgs: + r""" + Arguments that will end up in the final configuration go here. + + The precedence of the final configuration object will be: + - Defaults < Config file < Env < Cli args + + # Important + - Make sure all are [`Option`] so we can make it combinable with other + sources of configuration + """ + ... + +@typing.final +class DeprecatedArgs: + r""" + Deprecated arguments go here. + + # Important + - Make sure all are `hide = true` so it won't appear in `--help` + - Make sure all are [`Option`] so them being missing doesn't cause issues + - Edit [`Self::get_error`] to handle changes of new/removed args in here + """ + @property + def libp2p_port(self) -> typing.Optional[builtins.int]: ... + class FromSwarm: @typing.final class Connection(FromSwarm): @@ -31,6 +106,32 @@ class FromSwarm: ... +@typing.final +class LocatorArgs: + r""" + Arguments that are needed to resolve paths to files go here. + + This is needed for such things as resolving the path of the configuration `.toml` file, + therefore any args here cannot be specified by the configuration `.toml` file. + + By default, any path-like argument goes here, but can be moved to [`ConfigArgs`] if needed. + """ + @property + def exo_home(self) -> typing.Optional[pathlib.Path]: ... + @property + def config_file(self) -> typing.Optional[pathlib.Path]: ... + +@typing.final +class LocatorConfig: + @property + def exo_config_home(self) -> pathlib.Path: ... + @property + def exo_data_home(self) -> pathlib.Path: ... + @property + def exo_cache_home(self) -> pathlib.Path: ... + @property + def config_file(self) -> pathlib.Path: ... + @typing.final class NetworkingHandle: @staticmethod @@ -111,3 +212,12 @@ class PidfileError(builtins.Exception): def __repr__(self) -> builtins.str: ... def __str__(self) -> builtins.str: ... +@typing.final +class Verbosity(enum.Enum): + Off = ... + Error = ... + Warn = ... + Info = ... + Debug = ... + Trace = ... + diff --git a/rust/exo_rs/src/bin/stub_gen.rs b/rust/exo_rs/src/bin/stub_gen.rs index b1be28613..5238fc60d 100644 --- a/rust/exo_rs/src/bin/stub_gen.rs +++ b/rust/exo_rs/src/bin/stub_gen.rs @@ -1,12 +1,12 @@ use clap::Parser; -use exo_rs::config::cli::{CliArgs, EXO_VERSION}; +use exo_rs::config::cli::CliArgs; use pyo3_stub_gen::Result; fn main() -> Result<()> { let a = CliArgs::parse(); println!("{a:?}"); - println!("{:?} - {}", a.verbosity, EXO_VERSION); - return Ok(()); + println!("{:?}", a.verbosity); + // return Ok(()); env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init(); let stub = exo_rs::stub_info()?; diff --git a/rust/exo_rs/src/config/cli.rs b/rust/exo_rs/src/config/cli.rs index c7d98bd83..950ab48ed 100644 --- a/rust/exo_rs/src/config/cli.rs +++ b/rust/exo_rs/src/config/cli.rs @@ -1,12 +1,15 @@ use crate::config::locator::LocatorArgs; -use clap::{ArgAction, Parser, Subcommand, ValueEnum}; +use crate::version; +use clap::{ArgAction, Parser, ValueEnum}; +use pyo3::prelude::PyAnyMethods; +use pyo3::types::PyModule; +use pyo3::{pyclass, pymethods, PyResult, Python}; +use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods}; use serde::{Deserialize, Serialize}; +use std::ffi::OsString; -pub const EXO_VERSION: &str = match option_env!("EXO_PKG_VERSION") { - Some(v) => v, - None => env!("CARGO_PKG_VERSION"), -}; - +#[gen_stub_pyclass_enum] +#[pyclass(eq, eq_int, from_py_object)] #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ValueEnum)] #[serde(rename_all = "lowercase")] pub enum Verbosity { @@ -18,8 +21,10 @@ pub enum Verbosity { Trace, } +#[gen_stub_pyclass] +#[pyclass(from_py_object)] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Parser)] -#[command(name = "EXO", version = EXO_VERSION, about, long_about = None)] +#[command(name = "EXO", version = version::version(), about, long_about = None)] pub struct CliArgs { #[arg( short = 'v', @@ -29,6 +34,7 @@ pub struct CliArgs { value_name = "LEVEL", help = "Set the verbosity level" )] + #[pyo3(get)] pub verbosity: Verbosity, #[arg( @@ -37,6 +43,7 @@ pub struct CliArgs { action = ArgAction::SetTrue, help = "Force node to be master" )] + #[pyo3(get)] force_master: bool, #[arg( @@ -45,6 +52,7 @@ pub struct CliArgs { default_value_t = true, help = "Disable the API" )] + #[pyo3(get)] api_enabled: bool, #[arg( @@ -53,6 +61,7 @@ pub struct CliArgs { value_name = "PORT", help = "Port on which the API runs" )] + #[pyo3(get)] api_port: u16, #[arg( @@ -61,6 +70,7 @@ pub struct CliArgs { default_value_t = true, help = "Disable the worker" )] + #[pyo3(get)] worker_enabled: bool, #[arg( @@ -69,6 +79,7 @@ pub struct CliArgs { default_value_t = true, help = "Disable the download coordinator (node won't download models)" )] + #[pyo3(get)] downloads_enabled: bool, #[arg( @@ -76,6 +87,7 @@ pub struct CliArgs { action = ArgAction::SetTrue, help = "Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models" )] + #[pyo3(get)] offline: bool, #[arg( @@ -84,6 +96,7 @@ pub struct CliArgs { default_value_t = true, help = "Disable continuous batching, use sequential generation" )] + #[pyo3(get)] continuous_batching_enabled: bool, #[arg( @@ -91,6 +104,7 @@ pub struct CliArgs { action = ArgAction::SetTrue, help = "Run as a legacy SysV-style background daemon using double-fork daemonization" )] + #[pyo3(get)] legacy_daemon: bool, #[arg( @@ -99,15 +113,16 @@ pub struct CliArgs { value_name = "MULTIADDRS", help = "Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)" )] - #[deprecated] + #[pyo3(get)] bootstrap_peers: Option>, #[arg( long, - default_value_t = EXO_VERSION.to_string(), + default_value_t = version::version().to_string(), value_name = "STRING", help = "Discovery namespace, nodes with different namespaces will not connect." )] + #[pyo3(get)] namespace: String, #[arg( @@ -116,6 +131,7 @@ pub struct CliArgs { value_name = "PORT", help = "Fixed TCP port for zenoh to listen." )] + #[pyo3(get)] zenoh_port: u16, #[arg( @@ -124,6 +140,7 @@ pub struct CliArgs { value_name = "PORT", help = "Fixed UDP port for the discovery service." )] + #[pyo3(get)] discovery_port: u16, #[arg( @@ -131,18 +148,41 @@ pub struct CliArgs { value_name = "BOOL", help = "Force MLX FAST_SYNCH on/off (for JACCL backend); omit for auto" )] + #[pyo3(get)] fast_synch: Option, #[command(flatten)] + #[pyo3(get)] locator: LocatorArgs, #[command(flatten)] + #[pyo3(get)] config: ConfigArgs, #[command(flatten)] + #[pyo3(get)] deprecated: DeprecatedArgs, } +#[gen_stub_pymethods] +#[pymethods] +impl CliArgs { + #[staticmethod] + #[pyo3(name = "parse_from")] + pub fn py_parse_from(argv: Vec) -> Self { + CliArgs::parse_from(argv) + } + + #[staticmethod] + #[pyo3(name = "parse")] + pub fn py_parse(py: Python<'_>) -> PyResult { + // the correct CLI args to parse is `sys.argv`, because the original ones + // (i.e. `sys.orig_argv`) may contain extra arguments which would mess up parsing + let argv: Vec = PyModule::import(py, "sys")?.getattr("argv")?.extract()?; + Ok(CliArgs::parse_from(argv)) + } +} + /// Arguments that will end up in the final configuration go here. /// /// The precedence of the final configuration object will be: @@ -151,7 +191,10 @@ pub struct CliArgs { /// # Important /// - Make sure all are [`Option`] so we can make it combinable with other /// sources of configuration +#[gen_stub_pyclass] +#[pyclass(from_py_object)] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, clap::Args)] +#[command(about = None, long_about = None)] pub struct ConfigArgs {} /// Deprecated arguments go here. @@ -160,9 +203,13 @@ pub struct ConfigArgs {} /// - Make sure all are `hide = true` so it won't appear in `--help` /// - Make sure all are [`Option`] so them being missing doesn't cause issues /// - Edit [`Self::get_error`] to handle changes of new/removed args in here +#[gen_stub_pyclass] +#[pyclass(from_py_object)] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, clap::Args)] +#[command(about = None, long_about = None)] pub struct DeprecatedArgs { #[arg(long = "libp2p-port", hide = true)] + #[pyo3(get)] libp2p_port: Option, } @@ -184,7 +231,3 @@ impl DeprecatedArgs { } } } - -fn env_flag(name: &str) -> bool { - std::env::var(name).is_ok_and(|value| value.eq_ignore_ascii_case("true")) -} diff --git a/rust/exo_rs/src/config/locator.rs b/rust/exo_rs/src/config/locator.rs index d9e974796..d209d18b3 100644 --- a/rust/exo_rs/src/config/locator.rs +++ b/rust/exo_rs/src/config/locator.rs @@ -1,3 +1,5 @@ +use pyo3::pyclass; +use pyo3_stub_gen::derive::gen_stub_pyclass; use serde::{Deserialize, Serialize}; use std::path::PathBuf; use std::{fs, io}; @@ -8,7 +10,10 @@ use std::{fs, io}; /// therefore any args here cannot be specified by the configuration `.toml` file. /// /// By default, any path-like argument goes here, but can be moved to [`ConfigArgs`] if needed. +#[gen_stub_pyclass] +#[pyclass(from_py_object)] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, clap::Args)] +#[command(about = None, long_about = None)] pub struct LocatorArgs { #[arg( long, @@ -16,23 +21,33 @@ pub struct LocatorArgs { value_name = "PATH", help = "Path to Exo's home directory" )] + #[pyo3(get)] pub exo_home: Option, + #[arg( long, env = "EXO_CONFIG_FILE", value_name = "PATH", help = "Path to Exo's .toml config file" )] + #[pyo3(get)] pub config_file: Option, } +#[gen_stub_pyclass] +#[pyclass(from_py_object)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct LocatorConfig { // base XDG directories + #[pyo3(get)] pub exo_config_home: PathBuf, + #[pyo3(get)] pub exo_data_home: PathBuf, + #[pyo3(get)] pub exo_cache_home: PathBuf, // other + #[pyo3(get)] pub config_file: PathBuf, } diff --git a/rust/exo_rs/src/config/mod.rs b/rust/exo_rs/src/config/mod.rs index fc171177d..4ca52ca01 100644 --- a/rust/exo_rs/src/config/mod.rs +++ b/rust/exo_rs/src/config/mod.rs @@ -1,3 +1,19 @@ +use crate::config::cli::{CliArgs, ConfigArgs, DeprecatedArgs, Verbosity}; +use crate::config::locator::{LocatorArgs, LocatorConfig}; +use pyo3::prelude::{PyModule, PyModuleMethods}; +use pyo3::{Bound, PyResult}; + pub mod cli; pub mod defaults; pub mod locator; + +pub fn config_submodule(m: &Bound) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + Ok(()) +} diff --git a/rust/exo_rs/src/lib.rs b/rust/exo_rs/src/lib.rs index f7bfe64a1..557264857 100644 --- a/rust/exo_rs/src/lib.rs +++ b/rust/exo_rs/src/lib.rs @@ -9,10 +9,11 @@ pub mod config; mod networking; mod pidfile; +use crate::config::config_submodule; use crate::networking::networking_submodule; use crate::pidfile::pidfile_submodule; use pyo3::prelude::PyModule; -use pyo3::{Bound, PyResult, pymodule}; +use pyo3::{pymodule, Bound, PyResult}; use pyo3_stub_gen::define_stub_info_gatherer; /// Namespace for crate-wide extension traits/methods @@ -157,6 +158,48 @@ pub(crate) mod ext { } } +/// Resolving the version of the python project +pub(crate) mod version { + use pyo3::exceptions::PyRuntimeError; + use pyo3::prelude::PyAnyMethods; + use pyo3::types::PyModule; + use pyo3::{PyResult, Python}; + use std::env; + use std::sync::OnceLock; + + const DEFAULT_VERSION: &str = env!("CARGO_PKG_VERSION"); + static VERSION: OnceLock = OnceLock::new(); + + /// Returns either the configured version of Exo (once set by [`set_version_once`]) + /// or falls back to `CARGO_PKG_VERSION` if that hasn't been configured. + pub fn version() -> &'static str { + VERSION.get().map_or(DEFAULT_VERSION, String::as_str) + } + + /// First tries to find `EXO_PKG_VERSION` env-var, falls back to calling Python + /// `importlib.metadata.version("exo")` to resolve the version of Exo + pub fn set_version_once(py: Python<'_>) -> PyResult<()> { + let v = if let Ok(v) = env::var("EXO_PKG_VERSION") { + v + } else { + // essentially runs: + // ```python + // from importlib.metadata import version + // version("exo") + // ``` + PyModule::import(py, "importlib.metadata")? + .getattr("version")? + .call1(("exo",))? + .extract()? + }; + + // sets version only once + VERSION + .set(v) + .map_err(|_| PyRuntimeError::new_err("Cannot set exo_rs version twice".to_string())) + } +} + /// A Python module implemented in Rust. The name of this function must match /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to /// import the module. @@ -164,6 +207,11 @@ pub(crate) mod ext { fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> { // install logger pyo3_log::init(); + + // resolve version + version::set_version_once(m.py())?; + + // configure runtime let mut builder = tokio::runtime::Builder::new_multi_thread(); builder.enable_all(); pyo3_async_runtimes::tokio::init(builder); @@ -172,8 +220,8 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> { // work with maturin, where the types generate correctly, in the right folder, without // too many importing issues... pidfile_submodule(m)?; - // m.add_class::()?; networking_submodule(m)?; + config_submodule(m)?; // top-level constructs // TODO: ... diff --git a/rust/exo_rs/src/networking.rs b/rust/exo_rs/src/networking.rs index 0f21df203..6b1f82a25 100644 --- a/rust/exo_rs/src/networking.rs +++ b/rust/exo_rs/src/networking.rs @@ -9,7 +9,7 @@ use networking::{is_valid_zid, Session}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::PyBytes; -use pyo3::{pymethods, Bound, Py, PyAny, PyErr, PyResult, Python}; +use pyo3::{pymethods, Bound, Py, PyErr, PyResult, Python}; use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods}; use tokio::sync::{mpsc, oneshot, Mutex}; diff --git a/src/exo/__init__.py b/src/exo/__init__.py index f7d034e53..06cbb6944 100644 --- a/src/exo/__init__.py +++ b/src/exo/__init__.py @@ -1,3 +1,6 @@ +import os from importlib.metadata import version +# set __version__ and env-var __version__ = version("exo") +os.environ["EXO_PKG_VERSION"] = __version__ diff --git a/src/exo/main.py b/src/exo/main.py index f504785e6..016ad986e 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -10,7 +10,7 @@ from typing import Self import anyio from anyio.lowlevel import checkpoint as anyio_checkpoint from daemon import DaemonContext # pyright: ignore[reportMissingTypeStubs] -from exo_rs import Pidfile, PidfileError +from exo_rs import CliArgs, Pidfile, PidfileError from loguru import logger from pydantic import PositiveInt @@ -276,6 +276,8 @@ class Node: def main(): + # a = CliArgs.parse() + # Parse args first => --help or bad args don't require PID-locking args = Args.parse()