diff --git a/crates/uv-settings/src/lib.rs b/crates/uv-settings/src/lib.rs index 20d276f35a1..a1e358862ed 100644 --- a/crates/uv-settings/src/lib.rs +++ b/crates/uv-settings/src/lib.rs @@ -775,6 +775,8 @@ pub struct EnvironmentOptions { pub init_bare: EnvFlag, pub malware_check: EnvFlag, pub malware_check_url: Option, + #[cfg(unix)] + pub run_ulimit: Option, } impl EnvironmentOptions { @@ -909,6 +911,8 @@ impl EnvironmentOptions { }) }) .transpose()?, + #[cfg(unix)] + run_ulimit: parse_integer_environment_variable(EnvVars::UV_RUN_ULIMIT, None)?, }) } } diff --git a/crates/uv-static/src/env_vars.rs b/crates/uv-static/src/env_vars.rs index 89b0c77af6e..e83c0459166 100644 --- a/crates/uv-static/src/env_vars.rs +++ b/crates/uv-static/src/env_vars.rs @@ -1346,6 +1346,13 @@ impl EnvVars { #[attr_added_in("0.5.19")] pub const UV_GIT_LFS: &'static str = "UV_GIT_LFS"; + /// Sets the soft open-file descriptor limit for commands executed by `uv run`. + /// + /// The limit is applied after uv prepares the environment and immediately before the command + /// is spawned. The hard open-file descriptor limit remains unchanged. Only supported on Unix. + #[attr_added_in("0.12.2")] + pub const UV_RUN_ULIMIT: &'static str = "UV_RUN_ULIMIT"; + /// Number of times that `uv run` has been recursively invoked. Used to guard against infinite /// recursion, e.g., when `uv run`` is used in a script shebang. #[attr_hidden] diff --git a/crates/uv-unix/src/lib.rs b/crates/uv-unix/src/lib.rs index e276391a87d..fc40ae00e35 100644 --- a/crates/uv-unix/src/lib.rs +++ b/crates/uv-unix/src/lib.rs @@ -6,4 +6,4 @@ mod resource_limits; -pub use resource_limits::{OpenFileLimitError, adjust_open_file_limit}; +pub use resource_limits::{OpenFileLimitError, adjust_open_file_limit, set_open_file_limit}; diff --git a/crates/uv-unix/src/resource_limits.rs b/crates/uv-unix/src/resource_limits.rs index 1ea8a042113..e28b086e693 100644 --- a/crates/uv-unix/src/resource_limits.rs +++ b/crates/uv-unix/src/resource_limits.rs @@ -14,7 +14,7 @@ //! See: use nix::errno::Errno; -use nix::sys::resource::{Resource, getrlimit, rlim_t, setrlimit}; +use nix::sys::resource::{RLIM_INFINITY, Resource, getrlimit, rlim_t, setrlimit}; use thiserror::Error; /// Errors that can occur when adjusting resource limits. @@ -29,7 +29,10 @@ pub enum OpenFileLimitError { #[error("soft limit ({current}) already meets the target ({target})")] AlreadySufficient { current: u64, target: u64 }, - #[error("failed to raise open file limit from {current} to {target}: {}", source.desc())] + #[error("requested open file limit ({target}) exceeds the hard limit ({hard})")] + ExceedsHardLimit { target: u64, hard: rlim_t }, + + #[error("failed to set open file limit from {current} to {target}: {}", source.desc())] SetLimitFailed { current: u64, target: u64, @@ -90,9 +93,36 @@ pub fn adjust_open_file_limit() -> Result { // Safe because target <= MAX_NOFILE_LIMIT which fits in both i64 and u64. let target_rlim = target as rlim_t; + set_open_file_limit_to(soft, target, target_rlim, hard) +} + +/// Set the soft open-file descriptor limit while preserving the hard limit. +pub fn set_open_file_limit(target: u32) -> Result { + let (soft, hard) = + getrlimit(Resource::RLIMIT_NOFILE).map_err(OpenFileLimitError::GetLimitFailed)?; + let Some(soft) = rlim_t_to_u64(soft) else { + return Err(OpenFileLimitError::NegativeSoftLimit { value: soft }); + }; + + let target_rlim = rlim_t::from(target); + let target = u64::from(target); + if hard != RLIM_INFINITY && target_rlim > hard { + return Err(OpenFileLimitError::ExceedsHardLimit { target, hard }); + } + + set_open_file_limit_to(soft, target, target_rlim, hard) +} + +/// Update the soft open-file descriptor limit while preserving the hard limit. +fn set_open_file_limit_to( + current: u64, + target: u64, + target_rlim: rlim_t, + hard: rlim_t, +) -> Result { setrlimit(Resource::RLIMIT_NOFILE, target_rlim, hard).map_err(|err| { OpenFileLimitError::SetLimitFailed { - current: soft, + current, target, source: err, } diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index 0d2a6ff828b..c8bc6fb4b69 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -121,6 +121,7 @@ pub(crate) async fn run( preview: Preview, max_recursion_depth: u32, malware_settings: MalwareCheckSettings, + #[cfg(unix)] run_ulimit: Option, ) -> anyhow::Result { // Check if max recursion depth was exceeded. This most commonly happens // for scripts with a shebang line like `#!/usr/bin/env -S uv run`, so try @@ -1303,6 +1304,16 @@ pub(crate) async fn run( process.env(EnvVars::VIRTUAL_ENV, interpreter.sys_prefix().as_os_str()); } + #[cfg(unix)] + if let Some(limit) = run_ulimit { + uv_unix::set_open_file_limit(limit).with_context(|| { + format!( + "Failed to apply `{}` value `{limit}`", + EnvVars::UV_RUN_ULIMIT + ) + })?; + } + // Spawn and wait for completion // Standard input, output, and error streams are all inherited // TODO(zanieb): Throw a nicer error message if the command is not found diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index a8371114396..d940237c77f 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -2342,6 +2342,8 @@ async fn run_project( globals.preview, args.max_recursion_depth, args.malware_settings, + #[cfg(unix)] + args.run_ulimit, )) .await } diff --git a/crates/uv/src/settings.rs b/crates/uv/src/settings.rs index 0325d8bd8f7..8dbe35661e3 100644 --- a/crates/uv/src/settings.rs +++ b/crates/uv/src/settings.rs @@ -681,6 +681,8 @@ pub(crate) struct RunSettings { pub(crate) env_file: EnvFile, pub(crate) max_recursion_depth: u32, pub(crate) malware_settings: MalwareCheckSettings, + #[cfg(unix)] + pub(crate) run_ulimit: Option, } impl RunSettings { @@ -847,6 +849,8 @@ impl RunSettings { .combine(filesystem_install_mirrors), max_recursion_depth: max_recursion_depth.unwrap_or(Self::DEFAULT_MAX_RECURSION_DEPTH), malware_settings, + #[cfg(unix)] + run_ulimit: environment.run_ulimit, }) } } diff --git a/crates/uv/tests/it/resource_limits.rs b/crates/uv/tests/it/resource_limits.rs index 97163b4da3f..4faf8129c0b 100644 --- a/crates/uv/tests/it/resource_limits.rs +++ b/crates/uv/tests/it/resource_limits.rs @@ -30,3 +30,86 @@ fn adjust_open_file_limit() { True "); } + +#[test] +fn run_open_file_limit_override() { + let context = uv_test::test_context!("3.12"); + let python = &context.python_versions[0].1; + + let mut command = Command::new(get_bin!()); + command + .arg("run") + .arg("--no-project") + .arg("--") + .arg(python) + .arg("-c") + .arg( + "import resource; soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE); print(soft); print(hard > soft)", + ) + .current_dir(context.temp_dir.path()) + .env(EnvVars::UV_CACHE_DIR, context.cache_dir.path()) + .env(EnvVars::UV_PYTHON_DOWNLOADS, "never") + .env(EnvVars::UV_RUN_ULIMIT, "128"); + + uv_snapshot!(context.filters(), command, @r" + exit_code: 0 (success) + ----- stdout ----- + 128 + True + "); +} + +#[test] +fn run_open_file_limit_override_invalid() { + let context = uv_test::test_context!("3.12"); + let python = &context.python_versions[0].1; + + let mut command = Command::new(get_bin!()); + command + .arg("run") + .arg("--no-project") + .arg("--") + .arg(python) + .arg("-c") + .arg("pass") + .current_dir(context.temp_dir.path()) + .env(EnvVars::UV_CACHE_DIR, context.cache_dir.path()) + .env(EnvVars::UV_PYTHON_DOWNLOADS, "never") + .env(EnvVars::UV_RUN_ULIMIT, "invalid"); + + uv_snapshot!(context.filters(), command, @r" + exit_code: 2 (failure) + ----- stderr ----- + error: Failed to parse environment variable `UV_RUN_ULIMIT` with invalid value `invalid`: invalid digit found in string + "); +} + +#[test] +fn run_open_file_limit_override_exceeds_hard_limit() { + let context = uv_test::test_context!("3.12"); + let python = &context.python_versions[0].1; + + let mut command = Command::new("sh"); + command + .arg("-c") + .arg("ulimit -S -n 128; ulimit -H -n 128; exec \"$@\"") + .arg("sh") + .arg(get_bin!()) + .arg("run") + .arg("--no-project") + .arg("--") + .arg(python) + .arg("-c") + .arg("pass") + .current_dir(context.temp_dir.path()) + .env(EnvVars::UV_CACHE_DIR, context.cache_dir.path()) + .env(EnvVars::UV_PYTHON_DOWNLOADS, "never") + .env(EnvVars::UV_RUN_ULIMIT, "256"); + + uv_snapshot!(context.filters(), command, @r" + exit_code: 2 (failure) + ----- stderr ----- + error: Failed to apply `UV_RUN_ULIMIT` value `256` + Caused by: requested open file limit (256) exceeds the hard limit (128) + "); +}