import pickle
import subprocess
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any
[docs]
class ThirdPartyModel:
"""Wrapper for loading and calling a third-party plugin model.
Clones a Git repository into a local directory (if not already present) and
invokes a specified function from that repository in an isolated uv environment.
uv is used over conda because it provides precise lockfile-based reproducibility,
strict package version pinning, and per-project Python version isolation — ensuring
the plugin runs in exactly the environment its author intended.
"""
[docs]
def __init__(
self,
entry_point: str,
path: str | Path,
url: str | None = None,
branch: str | None = None,
shallow: bool = True,
):
"""
Initialize the third-party model.
Args:
entry_point: Module and function to call, in the form 'module.path:function'.
path: Path to the plugin repository. Cloned here if it does not exist and url is provided.
url: Git repository URL to clone (optionally prefixed with 'git+'). If None, path must already exist.
branch: Branch to clone. If None, clones the default branch.
shallow: If True, clones only the latest commit (no full history). Defaults to True.
Raises:
ValueError: If entry_point is not of the form 'module:function'.
FileNotFoundError: If path does not exist and no url is provided.
"""
try:
module, fn = entry_point.split(":", 1)
except ValueError as e:
raise ValueError("entry_point must be of the form 'module:function'.") from e
if not module or not fn:
raise ValueError("entry_point must be of the form 'module:function'.")
self.repo_dir = Path(path).resolve()
self.module = module
self.fn = fn
_check_tool("uv")
if not self.repo_dir.exists():
if url is None:
raise FileNotFoundError(f"'{self.repo_dir}' does not exist. Provide a url to clone it.")
_check_tool("git")
_clone_git_repository(self.repo_dir, url, branch, shallow)
[docs]
def __call__(self, *args, **kwargs) -> Any:
"""
Execute the plugin's function with the given arguments.
Args:
*args: Positional arguments passed to the plugin function.
**kwargs: Keyword arguments passed to the plugin function.
Returns:
The return value of the plugin function.
"""
return _run(self.repo_dir, self.module, self.fn, args, kwargs)
def _check_tool(name: str) -> None:
if subprocess.run([name, "--version"], capture_output=True).returncode != 0:
raise RuntimeError(f"'{name}' is not installed or not found on PATH.")
def _clone_git_repository(repo_dir: Path, repo_url: str, branch: str | None = None, shallow: bool = True):
if repo_dir.exists():
raise FileExistsError(f"'{repo_dir}' already exists.")
repo_dir.parent.mkdir(parents=True, exist_ok=True)
url = repo_url.removeprefix("git+")
clone_cmd = ["git", "clone", url, str(repo_dir)]
if branch:
clone_cmd += ["-b", branch, "--single-branch"]
if shallow:
clone_cmd += ["--depth", "1"]
subprocess.check_call(clone_cmd, stdout=subprocess.DEVNULL)
def _wrap_code(module: str, fn: str) -> str:
return (
"import pickle, sys;"
f"import {module};"
"args, kwargs = pickle.load(open(sys.argv[1],'rb'));"
f"result = {module}.{fn}(*args, **kwargs);"
"pickle.dump(result, open(sys.argv[2],'wb'))"
)
def _run(repo_dir: Path, module: str, fn: str, args: tuple, kwargs: dict) -> Any:
with TemporaryDirectory(prefix="seqme") as tmpdir:
in_path = Path(tmpdir) / "input.pkl"
out_path = Path(tmpdir) / "output.pkl"
with open(in_path, "wb") as f:
pickle.dump((args, kwargs), f)
code = _wrap_code(module, fn)
cmd = ["uv", "run", "--project", str(repo_dir), "python", "-c", code, str(in_path), str(out_path)]
try:
subprocess.run(cmd, capture_output=True, text=True, check=True)
except subprocess.CalledProcessError as e:
raise RuntimeError(f"Plugin subprocess failed:\n{e.stderr}") from e
with open(out_path, "rb") as f:
return pickle.load(f)
