feat: add configurable PGO dependency group (#3278)

guide/src/config.md

@@ -95,6 +95,8 @@ use-base-python = false
 # virtualenv with the instrumented wheel installed.
 # Example: "python -m pytest tests/benchmarks"
 pgo-command = "python -m pytest tests/benchmarks"
+# Dependency group installed before the PGO training command (defaults to "dev")
+pgo-dependency-group = "testing-extra"
 # Select which Cargo compile targets to build when the crate defines more than
 # one matching target. Each entry is an object with a required `name` (as in
 # Cargo.toml) and an optional `kind` (`bin`, `cdylib`, `dylib`, `lib`, `rlib`,
@@ -159,6 +161,25 @@ pgo-command = "python -m pytest tests/benchmarks"
 See the `--pgo` option under [Build](./distribution.md) for the overall
 three-phase flow.
 
+#### `pgo-dependency-group`
+
+Dependency group installed into the temporary PGO virtualenv before running
+`pgo-command`. It defaults to `dev`. Set it when the training workload uses
+a different PEP 735 dependency group:
+
+```toml
+[dependency-groups]
+testing-extra = ["pytest", "pytest-benchmark"]
+
+[tool.maturin]
+pgo-dependency-group = "testing-extra"
+```
+
+When uv is available, maturin exports the selected group before installing it
+so versions from the uv lockfile are respected without replacing the temporary
+PGO environment. Otherwise maturin installs the group with pip's `--group`
+support.
+
 #### `generate-ci`
 
 Defaults for `maturin generate-ci` (currently GitHub Actions). CLI flags that

maturin.schema.json

@@ -156,6 +156,13 @@
         "null"
       ]
     },
