From 27330f132eeae3e8560c1afd851a03cf17b0b23d Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Wed, 29 Jul 2026 15:04:24 -0400 Subject: [PATCH 1/2] [INITIAL] Infer direct mode from existing pull requests Move the direct-mode determination out of `__init__` into a new `_initialize_direct` method that runs after PR info has been prefetched, so it can consult the existing pull requests in the stack. Previously direct mode was resolved solely from the `--direct` / `--no-direct` flags, falling back to the presence of the `.github/ghstack_direct` marker file. Now, when the user hasn't passed an explicit flag and the marker doesn't force direct on, we inspect the base refs of the already-submitted PRs: a `gh///base` base ref indicates non-direct style, anything else indicates direct. If the stack mixes both styles we bail out with an error asking the user to pass `--direct` / `--no-direct` explicitly, rather than guessing. This lets an existing stack that was first submitted in direct mode keep being updated in direct mode without repeating the flag. Test plumbing changes to support this: - `direct` is now `Optional[bool]` throughout `test_prelude.py`; `None` means "infer". `check_global_github_invariants` only enforces the non-direct base-ref-reuse invariant when `direct is False`. - `gh_submit` gained a `**submit_kwargs` passthrough and a per-call `direct_opt` override so a test can submit one commit with an explicit mode and a later update with inference. - conftest now parametrizes on `direct` as before, but recognizes `*.unparametrized.py.test` scripts and runs them once with `direct=None`. Collection also matches on the `.py.test` filename suffix directly. New test `test/submit/infer_direct.unparametrized.py.test`: submit a commit with `direct_opt=True`, then amend and update without specifying a mode, exercising the inference path. [ghstack-poisoned] --- conftest.py | 22 +++++++++++--- src/ghstack/submit.py | 30 +++++++++++++------ src/ghstack/test_prelude.py | 17 ++++++----- .../infer_direct.unparametrized.py.test | 11 +++++++ 4 files changed, 60 insertions(+), 20 deletions(-) create mode 100644 test/submit/infer_direct.unparametrized.py.test diff --git a/conftest.py b/conftest.py index e9cae79..950e192 100644 --- a/conftest.py +++ b/conftest.py @@ -17,15 +17,29 @@ def pytest_collect_file(file_path: pathlib.Path, parent): # NB: script name must not end with py, due to doctest picking it # up in that case - if file_path.suffixes == [".py", ".test"]: + if file_path.name.endswith(".py.test"): return Script.from_parent(parent, path=file_path) class Script(pytest.File): def collect(self): - yield ScriptItem.from_parent(self, name="default", direct=False) - if self.path.parent.name in ["submit", "unlink", "log", "sync"]: - yield ScriptItem.from_parent(self, name="direct", direct=True) + if self.path.name.endswith(".unparametrized.py.test"): + direct_values = [None] + else: + direct_values = [False] + if self.path.parent.name in ["submit", "unlink", "log", "sync"]: + direct_values.append(True) + + for direct in direct_values: + yield ScriptItem.from_parent( + self, + name=( + "unparametrized" + if direct is None + else "direct" if direct else "default" + ), + direct=direct, + ) class ScriptItem(pytest.Item): diff --git a/src/ghstack/submit.py b/src/ghstack/submit.py index a649f03..8ecc6cb 100644 --- a/src/ghstack/submit.py +++ b/src/ghstack/submit.py @@ -516,21 +516,32 @@ async def initialize(self) -> None: object.__setattr__(self, "base", default_branch) - # Check if direct should be used, if the user didn't explicitly - # specify an option + # ~~~~~~~~~~~~~~~~~~~~~~~~ + # The main algorithm + + async def _initialize_direct(self, pr_info_cache: Dict[GitHubNumber, Any]) -> None: direct = self.direct_opt if direct is None: - direct_r = await self.sh.agit( + direct = await self.sh.agit( "cat-file", "-e", "HEAD:.github/ghstack_direct", exitcode=True ) - assert isinstance(direct_r, bool) - direct = direct_r - + assert isinstance(direct, bool) + if self.direct_opt is None and not direct: + styles = { + re.fullmatch(r"gh/[^/]+/[0-9]+/base", base_ref) is None + for pr_info in pr_info_cache.values() + if (base_ref := self._pr_ref_name(pr_info, "base")) is not None + } + if len(styles) > 1: + raise RuntimeError( + "Cannot infer ghstack submission style: the stack contains " + "both direct and non-direct pull requests. Pass --direct or " + "--no-direct explicitly." + ) + if styles: + direct = styles.pop() object.__setattr__(self, "direct", direct) - # ~~~~~~~~~~~~~~~~~~~~~~~~ - # The main algorithm - async def run(self) -> List[DiffMeta]: timer = _Timer() if _TIMING_ENABLED else None @@ -601,6 +612,7 @@ async def run(self) -> List[DiffMeta]: ) pr_info_cache = await self._prefetch_pr_info(commits_to_rebase) + await self._initialize_direct(pr_info_cache) if not self.no_fetch: await self._fetch_foreign_pr_refs(pr_info_cache.values()) diff --git a/src/ghstack/test_prelude.py b/src/ghstack/test_prelude.py index 69ea201..ca34a7c 100644 --- a/src/ghstack/test_prelude.py +++ b/src/ghstack/test_prelude.py @@ -122,9 +122,9 @@ class Context: github: ghstack.github.GitHubEndpoint upstream_sh: ghstack.shell.Shell sh: ghstack.shell.Shell - direct: bool + direct: Optional[bool] - def __init__(self, direct: bool) -> None: + def __init__(self, direct: Optional[bool]) -> None: # Set up a "parent" repository with an empty initial commit that we'll operate on upstream_dir = tempfile.mkdtemp() self.upstream_sh = ghstack.shell.Shell(cwd=upstream_dir, testing=True) @@ -154,7 +154,7 @@ def cleanup(self) -> None: onerror=handle_remove_read_only, ) - async def check_global_github_invariants(self, direct: bool) -> None: + async def check_global_github_invariants(self, direct: Optional[bool]) -> None: r = await self.github.graphql( """ query { @@ -177,7 +177,7 @@ async def check_global_github_invariants(self, direct: bool) -> None: continue # In direct mode, only head refs may not be reused; # base refs can be reused in octopus situations - if not direct: + if direct is False: assert pr["baseRefName"] not in seen_refs seen_refs.add(pr["baseRefName"]) assert pr["headRefName"] not in seen_refs @@ -200,7 +200,7 @@ async def init_test() -> Context: @contextlib.asynccontextmanager -async def scoped_test(direct: bool) -> AsyncIterator[None]: +async def scoped_test(direct: Optional[bool]) -> AsyncIterator[None]: global CTX assert CTX is None try: @@ -224,8 +224,10 @@ async def gh_submit( reviewer: Optional[str] = None, label: Optional[str] = None, automsg: Optional[str] = None, + **submit_kwargs: Any, ) -> List[ghstack.submit.DiffMeta]: self = CTX + direct_opt = submit_kwargs.pop("direct_opt", self.direct) r = await ghstack.submit.main( msg=msg, username="ezyang", @@ -236,7 +238,7 @@ async def gh_submit( repo_owner_opt="pytorch", repo_name_opt="pytorch", short=short, - direct_opt=self.direct, + direct_opt=direct_opt, no_skip=no_skip, github_url="github.com", remote_name="origin", @@ -247,6 +249,7 @@ async def gh_submit( reviewer=reviewer, label=label, automsg=automsg, + **submit_kwargs, ) await self.check_global_github_invariants(self.direct) return r @@ -451,7 +454,7 @@ async def assert_github_state(expect: str, *, skip: int = 0) -> None: assert_expected_inline(await dump_github(), expect, skip=skip + 1) -def is_direct() -> bool: +def is_direct() -> Optional[bool]: return CTX.direct diff --git a/test/submit/infer_direct.unparametrized.py.test b/test/submit/infer_direct.unparametrized.py.test new file mode 100644 index 0000000..e18b314 --- /dev/null +++ b/test/submit/infer_direct.unparametrized.py.test @@ -0,0 +1,11 @@ +from ghstack.test_prelude import * + +await init_test() + +await commit("A") +await gh_submit("Initial", direct_opt=True) + +await amend("A2") +await gh_submit("Update") + +ok() From 0825aafc829166d9e686281f859688fd03eb23d6 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Wed, 29 Jul 2026 15:29:17 -0400 Subject: [PATCH 2/2] [UPDATE] Infer direct mode from existing PRs' base ref styles This update reworks how the direct-mode inference is wired in. Rather than just renaming the `cat-file` result variable, `_initialize_direct` now keeps the original `direct_r`/`direct` computation and takes a `pr_info_cache` argument so it can consult existing pull requests. When `--direct`/`--no-direct` was not passed on the command line and the `.github/ghstack_direct` marker is absent, it now examines the base ref styles of the already-submitted PRs (matching against the `gh///base` pattern) to infer whether the stack was previously submitted in direct mode, instead of defaulting to non-direct. The call in `run()` is updated to pass the prefetched `pr_info_cache`. Also updates AGENTS.md to document the lint workflow via `uv run --frozen lintrunner --all-files` (and `lintrunner format` for fixes), replacing the older black/flake8 instructions. [ghstack-poisoned] --- AGENTS.md | 2 +- src/ghstack/submit.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index bd70ad5..b5c1916 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,7 +2,7 @@ ## Lint -Run `python -m black src/ghstack/` and `python -m flake8 src/ghstack/` to fix and check lint before committing. If lint fixes are needed after a commit, amend the commit rather than creating a separate lint fix commit. +Run `uv run --frozen lintrunner --all-files` to reproduce CI lint locally before committing. Invoke lintrunner through `uv run` unless the project virtual environment is active; otherwise its Python-based linters may use a system Python without the development dependencies. To apply formatter fixes, run `uv run --frozen lintrunner format --all-files`, then rerun the full lint command. If lint fixes are needed after a commit, amend the commit rather than creating a separate lint fix commit. ## Tests diff --git a/src/ghstack/submit.py b/src/ghstack/submit.py index 8ecc6cb..065e60d 100644 --- a/src/ghstack/submit.py +++ b/src/ghstack/submit.py @@ -522,10 +522,11 @@ async def initialize(self) -> None: async def _initialize_direct(self, pr_info_cache: Dict[GitHubNumber, Any]) -> None: direct = self.direct_opt if direct is None: - direct = await self.sh.agit( + direct_r = await self.sh.agit( "cat-file", "-e", "HEAD:.github/ghstack_direct", exitcode=True ) - assert isinstance(direct, bool) + assert isinstance(direct_r, bool) + direct = direct_r if self.direct_opt is None and not direct: styles = { re.fullmatch(r"gh/[^/]+/[0-9]+/base", base_ref) is None