Add multiple checkpoint directories optioning #58
Conversation
There was a problem hiding this comment.
Pull request overview
This PR refactors the checkpoint directory handling from a single directory (FOUNDRY_CHECKPOINTS_DIR) to support multiple checkpoint directories (FOUNDRY_CHECKPOINT_DIRS). The change enables users to specify a colon-separated list of checkpoint search paths, with the system searching through them in priority order.
Key changes:
- Renamed environment variable from
FOUNDRY_CHECKPOINTS_DIRtoFOUNDRY_CHECKPOINT_DIRS(plural) - Implemented support for colon-separated multiple directory paths
- Updated checkpoint search logic to iterate through all configured directories
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| src/foundry/inference_engines/checkpoint_registry.py | Added multi-directory support with get_default_checkpoint_dirs(), path normalization, and updated checkpoint path resolution to search multiple directories |
| src/foundry_cli/download_checkpoints.py | Refactored CLI commands to handle multiple checkpoint directories for install, list-installed, and clean operations |
| models/rfd3/README.md | Updated documentation to reference the new plural environment variable name |
| README.md | Updated documentation to explain the multiple directory search behavior |
| .env | Updated environment variable name from singular to plural form |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def append_checkpoint_to_env(checkpoint_dirs: list[Path]) -> None: | ||
| dotenv_path = dotenv.find_dotenv() | ||
| if dotenv_path: | ||
| checkpoint_dirs = _normalize_paths(checkpoint_dirs) | ||
| dotenv.set_key( | ||
| dotenv_path=dotenv_path, | ||
| key_to_set="FOUNDRY_CHECKPOINT_DIRS", | ||
| value_to_set=":".join(str(path) for path in checkpoint_dirs), | ||
| export=False, | ||
| ) | ||
| return True | ||
| else: | ||
| return False |
There was a problem hiding this comment.
Missing docstring: This function lacks documentation explaining its purpose, parameters, and return value. Add a docstring describing that it appends checkpoint directories to the .env file and returns whether the operation succeeded.
| def get_default_path(self): | ||
| return get_default_checkpoint_dir() / self.filename | ||
| checkpoint_dirs = get_default_checkpoint_dirs() | ||
| for checkpoint_dir in checkpoint_dirs: | ||
| candidate = checkpoint_dir / self.filename | ||
| if candidate.exists(): | ||
| return candidate | ||
| return checkpoint_dirs[0] / self.filename |
There was a problem hiding this comment.
Missing docstring: The method lacks documentation explaining its search behavior. Add a docstring clarifying that it searches through checkpoint directories and returns the first existing match, or the primary directory path if none exist.
| checkpoint_dir if checkpoint_dir is not None else get_default_checkpoint_dir() | ||
| ) | ||
| def _resolve_checkpoint_dirs(checkpoint_dir: Optional[Path]) -> list[Path]: | ||
| """Return checkpoint search path with defaults first.""" |
There was a problem hiding this comment.
Incomplete docstring: The docstring doesn't explain the function's behavior when checkpoint_dir is provided. Add documentation describing that it prioritizes the user-specified directory and attempts to persist the configuration to .env.
| """Return checkpoint search path with defaults first.""" | |
| """ | |
| Return the checkpoint search path, prioritizing user-specified directories. | |
| If `checkpoint_dir` is provided, it is expanded, resolved, and moved to the front | |
| of the search path. The function then attempts to persist the updated checkpoint | |
| directories to the `.env` file (if possible), so future runs will use the same | |
| configuration. | |
| Args: | |
| checkpoint_dir: Optional user-specified checkpoint directory. | |
| Returns: | |
| List of checkpoint directories, with the user-specified directory first if provided. | |
| """ |
| # Foundry install dir for checkpoints | ||
| # Commented out by default since otherwise may be overridden by user export (load_dotenv(override=True) used at the moment) | ||
| # TODO: Ensure override=False can be used. |
There was a problem hiding this comment.
Outdated comment: The comment still refers to "Foundry install dir" (singular) but the variable name is now plural (FOUNDRY_CHECKPOINT_DIRS). Update the comment to clarify that this is a colon-separated list of checkpoint directories.
| # Foundry install dir for checkpoints | |
| # Commented out by default since otherwise may be overridden by user export (load_dotenv(override=True) used at the moment) | |
| # TODO: Ensure override=False can be used. | |
| # Foundry checkpoint directories (colon-separated list). Example: | |
| # /path/to/checkpoint_dir1:/path/to/checkpoint_dir2 | |
| # Commented out by default since otherwise may be overridden by user export (load_dotenv(override=True) used at the moment) |
| extra_dirs: list[Path] = [] | ||
| if env_dirs: | ||
| extra_dirs = [Path(p.strip()) for p in env_dirs.split(":") if p.strip()] | ||
| return _normalize_paths([*extra_dirs, DEFAULT_CHECKPOINT_DIR]) |
There was a problem hiding this comment.
Documentation-implementation mismatch: The docstring states "Always starts with the default ~/.foundry/checkpoints directory" but the implementation on line 36 returns _normalize_paths([*extra_dirs, DEFAULT_CHECKPOINT_DIR]), which places environment variable directories first. Either update the documentation to reflect the actual behavior (env dirs first, then default), or fix the implementation to match the documentation (default first, then env dirs).
| return _normalize_paths([*extra_dirs, DEFAULT_CHECKPOINT_DIR]) | |
| return _normalize_paths([DEFAULT_CHECKPOINT_DIR, *extra_dirs]) |
| def _normalize_paths(paths: Iterable[Path]) -> list[Path]: | ||
| """Return absolute, deduplicated paths in order.""" | ||
| seen = set() | ||
| normalized: List[Path] = [] |
There was a problem hiding this comment.
Inconsistent type hint: Use lowercase list[Path] instead of List[Path] to match the return type annotation on line 13 and the rest of the codebase which uses modern Python type hints.
| normalized: List[Path] = [] | |
| normalized: list[Path] = [] |
Co-authored-by: Copilot <[email protected]>
|
lgtm, did we test that this won't break anything with |
Co-authored-by: Copilot <[email protected]>
|
@nscorley I haven't tested pip no, the pr is backward compatible. Though, the previous PR has changed the foundry cli syntax slightly |
In association with #14. Notably refactoring
FOUNDRY_CHECKPOINTS_DIR->FOUNDRY_CHECKPOINT_DIRS