Clear ephemeral overlays when running tools (#11141)

## Summary

This PR removes the ephemeral `.pth` overlay when using a cached
environment. This solution isn't _completely_ safe, since we could
remove the `.pth` file just as another process is starting the
environment... But that risk already exists today, since we could
_overwrite_ the `.pth` file just as another process is starting the
environment, so I think what I've added here is a strict improvement.

Ideally, we wouldn't write this file at all, and we'd instead somehow
(e.g.) pass a file to the interpreter to run at startup? Or find some
other solution that doesn't require poisoning the cache like this.

Closes https://github.com/astral-sh/uv/issues/11117.

# Test Plan

Ran through the great reproduction steps from the linked issue.

Before:

![Screenshot 2025-01-31 at 2 11
31 PM](https://github.com/user-attachments/assets/d36e1db5-27b1-483a-9ced-bec67bd7081d)

After:

![Screenshot 2025-01-31 at 2 11
39 PM](https://github.com/user-attachments/assets/1f963ce0-7903-4acd-9fd6-753374c31705)
This commit is contained in:
Charlie Marsh 2025-02-04 17:45:45 -05:00 committed by GitHub
parent 2fad82c735
commit f615e81ad5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 185 additions and 91 deletions

View File

@ -1,12 +1,5 @@
use tracing::debug; use tracing::debug;
use crate::commands::pip::loggers::{InstallLogger, ResolveLogger};
use crate::commands::project::install_target::InstallTarget;
use crate::commands::project::{
resolve_environment, sync_environment, EnvironmentSpecification, PlatformState, ProjectError,
};
use crate::printer::Printer;
use crate::settings::ResolverInstallerSettings;
use uv_cache::{Cache, CacheBucket}; use uv_cache::{Cache, CacheBucket};
use uv_cache_key::{cache_digest, hash_digest}; use uv_cache_key::{cache_digest, hash_digest};
use uv_client::Connectivity; use uv_client::Connectivity;
@ -17,6 +10,14 @@ use uv_distribution_types::{Name, Resolution};
use uv_python::{Interpreter, PythonEnvironment}; use uv_python::{Interpreter, PythonEnvironment};
use uv_resolver::Installable; use uv_resolver::Installable;
use crate::commands::pip::loggers::{InstallLogger, ResolveLogger};
use crate::commands::project::install_target::InstallTarget;
use crate::commands::project::{
resolve_environment, sync_environment, EnvironmentSpecification, PlatformState, ProjectError,
};
use crate::printer::Printer;
use crate::settings::ResolverInstallerSettings;
/// A [`PythonEnvironment`] stored in the cache. /// A [`PythonEnvironment`] stored in the cache.
#[derive(Debug)] #[derive(Debug)]
pub(crate) struct CachedEnvironment(PythonEnvironment); pub(crate) struct CachedEnvironment(PythonEnvironment);
@ -215,6 +216,36 @@ impl CachedEnvironment {
Ok(Self(PythonEnvironment::from_root(root, cache)?)) Ok(Self(PythonEnvironment::from_root(root, cache)?))
} }
/// Set the ephemeral overlay for a Python environment.
#[allow(clippy::result_large_err)]
pub(crate) fn set_overlay(&self, contents: impl AsRef<[u8]>) -> Result<(), ProjectError> {
let site_packages = self
.0
.site_packages()
.next()
.ok_or(ProjectError::NoSitePackages)?;
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
fs_err::write(overlay_path, contents)?;
Ok(())
}
/// Clear the ephemeral overlay for a Python environment, if it exists.
#[allow(clippy::result_large_err)]
pub(crate) fn clear_overlay(&self) -> Result<(), ProjectError> {
let site_packages = self
.0
.site_packages()
.next()
.ok_or(ProjectError::NoSitePackages)?;
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
match fs_err::remove_file(overlay_path) {
Ok(()) => (),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => (),
Err(err) => return Err(ProjectError::OverlayRemoval(err)),
}
Ok(())
}
/// Convert the [`CachedEnvironment`] into an [`Interpreter`]. /// Convert the [`CachedEnvironment`] into an [`Interpreter`].
pub(crate) fn into_interpreter(self) -> Interpreter { pub(crate) fn into_interpreter(self) -> Interpreter {
self.0.into_interpreter() self.0.into_interpreter()

View File

@ -176,6 +176,12 @@ pub(crate) enum ProjectError {
#[error("Failed to parse PEP 723 script metadata")] #[error("Failed to parse PEP 723 script metadata")]
Pep723ScriptTomlParse(#[source] toml::de::Error), Pep723ScriptTomlParse(#[source] toml::de::Error),
#[error("Failed to remove ephemeral overlay")]
OverlayRemoval(#[source] std::io::Error),
#[error("Failed to find `site-packages` directory for environment")]
NoSitePackages,
#[error(transparent)] #[error(transparent)]
DependencyGroup(#[from] DependencyGroupError), DependencyGroup(#[from] DependencyGroupError),

View File

@ -283,6 +283,9 @@ pub(crate) async fn run(
Err(err) => return Err(err.into()), Err(err) => return Err(err.into()),
}; };
// Clear any existing overlay.
environment.clear_overlay()?;
Some(environment.into_interpreter()) Some(environment.into_interpreter())
} else { } else {
// If no lockfile is found, warn against `--locked` and `--frozen`. // If no lockfile is found, warn against `--locked` and `--frozen`.
@ -426,6 +429,9 @@ pub(crate) async fn run(
Err(err) => return Err(err.into()), Err(err) => return Err(err.into()),
}; };
// Clear any existing overlay.
environment.clear_overlay()?;
Some(environment.into_interpreter()) Some(environment.into_interpreter())
} else { } else {
// Create a virtual environment. // Create a virtual environment.
@ -911,98 +917,76 @@ pub(crate) async fn run(
}; };
// If necessary, create an environment for the ephemeral requirements or command. // If necessary, create an environment for the ephemeral requirements or command.
let temp_dir; let ephemeral_env = match spec {
let ephemeral_env = if can_skip_ephemeral(spec.as_ref(), &base_interpreter, &settings) { None => None,
None Some(spec) if can_skip_ephemeral(&spec, &base_interpreter, &settings) => None,
} else { Some(spec) => {
debug!("Creating ephemeral environment"); debug!("Syncing ephemeral requirements");
Some(match spec.filter(|spec| !spec.is_empty()) { let result = CachedEnvironment::from_spec(
None => { EnvironmentSpecification::from(spec).with_lock(
// Create a virtual environment lock.as_ref()
temp_dir = cache.venv_dir()?; .map(|(lock, install_path)| (lock, install_path.as_ref())),
uv_virtualenv::create_venv( ),
temp_dir.path(), &base_interpreter,
base_interpreter.clone(), &settings,
uv_virtualenv::Prompt::None, &sync_state,
false, if show_resolution {
false, Box::new(DefaultResolveLogger)
false, } else {
false, Box::new(SummaryResolveLogger)
)? },
} if show_resolution {
Some(spec) => { Box::new(DefaultInstallLogger)
debug!("Syncing ephemeral requirements"); } else {
Box::new(SummaryInstallLogger)
},
installer_metadata,
connectivity,
concurrency,
native_tls,
allow_insecure_host,
cache,
printer,
preview,
)
.await;
let result = CachedEnvironment::from_spec( let environment = match result {
EnvironmentSpecification::from(spec).with_lock( Ok(resolution) => resolution,
lock.as_ref() Err(ProjectError::Operation(err)) => {
.map(|(lock, install_path)| (lock, install_path.as_ref())), return diagnostics::OperationDiagnostic::native_tls(native_tls)
), .with_context("`--with`")
&base_interpreter, .report(err)
&settings, .map_or(Ok(ExitStatus::Failure), |err| Err(err.into()))
&sync_state, }
if show_resolution { Err(err) => return Err(err.into()),
Box::new(DefaultResolveLogger) };
} else {
Box::new(SummaryResolveLogger)
},
if show_resolution {
Box::new(DefaultInstallLogger)
} else {
Box::new(SummaryInstallLogger)
},
installer_metadata,
connectivity,
concurrency,
native_tls,
allow_insecure_host,
cache,
printer,
preview,
)
.await;
let environment = match result { Some(environment)
Ok(resolution) => resolution, }
Err(ProjectError::Operation(err)) => {
return diagnostics::OperationDiagnostic::native_tls(native_tls)
.with_context("`--with`")
.report(err)
.map_or(Ok(ExitStatus::Failure), |err| Err(err.into()))
}
Err(err) => return Err(err.into()),
};
environment.into()
}
})
}; };
// If we're running in an ephemeral environment, add a path file to enable loading of // If we're running in an ephemeral environment, add a path file to enable loading of
// the base environment's site packages. Setting `PYTHONPATH` is insufficient, as it doesn't // the base environment's site packages. Setting `PYTHONPATH` is insufficient, as it doesn't
// resolve `.pth` files in the base environment. // resolve `.pth` files in the base environment.
// And `sitecustomize.py` would be an alternative but it can be shadowed by an existing such //
// `sitecustomize.py` would be an alternative, but it can be shadowed by an existing such
// module in the python installation. // module in the python installation.
if let Some(ephemeral_env) = ephemeral_env.as_ref() { if let Some(ephemeral_env) = ephemeral_env.as_ref() {
let ephemeral_site_packages = ephemeral_env let site_packages = base_interpreter
.site_packages() .site_packages()
.next() .next()
.ok_or_else(|| anyhow!("Ephemeral environment has no site packages directory"))?; .ok_or_else(|| ProjectError::NoSitePackages)?;
let base_site_packages = base_interpreter ephemeral_env.set_overlay(format!(
.site_packages() "import site; site.addsitedir(\"{}\")",
.next() site_packages.escape_for_python()
.ok_or_else(|| anyhow!("Base environment has no site packages directory"))?; ))?;
fs_err::write(
ephemeral_site_packages.join("_uv_ephemeral_overlay.pth"),
format!(
"import site; site.addsitedir(\"{}\")",
base_site_packages.escape_for_python()
),
)?;
} }
// Cast from `CachedEnvironment` to `PythonEnvironment`.
let ephemeral_env = ephemeral_env.map(PythonEnvironment::from);
// Determine the Python interpreter to use for the command, if necessary. // Determine the Python interpreter to use for the command, if necessary.
let interpreter = ephemeral_env let interpreter = ephemeral_env
.as_ref() .as_ref()
@ -1125,15 +1109,10 @@ pub(crate) async fn run(
/// Returns `true` if we can skip creating an additional ephemeral environment in `uv run`. /// Returns `true` if we can skip creating an additional ephemeral environment in `uv run`.
fn can_skip_ephemeral( fn can_skip_ephemeral(
spec: Option<&RequirementsSpecification>, spec: &RequirementsSpecification,
base_interpreter: &Interpreter, base_interpreter: &Interpreter,
settings: &ResolverInstallerSettings, settings: &ResolverInstallerSettings,
) -> bool { ) -> bool {
// No additional requirements.
let Some(spec) = spec.as_ref() else {
return true;
};
let Ok(site_packages) = SitePackages::from_interpreter(base_interpreter) else { let Ok(site_packages) = SitePackages::from_interpreter(base_interpreter) else {
return false; return false;
}; };

View File

@ -766,5 +766,8 @@ async fn get_or_create_environment(
}, },
}; };
// Clear any existing overlay.
environment.clear_overlay()?;
Ok((from, environment.into())) Ok((from, environment.into()))
} }

View File

@ -3928,11 +3928,86 @@ fn run_repeated() -> Result<()> {
uv_snapshot!( uv_snapshot!(
context.filters(), context.filters(),
context.tool_run().arg("--with").arg("typing-extensions").arg("python").arg("-c").arg("import typing_extensions; import iniconfig"), @r###" context.tool_run().arg("--with").arg("typing-extensions").arg("python").arg("-c").arg("import typing_extensions; import iniconfig"), @r###"
success: false
exit_code: 1
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Traceback (most recent call last):
File "<string>", line 1, in <module>
import typing_extensions; import iniconfig
^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'iniconfig'
"###);
Ok(())
}
/// See: <https://github.com/astral-sh/uv/issues/11117>
#[test]
fn run_without_overlay() -> Result<()> {
let context = TestContext::new_with_versions(&["3.13"]);
let pyproject_toml = context.temp_dir.child("pyproject.toml");
pyproject_toml.write_str(indoc! { r#"
[project]
name = "foo"
version = "1.0.0"
requires-python = ">=3.11, <4"
dependencies = ["iniconfig"]
"#
})?;
// Import `iniconfig` in the context of the project.
uv_snapshot!(
context.filters(),
context.run().arg("--with").arg("typing-extensions").arg("python").arg("-c").arg("import typing_extensions; import iniconfig"), @r###"
success: true success: true
exit_code: 0 exit_code: 0
----- stdout ----- ----- stdout -----
----- stderr ----- ----- stderr -----
Using CPython 3.13.[X] interpreter at: [PYTHON-3.13]
Creating virtual environment at: .venv
Resolved 2 packages in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ iniconfig==2.0.0
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ typing-extensions==4.10.0
"###);
// Import `iniconfig` in the context of a `tool run` command, which should fail.
uv_snapshot!(
context.filters(),
context.tool_run().arg("--with").arg("typing-extensions").arg("python").arg("-c").arg("import typing_extensions; import iniconfig"), @r###"
success: false
exit_code: 1
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Traceback (most recent call last):
File "<string>", line 1, in <module>
import typing_extensions; import iniconfig
^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'iniconfig'
"###);
// Re-running in the context of the project should reset the overlay.
uv_snapshot!(
context.filters(),
context.run().arg("--with").arg("typing-extensions").arg("python").arg("-c").arg("import typing_extensions; import iniconfig"), @r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
Resolved 2 packages in [TIME]
Audited 1 package in [TIME]
Resolved 1 package in [TIME] Resolved 1 package in [TIME]
"###); "###);