+    "pgo-dependency-group": {
+      "description": "Dependency group to install before PGO profile generation.\nDefaults to `dev`.",
+      "type": [
+        "string",
+        "null"
+      ]
+    },
     "profile": {
       "description": "Build artifacts with the specified Cargo profile",
       "type": [

src/pgo.rs

@@ -176,6 +176,7 @@ impl PgoContext {
         let host_target = Target::from_target_triple(None)?;
         let venv_python = host_target.get_venv_python(venv_path);
         let venv_bin_dir = host_target.get_venv_bin_dir(venv_path);
+        let project_dir = build_context.project.project_layout.project_root.as_path();
 
         // Install the instrumented wheel
         eprintln!("📦 Installing instrumented wheel into temporary venv...");
@@ -206,27 +207,64 @@ impl PgoContext {
             }
         }
 
-        // Install dev dependency group if present (pip only — uv doesn't support --group yet)
-        if uv.is_none() {
-            let has_dev_group = build_context
-                .project
-                .pyproject_toml
-                .as_ref()
-                .and_then(|p| p.dependency_groups.as_ref())
-                .is_some_and(|dg| dg.0.contains_key("dev"));
-            if has_dev_group {
+        // Install the configured PGO dependency group, defaulting to `dev`.
+        // With uv, export the group first so the project's lockfile is respected
+        // without syncing away the already-installed instrumented wheel.
+        let dependency_group = build_context
+            .project
+            .pyproject_toml
+            .as_ref()
+            .and_then(|p| p.pgo_dependency_group())
+            .unwrap_or("dev");
+        let has_dependency_group = build_context
+            .project
+            .pyproject_toml
+            .as_ref()
+            .and_then(|p| p.dependency_groups.as_ref())
+            .is_some_and(|dg| dg.0.contains_key(dependency_group));
+        if has_dependency_group {
+            if let Some((uv_path, uv_args)) = &uv {
+                let requirements_path = venv_path.join("pgo-requirements.txt");
+                debug!("Exporting {dependency_group} dependency group with uv");
+                let status = Command::new(uv_path)
+                    .args(uv_args.iter().copied())
+                    .args(["export", "--group", dependency_group, "--no-emit-workspace"])
+                    .args(["--no-hashes", "--no-header", "--no-annotate", "-o"])
+                    .arg(&requirements_path)
+                    .current_dir(project_dir)
+                    .status()
+                    .context("Failed to export PGO dependency group with uv")?;
+                if !status.success() {
+                    eprintln!(
+                        "⚠️  Warning: failed to export {dependency_group} dependency group with uv"
+                    );
+                } else {
+                    debug!("Installing {dependency_group} dependency group from uv export");
+                    let status = self.pip_install(
+                        &uv,
+                        &venv_python,
+                        &["-r"],
+                        &[requirements_path.as_path()],
+                    )?;
+                    if !status.success() {
+                        eprintln!(
+                            "⚠️  Warning: failed to install {dependency_group} dependency group"
+                        );
+                    }
+                }
+            } else {
                 let pyproject_group = format!(
-                    "{}:dev",
+                    "{}:{dependency_group}",
                     dunce::simplified(&build_context.project.pyproject_toml_path).display()
                 );
-                debug!("Installing dev dependency group");
+                debug!("Installing {dependency_group} dependency group");
                 let status = Command::new(&venv_python)
                     .args(["-m", "pip", "install", "--group", &pyproject_group])
                     .status()
-                    .context("Failed to install dev dependency group")?;
+                    .context("Failed to install PGO dependency group")?;
                 if !status.success() {
                     eprintln!(
-                        "⚠️  Warning: failed to install dev dependency group \
+                        "⚠️  Warning: failed to install {dependency_group} dependency group \
                          (pip >= 25.1 required for --group support)"
                     );
                 }
@@ -245,8 +283,6 @@ impl PgoContext {
         let sep = if host_target.is_windows() { ";" } else { ":" };
         let path_env = format!("{}{sep}{current_path}", venv_bin_dir.display());
 
-        let project_dir = build_context.project.project_layout.project_root.as_path();
-
         // Run through the system shell with the venv's bin dir prepended to PATH,
         // so that `python`, `pytest`, etc. resolve to the venv's copies.
         let mut cmd = if cfg!(windows) {

src/pyproject_toml.rs

@@ -434,6 +434,9 @@ pub struct ToolMaturin {
     /// Executed in a temporary virtualenv with the instrumented wheel installed.
     /// Example: `python -m pytest tests/benchmarks`
     pub pgo_command: Option<String>,
+    /// Dependency group to install before PGO profile generation.
+    /// Defaults to `dev`.
+    pub pgo_dependency_group: Option<String>,
     /// CI generation configuration
     pub generate_ci: Option<GenerateCIConfig>,
 }
@@ -600,6 +603,12 @@ impl PyProjectToml {
         self.maturin().and_then(|m| m.pgo_command.as_deref())
     }
 
+    /// Returns the PGO dependency group from `[tool.maturin]`
+    pub fn pgo_dependency_group(&self) -> Option<&str> {
+        self.maturin()
+            .and_then(|m| m.pgo_dependency_group.as_deref())
+    }
+
     /// Returns the value of `[tool.maturin.compatibility]` in pyproject.toml
     pub fn compatibility(&self) -> Option<CompatibilityTag> {
         self.maturin()?.compatibility
@@ -1311,6 +1320,7 @@ mod tests {
 
             [tool.maturin]
             pgo-command = "python -m pytest tests/benchmarks"
+            pgo-dependency-group = "testing-extra"
             "#,
         )
         .unwrap();
@@ -1319,6 +1329,7 @@ mod tests {
             pyproject.pgo_command(),
             Some("python -m pytest tests/benchmarks")
         );
+        assert_eq!(pyproject.pgo_dependency_group(), Some("testing-extra"));
     }
 
     #[test]
@@ -1339,5 +1350,6 @@ mod tests {
         .unwrap();
         let pyproject = PyProjectToml::new(pyproject_file).unwrap();
         assert_eq!(pyproject.pgo_command(), None);
+        assert_eq!(pyproject.pgo_dependency_group(), None);
     }
 }

test-crates/pyo3-mixed/pgo_check.py

@@ -0,0 +1,14 @@
+import importlib
+import sys
+from pathlib import Path
+
+project_root = Path(__file__).resolve().parent
+sys.path = [path for path in sys.path if path and Path(path).resolve() != project_root]
+
+attrs = importlib.import_module("attrs")
+pyo3_mixed = importlib.import_module("pyo3_mixed")
+
+
+assert attrs is not None
+assert pyo3_mixed.get_42() == 42
+pyo3_mixed.print_cli_args()

test-crates/pyo3-mixed/pyproject.toml

@@ -18,6 +18,10 @@ dynamic = ["version"]
 get_42 = "pyo3_mixed:get_42"
 print_cli_args = "pyo3_mixed:print_cli_args"
 
+[dependency-groups]
+pgo = ["attrs"]
+
 [tool.maturin]
 include = ["pyo3_mixed/assets/*"]
-pgo-command = "print_cli_args"
+pgo-dependency-group = "pgo"
+pgo-command = "python pgo_check.py"