from pathlib import Path import tempfile from typing import Optional from lmm.config import LMMConfig from lmm.cmd import run_sudo_cmd, run_cmd def compute_workdirs_location(mount: Path) -> Path: with open("/proc/mounts", "r") as f: mounts = f.read().split("\n") mounted_dirs = [ m for line in mounts if line != "" and (m := line.strip().split(" ")[1]) != "/" ] mounted_dirs.sort() closest_mount = None for m in mounted_dirs: if str(mount).startswith(m) and len(m) > len(closest_mount or ""): closest_mount = m continue return Path("/tmp") if closest_mount is None else Path(closest_mount) / ".temp" class OverlayFSMount: upper: Optional[Path] lower: list[Path | str] workdir: Optional[tempfile.TemporaryDirectory] mount_path: Path _mounted: bool cd: Optional[Path] def __init__( self, upper: Optional[Path], lower: list[Path | str], mount: Path, cd: Optional[Path] = None, ): self.upper = upper self.lower = lower self.mount_path = mount self._mounted = False self.cd = cd assert len(lower) < 500, "overlayfs only supports up-to 500 lower directories" def create_workdir(self) -> Path: # Ensure the temporary directory for workdirs exists workdir_location = compute_workdirs_location(self.mount_path) if not workdir_location.exists(): workdir_location.mkdir(parents=True, exist_ok=True) # Create the workdir self.workdir = tempfile.TemporaryDirectory( dir=workdir_location, ) print(f"Temporary workdir {self.workdir.name} created") return Path(self.workdir.name) def mount(self) -> Optional[Path]: lower = ":".join([str(l) for l in self.lower]) options = f"lowerdir={lower}" if self.upper is not None: workdir = self.create_workdir() options += f",upperdir={self.upper},workdir={workdir}" if LMMConfig.use_fuse: result = run_cmd( [ LMMConfig.overlayfs_command, "-o", options, str(self.mount_path), ], cd=self.cd, ) else: result = run_sudo_cmd( [ "mount", "-t", "overlay", "overlay", "-o", options, str(self.mount_path), ], cd=self.cd, ) print("Mount result:", result) if not result == 0: return None self._mounted = True return self.mount_path def unmount(self) -> bool: if self._mounted: print("Unmounting...") # Remove the mount if LMMConfig.use_fuse: returncode = run_cmd( [ LMMConfig.fusermount_command, "-u", str(self.mount_path), ] ) == 0 else: returncode = run_sudo_cmd( [ "umount", str(self.mount_path), ] ) if returncode != 0: print("Unmounting failed") return False # Remove the temporary workdir if self.workdir is not None: try: # TODO: This fails because the workdir is owned by root self.workdir.cleanup() except PermissionError: print(f"Failed to clean temporary directory {self.workdir.name}") self._mounted = False return True def __del__(self): self.unmount()