diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a117b64..5ca6a32 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -108,7 +108,7 @@ jobs: - name: Setup dependencies run: | uv sync - git submodule update --remote --rebase + git submodule update cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/ - name: Build CPU Docker image @@ -183,7 +183,7 @@ jobs: - name: Setup dependencies run: | uv sync - git submodule update --remote --rebase + git submodule update cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/ - name: Build and push GPU Docker image (main) @@ -236,7 +236,7 @@ jobs: - name: Setup dependencies run: | uv sync - git submodule update --remote --rebase + git submodule update cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/ - name: Build and push CPU Docker image (main) @@ -299,7 +299,7 @@ jobs: - name: Setup dependencies run: | uv sync - git submodule update --remote --rebase + git submodule update cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/ - name: Build and push GPU Docker image (prod) @@ -362,7 +362,7 @@ jobs: - name: Setup dependencies run: | uv sync - git submodule update --remote --rebase + git submodule update cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/ - name: Build and push CPU Docker image (prod) diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index b420327..61e631c 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -84,8 +84,8 @@ jobs: - name: Setup dependencies run: | uv sync - git submodule update --remote --rebase - cp tetra-rp/src/tetra_rp/protos/remote_execution.py . + git submodule update + cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/ - name: Build and push GPU Docker image uses: docker/build-push-action@v6 diff --git a/CLAUDE.md b/CLAUDE.md index b4a6152..59f7748 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -146,6 +146,7 @@ git submodule update --remote --rebase # Update tetra-rp to latest - `RUNPOD_ENDPOINT_ID`: Used for workspace isolation (automatically set by RunPod) - `HF_HUB_ENABLE_HF_TRANSFER`: Set to "1" in Dockerfile to enable accelerated HuggingFace downloads - `HF_TOKEN`: Optional authentication token for private/gated HuggingFace models +- `HF_HOME=/hf-cache`: HuggingFace cache location, set outside `/root/.cache` to exclude from volume sync - `DEBIAN_FRONTEND=noninteractive`: Set during system package installation - `UV_CACHE_DIR`: Package cache configuration - `VIRTUAL_ENV`: Virtual environment path configuration diff --git a/Dockerfile b/Dockerfile index 896aaab..0ef55b4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,10 +4,16 @@ WORKDIR /app # Enable HuggingFace transfer acceleration ENV HF_HUB_ENABLE_HF_TRANSFER=1 +# Relocate HuggingFace cache outside /root/.cache to exclude from volume sync +ENV HF_HOME=/hf-cache + +# Configure APT cache to persist under /root/.cache for volume sync +RUN mkdir -p /root/.cache/apt/archives/partial \ + && echo 'Dir::Cache "/root/.cache/apt";' > /etc/apt/apt.conf.d/01cache # Install system dependencies and uv RUN apt-get update && apt-get install -y --no-install-recommends \ - curl ca-certificates nala \ + build-essential curl ca-certificates nala \ && curl -LsSf https://astral.sh/uv/install.sh | sh \ && cp ~/.local/bin/uv /usr/local/bin/uv \ && chmod +x /usr/local/bin/uv \ diff --git a/Dockerfile-cpu b/Dockerfile-cpu index 6c71c69..fd7b823 100644 --- a/Dockerfile-cpu +++ b/Dockerfile-cpu @@ -2,9 +2,13 @@ FROM python:3.12-slim WORKDIR /app +# Configure APT cache to persist under /root/.cache for volume sync +RUN mkdir -p /root/.cache/apt/archives/partial \ + && echo 'Dir::Cache "/root/.cache/apt";' > /etc/apt/apt.conf.d/01cache + # Install system dependencies and uv RUN apt-get update && apt-get install -y --no-install-recommends \ - curl ca-certificates nala \ + build-essential curl ca-certificates nala \ && curl -LsSf https://astral.sh/uv/install.sh | sh \ && cp ~/.local/bin/uv /usr/local/bin/uv \ && chmod +x /usr/local/bin/uv \ diff --git a/Makefile b/Makefile index 1a9c5f0..2f84a0e 100644 --- a/Makefile +++ b/Makefile @@ -23,6 +23,7 @@ dev: # Install development dependencies update: # Upgrade all dependencies uv sync --upgrade --all-groups uv lock --upgrade + git submodule update --remote clean: # Remove build artifacts and cache files rm -rf dist build *.egg-info @@ -31,8 +32,7 @@ clean: # Remove build artifacts and cache files find . -type f -name "*.pkl" -delete setup: dev # Initialize project, sync deps, update submodules - git submodule init - git submodule update --remote --rebase + git submodule update --init --recursive cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/ build: # Build both GPU and CPU Docker images diff --git a/docs/System_Python_Runtime_Architecture.md b/docs/System_Python_Runtime_Architecture.md index c9dc667..91ff08a 100644 --- a/docs/System_Python_Runtime_Architecture.md +++ b/docs/System_Python_Runtime_Architecture.md @@ -71,7 +71,7 @@ graph TB ### System Installation Strategy ```python # Docker environment -command = ["uv", "pip", "install", "--system", "--no-cache"] + packages +command = ["uv", "pip", "install", "--system"] + packages # Local environment command = ["uv", "pip", "install", "--python-preference=managed"] + packages diff --git a/src/cache_sync_manager.py b/src/cache_sync_manager.py new file mode 100644 index 0000000..0726730 --- /dev/null +++ b/src/cache_sync_manager.py @@ -0,0 +1,376 @@ +import os +import logging +import asyncio +import tempfile +from datetime import datetime +from pathlib import Path +from typing import Optional +from constants import NAMESPACE, CACHE_DIR, VOLUME_CACHE_PATH +from subprocess_utils import run_logged_subprocess + + +class CacheSyncManager: + """Manages async fire-and-forget cache synchronization to network volume.""" + + def __init__(self): + self.logger = logging.getLogger(f"{NAMESPACE}.{__name__.split('.')[-1]}") + self._should_sync_cached: Optional[bool] = None + self._endpoint_id = os.environ.get("RUNPOD_ENDPOINT_ID") + self._baseline_time: Optional[float] = None + + @property + def _tarball_path(self) -> str: + """Get the path to the cache tarball for this endpoint.""" + return f"{VOLUME_CACHE_PATH}/cache-{self._endpoint_id}.tar" + + @property + def _hydration_marker_path(self) -> str: + """Get the path to the cache hydration marker file.""" + return f"{CACHE_DIR}/.cache-last-hydrated" + + def _cleanup_temp_file(self, path: str, description: str) -> None: + """Clean up a temporary file, logging any errors at debug level.""" + if os.path.exists(path): + try: + os.remove(path) + except Exception as e: + self.logger.debug(f"Failed to clean up {description}: {e}") + + def should_sync(self) -> bool: + """ + Determine if cache sync functionality is available. + + Checks if all prerequisites for cache synchronization are met: + - RUNPOD_ENDPOINT_ID is set + - Network volume is mounted + - Volume cache directory exists or can be created + + Result is cached after first check. + + Returns: + True if sync functionality is available, False otherwise + """ + if self._should_sync_cached is not None: + return self._should_sync_cached + + # Skip if no endpoint ID + if not self._endpoint_id: + self.logger.debug("No RUNPOD_ENDPOINT_ID set, skipping cache sync") + self._should_sync_cached = False + return False + + # Skip if volume not mounted + volume_root = os.path.dirname(VOLUME_CACHE_PATH) + if not os.path.exists(volume_root): + self.logger.debug(f"Volume {volume_root} not mounted, skipping cache sync") + self._should_sync_cached = False + return False + + # Ensure volume cache directory exists + try: + os.makedirs(VOLUME_CACHE_PATH, exist_ok=True) + except Exception as e: + self.logger.warning( + f"Failed to create volume cache directory {VOLUME_CACHE_PATH}: {e}" + ) + self._should_sync_cached = False + return False + + self._should_sync_cached = True + return True + + def mark_baseline(self) -> None: + """Mark baseline timestamp before installation.""" + if not self.should_sync(): + return + + try: + tarball_path = self._tarball_path + if os.path.exists(tarball_path): + # Subsequent run: use tarball mtime as baseline + self._baseline_time = os.path.getmtime(tarball_path) + baseline_source = "tarball" + else: + # First run: use current time as baseline + self._baseline_time = datetime.now().timestamp() + baseline_source = "current time" + + self.logger.debug( + f"Baseline ({baseline_source}): {datetime.fromtimestamp(self._baseline_time).strftime('%Y-%m-%d %H:%M:%S')}" + ) + except Exception as e: + self.logger.warning(f"Failed to mark cache baseline: {e}") + self._baseline_time = None + + async def sync_to_volume(self) -> None: + """Background worker to collect delta and create tarball.""" + if not self.should_sync() or not self._baseline_time: + return + + try: + baseline_time = self._baseline_time + tarball_path = self._tarball_path + tarball_exists = os.path.exists(tarball_path) + + self.logger.debug( + f"Sync cache to persist from {CACHE_DIR} to {tarball_path}" + ) + + # Format timestamp for find -newermt + baseline_str = datetime.fromtimestamp(baseline_time).strftime( + "%Y-%m-%d %H:%M:%S" + ) + + # Find files newer than baseline + find_result = await asyncio.to_thread( + run_logged_subprocess, + command=[ + "find", + CACHE_DIR, + "-newermt", + baseline_str, + "-type", + "f", + "-not", + "-path", + "*/refs/*", + "-not", + "-path", + "*/.no_exist/*", + "-not", + "-name", + ".cache-last-hydrated", + ], + logger=self.logger, + operation_name="Finding new cache files", + suppress_output=True, + ) + + if not find_result.success: + self.logger.warning(f"Failed to find cache delta: {find_result.error}") + return + + # Check if there are any new files + new_files = (find_result.stdout or "").strip() + if not new_files: + self.logger.debug("No new cache files to sync") + return + + # Log summary instead of full file list + file_count = len(new_files.split("\n")) + self.logger.debug(f"Found {file_count} new cache files to sync") + + # Monitor tarball size if it exists + if tarball_exists: + try: + tarball_size = os.path.getsize(tarball_path) + tarball_mb = tarball_size / (1024 * 1024) + self.logger.debug(f"Current tarball size: {tarball_mb:.1f}MB") + + # Check volume capacity and warn if tarball exceeds 50% + volume_root = os.path.dirname(VOLUME_CACHE_PATH) + stat = os.statvfs(volume_root) + volume_total = stat.f_blocks * stat.f_frsize + threshold = volume_total * 0.75 + + if tarball_size > threshold: + volume_total_gb = volume_total / (1024**3) + self.logger.warning( + f"Tarball size ({tarball_mb:.1f}MB) exceeds 75% of volume capacity ({volume_total_gb:.1f}GB)" + ) + except OSError as e: + self.logger.debug(f"Failed to check tarball size: {e}") + + # Write file list to temporary file + file_list_fd = tempfile.NamedTemporaryFile( + prefix=".cache-files-", dir="/tmp", delete=False, mode="w" + ) + file_list_path = file_list_fd.name + try: + file_list_fd.write(new_files) + file_list_fd.close() + except Exception as e: + self.logger.warning(f"Failed to write file list: {e}") + file_list_fd.close() + return + + # Always create tarball of new files first + new_tarball = f"{tarball_path}.new" + temp_tarball = f"{tarball_path}.tmp" + + try: + # Create tarball containing only new files + create_result = await asyncio.to_thread( + run_logged_subprocess, + command=["tar", "cf", new_tarball, "-T", file_list_path], + logger=self.logger, + operation_name="Creating tarball of new files", + ) + + if not create_result.success: + self.logger.warning( + f"Failed to create new files tarball: {create_result.error}" + ) + return + + if tarball_exists: + # Move existing tarball to temp location + move_to_temp_result = await asyncio.to_thread( + run_logged_subprocess, + command=["mv", tarball_path, temp_tarball], + logger=self.logger, + operation_name="Moving existing tarball to temp", + ) + + if not move_to_temp_result.success: + self.logger.warning( + f"Failed to move tarball to temp: {move_to_temp_result.error}" + ) + return + + # Concatenate new tarball into temp (faster than append) + concat_result = await asyncio.to_thread( + run_logged_subprocess, + command=["tar", "-A", "-f", temp_tarball, new_tarball], + logger=self.logger, + operation_name="Concatenating new files to tarball", + ) + + if not concat_result.success: + self.logger.warning( + f"Failed to concatenate tarball: {concat_result.error}" + ) + return + + # Atomically move temp to final location + rename_result = await asyncio.to_thread( + run_logged_subprocess, + command=["mv", temp_tarball, tarball_path], + logger=self.logger, + operation_name="Moving tarball to final location", + ) + + if rename_result.success: + self.logger.info( + f"Successfully concatenated cache tarball at {tarball_path}" + ) + self.mark_last_hydrated() + else: + self.logger.warning( + f"Failed to move tarball: {rename_result.error}" + ) + else: + # No existing tarball, just move new one to final location + rename_result = await asyncio.to_thread( + run_logged_subprocess, + command=["mv", new_tarball, tarball_path], + logger=self.logger, + operation_name="Moving tarball to final location", + ) + + if rename_result.success: + self.logger.info( + f"Successfully created cache tarball at {tarball_path}" + ) + self.mark_last_hydrated() + else: + self.logger.warning( + f"Failed to move tarball: {rename_result.error}" + ) + finally: + # Clean up temporary files + self._cleanup_temp_file(file_list_path, "file list") + self._cleanup_temp_file(new_tarball, "new files tarball") + self._cleanup_temp_file(temp_tarball, "temp tarball") + + except Exception as e: + self.logger.error(f"Unexpected error in cache sync: {e}", exc_info=True) + + def should_hydrate(self) -> bool: + """ + Check if cache hydration should run. + + Returns: + True if tarball exists and is newer than last hydration, False otherwise + """ + if not self.should_sync(): + return False + + tarball_path = self._tarball_path + if not os.path.exists(tarball_path): + self.logger.debug( + f"Tarball {tarball_path} does not exist, skipping hydration" + ) + return False + + # Check last hydrated marker + marker_path = self._hydration_marker_path + if not os.path.exists(marker_path): + self.logger.debug("No hydration marker found, hydration needed") + return True + + try: + tarball_mtime = os.path.getmtime(tarball_path) + marker_mtime = os.path.getmtime(marker_path) + + if tarball_mtime > marker_mtime: + self.logger.debug( + "Tarball is newer than last hydration, hydration needed" + ) + return True + else: + self.logger.debug( + "Tarball is older than last hydration, skipping hydration" + ) + return False + except Exception as e: + self.logger.warning(f"Failed to check hydration status: {e}") + return True + + def mark_last_hydrated(self) -> None: + """Mark timestamp of last hydration.""" + if not self.should_sync(): + return + + try: + Path(self._hydration_marker_path).touch() + self.logger.debug( + f"Marked cache last hydrated at {self._hydration_marker_path}" + ) + except Exception as e: + self.logger.warning(f"Failed to mark cache last hydrated: {e}") + + async def hydrate_from_volume(self) -> None: + """Extract tarball from volume to hydrate local cache.""" + if not self.should_hydrate(): + return + + try: + tarball_path = self._tarball_path + self.logger.debug(f"Hydrating cache from {tarball_path} to {CACHE_DIR}") + + # Ensure cache directory exists + try: + os.makedirs(CACHE_DIR, exist_ok=True) + except Exception as e: + self.logger.warning( + f"Failed to create cache directory {CACHE_DIR}: {e}" + ) + return + + # Extract tarball to cache directory + tar_result = await asyncio.to_thread( + run_logged_subprocess, + command=["tar", "xf", tarball_path, "-C", "/"], + logger=self.logger, + operation_name="Extracting cache tarball", + ) + + if tar_result.success: + self.logger.info(f"Successfully hydrated cache from {tarball_path}") + self.mark_last_hydrated() + else: + self.logger.warning(f"Failed to extract tarball: {tar_result.error}") + + except Exception as e: + self.logger.error(f"Unexpected error during hydration: {e}", exc_info=True) diff --git a/src/constants.py b/src/constants.py index 26ddace..858ffa7 100644 --- a/src/constants.py +++ b/src/constants.py @@ -17,3 +17,10 @@ "wget", ] """List of system packages that benefit from nala's accelerated installation.""" + +# Cache Sync Configuration +CACHE_DIR = "/root/.cache" +"""Directory containing package and model caches.""" + +VOLUME_CACHE_PATH = "/runpod-volume/.cache" +"""Network volume path for cache tarball storage.""" diff --git a/src/dependency_installer.py b/src/dependency_installer.py index ea307e6..85c018d 100644 --- a/src/dependency_installer.py +++ b/src/dependency_installer.py @@ -37,7 +37,7 @@ def install_dependencies( if self._is_docker_environment(): if accelerate_downloads: # Packages are installed to the system location where they can be imported - command = ["uv", "pip", "install", "--system", "--no-cache"] + packages + command = ["uv", "pip", "install", "--system"] + packages else: # Use full path to system python command = ["pip", "install"] + packages @@ -47,21 +47,13 @@ def install_dependencies( operation_name = f"Installing Python packages ({'accelerated' if accelerate_downloads else 'standard'})" - # Set environment variables to avoid UV cache issues in read-only volumes - env = None - if self._is_docker_environment(): - env = os.environ.copy() - # Disable UV cache completely in Docker environments - env["UV_NO_CACHE"] = "1" - env["UV_CACHE_DIR"] = "/tmp/uv-cache" # Use writable temp directory - try: return run_logged_subprocess( command=command, logger=self.logger, operation_name=operation_name, timeout=300, - env=env, + env=os.environ.copy(), ) except Exception as e: return FunctionResponse(success=False, error=str(e)) diff --git a/src/huggingface_cache.py b/src/huggingface_cache.py index df8b91a..9afc41f 100644 --- a/src/huggingface_cache.py +++ b/src/huggingface_cache.py @@ -61,6 +61,8 @@ def cache_model_download( stdout=f"Model {model_id} already cached (cache hit)", ) + self.logger.info(f"Started downloading model {model_id}") + # Get HF authentication token if available hf_token = os.environ.get("HF_TOKEN") @@ -93,7 +95,7 @@ def _is_model_cached(self, model_id: str, revision: str = "main") -> bool: Args: model_id: HuggingFace model identifier - revision: Model revision/branch + revision: Model revision/branch/commit hash Returns: True if model is cached, False otherwise @@ -102,9 +104,15 @@ def _is_model_cached(self, model_id: str, revision: str = "main") -> bool: cache_info = scan_cache_dir() for repo in cache_info.repos: if repo.repo_id == model_id: - # Check if the specific revision is cached + # If revision is "main", accept any cached version of the model + if revision == "main": + return len(repo.revisions) > 0 + + # Check for specific revision by commit hash for rev in repo.revisions: - if rev.commit_hash == revision or revision == "main": + if rev.commit_hash.startswith(revision) or revision.startswith( + rev.commit_hash + ): return True return False except CacheNotFound: diff --git a/src/remote_executor.py b/src/remote_executor.py index 3d0e5e1..39e1d03 100644 --- a/src/remote_executor.py +++ b/src/remote_executor.py @@ -7,6 +7,7 @@ from function_executor import FunctionExecutor from class_executor import ClassExecutor from log_streamer import start_log_streaming, stop_log_streaming, get_streamed_logs +from cache_sync_manager import CacheSyncManager from constants import NAMESPACE @@ -25,6 +26,7 @@ def __init__(self): self.function_executor = FunctionExecutor() self.class_executor = ClassExecutor() self.hf_cache = HuggingFaceCacheAhead() + self.cache_sync = CacheSyncManager() async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse: """ @@ -51,6 +53,18 @@ async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse: ) try: + # Hydrate cache from volume if needed (before any installations) + has_installations = ( + request.dependencies + or request.system_dependencies + or request.hf_models_to_cache + ) + if has_installations: + await self.cache_sync.hydrate_from_volume() + + # Mark cache baseline before installation + self.cache_sync.mark_baseline() + # Install dependencies if request.accelerate_downloads: # Run installations in parallel when acceleration is enabled @@ -77,6 +91,9 @@ async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse: dep_result.stdout = logs return dep_result + # cache sync after installation + await self.cache_sync.sync_to_volume() + # Route to appropriate execution method based on type execution_type = getattr(request, "execution_type", "function") @@ -175,16 +192,11 @@ async def _install_dependencies_sequential( # Cache-ahead HuggingFace models if requested (should not happen when acceleration disabled) if request.accelerate_downloads and request.hf_models_to_cache: for model_id in request.hf_models_to_cache: - self.logger.info(f"Cache-ahead HuggingFace model: {model_id}") cache_result = self.hf_cache.cache_model_download(model_id) if cache_result.success: - self.logger.info( - f"Successfully cached model {model_id}: {cache_result.stdout}" - ) + self.logger.info(cache_result.stdout) else: - self.logger.warning( - f"Failed to cache model {model_id}: {cache_result.error}" - ) + self.logger.warning(cache_result.error) # Install Python dependencies next if request.dependencies: diff --git a/src/subprocess_utils.py b/src/subprocess_utils.py index dd2cbc4..b2b3e16 100644 --- a/src/subprocess_utils.py +++ b/src/subprocess_utils.py @@ -22,6 +22,7 @@ def run_logged_subprocess( capture_output: bool = True, text: bool = True, env: Optional[dict[str, str]] = None, + suppress_output: bool = False, **popen_kwargs, ) -> FunctionResponse: """ @@ -38,6 +39,8 @@ def run_logged_subprocess( timeout: Timeout in seconds for subprocess execution capture_output: Whether to capture stdout/stderr text: Whether to return strings instead of bytes + env: Environment variables to pass to subprocess + suppress_output: If True, only log command execution, not output **popen_kwargs: Additional arguments passed to subprocess.Popen Returns: @@ -74,14 +77,15 @@ def run_logged_subprocess( logger.debug(f"{log_prefix}Error: {error_msg}") return FunctionResponse(success=False, error=error_msg) - # Log subprocess output - if stdout: - logger.debug(f"{log_prefix}Output: {stdout.strip()}") - if stderr: - if process.returncode == 0: - logger.debug(f"{log_prefix}Warnings: {stderr.strip()}") - else: - logger.debug(f"{log_prefix}Errors: {stderr.strip()}") + # Log subprocess output (unless suppressed) + if not suppress_output: + if stdout: + logger.debug(f"{log_prefix}Output: {stdout.strip()}") + if stderr: + if process.returncode == 0: + logger.debug(f"{log_prefix}Warnings: {stderr.strip()}") + else: + logger.debug(f"{log_prefix}Errors: {stderr.strip()}") # Return appropriate response based on exit code if process.returncode == 0: diff --git a/tests/unit/test_cache_sync_manager.py b/tests/unit/test_cache_sync_manager.py new file mode 100644 index 0000000..153c9d7 --- /dev/null +++ b/tests/unit/test_cache_sync_manager.py @@ -0,0 +1,516 @@ +import os +import pytest +from unittest.mock import patch +from pathlib import Path +from cache_sync_manager import CacheSyncManager +from remote_execution import FunctionResponse + + +@pytest.fixture +def cache_sync(): + """Create a CacheSyncManager instance for testing.""" + return CacheSyncManager() + + +@pytest.fixture +def mock_env(monkeypatch): + """Mock environment variables.""" + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "test-endpoint-123") + + +class TestShouldSync: + def test_should_sync_no_endpoint_id(self, cache_sync): + """Test that sync is skipped when RUNPOD_ENDPOINT_ID is not set.""" + with patch.dict(os.environ, {}, clear=True): + cache_sync_new = CacheSyncManager() + assert cache_sync_new.should_sync() is False + + def test_should_sync_volume_not_mounted(self, cache_sync, mock_env): + """Test that sync is skipped when /runpod-volume is not mounted.""" + with patch("os.path.exists") as mock_exists: + mock_exists.return_value = False + assert cache_sync.should_sync() is False + + def test_should_sync_success(self, mock_env): + """Test that sync proceeds when conditions are met.""" + # Create cache_sync after environment is set + cache_sync = CacheSyncManager() + + with ( + patch("os.path.exists") as mock_exists, + patch("os.makedirs") as mock_makedirs, + ): + + def exists_side_effect(path): + if path == "/runpod-volume": + return True + return False + + mock_exists.side_effect = exists_side_effect + assert cache_sync.should_sync() is True + mock_makedirs.assert_called_once_with( + "/runpod-volume/.cache", exist_ok=True + ) + + def test_should_sync_cached_result(self, mock_env): + """Test that should_sync caches its result.""" + # Create cache_sync after environment is set + cache_sync = CacheSyncManager() + + with patch("os.path.exists") as mock_exists, patch("os.makedirs"): + + def exists_side_effect(path): + if path == "/runpod-volume": + return True + return False + + mock_exists.side_effect = exists_side_effect + + # First call + result1 = cache_sync.should_sync() + # Second call should use cached result + result2 = cache_sync.should_sync() + + assert result1 is True + assert result2 is True + # os.path.exists should be called only once (cached on second call) + assert mock_exists.call_count <= 1 + + +class TestMarkBaseline: + def test_mark_baseline_skips_when_should_not_sync(self, cache_sync): + """Test that mark_baseline skips when should_sync returns False.""" + with patch.object(cache_sync, "should_sync", return_value=False): + cache_sync.mark_baseline() + assert cache_sync._baseline_time is None + + def test_mark_baseline_stores_timestamp(self, cache_sync, mock_env): + """Test that mark_baseline stores current timestamp.""" + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("cache_sync_manager.datetime") as mock_datetime, + ): + # Mock datetime.now().timestamp() + mock_now = mock_datetime.now.return_value + mock_now.timestamp.return_value = 1234567890.0 + + cache_sync.mark_baseline() + + assert cache_sync._baseline_time == 1234567890.0 + + def test_mark_baseline_handles_exception(self, cache_sync, mock_env): + """Test that mark_baseline handles exceptions gracefully.""" + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("cache_sync_manager.datetime") as mock_datetime, + ): + mock_datetime.now.side_effect = Exception("Time error") + + cache_sync.mark_baseline() + assert cache_sync._baseline_time is None + + +class TestSyncToVolumeAsync: + @pytest.mark.asyncio + async def test_sync_skips_when_should_not_sync(self, cache_sync): + """Test that sync_to_volume skips when should_sync returns False.""" + with ( + patch.object(cache_sync, "should_sync", return_value=False), + patch("asyncio.to_thread") as mock_to_thread, + ): + await cache_sync.sync_to_volume() + # Verify no subprocess operations were attempted + mock_to_thread.assert_not_called() + + +class TestCollectAndTarball: + @pytest.mark.asyncio + async def test_sync_to_volume_no_new_files(self, cache_sync, mock_env): + """Test that sync_to_volume handles no new files.""" + cache_sync._endpoint_id = "test-endpoint-123" + cache_sync._baseline_time = 1234567890.0 + + mock_find_result = FunctionResponse(success=True, stdout="") + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists", return_value=True), + patch("os.path.getmtime", return_value=1234567890.0), + patch( + "asyncio.to_thread", side_effect=[mock_find_result] + ) as mock_to_thread, + ): + await cache_sync.sync_to_volume() + + # Only find command should be called, tar should be skipped + assert mock_to_thread.call_count == 1 + + @pytest.mark.asyncio + async def test_sync_to_volume_success_new(self, cache_sync, mock_env): + """Test successful tarball creation when no tarball exists (uses baseline_time).""" + cache_sync._endpoint_id = "test-endpoint-123" + cache_sync._baseline_time = 1234567890.0 + + mock_find_result = FunctionResponse( + success=True, stdout="/root/.cache/file1\n/root/.cache/file2" + ) + mock_create_result = FunctionResponse(success=True, stdout="") + mock_mv_result = FunctionResponse(success=True, stdout="") + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch( + "asyncio.to_thread", + side_effect=[mock_find_result, mock_create_result, mock_mv_result], + ) as mock_to_thread, + patch("os.path.exists") as mock_exists, + patch("os.remove") as mock_remove, + patch("tempfile.NamedTemporaryFile") as mock_tempfile, + patch.object(cache_sync, "mark_last_hydrated") as mock_mark, + ): + # Mock tempfile for file list + mock_file_list = mock_tempfile.return_value + mock_file_list.name = "/tmp/.cache-files-abc123" + + # Tarball doesn't exist initially, but temp files exist + def exists_side_effect(path): + if path == "/runpod-volume/.cache/cache-test-endpoint-123.tar": + return False + elif path.endswith(".new") or path.endswith(".tmp"): + return True + elif path.startswith("/tmp/.cache-files-"): + return True + return False + + mock_exists.side_effect = exists_side_effect + + await cache_sync.sync_to_volume() + + # find, tar cf (create new), and mv should be called + assert mock_to_thread.call_count == 3 + # File list and temp files should be cleaned up (3 calls) + assert mock_remove.call_count == 3 + # mark_last_hydrated should be called after successful sync + mock_mark.assert_called_once() + + @pytest.mark.asyncio + async def test_sync_to_volume_success_append(self, cache_sync, mock_env): + """Test successful tarball concatenation when tarball already exists (uses baseline_time).""" + cache_sync._endpoint_id = "test-endpoint-123" + cache_sync._baseline_time = 1234567890.0 + + mock_find_result = FunctionResponse( + success=True, stdout="/root/.cache/file3\n/root/.cache/file4" + ) + mock_create_result = FunctionResponse(success=True, stdout="") + mock_mv_to_temp_result = FunctionResponse(success=True, stdout="") + mock_concat_result = FunctionResponse(success=True, stdout="") + mock_mv_to_final_result = FunctionResponse(success=True, stdout="") + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch( + "asyncio.to_thread", + side_effect=[ + mock_find_result, + mock_create_result, + mock_mv_to_temp_result, + mock_concat_result, + mock_mv_to_final_result, + ], + ) as mock_to_thread, + patch("os.path.exists") as mock_exists, + patch("os.remove") as mock_remove, + patch("tempfile.NamedTemporaryFile") as mock_tempfile, + patch.object(cache_sync, "mark_last_hydrated") as mock_mark, + ): + # Mock tempfile for file list + mock_file_list = mock_tempfile.return_value + mock_file_list.name = "/tmp/.cache-files-abc123" + + # Tarball exists initially, plus temp files + def exists_side_effect(path): + if path == "/runpod-volume/.cache/cache-test-endpoint-123.tar": + return True + elif path.endswith(".new") or path.endswith(".tmp"): + return True + elif path.startswith("/tmp/.cache-files-"): + return True + return False + + mock_exists.side_effect = exists_side_effect + + await cache_sync.sync_to_volume() + + # find, tar cf (create new), mv (to temp), tar -A (concat), mv (to final) + assert mock_to_thread.call_count == 5 + # File list and temp files should be cleaned up (3 calls) + assert mock_remove.call_count == 3 + # mark_last_hydrated should be called after successful sync + mock_mark.assert_called_once() + + @pytest.mark.asyncio + async def test_sync_to_volume_move_to_temp_failure(self, cache_sync, mock_env): + """Test handling of move to temp failure when concatenating.""" + cache_sync._endpoint_id = "test-endpoint-123" + cache_sync._baseline_time = 1234567890.0 + + mock_find_result = FunctionResponse( + success=True, stdout="/root/.cache/file3\n/root/.cache/file4" + ) + mock_create_result = FunctionResponse(success=True, stdout="") + mock_mv_to_temp_result = FunctionResponse( + success=False, error="Move to temp failed" + ) + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch( + "asyncio.to_thread", + side_effect=[ + mock_find_result, + mock_create_result, + mock_mv_to_temp_result, + ], + ) as mock_to_thread, + patch("os.path.exists") as mock_exists, + ): + # Tarball exists initially + def exists_side_effect(path): + if path == "/runpod-volume/.cache/cache-test-endpoint-123.tar": + return True + return False + + mock_exists.side_effect = exists_side_effect + + await cache_sync.sync_to_volume() + + # find, tar cf (create new), and mv to temp should be called + assert mock_to_thread.call_count == 3 + + @pytest.mark.asyncio + async def test_sync_to_volume_find_failure(self, cache_sync, mock_env): + """Test handling of find command failure.""" + cache_sync._endpoint_id = "test-endpoint-123" + cache_sync._baseline_time = 1234567890.0 + + mock_find_result = FunctionResponse(success=False, error="Find failed") + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists", return_value=True), + patch( + "asyncio.to_thread", side_effect=[mock_find_result] + ) as mock_to_thread, + ): + await cache_sync.sync_to_volume() + + # Only find should be attempted, tar should be skipped + assert mock_to_thread.call_count == 1 + + @pytest.mark.asyncio + async def test_sync_to_volume_handles_exception(self, cache_sync, mock_env): + """Test that sync_to_volume handles unexpected exceptions.""" + cache_sync._endpoint_id = "test-endpoint-123" + cache_sync._baseline_time = 1234567890.0 + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists", return_value=True), + patch("asyncio.to_thread", side_effect=Exception("Unexpected error")), + ): + # Should not raise exception + await cache_sync.sync_to_volume() + + +class TestShouldHydrate: + def test_should_hydrate_when_should_sync_false(self, cache_sync): + """Test that hydration skips when should_sync returns False.""" + with patch.object(cache_sync, "should_sync", return_value=False): + assert cache_sync.should_hydrate() is False + + def test_should_hydrate_when_no_tarball(self, cache_sync, mock_env): + """Test that hydration skips when tarball doesn't exist.""" + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists", return_value=False), + ): + assert cache_sync.should_hydrate() is False + + def test_should_hydrate_when_no_marker(self, cache_sync, mock_env): + """Test that hydration proceeds when no marker exists.""" + cache_sync._endpoint_id = "test-endpoint-123" + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists") as mock_exists, + ): + + def exists_side_effect(path): + if "cache-test-endpoint-123.tar" in path: + return True + return False + + mock_exists.side_effect = exists_side_effect + assert cache_sync.should_hydrate() is True + + def test_should_hydrate_when_tarball_newer(self, cache_sync, mock_env): + """Test that hydration proceeds when tarball is newer than marker.""" + cache_sync._endpoint_id = "test-endpoint-123" + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists", return_value=True), + patch("os.path.getmtime") as mock_getmtime, + ): + + def getmtime_side_effect(path): + if "cache-test-endpoint-123.tar" in path: + return 2000.0 + return 1000.0 + + mock_getmtime.side_effect = getmtime_side_effect + assert cache_sync.should_hydrate() is True + + def test_should_hydrate_when_tarball_older(self, cache_sync, mock_env): + """Test that hydration skips when tarball is older than marker.""" + cache_sync._endpoint_id = "test-endpoint-123" + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists", return_value=True), + patch("os.path.getmtime") as mock_getmtime, + ): + + def getmtime_side_effect(path): + if "cache-test-endpoint-123.tar" in path: + return 1000.0 + return 2000.0 + + mock_getmtime.side_effect = getmtime_side_effect + assert cache_sync.should_hydrate() is False + + def test_should_hydrate_handles_exception(self, cache_sync, mock_env): + """Test that should_hydrate handles exceptions gracefully.""" + cache_sync._endpoint_id = "test-endpoint-123" + + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch("os.path.exists", return_value=True), + patch("os.path.getmtime", side_effect=OSError("Permission denied")), + ): + # Should return True on exception (safe default) + assert cache_sync.should_hydrate() is True + + +class TestHydrateFromVolume: + @pytest.mark.asyncio + async def test_hydrate_skips_when_should_not_hydrate(self, cache_sync): + """Test that hydrate_from_volume skips when should_hydrate returns False.""" + with ( + patch.object(cache_sync, "should_hydrate", return_value=False), + patch("asyncio.to_thread") as mock_to_thread, + ): + await cache_sync.hydrate_from_volume() + mock_to_thread.assert_not_called() + + @pytest.mark.asyncio + async def test_hydrate_success(self, cache_sync, mock_env): + """Test successful cache hydration from tarball.""" + cache_sync._endpoint_id = "test-endpoint-123" + + mock_tar_result = FunctionResponse(success=True, stdout="") + + with ( + patch.object(cache_sync, "should_hydrate", return_value=True), + patch("os.makedirs") as mock_makedirs, + patch.object(Path, "glob", return_value=[]), + patch("asyncio.to_thread", return_value=mock_tar_result) as mock_to_thread, + patch.object(cache_sync, "mark_last_hydrated") as mock_mark, + ): + await cache_sync.hydrate_from_volume() + + # Cache dir should be created + mock_makedirs.assert_called_once_with("/root/.cache", exist_ok=True) + # Tar extraction should be called + assert mock_to_thread.call_count == 1 + # Hydration marker should be set + mock_mark.assert_called_once() + + @pytest.mark.asyncio + async def test_hydrate_tar_failure(self, cache_sync, mock_env): + """Test handling of tar extraction failure.""" + cache_sync._endpoint_id = "test-endpoint-123" + + mock_tar_result = FunctionResponse(success=False, error="Extraction failed") + + with ( + patch.object(cache_sync, "should_hydrate", return_value=True), + patch("os.makedirs"), + patch.object(Path, "glob", return_value=[]), + patch("asyncio.to_thread", return_value=mock_tar_result), + patch.object(cache_sync, "mark_last_hydrated") as mock_mark, + ): + await cache_sync.hydrate_from_volume() + + # Marker should NOT be set on failure + mock_mark.assert_not_called() + + @pytest.mark.asyncio + async def test_hydrate_makedirs_failure(self, cache_sync, mock_env): + """Test handling of cache directory creation failure.""" + cache_sync._endpoint_id = "test-endpoint-123" + + with ( + patch.object(cache_sync, "should_hydrate", return_value=True), + patch("os.makedirs", side_effect=OSError("Permission denied")), + patch("asyncio.to_thread") as mock_to_thread, + ): + await cache_sync.hydrate_from_volume() + + # Tar should not be attempted if mkdir fails + mock_to_thread.assert_not_called() + + @pytest.mark.asyncio + async def test_hydrate_handles_exception(self, cache_sync, mock_env): + """Test that hydrate_from_volume handles unexpected exceptions.""" + cache_sync._endpoint_id = "test-endpoint-123" + + with ( + patch.object(cache_sync, "should_hydrate", return_value=True), + patch("os.makedirs", side_effect=Exception("Unexpected error")), + ): + # Should not raise exception + await cache_sync.hydrate_from_volume() + + +class TestMarkLastHydrated: + def test_mark_last_hydrated_skips_when_should_not_sync(self, cache_sync): + """Test that mark_last_hydrated skips when should_sync returns False.""" + with ( + patch.object(cache_sync, "should_sync", return_value=False), + patch.object(Path, "touch") as mock_touch, + ): + cache_sync.mark_last_hydrated() + # Should not touch the file when should_sync is False + mock_touch.assert_not_called() + + def test_mark_last_hydrated_creates_marker(self, cache_sync, mock_env): + """Test that mark_last_hydrated creates a marker file.""" + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch.object(Path, "touch") as mock_touch, + ): + cache_sync.mark_last_hydrated() + + # Should touch the hydration marker path + mock_touch.assert_called_once() + + def test_mark_last_hydrated_handles_exception(self, cache_sync, mock_env): + """Test that mark_last_hydrated handles exceptions gracefully.""" + with ( + patch.object(cache_sync, "should_sync", return_value=True), + patch.object(Path, "touch", side_effect=OSError("Permission denied")), + ): + # Should not raise exception + cache_sync.mark_last_hydrated() diff --git a/tests/unit/test_remote_executor.py b/tests/unit/test_remote_executor.py index 9784df7..f801805 100644 --- a/tests/unit/test_remote_executor.py +++ b/tests/unit/test_remote_executor.py @@ -189,3 +189,119 @@ def test_component_attribute_exposure(self): # Test class executor attributes through component assert hasattr(self.executor.class_executor, "class_instances") assert hasattr(self.executor.class_executor, "instance_metadata") + + @pytest.mark.asyncio + async def test_hydration_before_installation_with_dependencies(self): + """Test that hydrate_from_volume is called before installations when there are dependencies.""" + request = FunctionRequest( + function_name="test_func", + function_code="def test_func(): return 'test'", + dependencies=["requests"], + args=[], + kwargs={}, + ) + + with ( + patch.object( + self.executor.cache_sync, + "hydrate_from_volume", + new_callable=AsyncMock, + ) as mock_hydrate, + patch.object( + self.executor.cache_sync, + "mark_baseline", + ) as mock_baseline, + patch.object( + self.executor.dependency_installer, + "install_dependencies_async", + new_callable=AsyncMock, + ) as mock_deps, + patch.object(self.executor.function_executor, "execute") as mock_execute, + patch.object( + self.executor.cache_sync, + "sync_to_volume", + new_callable=AsyncMock, + ) as mock_sync, + ): + from remote_execution import FunctionResponse + + mock_deps.return_value = FunctionResponse( + success=True, stdout="Deps installed" + ) + mock_execute.return_value = Mock(success=True, result="encoded_result") + + await self.executor.ExecuteFunction(request) + + # Verify hydration was called + mock_hydrate.assert_called_once() + # Verify baseline was marked + mock_baseline.assert_called_once() + # Verify sync was called after installation + mock_sync.assert_called_once() + + @pytest.mark.asyncio + async def test_no_hydration_without_dependencies(self): + """Test that hydrate_from_volume is not called when there are no dependencies.""" + request = FunctionRequest( + function_name="test_func", + function_code="def test_func(): return 'test'", + args=[], + kwargs={}, + ) + + with ( + patch.object( + self.executor.cache_sync, + "hydrate_from_volume", + new_callable=AsyncMock, + ) as mock_hydrate, + patch.object(self.executor.function_executor, "execute") as mock_execute, + ): + mock_execute.return_value = Mock(success=True, result="encoded_result") + + await self.executor.ExecuteFunction(request) + + # Verify hydration was NOT called (no dependencies) + mock_hydrate.assert_not_called() + + @pytest.mark.asyncio + async def test_hydration_with_hf_models(self): + """Test that hydrate_from_volume is called when hf_models_to_cache is present.""" + request = FunctionRequest( + function_name="test_func", + function_code="def test_func(): return 'test'", + hf_models_to_cache=["bert-base-uncased"], + args=[], + kwargs={}, + ) + + with ( + patch.object( + self.executor.cache_sync, + "hydrate_from_volume", + new_callable=AsyncMock, + ) as mock_hydrate, + patch.object( + self.executor.cache_sync, + "mark_baseline", + ) as mock_baseline, + patch.object( + self.executor.hf_cache, + "cache_model_download_async", + new_callable=AsyncMock, + ) as mock_hf_cache, + patch.object(self.executor.function_executor, "execute") as mock_execute, + ): + from remote_execution import FunctionResponse + + mock_hf_cache.return_value = FunctionResponse( + success=True, stdout="Model cached" + ) + mock_execute.return_value = Mock(success=True, result="encoded_result") + + await self.executor.ExecuteFunction(request) + + # Verify hydration was called (hf_models present) + mock_hydrate.assert_called_once() + # Verify baseline was marked + mock_baseline.assert_called_once()