diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e945e7e..a117b64 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -76,17 +76,26 @@ jobs:
- name: Check code style with ruff
run: make lint
- local-execution-test:
+ docker-test:
runs-on: ubuntu-latest
needs: [test, lint]
- strategy:
- matrix:
- image-type: [cpu]
steps:
+ - name: Clear Space
+ if: github.event_name == 'pull_request'
+ run: |
+ rm -rf /usr/share/dotnet
+ rm -rf /opt/ghc
+ rm -rf "/usr/local/share/boost"
+ rm -rf "$AGENT_TOOLSDIRECTORY"
+
- name: Checkout repository
uses: actions/checkout@v4
with:
submodules: recursive
+ fetch-depth: 0
+
+ - name: Set up QEMU
+ uses: docker/setup-qemu-action@v3
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
@@ -102,28 +111,27 @@ jobs:
git submodule update --remote --rebase
cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/
- - name: Build CPU Docker image for handler testing
- if: matrix.image-type == 'cpu'
+ - name: Build CPU Docker image
uses: docker/build-push-action@v6
with:
context: .
file: ./Dockerfile-cpu
platforms: linux/amd64
push: false
- tags: tetra-rp-cpu:local-test
+ tags: tetra-rp-cpu:test
cache-from: type=gha
cache-to: type=gha,mode=max
load: true
- name: Test CPU handler execution in Docker environment
- if: matrix.image-type == 'cpu'
run: |
echo "Testing CPU handler in Docker environment..."
- docker run --rm tetra-rp-cpu:local-test ./test-handler.sh
+ docker run --rm tetra-rp-cpu:test ./test-handler.sh
+
release:
runs-on: ubuntu-latest
- needs: [test, lint, local-execution-test]
+ needs: [test, lint, docker-test]
if: github.ref == 'refs/heads/main'
outputs:
release_created: ${{ steps.release.outputs.release_created }}
@@ -135,61 +143,10 @@ jobs:
release-type: python
token: ${{ secrets.GITHUB_TOKEN }}
- docker-pr:
- runs-on: ubuntu-latest
- needs: [test, lint, local-execution-test]
- if: github.event_name == 'pull_request'
- steps:
- - name: Clear Space
- run: |
- rm -rf /usr/share/dotnet
- rm -rf /opt/ghc
- rm -rf "/usr/local/share/boost"
- rm -rf "$AGENT_TOOLSDIRECTORY"
-
- - name: Checkout repository
- uses: actions/checkout@v4
- with:
- submodules: recursive
- fetch-depth: 0
-
- - name: Set up QEMU
- uses: docker/setup-qemu-action@v3
-
- - name: Set up Docker Buildx
- uses: docker/setup-buildx-action@v3
-
- - name: Set up uv
- uses: astral-sh/setup-uv@v4
- with:
- enable-cache: true
-
- - name: Setup dependencies
- run: |
- uv sync
- git submodule update --remote --rebase
- cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/
-
- - name: Build and test CPU Docker image (PR)
- uses: docker/build-push-action@v6
- with:
- context: .
- file: ./Dockerfile-cpu
- platforms: linux/amd64
- push: false
- tags: tetra-rp-cpu:pr-test
- cache-from: type=gha
- cache-to: type=gha,mode=max
- load: true
-
- - name: Test CPU Docker image with handler test suite
- run: |
- echo "Testing CPU image with handler test suite..."
- docker run --rm tetra-rp-cpu:pr-test ./test-handler.sh
docker-main-gpu:
runs-on: ubuntu-latest
- needs: [test, lint, local-execution-test, release]
+ needs: [test, lint, docker-test, release]
if: github.ref == 'refs/heads/main' && github.event_name == 'push' && !needs.release.outputs.release_created
steps:
- name: Clear Space
@@ -242,7 +199,7 @@ jobs:
docker-main-cpu:
runs-on: ubuntu-latest
- needs: [test, lint, local-execution-test, release]
+ needs: [test, lint, docker-test, release]
if: github.ref == 'refs/heads/main' && github.event_name == 'push' && !needs.release.outputs.release_created
steps:
- name: Clear Space
diff --git a/.gitignore b/.gitignore
index aa23d69..0c60b0d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -120,6 +120,7 @@ celerybeat.pid
*.sage.py
# Environments
+.envrc
.env
.venv
env/
diff --git a/CLAUDE.md b/CLAUDE.md
index 64afc6d..b4a6152 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -17,12 +17,12 @@ This is `worker-tetra`, a RunPod Serverless worker template that provides dynami
- **Function Executor** (`src/function_executor.py:12`): Handles individual function execution with full output capture (stdout, stderr, logs)
- **Class Executor** (`src/class_executor.py:14`): Manages class instantiation and method execution with instance persistence and metadata tracking
-### 2. Workspace & Environment Management (`src/workspace_manager.py:12`)
-- Local workspace configuration and directory management
-- Environment variable setup for execution contexts
-- Integration with HuggingFace accelerator for model downloads
-- Python path configuration for execution isolation
-- Change directory management for execution context
+### 2. HuggingFace Model Cache-Ahead (`src/huggingface_cache.py`)
+- **Model Pre-Caching**: Downloads HuggingFace models before user code execution
+- **Cache Validation**: Checks if models are already cached to avoid redundant downloads
+- **Authentication**: Supports HF_TOKEN for private/gated model access
+- **Transfer Acceleration**: Uses hf_transfer when HF_HUB_ENABLE_HF_TRANSFER=1 is set
+- **Transparent Caching**: User code references models without knowing they're pre-cached
### 3. Dependency Management System (`src/dependency_installer.py:14`)
- **Python Package Installation**: UV-based package management with environment-aware configuration (Docker vs local)
@@ -32,48 +32,41 @@ This is `worker-tetra`, a RunPod Serverless worker template that provides dynami
- **System Package Filtering**: Intelligent detection of system-available packages to avoid redundant installation
- **Universal Subprocess Integration**: All subprocess operations use centralized logging utility
-### 4. Download Acceleration Infrastructure
-- **Download Accelerator** (`src/download_accelerator.py:166`): HuggingFace transfer optimization using hf_transfer
-- **HuggingFace Integration** (`src/huggingface_accelerator.py`): Model caching and acceleration strategies
-- **Strategy Pattern**: Multiple download strategies (native HF, hf_transfer, factory-based selection)
-- **Performance Metrics**: Download speed tracking and optimization reporting
-
-### 5. Universal Subprocess Utility (`src/subprocess_utils.py`)
+### 4. Universal Subprocess Utility (`src/subprocess_utils.py`)
- **Centralized Subprocess Operations**: All subprocess calls use `run_logged_subprocess` for consistency
- **Automatic Logging Integration**: All subprocess output flows through log streamer at DEBUG level
- **Environment-Aware Execution**: Handles Docker vs local environment differences automatically
- **Standardized Error Handling**: Consistent FunctionResponse pattern for all subprocess operations
- **Timeout Management**: Configurable timeouts with proper cleanup on timeout/cancellation
-### 6. Serialization & Protocol Management
+### 5. Serialization & Protocol Management
- **Protocol Definitions** (`src/remote_execution.py:13`): Pydantic models for request/response with validation
- **Serialization Utils** (`src/serialization_utils.py`): CloudPickle-based data serialization for function arguments and results
- **Base Executor** (`src/base_executor.py`): Common execution interface and environment setup
-### 7. Tetra SDK Integration (`tetra-rp/` submodule)
+### 6. Tetra SDK Integration (`tetra-rp/` submodule)
- **Client Interface**: `@remote` decorator for marking functions for remote execution
- **Resource Management**: GPU/CPU configuration and provisioning through LiveServerless objects
- **Live Serverless**: Dynamic infrastructure provisioning with auto-scaling
- **Protocol Buffers**: Communication protocol definitions for distributed execution
-### 8. Testing Infrastructure (`tests/`)
+### 7. Testing Infrastructure (`tests/`)
- **Unit Tests** (`tests/unit/`): Component-level testing for individual modules with mocking
- **Integration Tests** (`tests/integration/`): End-to-end workflow testing with real execution
- **Test Fixtures** (`tests/conftest.py:1`): Shared test data, mock objects, and utility functions
-- **Handler Testing**: Local execution validation with JSON test files (`src/test_*.json`)
- - **Full Coverage**: All 14 handler tests pass with environment-aware dependency installation
+- **Handler Testing**: Local execution validation with JSON test files (`src/tests/`)
+ - **Full Coverage**: All handler tests pass with environment-aware dependency installation
- **Cross-Platform**: Works correctly in both Docker containers and local macOS/Linux environments
-### 9. Build & Deployment Pipeline
+### 8. Build & Deployment Pipeline
- **Docker Containerization**: GPU (`Dockerfile`) and CPU (`Dockerfile-cpu`) image builds
- **CI/CD Pipeline**: Automated testing, linting, and releases (`.github/workflows/`)
- **Quality Gates** (`Makefile:104`): Format checking, type checking, test coverage requirements
- **Release Management**: Automated semantic versioning and Docker Hub deployment
-### 10. Configuration & Constants
-- **Constants** (`src/constants.py`): System-wide configuration values and thresholds
-- **Environment Configuration**: RunPod API integration and workspace paths
-- **Performance Tuning**: Download acceleration thresholds and caching strategies
+### 9. Configuration & Constants
+- **Constants** (`src/constants.py`): System-wide configuration values (NAMESPACE, LARGE_SYSTEM_PACKAGES)
+- **Environment Configuration**: RunPod API integration and HuggingFace cache settings
## Architecture
@@ -97,12 +90,13 @@ This is `worker-tetra`, a RunPod Serverless worker template that provides dynami
### Key Patterns
1. **Remote Function Execution**: Functions decorated with `@remote` are automatically executed on RunPod GPU workers
-2. **Composition Pattern**: RemoteExecutor uses specialized components (WorkspaceManager, DependencyInstaller, Executors)
+2. **Composition Pattern**: RemoteExecutor uses specialized components (DependencyInstaller, HuggingFaceCacheAhead, Executors)
3. **Dynamic Dependency Management**: Dependencies specified in decorators are installed at runtime with differential updates
-4. **Universal Subprocess Operations**: All subprocess calls use centralized `run_logged_subprocess` for consistent logging and error handling
-5. **Environment-Aware Configuration**: Automatic Docker vs local environment detection for appropriate installation methods
-6. **Serialization**: Uses cloudpickle + base64 encoding for function arguments and results
-7. **Resource Configuration**: `LiveServerless` objects define GPU requirements, scaling, and worker configuration
+4. **HuggingFace Cache-Ahead**: Models specified in `hf_models_to_cache` are pre-downloaded before execution
+5. **Universal Subprocess Operations**: All subprocess calls use centralized `run_logged_subprocess` for consistent logging and error handling
+6. **Environment-Aware Configuration**: Automatic Docker vs local environment detection for appropriate installation methods
+7. **Serialization**: Uses cloudpickle + base64 encoding for function arguments and results
+8. **Resource Configuration**: `LiveServerless` objects define GPU requirements, scaling, and worker configuration
## Development Commands
@@ -150,6 +144,8 @@ git submodule update --remote --rebase # Update tetra-rp to latest
### Environment Variables
- `RUNPOD_API_KEY`: Required for RunPod Serverless integration
- `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
- `DEBIAN_FRONTEND=noninteractive`: Set during system package installation
- `UV_CACHE_DIR`: Package cache configuration
- `VIRTUAL_ENV`: Virtual environment path configuration
@@ -207,13 +203,12 @@ gpu_config = LiveServerless(
│ ├── remote_execution.py # Protocol definitions
│ ├── function_executor.py # Function execution with output capture
│ ├── class_executor.py # Class execution with persistence
-│ ├── workspace_manager.py # Workspace and environment management
│ ├── dependency_installer.py # Python and system dependency management
-│ ├── download_accelerator.py # HuggingFace download optimization
+│ ├── huggingface_cache.py # HuggingFace model cache-ahead system
│ ├── serialization_utils.py # CloudPickle serialization utilities
│ ├── base_executor.py # Common execution interface
│ ├── constants.py # System-wide configuration constants
-│ └── test_*.json # Local handler test files
+│ └── tests/ # Handler test JSON files
├── tests/ # Comprehensive test suite
│ ├── conftest.py # Shared test fixtures
│ ├── unit/ # Unit tests for individual components
diff --git a/Dockerfile b/Dockerfile
index 3112e5d..896aaab 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,6 +2,9 @@ FROM pytorch/pytorch:2.8.0-cuda12.8-cudnn9-runtime
WORKDIR /app
+# Enable HuggingFace transfer acceleration
+ENV HF_HUB_ENABLE_HF_TRANSFER=1
+
# Install system dependencies and uv
RUN apt-get update && apt-get install -y --no-install-recommends \
curl ca-certificates nala \
@@ -12,7 +15,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
&& rm -rf /var/lib/apt/lists/*
# Copy app code and install dependencies
-COPY README.md src/* pyproject.toml uv.lock ./
+COPY README.md pyproject.toml uv.lock ./
+COPY src/ ./
RUN uv export --format requirements-txt --no-dev --no-hashes > requirements.txt \
&& uv pip install --system -r requirements.txt
diff --git a/Dockerfile-cpu b/Dockerfile-cpu
index e628df3..6c71c69 100644
--- a/Dockerfile-cpu
+++ b/Dockerfile-cpu
@@ -12,7 +12,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
&& rm -rf /var/lib/apt/lists/*
# Copy app code and install dependencies
-COPY README.md src/* pyproject.toml uv.lock ./
+COPY README.md pyproject.toml uv.lock ./
+COPY src/ ./
RUN uv export --format requirements-txt --no-dev --no-hashes > requirements.txt \
&& uv pip install --system -r requirements.txt
diff --git a/docs/Centralized_Log_Streaming_System.md b/docs/Centralized_Log_Streaming_System.md
index 8a5a63c..bcc0918 100644
--- a/docs/Centralized_Log_Streaming_System.md
+++ b/docs/Centralized_Log_Streaming_System.md
@@ -35,14 +35,11 @@ sequenceDiagram
participant RE as RemoteExecutor
participant LS as LogStreamer
participant DI as DependencyInstaller
- participant WM as WorkspaceManager
C->>RE: Execute Function
RE->>LS: Start Log Streaming
RE->>DI: Install Dependencies
DI-->>LS: Log installation progress
- RE->>WM: Setup Workspace
- WM-->>LS: Log workspace operations
RE->>RE: Execute Function
RE-->>LS: Capture execution logs
LS->>RE: Streamed logs
diff --git a/docs/Endpoint Persistence.md b/docs/Endpoint Persistence.md
new file mode 100644
index 0000000..f492615
--- /dev/null
+++ b/docs/Endpoint Persistence.md
@@ -0,0 +1,70 @@
+# Endpoint Persistence using Network Volume and CDR
+
+## Definitions
+
+- Worker: docker container that relies on their local sandbox environment `/app` for all of its operations.
+
+- Volume: network volume attached to a Worker as provisioned by its parent Endpoint.
+
+- Workspace: designated environment residing in the network volume `/runpod-volume/runtimes/{endpoint_id}`. Serves as an Endpoint persistence disk.
+
+- CDR: continuous data replication daemon that ensures data is replicated to the network volume workspace with an optional "hydrate" function (data transfers from volume to container)
+
+## Logic
+
+- First container boots, and checks for volume presence and endpoint workspace. Create if not found.
+
+ 1. Container will proceed to download any system, python or HF pre-cache instructed from the remote decorator.
+
+ 2. Container runs its job.
+
+ 3. Container launches its own CDR daemon to monitor `/app` for changes and replicates `/app` to `/runpod-volume/runtimes/{endpoint_id}` as files are downloaded or changed.
+
+- Subsequent container boots, and checks for volume presence and endpoint workspace. Found.
+
+ 1. Container launches its own CDR daemon to hydrate its `/app` from the workspace and then watch `/app` for changes.
+
+ 2. Container completely skips downloading from the internet.
+
+ 3. Container runs its job.
+
+### Logic Flow
+```mermaid
+graph TD
+ A[Container Boot] --> B{Volume Present?}
+ B -->|No| C[Use Local /app Only]
+ B -->|Yes| D{Workspace Exists?}
+
+ D -->|No| E["Create Workspace
/runpod-volume/runtimes/{endpoint_id}"]
+ D -->|Yes| F[Workspace Found]
+
+ E --> G[Launch CDR Daemon
Monitor /app → Workspace]
+ F --> H{First Container?}
+
+ H -->|Yes| I[Launch CDR Daemon
Monitor /app → Workspace]
+ H -->|No| J[Launch CDR Daemon
Hydrate /app ← Workspace
Then Monitor /app →
+ Workspace]
+
+ G --> K[Download Dependencies
System + Python + HF]
+ I --> K
+ J --> L[Skip Downloads
Use Cached Data]
+
+ K --> M[CDR: Replicate Downloads
/app → Workspace]
+ L --> N[Execute Job]
+ M --> O[Execute Job]
+
+ O --> P[CDR: Continue Monitoring
/app → Workspace]
+ N --> Q[CDR: Continue Monitoring
/app → Workspace]
+
+ C --> R[Execute Job
No Persistence]
+
+ style A fill:#e1f5fe,color:#000000
+ style G fill:#f3e5f5,color:#000000
+ style I fill:#f3e5f5,color:#000000
+ style J fill:#fff3e0,color:#000000
+ style K fill:#e8f5e8,color:#000000
+ style L fill:#fff9c4,color:#000000
+ style M fill:#fce4ec,color:#000000
+ style P fill:#fce4ec,color:#000000
+ style Q fill:#fce4ec,color:#000000
+```
\ No newline at end of file
diff --git a/docs/System_Python_Runtime_Architecture.md b/docs/System_Python_Runtime_Architecture.md
index e5e26a3..c9dc667 100644
--- a/docs/System_Python_Runtime_Architecture.md
+++ b/docs/System_Python_Runtime_Architecture.md
@@ -19,7 +19,6 @@ graph TD
F --> H[Function Execution]
G --> H
- I[WorkspaceManager] --> C
J[DependencyInstaller] --> C
K[FunctionExecutor] --> C
```
@@ -43,7 +42,6 @@ flowchart LR
```mermaid
graph TB
A[handler.py] --> B[RemoteExecutor]
- B --> C[WorkspaceManager]
B --> D[DependencyInstaller]
B --> E[FunctionExecutor]
B --> F[ClassExecutor]
@@ -52,7 +50,6 @@ graph TB
G --> E
G --> F
- H[download_accelerator] --> D
I[serialization_utils] --> E
I --> F
```
diff --git a/src/base_executor.py b/src/base_executor.py
deleted file mode 100644
index 4e4c156..0000000
--- a/src/base_executor.py
+++ /dev/null
@@ -1,47 +0,0 @@
-"""Base executor class to ensure consistent patterns across all executors."""
-
-from abc import ABC, abstractmethod
-from remote_execution import FunctionRequest, FunctionResponse
-
-
-class BaseExecutor(ABC):
- """
- Base class for all executors to ensure consistent initialization and execution patterns.
-
- This class enforces that all executors:
- 1. Accept workspace_manager in constructor
- 2. Setup Python path before execution
- 3. Follow consistent error handling patterns
- """
-
- def __init__(self, workspace_manager):
- """
- Initialize executor with required workspace manager.
-
- Args:
- workspace_manager: WorkspaceManager instance for volume operations
- """
- if workspace_manager is None:
- raise ValueError("workspace_manager is required for all executors")
- self.workspace_manager = workspace_manager
-
- def _setup_execution_environment(self):
- """
- Setup execution environment including Python path.
-
- This method MUST be called before any code execution to ensure:
- - Volume-installed packages are available in sys.path
- - Workspace is properly configured
- """
- # Setup Python path for volume packages - CRITICAL for volume-installed dependencies
- self.workspace_manager.setup_python_path()
-
- @abstractmethod
- def execute(self, request: FunctionRequest) -> FunctionResponse:
- """
- Execute the request. Subclasses must implement this method.
-
- IMPORTANT: All implementations MUST call self._setup_execution_environment()
- before executing any user code.
- """
- pass
diff --git a/src/class_executor.py b/src/class_executor.py
index 0dc7fd5..9347ed7 100644
--- a/src/class_executor.py
+++ b/src/class_executor.py
@@ -6,22 +6,20 @@
from datetime import datetime
from typing import Dict, Any, Tuple
-from base_executor import BaseExecutor
from remote_execution import FunctionRequest, FunctionResponse
from serialization_utils import SerializationUtils
-class ClassExecutor(BaseExecutor):
+class ClassExecutor:
"""Handles execution of class methods with instance management."""
- def __init__(self, workspace_manager):
- super().__init__(workspace_manager)
+ def __init__(self):
# Instance registry for persistent class instances
self.class_instances: Dict[str, Any] = {}
self.instance_metadata: Dict[str, Dict[str, Any]] = {}
def execute(self, request: FunctionRequest) -> FunctionResponse:
- """Execute class method - required by BaseExecutor interface."""
+ """Execute class method."""
return self.execute_class_method(request)
def execute_class_method(self, request: FunctionRequest) -> FunctionResponse:
@@ -33,16 +31,13 @@ def execute_class_method(self, request: FunctionRequest) -> FunctionResponse:
log_io = io.StringIO()
with redirect_stdout(stdout_io), redirect_stderr(stderr_io):
- try:
- # Setup execution environment including Python path
- self._setup_execution_environment()
-
- # Setup logging
- log_handler = logging.StreamHandler(log_io)
- log_handler.setLevel(logging.DEBUG)
- logger = logging.getLogger()
- logger.addHandler(log_handler)
+ # Setup logging
+ log_handler = logging.StreamHandler(log_io)
+ log_handler.setLevel(logging.DEBUG)
+ logger = logging.getLogger()
+ logger.addHandler(log_handler)
+ try:
# Get or create class instance
instance, instance_id = self._get_or_create_instance(request)
diff --git a/src/constants.py b/src/constants.py
index 667327a..26ddace 100644
--- a/src/constants.py
+++ b/src/constants.py
@@ -1,74 +1,6 @@
-# RunPod Volume Paths
-RUNPOD_VOLUME_PATH = "/runpod-volume"
-"""Path to the RunPod persistent volume mount point."""
-
-DEFAULT_WORKSPACE_PATH = "/app"
-"""Default workspace path when no persistent volume is available."""
-
-# Directory Names
-VENV_DIR_NAME = ".venv"
-"""Name of the virtual environment directory."""
-
-UV_CACHE_DIR_NAME = ".uv-cache"
-"""Name of the UV package cache directory."""
-
-HF_CACHE_DIR_NAME = ".hf-cache"
-"""Name of the Hugging Face cache directory."""
-
-WORKSPACE_LOCK_FILE = ".initialization.lock"
-"""Name of the workspace initialization lock file."""
-
-RUNTIMES_DIR_NAME = "runtimes"
-"""Name of the runtimes directory containing per-endpoint workspaces."""
-
-# Download Acceleration Settings
-MIN_SIZE_FOR_ACCELERATION_MB = 10
-"""Minimum file size in MB to trigger download acceleration."""
-
-DOWNLOAD_TIMEOUT_SECONDS = 600
-"""Default timeout for download operations in seconds."""
-
-# New download accelerator settings
-HF_TRANSFER_ENABLED = True
-"""Enable hf_transfer for fresh HuggingFace downloads."""
-
-
-# Size Conversion Constants
-BYTES_PER_MB = 1024 * 1024
-"""Number of bytes in a megabyte."""
-
-MB_SIZE_THRESHOLD = 1 * BYTES_PER_MB
-"""Minimum file size threshold for considering acceleration (1MB)."""
-
-# HuggingFace Model Patterns
-LARGE_HF_MODEL_PATTERNS = [
- "albert-large",
- "albert-xlarge",
- "bart-large",
- "bert-large",
- "bert-base",
- "codegen",
- "diffusion",
- "distilbert-base",
- "falcon",
- "gpt",
- "hubert",
- "llama",
- "mistral",
- "mpt",
- "pegasus",
- "roberta-large",
- "roberta-base",
- "santacoder",
- "stable-diffusion",
- "t5",
- "vae",
- "wav2vec2",
- "whisper",
- "xlm-roberta",
- "xlnet",
-]
-"""List of HuggingFace model patterns that benefit from download acceleration."""
+# Logger Configuration
+NAMESPACE = "tetra"
+"""Application logger namespace for all components."""
# System Package Acceleration with Nala
LARGE_SYSTEM_PACKAGES = [
diff --git a/src/dependency_installer.py b/src/dependency_installer.py
index 1043734..ea307e6 100644
--- a/src/dependency_installer.py
+++ b/src/dependency_installer.py
@@ -5,18 +5,15 @@
from typing import List
from remote_execution import FunctionResponse
-from download_accelerator import DownloadAccelerator
-from constants import LARGE_SYSTEM_PACKAGES
+from constants import LARGE_SYSTEM_PACKAGES, NAMESPACE
from subprocess_utils import run_logged_subprocess
class DependencyInstaller:
"""Handles installation of system and Python dependencies."""
- def __init__(self, workspace_manager):
- self.workspace_manager = workspace_manager
- self.logger = logging.getLogger(f"worker_tetra.{__name__.split('.')[-1]}")
- self.download_accelerator = DownloadAccelerator(workspace_manager)
+ def __init__(self):
+ self.logger = logging.getLogger(f"{NAMESPACE}.{__name__.split('.')[-1]}")
self._nala_available = None # Cache nala availability check
self._is_docker = None # Cache Docker environment detection
diff --git a/src/download_accelerator.py b/src/download_accelerator.py
deleted file mode 100644
index 9f59385..0000000
--- a/src/download_accelerator.py
+++ /dev/null
@@ -1,266 +0,0 @@
-"""
-Download acceleration using hf_transfer for optimal HuggingFace model downloads.
-
-This module provides accelerated download capabilities optimized for HuggingFace models:
-- hf_transfer for accelerated downloads when available
-- hf_xet acceleration is automatically handled by HuggingFace Hub (huggingface_hub>=0.32.0)
-- Standard HF hub as reliable fallback
-"""
-
-import os
-import time
-import logging
-from dataclasses import dataclass
-from typing import Optional
-
-from remote_execution import FunctionResponse
-from constants import (
- MIN_SIZE_FOR_ACCELERATION_MB,
- HF_TRANSFER_ENABLED,
-)
-
-
-@dataclass
-class DownloadMetrics:
- """Performance metrics for download operations."""
-
- method: str
- file_size_bytes: int
- total_time_seconds: float
- average_speed_mbps: float
- success: bool
- error_message: Optional[str] = None
-
- @property
- def speed_mb_per_sec(self) -> float:
- """Convert to MB/s for easier reading."""
- return self.average_speed_mbps / 8.0
-
- @property
- def file_size_mb(self) -> float:
- """File size in megabytes."""
- return self.file_size_bytes / (1024 * 1024)
-
-
-class HfTransferDownloader:
- """HuggingFace Transfer downloader for fresh downloads."""
-
- def __init__(self):
- self.logger = logging.getLogger(__name__)
- self.hf_transfer_available = self._check_hf_transfer()
-
- def _check_hf_transfer(self) -> bool:
- """Check if hf_transfer is available."""
- import importlib.util
-
- if importlib.util.find_spec("hf_transfer") is not None:
- return HF_TRANSFER_ENABLED
- else:
- self.logger.debug("hf_transfer not available")
- return False
-
- def download(
- self,
- url: str,
- output_path: str,
- show_progress: bool = False,
- ) -> DownloadMetrics:
- """
- Download file using hf_transfer for maximum speed.
-
- Args:
- url: URL to download
- output_path: Local file path to save to
- show_progress: Whether to show real-time progress
-
- Returns:
- DownloadMetrics with performance data
- """
- if not self.hf_transfer_available:
- raise RuntimeError("hf_transfer not available")
-
- start_time = time.time()
-
- try:
- # Set HF_HUB_ENABLE_HF_TRANSFER environment variable
- env = os.environ.copy()
- env["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-
- # Add authentication if HF token is available
- hf_token = os.environ.get("HF_TOKEN")
- if hf_token:
- env["HF_TOKEN"] = hf_token
-
- # Use hf_transfer via huggingface_hub
- from huggingface_hub import hf_hub_download
-
- # Extract model_id and filename from URL
- # URL format: https://huggingface.co/{model_id}/resolve/{revision}/{filename}
- if "huggingface.co" in url and "/resolve/" in url:
- parts = url.replace("https://huggingface.co/", "").split("/resolve/")
- model_id = parts[0]
- revision_and_filename = parts[1].split("/", 1)
- revision = revision_and_filename[0]
- filename = revision_and_filename[1]
-
- # Create output directory
- os.makedirs(os.path.dirname(output_path), exist_ok=True)
-
- # Download using hf_hub_download with hf_transfer enabled
- downloaded_path = hf_hub_download(
- repo_id=model_id,
- filename=filename,
- revision=revision,
- cache_dir=os.path.dirname(output_path),
- local_dir=os.path.dirname(output_path),
- local_dir_use_symlinks=False,
- )
-
- # Move to expected location if needed
- if downloaded_path != output_path:
- import shutil
-
- shutil.move(downloaded_path, output_path)
-
- else:
- # Fallback to direct download for non-HF URLs
- raise ValueError("hf_transfer only supports HuggingFace URLs")
-
- end_time = time.time()
- file_size = (
- os.path.getsize(output_path) if os.path.exists(output_path) else 0
- )
- total_time = end_time - start_time
-
- if total_time > 0 and file_size > 0:
- bits_per_second = (file_size * 8) / total_time
- avg_speed = bits_per_second / (1024 * 1024)
- else:
- avg_speed = 0
-
- self.logger.info(
- f"Downloaded {file_size / (1024 * 1024):.1f}MB in {total_time:.1f}s "
- f"({avg_speed / 8:.1f} MB/s) using hf_transfer"
- )
-
- return DownloadMetrics(
- method="hf_transfer",
- file_size_bytes=file_size,
- total_time_seconds=total_time,
- average_speed_mbps=avg_speed,
- success=True,
- )
-
- except Exception as e:
- self.logger.error(f"hf_transfer download failed: {str(e)}")
- return DownloadMetrics(
- method="hf_transfer",
- file_size_bytes=0,
- total_time_seconds=time.time() - start_time,
- average_speed_mbps=0,
- success=False,
- error_message=str(e),
- )
-
-
-class DownloadAccelerator:
- """
- Main download acceleration coordinator using hf_transfer.
-
- Note: hf_xet acceleration is now automatically handled by HuggingFace Hub
- when using hf_hub_download() or snapshot_download() functions.
- """
-
- def __init__(self, workspace_manager=None):
- self.workspace_manager = workspace_manager
- self.logger = logging.getLogger(__name__)
- self.hf_transfer_downloader = HfTransferDownloader()
-
- def should_accelerate_download(
- self, url: str, estimated_size_mb: float = 0
- ) -> bool:
- """
- Determine if download should be accelerated.
-
- Args:
- url: Download URL
- estimated_size_mb: Estimated file size in MB
-
- Returns:
- True if download should be accelerated
- """
- # Only accelerate HuggingFace downloads with our new methods
- if "huggingface.co" not in url:
- return False
-
- if estimated_size_mb >= MIN_SIZE_FOR_ACCELERATION_MB:
- return True
-
- # For HuggingFace URLs, always try acceleration
- return True
-
- def is_file_cached(self, output_path: str) -> bool:
- """Check if file is already cached locally."""
- return os.path.exists(output_path) and os.path.getsize(output_path) > 0
-
- def download_with_fallback(
- self,
- url: str,
- output_path: str,
- estimated_size_mb: float = 0,
- show_progress: bool = False,
- ) -> FunctionResponse:
- """
- Download with HF optimization when applicable.
-
- Strategy:
- 1. Use hf_transfer for HF URLs when available and size warrants acceleration
- 2. Otherwise return failure - let HF's native download handling work
-
- Args:
- url: URL to download
- output_path: Local file path
- estimated_size_mb: Estimated size for acceleration decision
- show_progress: Whether to show progress
-
- Returns:
- FunctionResponse with download result
- """
- if not self.should_accelerate_download(url, estimated_size_mb):
- self.logger.info(
- f"Not accelerating download, letting HF handle natively: {url}"
- )
- return FunctionResponse(
- success=False,
- error="No acceleration available - defer to HF native handling",
- )
-
- # Strategy 1: Try hf_transfer (hf_xet is automatically used by HF Hub when available)
- if self.hf_transfer_downloader.hf_transfer_available:
- try:
- self.logger.info(f"Using hf_transfer for download: {url}")
- metrics = self.hf_transfer_downloader.download(
- url, output_path, show_progress=show_progress
- )
-
- if metrics.success:
- return FunctionResponse(
- success=True,
- stdout=f"Downloaded {metrics.file_size_mb:.1f}MB in {metrics.total_time_seconds:.1f}s "
- f"({metrics.speed_mb_per_sec:.1f} MB/s) using hf_transfer",
- )
- else:
- self.logger.warning(
- f"hf_transfer download failed: {metrics.error_message}"
- )
- except Exception as e:
- self.logger.warning(f"hf_transfer download failed: {e}")
-
- # No acceleration available - let HF handle natively
- self.logger.info(
- f"No acceleration available for {url}, deferring to HF native handling"
- )
- return FunctionResponse(
- success=False,
- error="Acceleration not available - defer to HF native handling",
- )
diff --git a/src/function_executor.py b/src/function_executor.py
index 02465d4..7f94690 100644
--- a/src/function_executor.py
+++ b/src/function_executor.py
@@ -4,17 +4,13 @@
from contextlib import redirect_stdout, redirect_stderr
from typing import Dict, Any
-from base_executor import BaseExecutor
from remote_execution import FunctionRequest, FunctionResponse
from serialization_utils import SerializationUtils
-class FunctionExecutor(BaseExecutor):
+class FunctionExecutor:
"""Handles execution of individual functions with output capture."""
- def __init__(self, workspace_manager):
- super().__init__(workspace_manager)
-
def execute(self, request: FunctionRequest) -> FunctionResponse:
"""
Execute a function with full output capture.
@@ -24,83 +20,70 @@ def execute(self, request: FunctionRequest) -> FunctionResponse:
Returns:
FunctionResponse object with execution result
"""
- original_cwd = self.workspace_manager.change_to_workspace()
-
stdout_io = io.StringIO()
stderr_io = io.StringIO()
log_io = io.StringIO()
- try:
- # Setup execution environment including Python path
- self._setup_execution_environment()
-
- # Capture all stdout, stderr, and logs
- with redirect_stdout(stdout_io), redirect_stderr(stderr_io):
- try:
- # Setup logging capture
- log_handler = logging.StreamHandler(log_io)
- log_handler.setLevel(logging.DEBUG)
- logger = logging.getLogger()
- logger.addHandler(log_handler)
-
- # Execute function code in namespace
- namespace: Dict[str, Any] = {}
- if request.function_code:
- exec(request.function_code, namespace)
-
- if request.function_name not in namespace:
- return FunctionResponse(
- success=False,
- result=f"Function '{request.function_name}' not found in the provided code",
- )
-
- func = namespace[request.function_name]
-
- # Deserialize arguments
- args = SerializationUtils.deserialize_args(request.args)
- kwargs = SerializationUtils.deserialize_kwargs(request.kwargs)
-
- # Execute the function
- result = func(*args, **kwargs)
-
- except Exception as e:
- # Combine output streams
- combined_output = (
- stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue()
- )
-
- # Capture full traceback
- traceback_str = traceback.format_exc()
- error_message = f"{str(e)}\n{traceback_str}"
-
+ # Capture all stdout, stderr, and logs
+ with redirect_stdout(stdout_io), redirect_stderr(stderr_io):
+ # Setup logging capture
+ log_handler = logging.StreamHandler(log_io)
+ log_handler.setLevel(logging.DEBUG)
+ logger = logging.getLogger()
+ logger.addHandler(log_handler)
+
+ try:
+ # Execute function code in namespace
+ namespace: Dict[str, Any] = {}
+ if request.function_code:
+ exec(request.function_code, namespace)
+
+ if request.function_name not in namespace:
return FunctionResponse(
success=False,
- error=error_message,
- stdout=combined_output,
+ result=f"Function '{request.function_name}' not found in the provided code",
)
- finally:
- # Clean up logging handler
- if "logger" in locals() and "log_handler" in locals():
- logger.removeHandler(log_handler)
-
- # Serialize result
- serialized_result = SerializationUtils.serialize_result(result)
-
- # Combine output streams
- combined_output = (
- stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue()
- )
-
- return FunctionResponse(
- success=True,
- result=serialized_result,
- stdout=combined_output,
- )
-
- finally:
- # Restore original working directory
- if original_cwd:
- import os
-
- os.chdir(original_cwd)
+ func = namespace[request.function_name]
+
+ # Deserialize arguments
+ args = SerializationUtils.deserialize_args(request.args)
+ kwargs = SerializationUtils.deserialize_kwargs(request.kwargs)
+
+ # Execute the function
+ result = func(*args, **kwargs)
+
+ except Exception as e:
+ # Combine output streams
+ combined_output = (
+ stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue()
+ )
+
+ # Capture full traceback
+ traceback_str = traceback.format_exc()
+ error_message = f"{str(e)}\n{traceback_str}"
+
+ return FunctionResponse(
+ success=False,
+ error=error_message,
+ stdout=combined_output,
+ )
+
+ finally:
+ # Clean up logging handler
+ if "logger" in locals() and "log_handler" in locals():
+ logger.removeHandler(log_handler)
+
+ # Serialize result
+ serialized_result = SerializationUtils.serialize_result(result)
+
+ # Combine output streams
+ combined_output = (
+ stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue()
+ )
+
+ return FunctionResponse(
+ success=True,
+ result=serialized_result,
+ stdout=combined_output,
+ )
diff --git a/src/hf_download_strategy.py b/src/hf_download_strategy.py
deleted file mode 100644
index d8e1df0..0000000
--- a/src/hf_download_strategy.py
+++ /dev/null
@@ -1,81 +0,0 @@
-"""
-HuggingFace download strategy interface.
-
-Provides pluggable download strategies for HuggingFace models to allow
-switching between different acceleration methods and benchmarking performance.
-"""
-
-from abc import ABC, abstractmethod
-from typing import Dict, Any
-from remote_execution import FunctionResponse
-
-
-class HFDownloadStrategy(ABC):
- """Abstract base class for HuggingFace download strategies."""
-
- @abstractmethod
- def download_model(self, model_id: str, revision: str = "main") -> FunctionResponse:
- """
- Download a HuggingFace model.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- FunctionResponse with download results
- """
- pass
-
- @abstractmethod
- def is_model_cached(self, model_id: str, revision: str = "main") -> bool:
- """
- Check if model is already cached.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- True if model appears to be cached
- """
- pass
-
- @abstractmethod
- def get_cache_info(self, model_id: str) -> Dict[str, Any]:
- """
- Get cache information for a model.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- Dictionary with cache information
- """
- pass
-
- @abstractmethod
- def should_accelerate(self, model_id: str) -> bool:
- """
- Determine if model should use acceleration.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- True if acceleration should be used
- """
- pass
-
- @abstractmethod
- def clear_model_cache(self, model_id: str) -> FunctionResponse:
- """
- Clear cache for a specific model.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- FunctionResponse with clearing result
- """
- pass
diff --git a/src/hf_downloader_native.py b/src/hf_downloader_native.py
deleted file mode 100644
index 4e1f630..0000000
--- a/src/hf_downloader_native.py
+++ /dev/null
@@ -1,175 +0,0 @@
-"""
-Native HuggingFace downloader strategy.
-
-This strategy implements the current simplified approach using HF Hub's
-native snapshot_download() with built-in acceleration support.
-"""
-
-import logging
-from typing import Dict, Any
-
-from huggingface_hub import HfApi, snapshot_download
-from remote_execution import FunctionResponse
-from hf_download_strategy import HFDownloadStrategy
-from constants import LARGE_HF_MODEL_PATTERNS, BYTES_PER_MB
-
-
-class NativeHFDownloader(HFDownloadStrategy):
- """Native HuggingFace downloader using HF Hub's built-in acceleration."""
-
- def __init__(self, workspace_manager):
- self.workspace_manager = workspace_manager
- self.logger = logging.getLogger(__name__)
- self.api = HfApi()
-
- # HF will automatically use HF_HOME environment variable set by workspace_manager
- # No need to manually manage cache directories
-
- def should_accelerate(self, model_id: str) -> bool:
- """
- Determine if model should be pre-cached.
- HF Hub automatically uses hf_transfer when available.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- True if model should be pre-cached
- """
- model_lower = model_id.lower()
- return any(pattern in model_lower for pattern in LARGE_HF_MODEL_PATTERNS)
-
- def download_model(self, model_id: str, revision: str = "main") -> FunctionResponse:
- """
- Pre-download HuggingFace model using HF Hub's native caching.
-
- This method downloads the complete model snapshot to HF's standard cache
- location, leveraging hf_transfer when available.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- FunctionResponse with download results
- """
- if not self.should_accelerate(model_id):
- return FunctionResponse(
- success=True, stdout=f"Model {model_id} does not require pre-caching"
- )
-
- self.logger.info(f"Pre-caching model: {model_id}")
-
- try:
- # Use HF Hub's native snapshot download with acceleration
- snapshot_path = snapshot_download(
- repo_id=model_id,
- revision=revision,
- # HF automatically uses HF_HOME/HF_HUB_CACHE from environment
- # and applies hf_transfer acceleration when available
- )
-
- return FunctionResponse(
- success=True,
- stdout=f"Successfully pre-cached model {model_id} to {snapshot_path}",
- )
-
- except Exception as e:
- return FunctionResponse(
- success=False,
- error=f"Failed to pre-cache model {model_id}: {str(e)}",
- )
-
- def is_model_cached(self, model_id: str, revision: str = "main") -> bool:
- """
- Check if model is already cached using HF Hub's cache utilities.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- True if model appears to be cached
- """
- try:
- from huggingface_hub import try_to_load_from_cache
-
- # Check for common model files that indicate a cached model
- key_files = ["config.json", "pytorch_model.bin", "model.safetensors"]
-
- for filename in key_files:
- cached_path = try_to_load_from_cache(
- repo_id=model_id, filename=filename, revision=revision
- )
- if cached_path is not None: # Found cached file
- return True
-
- return False
- except Exception:
- return False
-
- def get_cache_info(self, model_id: str) -> Dict[str, Any]:
- """
- Get cache information for a model using HF Hub utilities.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- Dictionary with cache information
- """
- try:
- from huggingface_hub import scan_cache_dir
-
- cache_info = scan_cache_dir()
-
- # Find our specific model in the cache
- for repo in cache_info.repos:
- if repo.repo_id == model_id:
- return {
- "cached": True,
- "cache_size_mb": repo.size_on_disk / BYTES_PER_MB,
- "file_count": len(list(repo.revisions)[0].files)
- if repo.revisions
- else 0,
- "cache_path": str(repo.repo_path),
- }
-
- return {"cached": False, "cache_size_mb": 0, "file_count": 0}
-
- except Exception:
- return {"cached": False, "cache_size_mb": 0, "file_count": 0}
-
- def clear_model_cache(self, model_id: str) -> FunctionResponse:
- """
- Clear cache for a specific model using HF Hub utilities.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- FunctionResponse with clearing result
- """
- try:
- from huggingface_hub import scan_cache_dir
-
- cache_info = scan_cache_dir()
-
- # Find and delete our specific model
- for repo in cache_info.repos:
- if repo.repo_id == model_id:
- delete_strategy = cache_info.delete_revisions(repo.repo_id)
- delete_strategy.execute()
-
- return FunctionResponse(
- success=True, stdout=f"Cleared cache for model {model_id}"
- )
-
- return FunctionResponse(
- success=True, stdout=f"No cache found for model {model_id}"
- )
-
- except Exception as e:
- return FunctionResponse(
- success=False, error=f"Failed to clear cache for {model_id}: {str(e)}"
- )
diff --git a/src/hf_downloader_tetra.py b/src/hf_downloader_tetra.py
deleted file mode 100644
index d9fa6ab..0000000
--- a/src/hf_downloader_tetra.py
+++ /dev/null
@@ -1,270 +0,0 @@
-"""
-Tetra HuggingFace downloader strategy.
-
-This strategy implements a custom acceleration logic with
-manual file enumeration and file-by-file downloads using
-hf_transfer and custom acceleration methods.
-"""
-
-import logging
-from typing import Dict, List, Any
-from pathlib import Path
-
-from huggingface_hub import HfApi
-from remote_execution import FunctionResponse
-from hf_download_strategy import HFDownloadStrategy
-from download_accelerator import DownloadAccelerator
-from constants import LARGE_HF_MODEL_PATTERNS, BYTES_PER_MB, MB_SIZE_THRESHOLD
-
-
-class TetraHFDownloader(HFDownloadStrategy):
- """Custom Tetra HuggingFace downloader with manual acceleration logic."""
-
- def __init__(self, workspace_manager):
- self.workspace_manager = workspace_manager
- self.logger = logging.getLogger(__name__)
- self.download_accelerator = DownloadAccelerator(workspace_manager)
- self.api = HfApi()
-
- # Use workspace manager's HF cache if available
- if workspace_manager and workspace_manager.hf_cache_path:
- self.cache_dir = Path(workspace_manager.hf_cache_path)
- else:
- self.cache_dir = Path.home() / ".cache" / "huggingface"
-
- self.cache_dir.mkdir(parents=True, exist_ok=True)
-
- def get_model_files(
- self, model_id: str, revision: str = "main"
- ) -> List[Dict[str, Any]]:
- """
- Get list of files for a HuggingFace model using the HF Hub API.
-
- Args:
- model_id: HuggingFace model identifier (e.g., 'gpt2', 'microsoft/DialoGPT-medium')
- revision: Model revision/branch (default: 'main')
-
- Returns:
- List of file information dictionaries
- """
- try:
- # Use HF Hub's native API instead of manual requests
- repo_info = self.api.repo_info(model_id, revision=revision)
-
- files = []
- if repo_info.siblings:
- for sibling in repo_info.siblings:
- if sibling.rfilename: # Only include actual files
- files.append(
- {
- "path": sibling.rfilename,
- "size": getattr(sibling, "size", 0) or 0,
- "url": f"https://huggingface.co/{model_id}/resolve/{revision}/{sibling.rfilename}",
- }
- )
-
- return files
-
- except Exception as e:
- self.logger.warning(f"Could not fetch model file list for {model_id}: {e}")
- return []
-
- def should_accelerate(self, model_id: str) -> bool:
- """
- Determine if model downloads should be accelerated.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- True if acceleration should be used
- """
- # Check if hf_transfer is available
- has_hf_transfer = (
- self.download_accelerator.hf_transfer_downloader.hf_transfer_available
- )
-
- if not has_hf_transfer:
- return False
-
- model_lower = model_id.lower()
- return any(pattern in model_lower for pattern in LARGE_HF_MODEL_PATTERNS)
-
- def download_model(self, model_id: str, revision: str = "main") -> FunctionResponse:
- """
- Download HuggingFace model files using Tetra's custom acceleration.
-
- This method downloads model files to the cache before transformers tries to access them,
- using hf_transfer or custom acceleration for optimized downloads.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- FunctionResponse with download results
- """
- if not self.should_accelerate(model_id):
- return FunctionResponse(
- success=True, stdout=f"Model {model_id} does not require acceleration"
- )
-
- self.logger.info(f"Accelerating model download: {model_id}")
-
- # Get model file list
- files = self.get_model_files(model_id, revision)
- if not files:
- return FunctionResponse(
- success=False, error=f"Could not get file list for model {model_id}"
- )
-
- # Filter for main model files (ignore small config files)
- large_files = [f for f in files if f["size"] > MB_SIZE_THRESHOLD]
-
- if not large_files:
- return FunctionResponse(
- success=True, stdout=f"No large files found for model {model_id}"
- )
-
- self.logger.info(
- f"Found {len(large_files)} large files to download for {model_id}"
- )
-
- # Create model-specific cache directory
- model_cache_dir = self.cache_dir / "transformers" / model_id.replace("/", "--")
- model_cache_dir.mkdir(parents=True, exist_ok=True)
-
- successful_downloads = 0
- total_size = sum(f["size"] for f in large_files)
-
- for file_info in large_files:
- file_path = model_cache_dir / file_info["path"]
- file_path.parent.mkdir(parents=True, exist_ok=True)
-
- # Skip if file already exists and is correct size
- if file_path.exists() and file_path.stat().st_size == file_info["size"]:
- self.logger.info(f"✓ {file_info['path']} (cached)")
- successful_downloads += 1
- continue
-
- try:
- file_size_mb = file_info["size"] / BYTES_PER_MB
- self.logger.info(
- f"Downloading {file_info['path']} ({file_size_mb:.1f}MB)..."
- )
-
- # Use download accelerator
- result = self.download_accelerator.download_with_fallback(
- file_info["url"],
- str(file_path),
- estimated_size_mb=file_size_mb,
- show_progress=True,
- )
-
- if result.success:
- successful_downloads += 1
- self.logger.info(f"✓ {file_info['path']} downloaded successfully")
- else:
- self.logger.error(f"✗ {file_info['path']} failed: {result.error}")
-
- except Exception as e:
- self.logger.error(
- f"✗ {file_info['path']} failed with exception: {str(e)}"
- )
-
- success = successful_downloads == len(large_files)
-
- if success:
- return FunctionResponse(
- success=True,
- stdout=f"Successfully pre-downloaded {successful_downloads} files "
- f"({total_size / BYTES_PER_MB:.1f}MB) for model {model_id}",
- )
- else:
- return FunctionResponse(
- success=False,
- error=f"Failed to download {len(large_files) - successful_downloads} files for {model_id}",
- stdout=f"Downloaded {successful_downloads}/{len(large_files)} files",
- )
-
- def is_model_cached(self, model_id: str, revision: str = "main") -> bool:
- """
- Check if model is already cached.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- True if model appears to be cached
- """
- model_cache_dir = self.cache_dir / "transformers" / model_id.replace("/", "--")
-
- if not model_cache_dir.exists():
- return False
-
- # Check if there are any model files
- model_files = list(model_cache_dir.glob("**/*.bin")) + list(
- model_cache_dir.glob("**/*.safetensors")
- )
- return len(model_files) > 0
-
- def get_cache_info(self, model_id: str) -> Dict[str, Any]:
- """
- Get cache information for a model.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- Dictionary with cache information
- """
- model_cache_dir = self.cache_dir / "transformers" / model_id.replace("/", "--")
-
- if not model_cache_dir.exists():
- return {"cached": False, "cache_size_mb": 0, "file_count": 0}
-
- total_size = 0
- file_count = 0
-
- for file_path in model_cache_dir.rglob("*"):
- if file_path.is_file():
- total_size += file_path.stat().st_size
- file_count += 1
-
- return {
- "cached": file_count > 0,
- "cache_size_mb": total_size / BYTES_PER_MB,
- "file_count": file_count,
- "cache_path": str(model_cache_dir),
- }
-
- def clear_model_cache(self, model_id: str) -> FunctionResponse:
- """
- Clear cache for a specific model.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- FunctionResponse with clearing result
- """
- model_cache_dir = self.cache_dir / "transformers" / model_id.replace("/", "--")
-
- if not model_cache_dir.exists():
- return FunctionResponse(
- success=True, stdout=f"No cache found for model {model_id}"
- )
-
- try:
- import shutil
-
- shutil.rmtree(model_cache_dir)
-
- return FunctionResponse(
- success=True, stdout=f"Cleared cache for model {model_id}"
- )
- except Exception as e:
- return FunctionResponse(
- success=False, error=f"Failed to clear cache for {model_id}: {str(e)}"
- )
diff --git a/src/hf_strategy_factory.py b/src/hf_strategy_factory.py
deleted file mode 100644
index 1ce81de..0000000
--- a/src/hf_strategy_factory.py
+++ /dev/null
@@ -1,119 +0,0 @@
-"""
-HuggingFace download strategy factory.
-
-Provides configuration system for switching between different HF download strategies
-and creating the appropriate downloader instance based on environment variables.
-"""
-
-import os
-import logging
-from typing import Optional, Dict, Any
-
-from hf_download_strategy import HFDownloadStrategy
-from hf_downloader_tetra import TetraHFDownloader
-from hf_downloader_native import NativeHFDownloader
-
-
-class HFStrategyFactory:
- """Factory for creating HF download strategy instances."""
-
- # Environment variable name
- STRATEGY_ENV_VAR = "HF_DOWNLOAD_STRATEGY"
-
- # Available strategy names
- TETRA_STRATEGY = "tetra"
- NATIVE_STRATEGY = "native"
-
- # Default strategy
- DEFAULT_STRATEGY = TETRA_STRATEGY
-
- @classmethod
- def get_available_strategies(cls) -> list[str]:
- """Get list of available strategy names."""
- return [cls.TETRA_STRATEGY, cls.NATIVE_STRATEGY]
-
- @classmethod
- def get_configured_strategy(cls) -> str:
- """
- Get the configured strategy name from environment variables.
-
- Returns:
- Strategy name (defaults to native if not configured)
- """
- strategy = os.environ.get(cls.STRATEGY_ENV_VAR, cls.DEFAULT_STRATEGY).lower()
-
- # Validate strategy
- if strategy not in cls.get_available_strategies():
- logger = logging.getLogger(__name__)
- logger.warning(
- f"Unknown HF download strategy '{strategy}', falling back to '{cls.DEFAULT_STRATEGY}'"
- )
- return cls.DEFAULT_STRATEGY
-
- return strategy
-
- @classmethod
- def create_strategy(
- cls, workspace_manager, strategy: Optional[str] = None
- ) -> HFDownloadStrategy:
- """
- Create HF download strategy instance.
-
- Args:
- workspace_manager: Workspace manager instance
- strategy: Optional strategy override (defaults to environment configuration)
-
- Returns:
- HFDownloadStrategy instance
- """
- if strategy is None:
- strategy = cls.get_configured_strategy()
-
- logger = logging.getLogger(__name__)
- logger.info(f"Creating HF download strategy: {strategy}")
-
- if strategy == cls.TETRA_STRATEGY:
- return TetraHFDownloader(workspace_manager)
- elif strategy == cls.NATIVE_STRATEGY:
- return NativeHFDownloader(workspace_manager)
- else:
- # Fallback to native
- logger.warning(f"Unknown strategy '{strategy}', using native")
- return NativeHFDownloader(workspace_manager)
-
- @classmethod
- def set_strategy(cls, strategy: str) -> None:
- """
- Set the HF download strategy via environment variable.
-
- Args:
- strategy: Strategy name to set
- """
- if strategy not in cls.get_available_strategies():
- raise ValueError(
- f"Invalid strategy '{strategy}'. Available: {cls.get_available_strategies()}"
- )
-
- os.environ[cls.STRATEGY_ENV_VAR] = strategy
-
- logger = logging.getLogger(__name__)
- logger.info(f"Set HF download strategy to: {strategy}")
-
- @classmethod
- def get_strategy_info(cls) -> Dict[str, Any]:
- """
- Get information about the current strategy configuration.
-
- Returns:
- Dictionary with strategy configuration info
- """
- current_strategy = cls.get_configured_strategy()
- env_value = os.environ.get(cls.STRATEGY_ENV_VAR, "not set")
-
- return {
- "current_strategy": current_strategy,
- "environment_variable": cls.STRATEGY_ENV_VAR,
- "environment_value": env_value,
- "default_strategy": cls.DEFAULT_STRATEGY,
- "available_strategies": cls.get_available_strategies(),
- }
diff --git a/src/huggingface_accelerator.py b/src/huggingface_accelerator.py
deleted file mode 100644
index 2f2b2ad..0000000
--- a/src/huggingface_accelerator.py
+++ /dev/null
@@ -1,150 +0,0 @@
-"""
-HuggingFace model download acceleration.
-
-This module provides accelerated downloads for HuggingFace models and datasets,
-integrating with the existing volume workspace caching system using pluggable
-download strategies.
-"""
-
-import logging
-from typing import Dict, List, Any
-
-from huggingface_hub import HfApi
-from remote_execution import FunctionResponse
-from hf_strategy_factory import HFStrategyFactory
-from hf_download_strategy import HFDownloadStrategy
-
-
-class HuggingFaceAccelerator:
- """Accelerated downloads for HuggingFace models and files using pluggable strategies."""
-
- def __init__(self, workspace_manager):
- self.workspace_manager = workspace_manager
- self.logger = logging.getLogger(__name__)
- self.api = HfApi()
-
- # Create the configured download strategy
- self.strategy: HFDownloadStrategy = HFStrategyFactory.create_strategy(
- workspace_manager
- )
-
- def get_model_files(
- self, model_id: str, revision: str = "main"
- ) -> List[Dict[str, Any]]:
- """
- Get list of files for a HuggingFace model using the HF Hub API.
-
- Args:
- model_id: HuggingFace model identifier (e.g., 'gpt2', 'microsoft/DialoGPT-medium')
- revision: Model revision/branch (default: 'main')
-
- Returns:
- List of file information dictionaries
- """
- try:
- # Use HF Hub's native API instead of manual requests
- repo_info = self.api.repo_info(model_id, revision=revision)
-
- files = []
- if repo_info.siblings:
- for sibling in repo_info.siblings:
- if sibling.rfilename: # Only include actual files
- files.append(
- {
- "path": sibling.rfilename,
- "size": getattr(sibling, "size", 0) or 0,
- "url": f"https://huggingface.co/{model_id}/resolve/{revision}/{sibling.rfilename}",
- }
- )
-
- return files
-
- except Exception as e:
- self.logger.warning(f"Could not fetch model file list for {model_id}: {e}")
- return []
-
- def should_accelerate_model(self, model_id: str) -> bool:
- """
- Determine if model should be pre-cached using the configured strategy.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- True if model should be pre-cached
- """
- return self.strategy.should_accelerate(model_id)
-
- def accelerate_model_download(
- self, model_id: str, revision: str = "main"
- ) -> FunctionResponse:
- """
- Pre-download HuggingFace model using the configured download strategy.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- FunctionResponse with download results
- """
- return self.strategy.download_model(model_id, revision)
-
- def is_model_cached(self, model_id: str, revision: str = "main") -> bool:
- """
- Check if model is already cached using the configured strategy.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- True if model appears to be cached
- """
- return self.strategy.is_model_cached(model_id, revision)
-
- def get_cache_info(self, model_id: str) -> Dict[str, Any]:
- """
- Get cache information for a model using the configured strategy.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- Dictionary with cache information
- """
- return self.strategy.get_cache_info(model_id)
-
- def clear_model_cache(self, model_id: str) -> FunctionResponse:
- """
- Clear cache for a specific model using the configured strategy.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- FunctionResponse with clearing result
- """
- return self.strategy.clear_model_cache(model_id)
-
- def get_strategy_info(self) -> Dict[str, Any]:
- """
- Get information about the current download strategy.
-
- Returns:
- Dictionary with strategy information
- """
- strategy_info = HFStrategyFactory.get_strategy_info()
- strategy_info["strategy_instance"] = type(self.strategy).__name__
- return strategy_info
-
- def set_strategy(self, strategy: str) -> None:
- """
- Change the download strategy (creates new strategy instance).
-
- Args:
- strategy: Strategy name ("tetra" or "native")
- """
- HFStrategyFactory.set_strategy(strategy)
- self.strategy = HFStrategyFactory.create_strategy(self.workspace_manager)
- self.logger.info(f"Switched to {strategy} download strategy")
diff --git a/src/huggingface_cache.py b/src/huggingface_cache.py
new file mode 100644
index 0000000..df8b91a
--- /dev/null
+++ b/src/huggingface_cache.py
@@ -0,0 +1,118 @@
+"""
+HuggingFace model download caching.
+
+This module provides cache-ahead downloads for HuggingFace models and datasets.
+"""
+
+import os
+import asyncio
+import logging
+
+from huggingface_hub import snapshot_download, scan_cache_dir
+from huggingface_hub.errors import CacheNotFound
+from remote_execution import FunctionResponse
+
+
+class HuggingFaceCacheAhead:
+ """Cache-ahead downloads for HuggingFace models and files."""
+
+ def __init__(self):
+ self.logger = logging.getLogger(__name__)
+
+ async def cache_model_download_async(
+ self, model_id: str, revision: str = "main"
+ ) -> FunctionResponse:
+ """
+ Async wrapper for pre-downloading HuggingFace models.
+
+ Args:
+ model_id: HuggingFace model identifier
+ revision: Model revision/branch
+
+ Returns:
+ FunctionResponse with download results
+ """
+ return await asyncio.to_thread(self.cache_model_download, model_id, revision)
+
+ def cache_model_download(
+ self, model_id: str, revision: str = "main"
+ ) -> FunctionResponse:
+ """
+ Pre-download HuggingFace model using HF Hub's native caching.
+
+ This method downloads the complete model snapshot to HF's standard cache
+ location. HF Hub automatically uses hf_transfer/hf_xet acceleration when
+ HF_HUB_ENABLE_HF_TRANSFER=1 is set in the environment.
+
+ Args:
+ model_id: HuggingFace model identifier
+ revision: Model revision/branch
+
+ Returns:
+ FunctionResponse with download results
+ """
+ try:
+ # Check if model is already cached
+ cache_hit = self._is_model_cached(model_id, revision)
+ if cache_hit:
+ self.logger.info(f"Model {model_id} already cached, skipping download")
+ return FunctionResponse(
+ success=True,
+ stdout=f"Model {model_id} already cached (cache hit)",
+ )
+
+ # Get HF authentication token if available
+ hf_token = os.environ.get("HF_TOKEN")
+
+ # Use HF Hub's native snapshot download with acceleration
+ snapshot_path = snapshot_download(
+ repo_id=model_id,
+ revision=revision,
+ token=hf_token,
+ # HF automatically uses HF_HOME/HF_HUB_CACHE from environment
+ # and applies hf_transfer acceleration when available
+ )
+
+ success_message = f"Successfully cached model {model_id} to {snapshot_path}"
+ self.logger.info(success_message)
+
+ return FunctionResponse(
+ success=True,
+ stdout=success_message,
+ )
+
+ except Exception as e:
+ return FunctionResponse(
+ success=False,
+ error=f"Failed to cache-ahead model {model_id}: {str(e)}",
+ )
+
+ def _is_model_cached(self, model_id: str, revision: str = "main") -> bool:
+ """
+ Check if a model is already cached locally.
+
+ Args:
+ model_id: HuggingFace model identifier
+ revision: Model revision/branch
+
+ Returns:
+ True if model is cached, False otherwise
+ """
+ try:
+ cache_info = scan_cache_dir()
+ for repo in cache_info.repos:
+ if repo.repo_id == model_id:
+ # Check if the specific revision is cached
+ for rev in repo.revisions:
+ if rev.commit_hash == revision or revision == "main":
+ return True
+ return False
+ except CacheNotFound:
+ # Cache directory doesn't exist yet - this is expected on first use
+ self.logger.debug(
+ f"Cache directory not found for {model_id}, will be created on download"
+ )
+ return False
+ except Exception as e:
+ self.logger.debug(f"Cache check failed for {model_id}: {e}")
+ return False
diff --git a/src/logger.py b/src/logger.py
index 51c4118..8f0e72f 100644
--- a/src/logger.py
+++ b/src/logger.py
@@ -9,9 +9,6 @@
import sys
from typing import Union, Optional
-# Application logger namespace
-APP_LOGGER_NAME = "tetra"
-
def get_log_level() -> int:
"""Get log level from environment variable, defaulting to INFO."""
@@ -34,7 +31,7 @@ def setup_logging(
) -> None:
"""
Setup logging configuration for worker-tetra.
- Only shows DEBUG logs from worker_tetra namespace when LOG_LEVEL=DEBUG.
+ Only shows DEBUG logs from tetra namespace when LOG_LEVEL=DEBUG.
Args:
level: Log level (defaults to LOG_LEVEL env var or INFO)
diff --git a/src/remote_executor.py b/src/remote_executor.py
index bf18ea9..3d0e5e1 100644
--- a/src/remote_executor.py
+++ b/src/remote_executor.py
@@ -1,12 +1,13 @@
import logging
import asyncio
from typing import List, Any
+from huggingface_cache import HuggingFaceCacheAhead
from remote_execution import FunctionRequest, FunctionResponse, RemoteExecutorStub
-from workspace_manager import WorkspaceManager
from dependency_installer import DependencyInstaller
from function_executor import FunctionExecutor
from class_executor import ClassExecutor
from log_streamer import start_log_streaming, stop_log_streaming, get_streamed_logs
+from constants import NAMESPACE
class RemoteExecutor(RemoteExecutorStub):
@@ -17,13 +18,13 @@ class RemoteExecutor(RemoteExecutorStub):
def __init__(self):
super().__init__()
- self.logger = logging.getLogger(f"worker_tetra.{__name__.split('.')[-1]}")
+ self.logger = logging.getLogger(f"{NAMESPACE}.{__name__.split('.')[-1]}")
# Initialize components using composition
- self.workspace_manager = WorkspaceManager()
- self.dependency_installer = DependencyInstaller(self.workspace_manager)
- self.function_executor = FunctionExecutor(self.workspace_manager)
- self.class_executor = ClassExecutor(self.workspace_manager)
+ self.dependency_installer = DependencyInstaller()
+ self.function_executor = FunctionExecutor()
+ self.class_executor = ClassExecutor()
+ self.hf_cache = HuggingFaceCacheAhead()
async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse:
"""
@@ -50,22 +51,7 @@ async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse:
)
try:
- # Initialize workspace if using volume
- if self.workspace_manager.has_runpod_volume:
- workspace_init = self.workspace_manager.initialize_workspace()
- if not workspace_init.success:
- # Add any buffered logs to the failed response
- logs = get_streamed_logs(clear_buffer=True)
- if logs:
- if workspace_init.stdout:
- workspace_init.stdout += "\n" + logs
- else:
- workspace_init.stdout = logs
- return workspace_init
- if workspace_init.stdout:
- self.logger.info(workspace_init.stdout)
-
- # Install dependencies and cache models
+ # Install dependencies
if request.accelerate_downloads:
# Run installations in parallel when acceleration is enabled
dep_result = await self._install_dependencies_parallel(request)
@@ -100,9 +86,6 @@ async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse:
else:
result = self.function_executor.execute(request)
- # Add acceleration summary to the result
- self._log_acceleration_summary(request, result)
-
# Add all captured system logs to the result
system_logs = get_streamed_logs(clear_buffer=True)
if system_logs:
@@ -117,68 +100,6 @@ async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse:
# Always stop log streaming to clean up
stop_log_streaming()
- def _log_acceleration_summary(
- self, request: FunctionRequest, result: FunctionResponse
- ):
- """Log acceleration impact summary for performance visibility."""
- if not hasattr(self.dependency_installer, "download_accelerator"):
- return
-
- acceleration_enabled = request.accelerate_downloads
- has_volume = self.workspace_manager.has_runpod_volume
- hf_transfer_available = self.dependency_installer.download_accelerator.hf_transfer_downloader.hf_transfer_available
- nala_available = self.dependency_installer._check_nala_available()
-
- # Build summary message
- summary_parts = []
-
- if acceleration_enabled:
- summary_parts.append("✓ Download acceleration ENABLED")
-
- if has_volume:
- summary_parts.append(
- f"✓ Volume workspace: {self.workspace_manager.workspace_path}"
- )
- summary_parts.append("✓ Persistent caching enabled")
- else:
- summary_parts.append("ℹ No persistent volume - using temporary cache")
-
- # System package acceleration status
- if request.system_dependencies:
- large_system_packages = (
- self.dependency_installer._identify_large_system_packages(
- request.system_dependencies
- )
- )
- if large_system_packages and nala_available:
- summary_parts.append(
- f"✓ System packages with nala: {len(large_system_packages)}"
- )
- elif request.system_dependencies:
- summary_parts.append("→ System packages using standard apt-get")
-
- if request.hf_models_to_cache:
- summary_parts.append(
- f"✓ HF models pre-cached: {len(request.hf_models_to_cache)}"
- )
-
- elif acceleration_enabled and not (hf_transfer_available or nala_available):
- summary_parts.append(
- "⚠ Download acceleration REQUESTED but no accelerators available"
- )
- summary_parts.append("→ Using standard downloads")
-
- elif not acceleration_enabled:
- summary_parts.append("- Download acceleration DISABLED")
- summary_parts.append("→ Using standard downloads")
-
- # Log the summary
- if summary_parts:
- self.logger.debug("=== DOWNLOAD ACCELERATION SUMMARY ===")
- for part in summary_parts:
- self.logger.debug(part)
- self.logger.debug("=====================================")
-
async def _install_dependencies_parallel(
self, request: FunctionRequest
) -> FunctionResponse:
@@ -210,10 +131,10 @@ async def _install_dependencies_parallel(
tasks.append(task)
task_names.append("python_dependencies")
- # Add HF model caching tasks
+ # Add HF model cache-ahead tasks
if request.hf_models_to_cache:
for model_id in request.hf_models_to_cache:
- task = self.workspace_manager.accelerate_model_download_async(model_id)
+ task = self.hf_cache.cache_model_download_async(model_id)
tasks.append(task)
task_names.append(f"hf_model_{model_id}")
@@ -251,13 +172,11 @@ async def _install_dependencies_sequential(
return sys_installed
self.logger.info(sys_installed.stdout)
- # Pre-cache HuggingFace models if requested (should not happen when acceleration disabled)
+ # 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"Pre-caching HuggingFace model: {model_id}")
- cache_result = self.workspace_manager.accelerate_model_download(
- model_id
- )
+ 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}"
@@ -330,16 +249,10 @@ def _process_parallel_results(
stdout=f"Parallel installation: {success_count}/{len(results)} tasks succeeded\n"
+ "\n".join(stdout_parts),
)
- else:
- # All tasks succeeded
- return FunctionResponse(
- success=True,
- stdout=f"Parallel installation: {success_count}/{len(results)} tasks completed successfully\n"
- + "\n".join(stdout_parts),
- )
- # All tasks succeeded
- return FunctionResponse(
- success=True,
- stdout=f"Parallel installation: {success_count}/{len(results)} tasks completed successfully\n"
- + "\n".join(stdout_parts),
- )
+
+ # All tasks succeeded
+ return FunctionResponse(
+ success=True,
+ stdout=f"Parallel installation: {success_count}/{len(results)} tasks completed successfully\n"
+ + "\n".join(stdout_parts),
+ )
diff --git a/src/test-handler.sh b/src/test-handler.sh
index e7c9ee3..1304a8d 100755
--- a/src/test-handler.sh
+++ b/src/test-handler.sh
@@ -6,7 +6,7 @@ failed_tests=""
test_count=0
passed_count=0
-for test_file in test_*.json; do
+for test_file in tests/test_*.json; do
if [ ! -f "$test_file" ]; then
echo "No test_*.json files found"
exit 1
diff --git a/src/test_class_custom_method.json b/src/tests/test_class_custom_method.json
similarity index 100%
rename from src/test_class_custom_method.json
rename to src/tests/test_class_custom_method.json
diff --git a/src/test_class_input.json b/src/tests/test_class_input.json
similarity index 100%
rename from src/test_class_input.json
rename to src/tests/test_class_input.json
diff --git a/src/test_class_persistence.json b/src/tests/test_class_persistence.json
similarity index 100%
rename from src/test_class_persistence.json
rename to src/tests/test_class_persistence.json
diff --git a/src/test_dependencies.json b/src/tests/test_dependencies.json
similarity index 100%
rename from src/test_dependencies.json
rename to src/tests/test_dependencies.json
diff --git a/src/test_function_args.json b/src/tests/test_function_args.json
similarity index 100%
rename from src/test_function_args.json
rename to src/tests/test_function_args.json
diff --git a/src/test_hf_accelerated_input.json b/src/tests/test_hf_accelerated_input.json
similarity index 100%
rename from src/test_hf_accelerated_input.json
rename to src/tests/test_hf_accelerated_input.json
diff --git a/src/test_input.json b/src/tests/test_input.json
similarity index 100%
rename from src/test_input.json
rename to src/tests/test_input.json
diff --git a/src/test_installed_packages.json b/src/tests/test_installed_packages.json
similarity index 100%
rename from src/test_installed_packages.json
rename to src/tests/test_installed_packages.json
diff --git a/src/test_log_streaming.json b/src/tests/test_log_streaming.json
similarity index 100%
rename from src/test_log_streaming.json
rename to src/tests/test_log_streaming.json
diff --git a/src/test_mixed_dependencies.json b/src/tests/test_mixed_dependencies.json
similarity index 100%
rename from src/test_mixed_dependencies.json
rename to src/tests/test_mixed_dependencies.json
diff --git a/src/test_pip_package_access.json b/src/tests/test_pip_package_access.json
similarity index 100%
rename from src/test_pip_package_access.json
rename to src/tests/test_pip_package_access.json
diff --git a/src/test_runpod_import.json b/src/tests/test_runpod_import.json
similarity index 100%
rename from src/test_runpod_import.json
rename to src/tests/test_runpod_import.json
diff --git a/src/test_system_dependencies.json b/src/tests/test_system_dependencies.json
similarity index 100%
rename from src/test_system_dependencies.json
rename to src/tests/test_system_dependencies.json
diff --git a/src/test_uv_no_acceleration.json b/src/tests/test_uv_no_acceleration.json
similarity index 100%
rename from src/test_uv_no_acceleration.json
rename to src/tests/test_uv_no_acceleration.json
diff --git a/src/workspace_manager.py b/src/workspace_manager.py
deleted file mode 100644
index a3db7fb..0000000
--- a/src/workspace_manager.py
+++ /dev/null
@@ -1,423 +0,0 @@
-import os
-import fcntl
-import time
-import logging
-import asyncio
-from typing import Optional, TYPE_CHECKING, Any, Dict
-
-if TYPE_CHECKING:
- from huggingface_accelerator import HuggingFaceAccelerator
-
-from remote_execution import FunctionResponse
-from subprocess_utils import run_logged_subprocess
-from constants import (
- RUNPOD_VOLUME_PATH,
- DEFAULT_WORKSPACE_PATH,
- VENV_DIR_NAME,
- UV_CACHE_DIR_NAME,
- HF_CACHE_DIR_NAME,
- WORKSPACE_LOCK_FILE,
- RUNTIMES_DIR_NAME,
-)
-
-
-class WorkspaceManager:
- """Manages RunPod volume workspace initialization and configuration."""
-
- venv_path: Optional[str]
- cache_path: Optional[str]
- hf_cache_path: Optional[str]
-
- def __init__(self) -> None:
- self.logger = logging.getLogger(f"worker_tetra.{__name__.split('.')[-1]}")
- self.has_runpod_volume = os.path.exists(RUNPOD_VOLUME_PATH)
- self.endpoint_id = os.environ.get("RUNPOD_ENDPOINT_ID", "default")
-
- # Set up workspace paths
- if self.has_runpod_volume:
- # Endpoint-specific workspace: /runpod-volume/runtimes/{endpoint_id}
- self.workspace_path = os.path.join(
- RUNPOD_VOLUME_PATH, RUNTIMES_DIR_NAME, self.endpoint_id
- )
- self.venv_path = os.path.join(self.workspace_path, VENV_DIR_NAME)
- # Shared caches at volume root for all endpoints
- self.cache_path = os.path.join(RUNPOD_VOLUME_PATH, UV_CACHE_DIR_NAME)
- self.hf_cache_path = os.path.join(RUNPOD_VOLUME_PATH, HF_CACHE_DIR_NAME)
- else:
- # Fallback to container workspace
- self.workspace_path = DEFAULT_WORKSPACE_PATH
- self.venv_path = None
- self.cache_path = None
- self.hf_cache_path = None
-
- # Initialize HuggingFace accelerator after paths are set
- self._hf_accelerator: Optional[HuggingFaceAccelerator] = None
-
- if self.has_runpod_volume:
- self._configure_uv_cache()
- self._configure_huggingface_cache()
- self._configure_volume_environment()
-
- def _configure_uv_cache(self):
- """Configure uv to use the shared volume cache."""
- if self.cache_path:
- os.environ["UV_CACHE_DIR"] = self.cache_path
-
- def _configure_huggingface_cache(self):
- """Configure Hugging Face to use the shared volume cache."""
- if self.hf_cache_path:
- # Ensure HF cache directory exists
- os.makedirs(self.hf_cache_path, exist_ok=True)
-
- # Set main HF cache directory - HF will automatically create subdirectories
- os.environ["HF_HOME"] = self.hf_cache_path
-
- # HF automatically creates and manages these subdirectories:
- # - hub/ (for model downloads and cache)
- # - transformers/ (legacy, but still used by some components)
- # - datasets/ (for HF datasets)
- # Let HF handle the hierarchy instead of forcing specific paths
-
- def _configure_volume_environment(self):
- """Configure environment variables for volume usage."""
- if self.venv_path:
- os.environ["VIRTUAL_ENV"] = self.venv_path
- venv_bin = os.path.join(self.venv_path, "bin")
- current_path = os.environ.get("PATH", "")
- os.environ["PATH"] = f"{venv_bin}:{current_path}"
-
- # Additional environment variables that may help with subprocess execution
- venv_lib_python = os.path.join(
- self.venv_path, "lib", "python*", "site-packages"
- )
- import glob
-
- site_packages_dirs = glob.glob(venv_lib_python)
- if site_packages_dirs:
- # Set PYTHONPATH to include volume site-packages
- current_pythonpath = os.environ.get("PYTHONPATH", "")
- new_pythonpath = ":".join(site_packages_dirs)
- if current_pythonpath:
- os.environ["PYTHONPATH"] = f"{new_pythonpath}:{current_pythonpath}"
- else:
- os.environ["PYTHONPATH"] = new_pythonpath
- self.logger.info(
- f"Set PYTHONPATH to include volume packages: {new_pythonpath}"
- )
-
- def initialize_workspace(self, timeout: int = 30) -> FunctionResponse:
- """
- Initialize the RunPod volume workspace with virtual environment.
-
- Args:
- timeout: Maximum time to wait for workspace initialization
-
- Returns:
- FunctionResponse: Success or failure of initialization
- """
- if not self.has_runpod_volume:
- return FunctionResponse(
- success=True, stdout="No volume available, using container workspace"
- )
-
- # Check if workspace is already initialized and functional
- if self.venv_path and os.path.exists(self.venv_path):
- validation_result = self._validate_virtual_environment()
- if validation_result.success:
- return FunctionResponse(
- success=True, stdout="Workspace already initialized"
- )
- else:
- # Virtual environment exists but is broken, recreate it
- self.logger.warning(
- f"Virtual environment validation failed: {validation_result.error}"
- )
- self.logger.info("Recreating virtual environment...")
- self._remove_broken_virtual_environment()
-
- # Use file-based locking for concurrent initialization
- lock_file = os.path.join(self.workspace_path, WORKSPACE_LOCK_FILE)
-
- try:
- # Ensure workspace directory exists
- os.makedirs(self.workspace_path, exist_ok=True)
-
- with open(lock_file, "w") as lock:
- try:
- # Try to acquire exclusive lock with timeout
- fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
- except (BlockingIOError, OSError):
- # Lock not available, wait for initialization by another worker
- start_time = time.time()
- while time.time() - start_time < timeout:
- if self.venv_path and os.path.exists(self.venv_path):
- validation_result = self._validate_virtual_environment()
- if validation_result.success:
- return FunctionResponse(
- success=True,
- stdout="Workspace initialized by another worker",
- )
- time.sleep(0.5)
-
- return FunctionResponse(
- success=False, error="Workspace initialization timeout"
- )
-
- # We have the lock, initialize the workspace
- return self._create_virtual_environment()
-
- except Exception as e:
- return FunctionResponse(
- success=False, error=f"Failed to initialize workspace: {str(e)}"
- )
- finally:
- # Clean up lock file
- try:
- if os.path.exists(lock_file):
- os.remove(lock_file)
- except OSError:
- pass
-
- def _create_virtual_environment(self) -> FunctionResponse:
- """Create virtual environment in the volume."""
- if not self.venv_path:
- return FunctionResponse(
- success=False, error="Virtual environment path not configured"
- )
-
- result = run_logged_subprocess(
- command=["uv", "venv", self.venv_path],
- logger=self.logger,
- operation_name="Creating virtual environment",
- )
-
- if not result.success:
- return FunctionResponse(
- success=False,
- error="Failed to create virtual environment",
- stdout=result.error,
- )
- else:
- # Create symlink from /app/.venv to volume venv for libraries that hardcode /app/.venv
- self._create_app_venv_symlink()
- return FunctionResponse(success=True, stdout=result.stdout)
-
- def _create_app_venv_symlink(self):
- """
- Create symlink from /app/.venv to volume virtual environment.
-
- This ensures libraries that hardcode /app/.venv (like vLLM) use the volume's venv
- instead of the container's default venv.
- """
- if not self.venv_path:
- return
-
- app_venv_path = "/app/.venv"
-
- try:
- # Remove existing /app/.venv if it exists (file, dir, or broken symlink)
- if os.path.exists(app_venv_path) or os.path.islink(app_venv_path):
- if os.path.isdir(app_venv_path) and not os.path.islink(app_venv_path):
- # It's a real directory, remove recursively
- import shutil
-
- shutil.rmtree(app_venv_path)
- else:
- # It's a file or symlink, remove it
- os.remove(app_venv_path)
-
- # Create symlink to volume venv
- os.symlink(self.venv_path, app_venv_path)
- self.logger.info(f"Created symlink: {app_venv_path} -> {self.venv_path}")
-
- except Exception as e:
- # Log error but don't fail workspace initialization
- self.logger.warning(f"Failed to create /app/.venv symlink: {str(e)}")
-
- def _remove_app_venv_symlink(self):
- """
- Remove /app/.venv symlink if it points to our virtual environment.
- """
- app_venv_path = "/app/.venv"
-
- try:
- if os.path.islink(app_venv_path):
- # Check if it points to our venv
- link_target = os.readlink(app_venv_path)
- if link_target == self.venv_path:
- os.remove(app_venv_path)
- self.logger.info(f"Removed symlink: {app_venv_path}")
- except Exception as e:
- self.logger.warning(f"Failed to remove /app/.venv symlink: {str(e)}")
-
- def change_to_workspace(self) -> Optional[str]:
- """
- Change to workspace directory and return original working directory.
-
- Returns:
- Original working directory path if changed, None otherwise
- """
- if self.has_runpod_volume:
- original_cwd = os.getcwd()
- os.chdir(self.workspace_path)
- return original_cwd
- return None
-
- def setup_python_path(self):
- """Add virtual environment packages to Python path if available."""
- if self.has_runpod_volume and self.venv_path and os.path.exists(self.venv_path):
- # Validate venv before using it
- validation_result = self._validate_virtual_environment()
- if not validation_result.success:
- self.logger.warning(
- f"Virtual environment is invalid: {validation_result.error}"
- )
- return
- import glob
- import sys
-
- site_packages = glob.glob(
- os.path.join(self.venv_path, "lib", "python*", "site-packages")
- )
- for site_package_path in site_packages:
- if site_package_path not in sys.path:
- sys.path.insert(0, site_package_path)
-
- def _validate_virtual_environment(self) -> FunctionResponse:
- """
- Validate that the virtual environment is functional.
-
- Returns:
- FunctionResponse indicating if the venv is valid
- """
- if not self.venv_path or not os.path.exists(self.venv_path):
- return FunctionResponse(
- success=False, error="Virtual environment does not exist"
- )
-
- python_exe = os.path.join(self.venv_path, "bin", "python3")
-
- # Check if Python executable exists and is not a broken symlink
- if not os.path.exists(python_exe):
- return FunctionResponse(
- success=False, error=f"Python executable not found at {python_exe}"
- )
-
- # Check if it's a broken symlink (need to resolve the full path)
- if os.path.islink(python_exe):
- try:
- # Use os.path.realpath to resolve the full symlink chain
- resolved_path = os.path.realpath(python_exe)
- if not os.path.exists(resolved_path):
- return FunctionResponse(
- success=False,
- error=f"Broken symlink at {python_exe}, underlying Python interpreter removed",
- )
- except (OSError, ValueError) as e:
- return FunctionResponse(
- success=False,
- error=f"Error resolving symlink at {python_exe}: {str(e)}",
- )
-
- # Try to execute a simple Python command to verify functionality
- result = run_logged_subprocess(
- command=[python_exe, "-c", "import sys; print(sys.version)"],
- logger=self.logger,
- operation_name="Validating Python interpreter",
- timeout=10,
- )
-
- if not result.success:
- return FunctionResponse(
- success=False,
- error=f"Python interpreter failed to execute: {result.error}",
- )
-
- return FunctionResponse(
- success=True, stdout="Virtual environment is functional"
- )
-
- def _remove_broken_virtual_environment(self):
- """Remove broken virtual environment directory and associated symlink."""
- if self.venv_path and os.path.exists(self.venv_path):
- import shutil
-
- try:
- shutil.rmtree(self.venv_path)
- self.logger.info(
- f"Removed broken virtual environment at {self.venv_path}"
- )
-
- # Also remove the /app/.venv symlink if it points to this venv
- self._remove_app_venv_symlink()
-
- except Exception as e:
- self.logger.error(
- f"Error removing broken virtual environment: {str(e)}"
- )
-
- @property
- def hf_accelerator(self) -> "HuggingFaceAccelerator":
- """Lazy-loaded HuggingFace accelerator."""
- if self._hf_accelerator is None:
- from huggingface_accelerator import HuggingFaceAccelerator
-
- self._hf_accelerator = HuggingFaceAccelerator(self)
- return self._hf_accelerator
-
- def accelerate_model_download(
- self, model_id: str, revision: str = "main"
- ) -> FunctionResponse:
- """
- Pre-download HuggingFace model using acceleration if beneficial.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- FunctionResponse with download result
- """
- return self.hf_accelerator.accelerate_model_download(model_id, revision)
-
- async def accelerate_model_download_async(
- self, model_id: str, revision: str = "main"
- ) -> FunctionResponse:
- """
- Async wrapper for HuggingFace model download acceleration.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- FunctionResponse with download result
- """
- return await asyncio.to_thread(
- self.accelerate_model_download, model_id, revision
- )
-
- def is_model_cached(self, model_id: str, revision: str = "main") -> bool:
- """
- Check if a HuggingFace model is cached.
-
- Args:
- model_id: HuggingFace model identifier
- revision: Model revision/branch
-
- Returns:
- True if model is cached
- """
- return self.hf_accelerator.is_model_cached(model_id, revision)
-
- def get_model_cache_info(self, model_id: str) -> Dict[str, Any]:
- """
- Get cache information for a HuggingFace model.
-
- Args:
- model_id: HuggingFace model identifier
-
- Returns:
- Dictionary with cache information
- """
- return self.hf_accelerator.get_cache_info(model_id)
diff --git a/tests/integration/test_download_acceleration_integration.py b/tests/integration/test_download_acceleration_integration.py
deleted file mode 100644
index 037d0ac..0000000
--- a/tests/integration/test_download_acceleration_integration.py
+++ /dev/null
@@ -1,363 +0,0 @@
-"""
-Integration tests for download acceleration functionality using hf_transfer.
-"""
-
-import pytest
-import tempfile
-import shutil
-from pathlib import Path
-from unittest.mock import Mock, patch, AsyncMock
-
-from src.download_accelerator import (
- DownloadAccelerator,
- HfTransferDownloader,
-)
-from src.huggingface_accelerator import HuggingFaceAccelerator
-from src.dependency_installer import DependencyInstaller
-from src.workspace_manager import WorkspaceManager
-from src.remote_executor import RemoteExecutor
-from src.remote_execution import FunctionRequest
-
-
-class TestDownloadAccelerationIntegration:
- """Integration tests for download acceleration components."""
-
- def setup_method(self):
- """Set up test environment."""
- self.temp_dir = Path(tempfile.mkdtemp())
- self.mock_workspace_manager = Mock(spec=WorkspaceManager)
- self.mock_workspace_manager.has_runpod_volume = True
- self.mock_workspace_manager.hf_cache_path = str(self.temp_dir / ".hf-cache")
- self.mock_workspace_manager.workspace_path = str(self.temp_dir)
- self.mock_workspace_manager.venv_path = str(self.temp_dir / ".venv")
-
- def teardown_method(self):
- """Clean up test environment."""
- shutil.rmtree(self.temp_dir, ignore_errors=True)
-
- @patch("src.download_accelerator.HF_TRANSFER_ENABLED", True)
- def test_hf_transfer_availability_detection(self):
- """Test detection of hf_transfer availability."""
- with patch("importlib.util.find_spec") as mock_find_spec:
- # Test when hf_transfer is available
- mock_find_spec.return_value = Mock() # Not None means available
- downloader = HfTransferDownloader()
- assert downloader.hf_transfer_available is True
-
- # Test when hf_transfer is not available
- mock_find_spec.return_value = None # None means not available
- downloader = HfTransferDownloader()
- assert downloader.hf_transfer_available is False
-
- def test_download_accelerator_decision_logic(self):
- """Test when acceleration should be used."""
- accelerator = DownloadAccelerator(self.mock_workspace_manager)
-
- # Mock hf_transfer as available
- accelerator.hf_transfer_downloader.hf_transfer_available = True
-
- # Should accelerate large HuggingFace files
- assert (
- accelerator.should_accelerate_download(
- "https://huggingface.co/model/resolve/main/large.bin", 50.0
- )
- is True
- )
-
- # Should accelerate HuggingFace URLs regardless of size
- assert (
- accelerator.should_accelerate_download(
- "https://huggingface.co/model/resolve/main/file", 5.0
- )
- is True
- )
-
- # Should not accelerate non-HF files
- assert (
- accelerator.should_accelerate_download("http://example.com/large.bin", 50.0)
- is False
- )
- assert (
- accelerator.should_accelerate_download("http://example.com/small.txt", 1.0)
- is False
- )
-
- @patch("src.huggingface_accelerator.HfApi.repo_info")
- def test_hf_model_file_fetching(self, mock_repo_info):
- """Test fetching HuggingFace model file information."""
- # Mock successful API response using HF Hub's native API
- from unittest.mock import Mock
-
- mock_repo_info_obj = Mock()
- mock_repo_info_obj.siblings = [
- Mock(rfilename="pytorch_model.bin", size=500 * 1024 * 1024), # 500MB
- Mock(rfilename="config.json", size=1024), # 1KB
- ]
- mock_repo_info.return_value = mock_repo_info_obj
-
- accelerator = HuggingFaceAccelerator(self.mock_workspace_manager)
- files = accelerator.get_model_files("gpt2")
-
- assert len(files) == 2
- assert files[0]["path"] == "pytorch_model.bin"
- assert files[0]["size"] == 500 * 1024 * 1024
- assert "huggingface.co/gpt2/resolve/main/pytorch_model.bin" in files[0]["url"]
-
- def test_hf_model_acceleration_decision(self):
- """Test when HuggingFace models should be pre-cached."""
- accelerator = HuggingFaceAccelerator(self.mock_workspace_manager)
-
- # Should pre-cache known large models (HF handles acceleration automatically)
- assert accelerator.should_accelerate_model("gpt2") is True
- assert accelerator.should_accelerate_model("bert-base-uncased") is True
- assert accelerator.should_accelerate_model("microsoft/DialoGPT-medium") is True
- assert accelerator.should_accelerate_model("stable-diffusion-v1-5") is True
-
- # Should not pre-cache unknown/small models
- assert accelerator.should_accelerate_model("unknown/tiny-model") is False
-
- @patch("src.workspace_manager.WorkspaceManager.__init__")
- def test_remote_executor_with_acceleration(self, mock_workspace_init):
- """Test RemoteExecutor integration with download acceleration."""
- # Mock workspace manager
- mock_workspace_init.return_value = None
-
- executor = RemoteExecutor()
- executor.workspace_manager = self.mock_workspace_manager
- executor.workspace_manager.has_runpod_volume = True
- executor.workspace_manager.initialize_workspace = Mock(
- return_value=Mock(success=True)
- )
- executor.workspace_manager.accelerate_model_download = Mock(
- return_value=Mock(success=True, stdout="Model cached successfully")
- )
-
- # Mock dependency installer
- executor.dependency_installer = Mock()
- executor.dependency_installer.install_system_dependencies = Mock(
- return_value=Mock(success=True, stdout="System deps installed")
- )
- executor.dependency_installer.install_dependencies_async = AsyncMock(
- return_value=Mock(success=True, stdout="Python deps installed")
- )
- executor.workspace_manager.accelerate_model_download_async = AsyncMock(
- return_value=Mock(success=True, stdout="Model cached")
- )
- executor.dependency_installer._identify_large_packages = Mock(
- return_value=["torch", "transformers"]
- )
- executor.dependency_installer.download_accelerator = Mock()
- executor.dependency_installer.download_accelerator.hf_transfer_downloader = (
- Mock()
- )
- executor.dependency_installer.download_accelerator.hf_transfer_downloader.hf_transfer_available = True
-
- # Mock executors
- executor.function_executor = Mock()
- executor.function_executor.execute = Mock(
- return_value=Mock(success=True, result="Function executed")
- )
-
- # Create request with acceleration enabled
- request = FunctionRequest(
- function_name="test_function",
- function_code="def test_function(): return 'test'",
- dependencies=["torch", "transformers"],
- accelerate_downloads=True,
- hf_models_to_cache=["gpt2", "bert-base-uncased"],
- )
-
- # Execute function
- import asyncio
-
- asyncio.run(executor.ExecuteFunction(request))
-
- # Verify model caching was attempted (async method is called)
- assert (
- executor.workspace_manager.accelerate_model_download_async.call_count == 2
- )
- executor.workspace_manager.accelerate_model_download_async.assert_any_call(
- "gpt2"
- )
- executor.workspace_manager.accelerate_model_download_async.assert_any_call(
- "bert-base-uncased"
- )
-
- # Verify dependencies were installed with acceleration enabled (async method)
- executor.dependency_installer.install_dependencies_async.assert_called_once_with(
- ["torch", "transformers"], True
- )
-
- @patch.dict("os.environ", {"HF_TOKEN": "test_token"})
- def test_hf_token_authentication(self):
- """Test that HF_TOKEN is properly used for authentication."""
- downloader = HfTransferDownloader()
- # Test that downloader correctly checks for availability
- # Since hf_transfer may not be installed, this will be False
- # and that's expected behavior
- assert isinstance(downloader.hf_transfer_available, bool)
-
- def test_strategy_selection_logic(self):
- """Test the download strategy selection logic."""
- accelerator = DownloadAccelerator(self.mock_workspace_manager)
- accelerator.hf_transfer_downloader.hf_transfer_available = True
-
- # Test file caching detection
- non_existent_file = str(self.temp_dir / "non_existent.bin")
- existing_file = str(self.temp_dir / "existing.bin")
-
- # Create existing file
- Path(existing_file).write_bytes(b"existing data")
-
- assert accelerator.is_file_cached(non_existent_file) is False
- assert accelerator.is_file_cached(existing_file) is True
-
- def test_fallback_behavior_without_accelerators(self):
- """Test graceful fallback when accelerators are not available."""
- accelerator = DownloadAccelerator(self.mock_workspace_manager)
- accelerator.hf_transfer_downloader.hf_transfer_available = False
-
- # With new logic, when acceleration is not available, we defer to HF native handling
- result = accelerator.download_with_fallback(
- "https://huggingface.co/gpt2/resolve/main/file.bin",
- str(self.temp_dir / "file.bin"),
- )
-
- # Should return failure and defer to HF native handling
- assert result.success is False
- assert "defer to HF native handling" in result.error
-
- @patch("src.dependency_installer.run_logged_subprocess")
- def test_dependency_installation_without_acceleration(self, mock_subprocess):
- """Test that packages install normally without aria2c acceleration."""
- # Mock successful installation
- from remote_execution import FunctionResponse
-
- mock_subprocess.return_value = FunctionResponse(
- success=True, stdout="Installed successfully"
- )
-
- installer = DependencyInstaller(self.mock_workspace_manager)
-
- # Install packages
- packages = ["torch==2.0.0", "transformers>=4.20.0"]
- result = installer.install_dependencies(packages)
-
- assert result.success is True
-
- # Verify the installation was called
- mock_subprocess.assert_called_once()
-
- @patch("src.hf_downloader_tetra.DownloadAccelerator")
- def test_model_cache_management(self, mock_download_accelerator):
- """Test model cache information and management using tetra strategy."""
- accelerator = HuggingFaceAccelerator(self.mock_workspace_manager)
-
- # Test cache info for non-existent model
- cache_info = accelerator.get_cache_info("non-existent-model")
- assert cache_info["cached"] is False
- assert cache_info["cache_size_mb"] == 0
- assert cache_info["file_count"] == 0
-
- # Create mock cache files for existing model
- model_cache_dir = self.temp_dir / ".hf-cache" / "transformers" / "gpt2"
- model_cache_dir.mkdir(parents=True, exist_ok=True)
-
- # Create mock model files
- config_file = model_cache_dir / "config.json"
- model_file = model_cache_dir / "pytorch_model.bin"
-
- config_file.write_text('{"model_type": "gpt2"}') # ~25 bytes
- model_file.write_bytes(b"0" * (150 * 1024 * 1024)) # 150MB of zeros
-
- # Test cache info for cached model
- cache_info = accelerator.get_cache_info("gpt2")
- assert cache_info["cached"] is True
- assert (
- abs(cache_info["cache_size_mb"] - 150.0) < 0.1
- ) # Allow for small differences
- assert cache_info["file_count"] == 2
-
- # Test cache clearing
- result = accelerator.clear_model_cache("gpt2")
- assert result.success is True
- assert not model_cache_dir.exists()
-
-
-class TestDownloadAccelerationErrorHandling:
- """Test error handling and edge cases in download acceleration."""
-
- def setup_method(self):
- """Set up test environment."""
- self.temp_dir = Path(tempfile.mkdtemp())
-
- def teardown_method(self):
- """Clean up test environment."""
- shutil.rmtree(self.temp_dir, ignore_errors=True)
-
- def test_hf_transfer_download_failure_fallback(self):
- """Test fallback to standard download when hf_transfer fails."""
- downloader = HfTransferDownloader()
-
- # Test that unavailable downloader raises error
- if not downloader.hf_transfer_available:
- try:
- result = downloader.download(
- "https://huggingface.co/gpt2/resolve/main/file.bin",
- str(self.temp_dir / "file.bin"),
- )
- assert not result.success
- except RuntimeError as e:
- assert "hf_transfer not available" in str(e)
-
- @patch("src.huggingface_accelerator.HfApi.repo_info")
- def test_hf_api_failure_handling(self, mock_repo_info):
- """Test handling of HuggingFace API failures."""
- # Mock API failure
- mock_repo_info.side_effect = Exception("API error")
-
- accelerator = HuggingFaceAccelerator(None)
- files = accelerator.get_model_files("gpt2")
-
- # Should return empty list on failure
- assert files == []
-
- def test_invalid_model_acceleration(self):
- """Test acceleration with invalid model specifications."""
- mock_workspace = Mock()
- mock_workspace.has_runpod_volume = True
- mock_workspace.hf_cache_path = str(self.temp_dir)
-
- accelerator = HuggingFaceAccelerator(mock_workspace)
-
- # Test with empty model ID - should return success but indicate no pre-caching needed
- result = accelerator.accelerate_model_download("")
- assert result.success is True
- assert result.stdout is not None
- assert "does not require acceleration" in result.stdout
-
- def test_non_hf_url_handling(self):
- """Test handling of non-HuggingFace URLs."""
- downloader = HfTransferDownloader()
-
- # Test error handling for non-HF URLs when downloader is available
- if downloader.hf_transfer_available:
- result = downloader.download(
- "http://example.com/file.bin", str(self.temp_dir / "file.bin")
- )
- assert result.success is False
- assert result.error_message is not None
- assert "only supports HuggingFace URLs" in result.error_message
- else:
- # When not available, should raise RuntimeError
- try:
- result = downloader.download(
- "http://example.com/file.bin", str(self.temp_dir / "file.bin")
- )
- assert not result.success
- except RuntimeError as e:
- assert "hf_transfer not available" in str(e)
-
-
-if __name__ == "__main__":
- pytest.main([__file__])
diff --git a/tests/integration/test_handler_integration.py b/tests/integration/test_handler_integration.py
index f12bc4b..6b7fdc6 100644
--- a/tests/integration/test_handler_integration.py
+++ b/tests/integration/test_handler_integration.py
@@ -1,8 +1,10 @@
+import os
import pytest
import json
import base64
import cloudpickle
from pathlib import Path
+from unittest.mock import patch
from handler import handler, RemoteExecutor
from remote_execution import FunctionRequest
@@ -13,7 +15,7 @@ class TestHandlerIntegration:
def setup_method(self):
"""Setup for each test method."""
- self.test_data_dir = Path(__file__).parent.parent.parent / "src"
+ self.test_data_dir = Path(__file__).parent.parent.parent / "src" / "tests"
self.test_input_file = self.test_data_dir / "test_input.json"
self.test_class_input_file = self.test_data_dir / "test_class_input.json"
@@ -462,3 +464,268 @@ def process_data(data):
assert decoded_result["sum"] == 15
assert decoded_result["name"] == "test"
assert decoded_result["processed"] is True
+
+ @pytest.mark.integration
+ @pytest.mark.asyncio
+ async def test_hf_cache_ahead_basic(self):
+ """Test basic HuggingFace model cache-ahead functionality."""
+ event = {
+ "input": {
+ "function_name": "test_model_usage",
+ "function_code": """
+def test_model_usage():
+ from transformers import AutoTokenizer
+
+ # Use the pre-cached model
+ tokenizer = AutoTokenizer.from_pretrained("gpt2")
+ result = tokenizer("Hello world")
+
+ return {
+ "tokens": len(result["input_ids"]),
+ "model": "gpt2"
+ }
+""",
+ "dependencies": ["transformers"],
+ "hf_models_to_cache": ["gpt2"],
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ result = await handler(event)
+
+ assert result["success"] is True
+ decoded_result = cloudpickle.loads(base64.b64decode(result["result"]))
+ assert decoded_result["model"] == "gpt2"
+ assert decoded_result["tokens"] > 0
+
+ @pytest.mark.integration
+ @pytest.mark.asyncio
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ async def test_hf_cache_hit_scenario(self, mock_is_cached):
+ """Test cache hit detection prevents redundant downloads."""
+ # First call: simulate cache miss
+ mock_is_cached.return_value = False
+
+ event = {
+ "input": {
+ "function_name": "first_call",
+ "function_code": """
+def first_call():
+ return "first"
+""",
+ "hf_models_to_cache": ["gpt2"],
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ result1 = await handler(event)
+ assert result1["success"] is True
+
+ # Second call: simulate cache hit
+ mock_is_cached.return_value = True
+
+ result2 = await handler(event)
+ assert result2["success"] is True
+ # Verify cache hit message in stdout
+ assert "already cached" in result2["stdout"] or "cache hit" in result2["stdout"]
+
+ @pytest.mark.integration
+ @pytest.mark.asyncio
+ async def test_hf_multiple_models_parallel(self):
+ """Test downloading multiple HF models in parallel."""
+ event = {
+ "input": {
+ "function_name": "multi_model_test",
+ "function_code": """
+def multi_model_test():
+ return "models cached"
+""",
+ "hf_models_to_cache": ["gpt2", "distilbert-base-uncased"],
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ result = await handler(event)
+
+ assert result["success"] is True
+ # Both models should be mentioned in stdout
+ stdout_lower = result["stdout"].lower()
+ assert "gpt2" in stdout_lower or "model" in stdout_lower
+
+ @pytest.mark.integration
+ @pytest.mark.asyncio
+ @patch.dict(os.environ, {"HF_TOKEN": "test_token_value"})
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ async def test_hf_authentication_with_token(
+ self, mock_is_cached, mock_snapshot_download
+ ):
+ """Test HF_TOKEN is used for authentication."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.return_value = "/cache/path/private-model"
+
+ event = {
+ "input": {
+ "function_name": "private_model_test",
+ "function_code": """
+def private_model_test():
+ return "authenticated"
+""",
+ "hf_models_to_cache": ["private/model"],
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ result = await handler(event)
+
+ assert result["success"] is True
+ # Verify token was passed to snapshot_download
+ mock_snapshot_download.assert_called()
+ call_kwargs = mock_snapshot_download.call_args.kwargs
+ assert call_kwargs["token"] == "test_token_value"
+
+ @pytest.mark.integration
+ @pytest.mark.asyncio
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ async def test_hf_cache_failure_continues_execution(
+ self, mock_is_cached, mock_snapshot_download
+ ):
+ """Test that cache failures don't block function execution."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.side_effect = Exception("Network error")
+
+ event = {
+ "input": {
+ "function_name": "resilient_test",
+ "function_code": """
+def resilient_test():
+ return "execution continues despite cache failure"
+""",
+ "hf_models_to_cache": ["invalid-model"],
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ result = await handler(event)
+
+ # Execution should fail at dependency installation stage
+ # since cache-ahead failed
+ assert result["success"] is False
+ assert (
+ "Network error" in result["error"] or "Failed to cache" in result["error"]
+ )
+
+ @pytest.mark.integration
+ @pytest.mark.asyncio
+ async def test_hf_cache_with_custom_revision(self):
+ """Test caching specific model revisions."""
+ event = {
+ "input": {
+ "function_name": "revision_test",
+ "function_code": """
+def revision_test():
+ return "revision cached"
+""",
+ "hf_models_to_cache": ["gpt2"], # Default revision (main)
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ result = await handler(event)
+
+ assert result["success"] is True
+
+ @pytest.mark.integration
+ @pytest.mark.asyncio
+ async def test_hf_cache_ahead_prevents_redownload(self):
+ """
+ Test that cache-ahead prevents re-downloads when user code references the model.
+
+ This is the critical test: User's code doesn't know model was pre-cached,
+ but should use cached version instead of downloading from network.
+ """
+ import time
+
+ # Step 1: Cache-ahead the model WITHOUT using it
+ cache_event = {
+ "input": {
+ "function_name": "just_cache",
+ "function_code": """
+def just_cache():
+ return "model should be cached now"
+""",
+ "hf_models_to_cache": ["gpt2"],
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ cache_result = await handler(cache_event)
+ assert cache_result["success"] is True
+
+ # Step 2: User code loads the model (doesn't know it's cached)
+ # If cache works, this should be FAST (< 1 second)
+ # If cache fails, this downloads from network (> 5 seconds)
+ use_event = {
+ "input": {
+ "function_name": "use_cached_model",
+ "function_code": """
+def use_cached_model():
+ import time
+ from transformers import AutoTokenizer
+
+ start = time.time()
+
+ # User code doesn't know model was pre-cached
+ tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+ load_time = time.time() - start
+
+ return {
+ "load_time_seconds": round(load_time, 2),
+ "vocab_size": tokenizer.vocab_size,
+ "cache_was_used": load_time < 5.0 # Cache hit should be < 5s
+ }
+""",
+ "dependencies": ["transformers"],
+ "accelerate_downloads": True,
+ "args": [],
+ "kwargs": {},
+ }
+ }
+
+ start_time = time.time()
+ use_result = await handler(use_event)
+ total_time = time.time() - start_time
+
+ assert use_result["success"] is True
+ decoded_result = cloudpickle.loads(base64.b64decode(use_result["result"]))
+
+ # Verify model loaded successfully
+ assert decoded_result["vocab_size"] == 50257 # GPT2 vocab size
+
+ # Critical assertion: model load was fast (cache hit)
+ assert decoded_result["cache_was_used"] is True, (
+ f"Model took {decoded_result['load_time_seconds']}s to load. "
+ f"Expected < 5s (cache hit), suggesting re-download occurred."
+ )
+
+ # Total execution should also be fast
+ assert total_time < 30, (
+ f"Total execution took {total_time:.2f}s. "
+ f"Should be < 30s with cached model."
+ )
diff --git a/tests/integration/test_hf_strategy_integration.py b/tests/integration/test_hf_strategy_integration.py
deleted file mode 100644
index dd07bcf..0000000
--- a/tests/integration/test_hf_strategy_integration.py
+++ /dev/null
@@ -1,162 +0,0 @@
-"""
-Integration tests for HuggingFace download strategy system.
-"""
-
-import os
-import pytest
-from unittest.mock import Mock, patch
-
-from src.huggingface_accelerator import HuggingFaceAccelerator
-from src.hf_strategy_factory import HFStrategyFactory
-from hf_downloader_tetra import TetraHFDownloader
-from hf_downloader_native import NativeHFDownloader
-
-
-@pytest.fixture
-def mock_workspace_manager():
- """Mock workspace manager for integration tests."""
- workspace_manager = Mock()
- workspace_manager.hf_cache_path = "/tmp/test_cache"
- return workspace_manager
-
-
-class TestHuggingFaceAcceleratorIntegration:
- """Integration tests for HuggingFaceAccelerator with strategy pattern."""
-
- def test_accelerator_uses_configured_strategy(self, mock_workspace_manager):
- """Test that accelerator uses the configured strategy."""
- # Set environment to use tetra strategy
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "tetra"
-
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- accelerator = HuggingFaceAccelerator(mock_workspace_manager)
- assert isinstance(accelerator.strategy, TetraHFDownloader)
-
- def test_accelerator_strategy_delegation(self, mock_workspace_manager):
- """Test that accelerator properly delegates to strategy methods."""
- # Set to native strategy for simpler testing
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "native"
-
- accelerator = HuggingFaceAccelerator(mock_workspace_manager)
-
- # Mock the strategy methods
- accelerator.strategy.should_accelerate = Mock(return_value=True)
- accelerator.strategy.download_model = Mock(return_value=Mock(success=True))
- accelerator.strategy.is_model_cached = Mock(return_value=False)
- accelerator.strategy.get_cache_info = Mock(return_value={"cached": False})
- accelerator.strategy.clear_model_cache = Mock(return_value=Mock(success=True))
-
- # Test delegation
- assert accelerator.should_accelerate_model("gpt2")
- accelerator.strategy.should_accelerate.assert_called_once_with("gpt2")
-
- accelerator.accelerate_model_download("gpt2", "main")
- accelerator.strategy.download_model.assert_called_once_with("gpt2", "main")
-
- assert not accelerator.is_model_cached("gpt2", "main")
- accelerator.strategy.is_model_cached.assert_called_once_with("gpt2", "main")
-
- cache_info = accelerator.get_cache_info("gpt2")
- assert cache_info == {"cached": False}
- accelerator.strategy.get_cache_info.assert_called_once_with("gpt2")
-
- accelerator.clear_model_cache("gpt2")
- accelerator.strategy.clear_model_cache.assert_called_once_with("gpt2")
-
- def test_accelerator_strategy_switching(self, mock_workspace_manager):
- """Test runtime strategy switching."""
- # Start with native strategy
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "native"
-
- accelerator = HuggingFaceAccelerator(mock_workspace_manager)
- assert isinstance(accelerator.strategy, NativeHFDownloader)
-
- # Switch to tetra strategy
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- accelerator.set_strategy("tetra")
- assert isinstance(accelerator.strategy, TetraHFDownloader)
-
- # Check environment was updated
- assert os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] == "tetra"
-
- def test_accelerator_get_strategy_info(self, mock_workspace_manager):
- """Test getting strategy information from accelerator."""
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "native"
-
- accelerator = HuggingFaceAccelerator(mock_workspace_manager)
- info = accelerator.get_strategy_info()
-
- assert info["current_strategy"] == "native"
- assert info["strategy_instance"] == "NativeHFDownloader"
- assert info["environment_variable"] == HFStrategyFactory.STRATEGY_ENV_VAR
-
-
-class TestStrategyEnvironmentIntegration:
- """Test environment variable integration across the system."""
-
- def test_strategy_persistence_across_instances(self, mock_workspace_manager):
- """Test that strategy setting persists across new instances."""
- # Set strategy
- HFStrategyFactory.set_strategy("tetra")
-
- # Create first instance
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- accelerator1 = HuggingFaceAccelerator(mock_workspace_manager)
- assert isinstance(accelerator1.strategy, TetraHFDownloader)
-
- # Create second instance - should use same strategy
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- accelerator2 = HuggingFaceAccelerator(mock_workspace_manager)
- assert isinstance(accelerator2.strategy, TetraHFDownloader)
-
- def test_invalid_strategy_fallback(self, mock_workspace_manager):
- """Test fallback behavior with invalid strategy."""
- # Set invalid strategy
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "invalid_strategy"
-
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- accelerator = HuggingFaceAccelerator(mock_workspace_manager)
- # Should fallback to tetra (default)
- assert isinstance(accelerator.strategy, TetraHFDownloader)
-
- def test_no_env_var_uses_default(self, mock_workspace_manager):
- """Test default strategy when no environment variable is set."""
- # Clear environment variable
- if HFStrategyFactory.STRATEGY_ENV_VAR in os.environ:
- del os.environ[HFStrategyFactory.STRATEGY_ENV_VAR]
-
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- accelerator = HuggingFaceAccelerator(mock_workspace_manager)
- # Should use default (tetra)
- assert isinstance(accelerator.strategy, TetraHFDownloader)
-
-
-class TestWorkspaceManagerIntegration:
- """Test integration with workspace manager."""
-
- def test_strategy_uses_workspace_cache_path(self):
- """Test that strategies use workspace manager's cache path."""
- import tempfile
-
- with tempfile.TemporaryDirectory() as temp_dir:
- workspace_manager = Mock()
- workspace_manager.hf_cache_path = temp_dir
-
- # Test tetra strategy
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- tetra_strategy = TetraHFDownloader(workspace_manager)
- assert str(tetra_strategy.cache_dir) == temp_dir
-
- # Test native strategy (doesn't use cache_dir directly but should store workspace_manager)
- native_strategy = NativeHFDownloader(workspace_manager)
- assert native_strategy.workspace_manager == workspace_manager
-
- def test_strategy_with_no_cache_path(self):
- """Test strategy behavior when workspace manager has no cache path."""
- workspace_manager = Mock()
- workspace_manager.hf_cache_path = None
-
- with patch("src.hf_downloader_tetra.DownloadAccelerator"):
- tetra_strategy = TetraHFDownloader(workspace_manager)
- # Should fall back to default cache location
- assert "huggingface" in str(tetra_strategy.cache_dir)
diff --git a/tests/integration/test_python_path_integration.py b/tests/integration/test_python_path_integration.py
deleted file mode 100644
index 5051b24..0000000
--- a/tests/integration/test_python_path_integration.py
+++ /dev/null
@@ -1,152 +0,0 @@
-"""Integration tests to ensure Python path setup works for both function and class execution."""
-
-from unittest.mock import patch, MagicMock
-
-from class_executor import ClassExecutor
-from function_executor import FunctionExecutor
-from remote_execution import FunctionRequest
-from workspace_manager import WorkspaceManager
-
-
-class TestPythonPathIntegration:
- """Test Python path setup for volume-installed packages."""
-
- def setup_method(self):
- """Setup test environment."""
- self.mock_workspace_manager = MagicMock(spec=WorkspaceManager)
- self.mock_workspace_manager.has_runpod_volume = True
- self.mock_workspace_manager.venv_path = "/runpod-volume/runtimes/test/.venv"
-
- self.function_executor = FunctionExecutor(self.mock_workspace_manager)
- self.class_executor = ClassExecutor(self.mock_workspace_manager)
-
- def test_function_executor_calls_setup_python_path(self):
- """Test that FunctionExecutor calls setup_python_path before execution."""
- # Create a simple function request
- request = FunctionRequest(
- function_code="def test_func(): return 'success'",
- function_name="test_func",
- args=[],
- kwargs={},
- dependencies=[],
- system_dependencies=[],
- )
-
- # Mock the setup_python_path method
- with patch.object(
- self.mock_workspace_manager, "setup_python_path"
- ) as mock_setup:
- with patch.object(
- self.mock_workspace_manager, "change_to_workspace", return_value=None
- ):
- self.function_executor.execute(request)
- mock_setup.assert_called_once()
-
- def test_class_executor_calls_setup_python_path(self):
- """Test that ClassExecutor calls setup_python_path before execution."""
- # Create a simple class request
- request = FunctionRequest(
- class_code="class TestClass:\n def __call__(self): return 'success'",
- class_name="TestClass",
- execution_type="class",
- args=[],
- kwargs={},
- dependencies=[],
- system_dependencies=[],
- )
-
- # Mock the setup_python_path method
- with patch.object(
- self.mock_workspace_manager, "setup_python_path"
- ) as mock_setup:
- self.class_executor.execute_class_method(request)
- mock_setup.assert_called_once()
-
- def test_volume_package_import_simulation(self):
- """Test simulation of importing a volume-installed package."""
- # This test simulates the scenario where a package is installed in the volume
- # and needs to be available during class instantiation
-
- class_code = """
-class VolumePackageUser:
- def __init__(self):
- # This would normally fail if setup_python_path() wasn't called
- import sys
- self.paths = sys.path
-
- def get_paths(self):
- return self.paths
-"""
-
- request = FunctionRequest(
- class_code=class_code,
- class_name="VolumePackageUser",
- method_name="get_paths",
- execution_type="class",
- args=[],
- kwargs={},
- dependencies=[],
- system_dependencies=[],
- )
-
- # Mock setup_python_path to add a fake volume path
- def mock_setup_python_path():
- import sys
-
- fake_volume_path = (
- "/runpod-volume/runtimes/test/.venv/lib/python3.12/site-packages"
- )
- if fake_volume_path not in sys.path:
- sys.path.insert(0, fake_volume_path)
-
- with patch.object(
- self.mock_workspace_manager,
- "setup_python_path",
- side_effect=mock_setup_python_path,
- ):
- result = self.class_executor.execute_class_method(request)
-
- # Verify execution succeeded
- assert result.success is True
-
- # Verify the fake volume path was added to sys.path
- import sys
-
- fake_volume_path = (
- "/runpod-volume/runtimes/test/.venv/lib/python3.12/site-packages"
- )
- assert fake_volume_path in sys.path
-
- def test_base_executor_enforces_workspace_manager(self):
- """Test that BaseExecutor enforces workspace_manager requirement."""
- from base_executor import BaseExecutor
-
- # Test that BaseExecutor requires workspace_manager
- try:
-
- class TestExecutor(BaseExecutor):
- def execute(self, request):
- return None
-
- # This should raise ValueError
- TestExecutor(None)
- assert False, "Should have raised ValueError for None workspace_manager"
- except ValueError as e:
- assert "workspace_manager is required" in str(e)
-
- def test_setup_execution_environment_called_by_base_class(self):
- """Test that _setup_execution_environment is properly called."""
- from base_executor import BaseExecutor
-
- class TestExecutor(BaseExecutor):
- def execute(self, request):
- self._setup_execution_environment()
- return "executed"
-
- mock_workspace = MagicMock()
- executor = TestExecutor(mock_workspace)
-
- executor.execute(None)
-
- # Verify setup_python_path was called
- mock_workspace.setup_python_path.assert_called_once()
diff --git a/tests/integration/test_runpod_volume_integration.py b/tests/integration/test_runpod_volume_integration.py
deleted file mode 100644
index 6423478..0000000
--- a/tests/integration/test_runpod_volume_integration.py
+++ /dev/null
@@ -1,587 +0,0 @@
-"""Integration tests for RunPod volume workspace functionality."""
-
-import asyncio
-import base64
-import cloudpickle
-import threading
-from unittest.mock import Mock, patch, MagicMock
-
-from src.handler import RemoteExecutor, handler
-from remote_execution import FunctionResponse
-from src.constants import RUNPOD_VOLUME_PATH, VENV_DIR_NAME, RUNTIMES_DIR_NAME
-
-
-class TestFullWorkflowWithVolume:
- """Test complete request workflows with volume integration."""
-
- def setup_method(self):
- # Patch subprocess.run globally for all tests in this class
- class ContextManagerMock(MagicMock):
- def __enter__(self):
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- pass
-
- self.subprocess_run_patcher = patch("subprocess.run", new=ContextManagerMock())
- self.subprocess_run_patcher.start()
-
- def teardown_method(self):
- self.subprocess_run_patcher.stop()
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("workspace_manager.run_logged_subprocess")
- @patch("dependency_installer.run_logged_subprocess")
- @patch("dependency_installer.DependencyInstaller._is_docker_environment")
- @patch("os.chdir")
- @patch("glob.glob")
- async def test_full_workflow_with_volume(
- self,
- mock_glob,
- mock_chdir,
- mock_is_docker,
- mock_dependency_subprocess,
- mock_workspace_subprocess,
- mock_exists,
- mock_validate,
- mock_makedirs,
- ):
- """Test complete workflow from handler to execution with volume."""
- # Mock Docker environment detection to return True (simulating Docker container)
- mock_is_docker.return_value = True
-
- # Mock volume exists with endpoint-specific workspace
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- expected_venv = f"{expected_workspace}/{VENV_DIR_NAME}"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_workspace,
- expected_venv,
- ]
-
- # Mock glob for site-packages in endpoint-specific workspace
- mock_glob.return_value = [f"{expected_venv}/lib/python3.12/site-packages"]
-
- # Mock virtual environment validation
- mock_validate.return_value = FunctionResponse(success=True, stdout="Valid venv")
-
- # Mock successful dependency installation
- mock_dependency_subprocess.return_value = FunctionResponse(
- success=True, stdout="Successfully installed numpy"
- )
- mock_workspace_subprocess.return_value = FunctionResponse(
- success=True, stdout="Virtual environment created"
- )
-
- # Mock numpy module
- with patch.dict("sys.modules", {"numpy": Mock(__version__="1.21.0")}):
- # Complete request with dependencies and function
- event = {
- "input": {
- "function_name": "numpy_test",
- "function_code": """
-def numpy_test():
- import numpy as np
- return f"NumPy version: {np.__version__}"
-""",
- "args": [],
- "kwargs": {},
- "dependencies": ["numpy==1.21.0"],
- }
- }
-
- # This will fail until full integration is implemented
- result = await handler(event)
-
- assert result["success"] is True
- assert "error" not in result or result["error"] is None
-
- # Should have changed to endpoint-specific workspace directory
- chdir_calls = [call[0][0] for call in mock_chdir.call_args_list]
- assert expected_workspace in chdir_calls
-
- # Should have installed dependencies
- assert mock_dependency_subprocess.called
-
- @patch("os.makedirs")
- @patch("platform.system")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("dependency_installer.run_logged_subprocess")
- @patch("os.chdir")
- @patch("glob.glob")
- async def test_workflow_with_system_dependencies(
- self,
- mock_glob,
- mock_chdir,
- mock_subprocess,
- mock_exists,
- mock_validate,
- mock_platform,
- mock_makedirs,
- ):
- """Test workflow that includes both system and Python dependencies."""
- # Mock platform to return Linux to enable system dependency installation
- mock_platform.return_value = "Linux"
-
- # Mock volume exists with endpoint-specific workspace
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- expected_venv = f"{expected_workspace}/{VENV_DIR_NAME}"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_workspace,
- expected_venv,
- ]
-
- # Mock glob for site-packages in endpoint-specific workspace
- mock_glob.return_value = [f"{expected_venv}/lib/python3.12/site-packages"]
-
- # Mock virtual environment validation
- mock_validate.return_value = FunctionResponse(success=True, stdout="Valid venv")
-
- # Mock apt-get update and install
- apt_update_process = Mock()
- apt_update_process.returncode = 0
- apt_update_process.communicate.return_value = (b"Package lists updated", b"")
-
- apt_install_process = Mock()
- apt_install_process.returncode = 0
- apt_install_process.communicate.return_value = (
- b"System packages installed",
- b"",
- )
-
- # Mock uv pip list (for _get_installed_packages)
- pip_list_process = Mock()
- pip_list_process.returncode = 0
- pip_list_process.communicate.return_value = (
- b"", # No packages installed yet
- b"",
- )
-
- # Mock uv pip install
- pip_install_process = Mock()
- pip_install_process.returncode = 0
- pip_install_process.communicate.return_value = (
- b"Python packages installed",
- b"",
- )
-
- # Mock subprocess calls in order:
- # 1. which nala (system package acceleration check)
- # 2. apt-get update
- # 3. apt-get install
- # 4. uv pip list (get installed packages)
- # 5. uv pip install
- nala_check_process = Mock()
- nala_check_process.returncode = 1 # nala not available
- nala_check_process.communicate.return_value = (b"", b"which: nala: not found")
-
- # Create a function that returns appropriate mock based on the command
- def popen_side_effect(*args, **kwargs):
- cmd = args[0]
- if "nala" in str(cmd) or "which" in str(cmd):
- return nala_check_process
- elif "apt-get" in str(cmd) and "update" in str(cmd):
- return apt_update_process
- elif "apt-get" in str(cmd) and "install" in str(cmd):
- return apt_install_process
- elif "uv" in str(cmd) and "list" in str(cmd):
- return pip_list_process
- elif "uv" in str(cmd) and "install" in str(cmd):
- return pip_install_process
- else:
- # Return a generic successful process for any other calls
- generic_process = Mock()
- generic_process.returncode = 0
- generic_process.communicate.return_value = (b"", b"")
- return generic_process
-
- # Simplified mocking: just return success for all subprocess calls
- mock_subprocess.return_value = FunctionResponse(
- success=True, stdout="Successfully installed packages"
- )
-
- # Mock subprocess.run for the test function
- mock_run_result = Mock()
- mock_run_result.stdout = "/usr/bin/curl"
-
- with patch("subprocess.run", return_value=mock_run_result):
- with patch.dict("sys.modules", {"requests": Mock(__version__="2.25.1")}):
- event = {
- "input": {
- "function_name": "system_test",
- "function_code": """
-def system_test():
- import subprocess
- result = subprocess.run(['which', 'wget'], capture_output=True, text=True)
- return result.stdout.strip()
-""",
- "args": [],
- "kwargs": {},
- "system_dependencies": ["wget"],
- "dependencies": ["requests==2.25.1"],
- }
- }
-
- # This will fail until system dependency integration is implemented
- result = await handler(event)
-
- assert result["success"] is True
-
- # Should have called subprocess utility for dependency installation
- assert mock_subprocess.called
-
-
-class TestConcurrentRequests:
- """Test realistic concurrent access scenarios."""
-
- def setup_method(self):
- # Patch subprocess.run globally for all tests in this class
- class ContextManagerMock(MagicMock):
- def __enter__(self):
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- pass
-
- self.subprocess_run_patcher = patch("subprocess.run", new=ContextManagerMock())
- self.subprocess_run_patcher.start()
-
- def teardown_method(self):
- self.subprocess_run_patcher.stop()
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("dependency_installer.run_logged_subprocess")
- @patch("fcntl.flock")
- @patch("os.chdir")
- @patch("glob.glob")
- async def test_multiple_concurrent_requests(
- self,
- mock_glob,
- mock_chdir,
- mock_flock,
- mock_subprocess,
- mock_exists,
- mock_validate,
- mock_makedirs,
- ):
- """Test multiple concurrent requests to the same endpoint."""
- # Mock volume exists with endpoint-specific workspace
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- expected_venv = f"{expected_workspace}/{VENV_DIR_NAME}"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_workspace,
- expected_venv,
- ]
-
- # Mock glob for site-packages in endpoint-specific workspace
- mock_glob.return_value = [f"{expected_venv}/lib/python3.12/site-packages"]
-
- # Mock virtual environment validation
- mock_validate.return_value = FunctionResponse(success=True, stdout="Valid venv")
-
- # Mock successful installations
- mock_subprocess.return_value = FunctionResponse(
- success=True, stdout="Installation complete"
- )
-
- # Mock the time module
- mock_time = Mock()
- mock_time.sleep = Mock()
-
- with patch.dict(
- "sys.modules", {"time": mock_time, "numpy": Mock(__version__="1.21.0")}
- ):
-
- async def make_request(request_id):
- event = {
- "input": {
- "function_name": "concurrent_test",
- "function_code": f"""
-def concurrent_test():
- import time
- time.sleep(0.1) # Simulate some work
- return "Request {request_id} completed"
-""",
- "args": [],
- "kwargs": {},
- "dependencies": ["numpy==1.21.0"],
- }
- }
- return await handler(event)
-
- # Start 5 concurrent requests
- tasks = [make_request(i) for i in range(5)]
-
- # This will fail until concurrent safety is implemented
- results = await asyncio.gather(*tasks)
-
- # All requests should succeed
- for i, result in enumerate(results):
- assert result["success"] is True
- decoded_result = cloudpickle.loads(base64.b64decode(result["result"]))
- assert f"Request {i} completed" in decoded_result
-
- # Since workspace is already initialized, flock might not be called
- # Just verify that all requests succeeded
- assert len(results) == 5
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("dependency_installer.run_logged_subprocess")
- def test_concurrent_dependency_installation(
- self, mock_subprocess, mock_exists, mock_validate, mock_makedirs
- ):
- """Test that concurrent dependency installations don't conflict."""
- # Mock volume exists with endpoint-specific workspace
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- expected_venv = f"{expected_workspace}/{VENV_DIR_NAME}"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_workspace,
- expected_venv,
- ]
-
- # Track installation calls
- install_calls = []
-
- def track_subprocess(command, *args, **kwargs):
- # Track subprocess calls for verification
- if (
- isinstance(command, list)
- and "uv" in str(command)
- and "pip" in str(command)
- ):
- install_calls.append(command)
- return FunctionResponse(success=True, stdout="Installation complete")
-
- mock_subprocess.side_effect = track_subprocess
-
- def install_deps(executor, packages):
- return executor.dependency_installer.install_dependencies(packages)
-
- # Create multiple executors trying to install different packages
- executors = [RemoteExecutor() for _ in range(3)]
- package_sets = [["numpy==1.21.0"], ["pandas==1.3.0"], ["scipy==1.7.0"]]
-
- threads = []
- results = []
-
- for executor, packages in zip(executors, package_sets):
- thread = threading.Thread(
- target=lambda e=executor, p=packages: results.append(install_deps(e, p))
- )
- threads.append(thread)
- thread.start()
-
- for thread in threads:
- thread.join()
-
- # This will fail until concurrent installation safety is implemented
- assert len(results) == 3
- assert all(result.success for result in results)
-
- # Should have made installation calls for all packages
- all_packages = ["numpy==1.21.0", "pandas==1.3.0", "scipy==1.7.0"]
- for package in all_packages:
- assert any(package in " ".join(call) for call in install_calls)
-
-
-class TestMixedExecution:
- """Test mixed volume and non-volume execution scenarios."""
-
- def setup_method(self):
- # Patch subprocess.run globally for all tests in this class
- class ContextManagerMock(MagicMock):
- def __enter__(self):
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- pass
-
- self.subprocess_run_patcher = patch("subprocess.run", new=ContextManagerMock())
- self.subprocess_run_patcher.start()
-
- def teardown_method(self):
- self.subprocess_run_patcher.stop()
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("os.chdir")
- async def test_mixed_volume_and_non_volume_execution(
- self, mock_chdir, mock_exists, mock_validate, mock_makedirs
- ):
- """Test that handlers work both with and without volumes."""
- # First request - no volume available
- mock_exists.return_value = False
-
- event_no_volume = {
- "input": {
- "function_name": "simple_test",
- "function_code": "def simple_test():\n return 'no volume'",
- "args": [],
- "kwargs": {},
- }
- }
-
- result_no_volume = await handler(event_no_volume)
- assert result_no_volume["success"] is True
-
- # Second request - volume becomes available
- # Mock volume exists with endpoint-specific workspace
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- expected_venv = f"{expected_workspace}/{VENV_DIR_NAME}"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_workspace,
- expected_venv,
- ]
-
- event_with_volume = {
- "input": {
- "function_name": "volume_test",
- "function_code": "def volume_test():\n return 'with volume'",
- "args": [],
- "kwargs": {},
- }
- }
-
- # This will fail until mixed execution is properly handled
- result_with_volume = await handler(event_with_volume)
- assert result_with_volume["success"] is True
- chdir_calls = [call[0][0] for call in mock_chdir.call_args_list]
- # Should change to endpoint-specific workspace, not just volume root
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- assert expected_workspace in chdir_calls
-
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("dependency_installer.run_logged_subprocess")
- @patch("os.makedirs")
- @patch("builtins.open")
- async def test_fallback_on_volume_initialization_failure(
- self, mock_open, mock_makedirs, mock_subprocess, mock_exists, mock_validate
- ):
- """Test graceful fallback when volume initialization fails."""
- mock_exists.side_effect = (
- lambda path: path == RUNPOD_VOLUME_PATH
- ) # Volume exists but venv doesn't exist
-
- # Mock file operations
- mock_file = MagicMock()
- mock_file.fileno.return_value = 3
- mock_open.return_value.__enter__.return_value = mock_file
-
- mock_subprocess.return_value = FunctionResponse(
- success=False, error="Failed to create venv"
- )
-
- event = {
- "input": {
- "function_name": "fallback_test",
- "function_code": "def fallback_test():\n return 'fallback execution'",
- "args": [],
- "kwargs": {},
- "dependencies": ["numpy==1.21.0"],
- }
- }
-
- # This will fail until fallback mechanism is implemented
- result = await handler(event)
-
- # Should fail because venv creation failed and no fallback implemented yet
- assert result["success"] is False
- assert "failed to create virtual environment" in result.get("error", "").lower()
-
-
-class TestErrorHandlingIntegration:
- """Test error handling in integrated volume scenarios."""
-
- def setup_method(self):
- # Patch subprocess.run globally for all tests in this class
- class ContextManagerMock(MagicMock):
- def __enter__(self):
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- pass
-
- self.subprocess_run_patcher = patch("subprocess.run", new=ContextManagerMock())
- self.subprocess_run_patcher.start()
-
- def teardown_method(self):
- self.subprocess_run_patcher.stop()
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("dependency_installer.run_logged_subprocess")
- async def test_dependency_installation_failure_with_volume(
- self, mock_subprocess, mock_exists, mock_validate, mock_makedirs
- ):
- """Test proper error handling when dependency installation fails in volume."""
- # Mock volume exists with endpoint-specific workspace
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- expected_venv = f"{expected_workspace}/{VENV_DIR_NAME}"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_workspace,
- expected_venv,
- ]
-
- # Mock failed dependency installation
- mock_subprocess.return_value = FunctionResponse(
- success=False, error="Package not found: nonexistent-package"
- )
-
- event = {
- "input": {
- "function_name": "test_func",
- "function_code": "def test_func():\n return 'should not execute'",
- "args": [],
- "kwargs": {},
- "dependencies": ["nonexistent-package==999.999.999"],
- }
- }
-
- result = await handler(event)
-
- assert result["success"] is False
- assert "package not found" in result.get("error", "").lower()
- # Function should not have been executed
- assert "result" not in result or result["result"] is None
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("os.chdir")
- async def test_volume_permission_error_handling(
- self, mock_chdir, mock_exists, mock_validate, mock_makedirs
- ):
- """Test handling of permission errors when accessing volume."""
- mock_exists.return_value = True
- mock_chdir.side_effect = PermissionError("Permission denied")
-
- event = {
- "input": {
- "function_name": "permission_test",
- "function_code": "def permission_test():\n return 'test'",
- "args": [],
- "kwargs": {},
- }
- }
-
- # This will fail until permission error handling is implemented
- result = await handler(event)
-
- # Should handle permission error gracefully
- assert result["success"] is False
- assert "permission denied" in result.get("error", "").lower()
diff --git a/tests/unit/test_class_executor.py b/tests/unit/test_class_executor.py
index fdf8dd9..727231d 100644
--- a/tests/unit/test_class_executor.py
+++ b/tests/unit/test_class_executor.py
@@ -3,7 +3,6 @@
import base64
import cloudpickle
from datetime import datetime
-from unittest.mock import MagicMock
from class_executor import ClassExecutor
from remote_execution import FunctionRequest
@@ -14,8 +13,7 @@ class TestClassExecution:
def setup_method(self):
"""Setup for each test method."""
- mock_workspace_manager = MagicMock()
- self.executor = ClassExecutor(mock_workspace_manager)
+ self.executor = ClassExecutor()
def encode_args(self, *args):
"""Helper to encode arguments."""
@@ -159,8 +157,7 @@ class TestInstanceManagement:
def setup_method(self):
"""Setup for each test method."""
- mock_workspace_manager = MagicMock()
- self.executor = ClassExecutor(mock_workspace_manager)
+ self.executor = ClassExecutor()
def encode_args(self, *args):
"""Helper to encode arguments."""
diff --git a/tests/unit/test_dependency_installer.py b/tests/unit/test_dependency_installer.py
index edb71bd..e30a81c 100644
--- a/tests/unit/test_dependency_installer.py
+++ b/tests/unit/test_dependency_installer.py
@@ -1,9 +1,8 @@
"""Tests for DependencyInstaller component."""
-from unittest.mock import Mock, patch
+from unittest.mock import patch
from dependency_installer import DependencyInstaller
-from workspace_manager import WorkspaceManager
from remote_execution import FunctionResponse
@@ -12,8 +11,7 @@ class TestSystemDependencies:
def setup_method(self):
"""Setup for each test method."""
- self.workspace_manager = Mock(spec=WorkspaceManager)
- self.installer = DependencyInstaller(self.workspace_manager)
+ self.installer = DependencyInstaller()
@patch("platform.system")
@patch("dependency_installer.run_logged_subprocess")
@@ -67,8 +65,7 @@ class TestSystemPackageAcceleration:
def setup_method(self):
"""Setup for each test method."""
- self.workspace_manager = Mock(spec=WorkspaceManager)
- self.installer = DependencyInstaller(self.workspace_manager)
+ self.installer = DependencyInstaller()
@patch("dependency_installer.run_logged_subprocess")
def test_nala_availability_check_available(self, mock_subprocess):
@@ -207,10 +204,7 @@ class TestPythonDependencies:
def setup_method(self):
"""Setup for each test method."""
- self.workspace_manager = Mock(spec=WorkspaceManager)
- self.workspace_manager.has_runpod_volume = False
- self.workspace_manager.cache_path = None
- self.installer = DependencyInstaller(self.workspace_manager)
+ self.installer = DependencyInstaller()
@patch("dependency_installer.run_logged_subprocess")
def test_install_dependencies_success(self, mock_subprocess):
diff --git a/tests/unit/test_function_executor.py b/tests/unit/test_function_executor.py
index 0779f6c..815e326 100644
--- a/tests/unit/test_function_executor.py
+++ b/tests/unit/test_function_executor.py
@@ -2,10 +2,8 @@
import base64
import cloudpickle
-from unittest.mock import Mock, patch
from function_executor import FunctionExecutor
-from workspace_manager import WorkspaceManager
from remote_execution import FunctionRequest
@@ -14,9 +12,7 @@ class TestFunctionExecution:
def setup_method(self):
"""Setup for each test method."""
- self.workspace_manager = Mock(spec=WorkspaceManager)
- self.workspace_manager.change_to_workspace.return_value = None
- self.executor = FunctionExecutor(self.workspace_manager)
+ self.executor = FunctionExecutor()
def encode_args(self, *args):
"""Helper to encode arguments."""
@@ -134,39 +130,15 @@ def output_func():
assert "log message" in response.stdout
-class TestWorkspaceIntegration:
- """Test integration with workspace manager."""
+class TestErrorHandling:
+ """Test error handling in function execution."""
def setup_method(self):
"""Setup for each test method."""
- self.workspace_manager = Mock(spec=WorkspaceManager)
- self.executor = FunctionExecutor(self.workspace_manager)
-
- def test_execute_function_in_workspace(self):
- """Test that function execution uses workspace directory."""
- self.workspace_manager.change_to_workspace.return_value = "/original"
-
- request = FunctionRequest(
- function_name="test_func",
- function_code="def test_func():\n return 'test'",
- args=[],
- kwargs={},
- )
-
- with patch("os.chdir") as mock_chdir:
- self.executor.execute(request)
-
- # Verify workspace methods were called
- self.workspace_manager.change_to_workspace.assert_called_once()
- self.workspace_manager.setup_python_path.assert_called_once()
-
- # Verify directory was restored
- mock_chdir.assert_called_once_with("/original")
-
- def test_execute_function_workspace_restoration_on_error(self):
- """Test that workspace directory is restored even on error."""
- self.workspace_manager.change_to_workspace.return_value = "/original"
+ self.executor = FunctionExecutor()
+ def test_execute_function_handles_errors(self):
+ """Test that function execution properly handles errors."""
request = FunctionRequest(
function_name="error_func",
function_code="def error_func():\n raise Exception('test error')",
@@ -174,9 +146,8 @@ def test_execute_function_workspace_restoration_on_error(self):
kwargs={},
)
- with patch("os.chdir") as mock_chdir:
- response = self.executor.execute(request)
+ response = self.executor.execute(request)
- # Verify directory was restored even after error
- mock_chdir.assert_called_once_with("/original")
+ # Verify error was captured
assert response.success is False
+ assert "test error" in response.error
diff --git a/tests/unit/test_hf_download_strategies.py b/tests/unit/test_hf_download_strategies.py
deleted file mode 100644
index 898ab17..0000000
--- a/tests/unit/test_hf_download_strategies.py
+++ /dev/null
@@ -1,260 +0,0 @@
-"""
-Unit tests for HuggingFace download strategies.
-"""
-
-import os
-import pytest
-from unittest.mock import Mock, patch
-
-from src.hf_downloader_tetra import TetraHFDownloader
-from src.hf_downloader_native import NativeHFDownloader
-from src.hf_strategy_factory import HFStrategyFactory
-from src.remote_execution import FunctionResponse
-
-
-@pytest.fixture
-def mock_workspace_manager():
- """Mock workspace manager."""
- workspace_manager = Mock()
- workspace_manager.hf_cache_path = "/tmp/test_cache"
- return workspace_manager
-
-
-@pytest.fixture
-def mock_download_accelerator():
- """Mock download accelerator."""
- accelerator = Mock()
- accelerator.hf_transfer_downloader = Mock()
- accelerator.hf_transfer_downloader.hf_transfer_available = True
- return accelerator
-
-
-class TestHFStrategyFactory:
- """Tests for HF strategy factory."""
-
- def test_get_available_strategies(self):
- """Test getting available strategies."""
- strategies = HFStrategyFactory.get_available_strategies()
- assert HFStrategyFactory.TETRA_STRATEGY in strategies
- assert HFStrategyFactory.NATIVE_STRATEGY in strategies
-
- def test_get_configured_strategy_default(self):
- """Test default strategy when no env var set."""
- # Clear environment variable
- if HFStrategyFactory.STRATEGY_ENV_VAR in os.environ:
- del os.environ[HFStrategyFactory.STRATEGY_ENV_VAR]
-
- strategy = HFStrategyFactory.get_configured_strategy()
- assert strategy == HFStrategyFactory.DEFAULT_STRATEGY
-
- def test_get_configured_strategy_from_env(self):
- """Test getting strategy from environment variable."""
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "tetra"
- strategy = HFStrategyFactory.get_configured_strategy()
- assert strategy == "tetra"
-
- def test_get_configured_strategy_invalid_fallback(self):
- """Test fallback to default for invalid strategy."""
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "invalid_strategy"
- strategy = HFStrategyFactory.get_configured_strategy()
- assert strategy == HFStrategyFactory.DEFAULT_STRATEGY
-
- def test_create_tetra_strategy(self, mock_workspace_manager):
- """Test creating tetra strategy."""
- with patch("src.hf_strategy_factory.TetraHFDownloader") as mock_tetra:
- mock_instance = Mock()
- mock_tetra.return_value = mock_instance
-
- strategy = HFStrategyFactory.create_strategy(
- mock_workspace_manager, HFStrategyFactory.TETRA_STRATEGY
- )
-
- mock_tetra.assert_called_once_with(mock_workspace_manager)
- assert strategy == mock_instance
-
- def test_create_native_strategy(self, mock_workspace_manager):
- """Test creating native strategy."""
- with patch("src.hf_strategy_factory.NativeHFDownloader") as mock_native:
- mock_instance = Mock()
- mock_native.return_value = mock_instance
-
- strategy = HFStrategyFactory.create_strategy(
- mock_workspace_manager, HFStrategyFactory.NATIVE_STRATEGY
- )
-
- mock_native.assert_called_once_with(mock_workspace_manager)
- assert strategy == mock_instance
-
- def test_set_strategy(self):
- """Test setting strategy environment variable."""
- HFStrategyFactory.set_strategy("tetra")
- assert os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] == "tetra"
-
- def test_set_strategy_invalid(self):
- """Test setting invalid strategy raises error."""
- with pytest.raises(ValueError):
- HFStrategyFactory.set_strategy("invalid_strategy")
-
- def test_get_strategy_info(self):
- """Test getting strategy information."""
- os.environ[HFStrategyFactory.STRATEGY_ENV_VAR] = "tetra"
-
- info = HFStrategyFactory.get_strategy_info()
-
- assert info["current_strategy"] == "tetra"
- assert info["environment_variable"] == HFStrategyFactory.STRATEGY_ENV_VAR
- assert info["environment_value"] == "tetra"
- assert info["default_strategy"] == HFStrategyFactory.DEFAULT_STRATEGY
- assert "tetra" in info["available_strategies"]
- assert "native" in info["available_strategies"]
-
-
-class TestTetraHFDownloader:
- """Tests for Tetra HF downloader strategy."""
-
- def test_init(self, mock_workspace_manager):
- """Test TetraHFDownloader initialization."""
- with patch(
- "src.hf_downloader_tetra.DownloadAccelerator"
- ) as mock_accelerator_class:
- downloader = TetraHFDownloader(mock_workspace_manager)
-
- assert downloader.workspace_manager == mock_workspace_manager
- mock_accelerator_class.assert_called_once_with(mock_workspace_manager)
-
- def test_should_accelerate_with_hf_transfer(self, mock_workspace_manager):
- """Test should_accelerate when hf_transfer is available."""
- with patch(
- "src.hf_downloader_tetra.DownloadAccelerator"
- ) as mock_accelerator_class:
- mock_accelerator = Mock()
- mock_accelerator.hf_transfer_downloader.hf_transfer_available = True
- mock_accelerator_class.return_value = mock_accelerator
-
- downloader = TetraHFDownloader(mock_workspace_manager)
-
- # Should accelerate large models
- assert downloader.should_accelerate("gpt-3.5-turbo")
- assert downloader.should_accelerate("llama")
-
- # Should not accelerate small models
- assert not downloader.should_accelerate("prajjwal1/bert-tiny")
-
- def test_should_accelerate_without_hf_transfer(self, mock_workspace_manager):
- """Test should_accelerate when hf_transfer is not available."""
- with patch(
- "src.hf_downloader_tetra.DownloadAccelerator"
- ) as mock_accelerator_class:
- mock_accelerator = Mock()
- mock_accelerator.hf_transfer_downloader.hf_transfer_available = False
- mock_accelerator_class.return_value = mock_accelerator
-
- downloader = TetraHFDownloader(mock_workspace_manager)
-
- # Should not accelerate any models without hf_transfer
- assert not downloader.should_accelerate("gpt-3.5-turbo")
- assert not downloader.should_accelerate("llama")
-
- @patch("src.hf_downloader_tetra.Path.mkdir")
- def test_download_model_success(self, mock_mkdir, mock_workspace_manager):
- """Test successful model download."""
- with patch(
- "src.hf_downloader_tetra.DownloadAccelerator"
- ) as mock_accelerator_class:
- mock_accelerator = Mock()
- mock_accelerator.hf_transfer_downloader.hf_transfer_available = True
- mock_accelerator_class.return_value = mock_accelerator
-
- downloader = TetraHFDownloader(mock_workspace_manager)
-
- # Mock get_model_files to return test files
- downloader.get_model_files = Mock(
- return_value=[
- {
- "path": "pytorch_model.bin",
- "size": 100 * 1024 * 1024,
- "url": "https://test.com/file",
- }
- ]
- )
-
- # Mock download_with_fallback to succeed
- mock_accelerator.download_with_fallback.return_value = FunctionResponse(
- success=True
- )
-
- result = downloader.download_model("gpt2")
-
- assert result.success
- assert "Successfully pre-downloaded" in result.stdout
-
- def test_download_model_no_acceleration_needed(self, mock_workspace_manager):
- """Test download when no acceleration is needed."""
- with patch(
- "src.hf_downloader_tetra.DownloadAccelerator"
- ) as mock_accelerator_class:
- mock_accelerator = Mock()
- mock_accelerator.hf_transfer_downloader.hf_transfer_available = False
- mock_accelerator_class.return_value = mock_accelerator
-
- downloader = TetraHFDownloader(mock_workspace_manager)
-
- result = downloader.download_model("prajjwal1/bert-tiny")
-
- assert result.success
- assert "does not require acceleration" in result.stdout
-
-
-class TestNativeHFDownloader:
- """Tests for Native HF downloader strategy."""
-
- def test_init(self, mock_workspace_manager):
- """Test NativeHFDownloader initialization."""
- downloader = NativeHFDownloader(mock_workspace_manager)
- assert downloader.workspace_manager == mock_workspace_manager
-
- def test_should_accelerate(self, mock_workspace_manager):
- """Test should_accelerate logic."""
- downloader = NativeHFDownloader(mock_workspace_manager)
-
- # Should accelerate large models
- assert downloader.should_accelerate("gpt-3.5-turbo")
- assert downloader.should_accelerate("llama")
-
- # Should not accelerate small models
- assert not downloader.should_accelerate("prajjwal1/bert-tiny")
-
- @patch("src.hf_downloader_native.snapshot_download")
- def test_download_model_success(
- self, mock_snapshot_download, mock_workspace_manager
- ):
- """Test successful model download."""
- mock_snapshot_download.return_value = "/cache/models/gpt2"
-
- downloader = NativeHFDownloader(mock_workspace_manager)
- result = downloader.download_model("gpt2")
-
- assert result.success
- assert "Successfully pre-cached model gpt2" in result.stdout
- mock_snapshot_download.assert_called_once_with(repo_id="gpt2", revision="main")
-
- @patch("src.hf_downloader_native.snapshot_download")
- def test_download_model_failure(
- self, mock_snapshot_download, mock_workspace_manager
- ):
- """Test failed model download."""
- mock_snapshot_download.side_effect = Exception("Download failed")
-
- downloader = NativeHFDownloader(mock_workspace_manager)
- result = downloader.download_model("gpt2")
-
- assert not result.success
- assert "Failed to pre-cache model gpt2" in result.error
-
- def test_download_model_no_acceleration_needed(self, mock_workspace_manager):
- """Test download when no acceleration is needed."""
- downloader = NativeHFDownloader(mock_workspace_manager)
- result = downloader.download_model("prajjwal1/bert-tiny")
-
- assert result.success
- assert "does not require pre-caching" in result.stdout
diff --git a/tests/unit/test_huggingface_cache.py b/tests/unit/test_huggingface_cache.py
new file mode 100644
index 0000000..1106daf
--- /dev/null
+++ b/tests/unit/test_huggingface_cache.py
@@ -0,0 +1,172 @@
+"""Tests for HuggingFaceCacheAhead component."""
+
+import os
+from unittest.mock import patch, Mock
+import pytest
+
+from huggingface_cache import HuggingFaceCacheAhead
+from remote_execution import FunctionResponse
+
+
+class TestHuggingFaceCacheAhead:
+ """Test HuggingFace cache-ahead functionality."""
+
+ def setup_method(self):
+ """Setup for each test method."""
+ self.hf_cache = HuggingFaceCacheAhead()
+
+ @patch("huggingface_cache.scan_cache_dir")
+ def test_is_model_cached_returns_true_when_cached(self, mock_scan):
+ """Test cache detection returns True when model is cached."""
+ # Mock cache info with our model present
+ mock_repo = Mock()
+ mock_repo.repo_id = "gpt2"
+ mock_rev = Mock()
+ mock_rev.commit_hash = "main"
+ mock_repo.revisions = [mock_rev]
+
+ mock_cache_info = Mock()
+ mock_cache_info.repos = [mock_repo]
+ mock_scan.return_value = mock_cache_info
+
+ result = self.hf_cache._is_model_cached("gpt2", "main")
+ assert result is True
+
+ @patch("huggingface_cache.scan_cache_dir")
+ def test_is_model_cached_returns_false_when_not_cached(self, mock_scan):
+ """Test cache detection returns False when model is not cached."""
+ # Mock empty cache
+ mock_cache_info = Mock()
+ mock_cache_info.repos = []
+ mock_scan.return_value = mock_cache_info
+
+ result = self.hf_cache._is_model_cached("gpt2", "main")
+ assert result is False
+
+ @patch("huggingface_cache.scan_cache_dir")
+ def test_is_model_cached_returns_false_on_error(self, mock_scan):
+ """Test cache detection returns False when cache check fails."""
+ mock_scan.side_effect = Exception("Cache error")
+
+ result = self.hf_cache._is_model_cached("gpt2", "main")
+ assert result is False
+
+ @patch.dict(os.environ, {"HF_TOKEN": "test_token_12345"})
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ def test_cache_model_download_uses_hf_token(
+ self, mock_is_cached, mock_snapshot_download
+ ):
+ """Test that HF_TOKEN is passed to snapshot_download when present."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.return_value = "/cache/path/gpt2"
+
+ result = self.hf_cache.cache_model_download("gpt2")
+
+ assert result.success is True
+ mock_snapshot_download.assert_called_once_with(
+ repo_id="gpt2", revision="main", token="test_token_12345"
+ )
+
+ @patch.dict(os.environ, {}, clear=True)
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ def test_cache_model_download_without_token(
+ self, mock_is_cached, mock_snapshot_download
+ ):
+ """Test that None is passed when HF_TOKEN is not present."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.return_value = "/cache/path/gpt2"
+
+ result = self.hf_cache.cache_model_download("gpt2")
+
+ assert result.success is True
+ mock_snapshot_download.assert_called_once_with(
+ repo_id="gpt2", revision="main", token=None
+ )
+
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ def test_cache_model_download_skips_when_cached(self, mock_is_cached):
+ """Test that download is skipped when model is already cached."""
+ mock_is_cached.return_value = True
+
+ result = self.hf_cache.cache_model_download("gpt2")
+
+ assert result.success is True
+ assert "already cached" in result.stdout
+ assert "cache hit" in result.stdout
+
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ def test_cache_model_download_success(self, mock_is_cached, mock_snapshot_download):
+ """Test successful model download."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.return_value = "/cache/path/gpt2"
+
+ result = self.hf_cache.cache_model_download("gpt2")
+
+ assert result.success is True
+ assert "gpt2" in result.stdout
+ assert "/cache/path/gpt2" in result.stdout
+
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ def test_cache_model_download_handles_network_error(
+ self, mock_is_cached, mock_snapshot_download
+ ):
+ """Test error handling for network failures."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.side_effect = Exception("Network error")
+
+ result = self.hf_cache.cache_model_download("gpt2")
+
+ assert result.success is False
+ assert "Failed to cache-ahead" in result.error
+ assert "gpt2" in result.error
+
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ def test_cache_model_download_handles_invalid_model(
+ self, mock_is_cached, mock_snapshot_download
+ ):
+ """Test error handling for invalid model IDs."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.side_effect = Exception("Model not found")
+
+ result = self.hf_cache.cache_model_download("invalid-model-id")
+
+ assert result.success is False
+ assert "Failed to cache-ahead" in result.error
+
+ @pytest.mark.asyncio
+ @patch("huggingface_cache.asyncio.to_thread")
+ async def test_cache_model_download_async_delegates_to_sync(self, mock_to_thread):
+ """Test async wrapper properly delegates to sync method."""
+ mock_response = FunctionResponse(success=True, stdout="Model cached")
+ mock_to_thread.return_value = mock_response
+
+ result = await self.hf_cache.cache_model_download_async("gpt2", "v1.0")
+
+ mock_to_thread.assert_called_once()
+ call_args = mock_to_thread.call_args
+ assert call_args[0][0] == self.hf_cache.cache_model_download
+ assert call_args[0][1] == "gpt2"
+ assert call_args[0][2] == "v1.0"
+ assert result == mock_response
+
+ @patch.dict(os.environ, {}, clear=True)
+ @patch("huggingface_cache.snapshot_download")
+ @patch("huggingface_cache.HuggingFaceCacheAhead._is_model_cached")
+ def test_cache_model_download_with_custom_revision(
+ self, mock_is_cached, mock_snapshot_download
+ ):
+ """Test downloading specific model revision."""
+ mock_is_cached.return_value = False
+ mock_snapshot_download.return_value = "/cache/path/gpt2-v2"
+
+ result = self.hf_cache.cache_model_download("gpt2", revision="v2.0")
+
+ assert result.success is True
+ mock_snapshot_download.assert_called_once_with(
+ repo_id="gpt2", revision="v2.0", token=None
+ )
diff --git a/tests/unit/test_remote_executor.py b/tests/unit/test_remote_executor.py
index 632423b..9784df7 100644
--- a/tests/unit/test_remote_executor.py
+++ b/tests/unit/test_remote_executor.py
@@ -30,13 +30,11 @@ def encode_kwargs(self, **kwargs):
def test_executor_composition_initialization(self):
"""Test RemoteExecutor initializes all component dependencies correctly."""
# Test that all components are created
- assert hasattr(self.executor, "workspace_manager")
assert hasattr(self.executor, "dependency_installer")
assert hasattr(self.executor, "function_executor")
assert hasattr(self.executor, "class_executor")
# Test that components are properly initialized
- assert self.executor.workspace_manager is not None
assert self.executor.dependency_installer is not None
assert self.executor.function_executor is not None
assert self.executor.class_executor is not None
@@ -52,19 +50,13 @@ async def test_execute_function_orchestration_success(self):
)
# Mock component methods to verify orchestration
- with patch.object(
- self.executor.workspace_manager, "initialize_workspace"
- ) as mock_init:
- with patch.object(
- self.executor.function_executor, "execute"
- ) as mock_execute:
- mock_init.return_value = Mock(success=True, stdout="Workspace ready")
- mock_execute.return_value = Mock(success=True, result="encoded_result")
+ with 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)
+ await self.executor.ExecuteFunction(request)
- # Verify function executor was called
- mock_execute.assert_called_once_with(request)
+ # Verify function executor was called
+ mock_execute.assert_called_once_with(request)
@pytest.mark.asyncio
async def test_execute_function_orchestration_class(self):
@@ -78,20 +70,16 @@ async def test_execute_function_orchestration_class(self):
)
with patch.object(
- self.executor.workspace_manager, "initialize_workspace"
- ) as mock_init:
- with patch.object(
- self.executor.class_executor, "execute_class_method"
- ) as mock_class_execute:
- mock_init.return_value = Mock(success=True, stdout="Workspace ready")
- mock_class_execute.return_value = Mock(
- success=True, result="encoded_result"
- )
+ self.executor.class_executor, "execute_class_method"
+ ) as mock_class_execute:
+ mock_class_execute.return_value = Mock(
+ success=True, result="encoded_result"
+ )
- await self.executor.ExecuteFunction(request)
+ await self.executor.ExecuteFunction(request)
- # Verify class executor was called
- mock_class_execute.assert_called_once_with(request)
+ # Verify class executor was called
+ mock_class_execute.assert_called_once_with(request)
@pytest.mark.asyncio
async def test_execute_function_with_dependencies_orchestration(self):
@@ -106,77 +94,37 @@ async def test_execute_function_with_dependencies_orchestration(self):
)
with patch.object(
- self.executor.workspace_manager, "initialize_workspace"
- ) as mock_init:
+ self.executor.dependency_installer,
+ "install_system_dependencies_async",
+ new_callable=AsyncMock,
+ ) as mock_sys_deps_async:
with patch.object(
self.executor.dependency_installer,
- "install_system_dependencies_async",
+ "install_dependencies_async",
new_callable=AsyncMock,
- ) as mock_sys_deps_async:
- with patch.object(
- self.executor.dependency_installer,
- "install_dependencies_async",
- new_callable=AsyncMock,
- ) as mock_py_deps_async:
- with patch.object(
- self.executor.function_executor, "execute"
- ) as mock_execute:
- # Setup successful responses
- mock_init.return_value = Mock(
- success=True, stdout="Workspace ready"
- )
-
- # Mock async methods with proper FunctionResponse returns
- from remote_execution import FunctionResponse
-
- mock_sys_deps_async.return_value = FunctionResponse(
- success=True, stdout="System deps installed"
- )
- mock_py_deps_async.return_value = FunctionResponse(
- success=True, stdout="Python deps installed"
- )
- mock_execute.return_value = Mock(
- success=True, result="encoded_result"
- )
-
- await self.executor.ExecuteFunction(request)
-
- # Verify all components were called in correct order
- mock_sys_deps_async.assert_called_once_with(["curl"], True)
- mock_py_deps_async.assert_called_once_with(["requests"], True)
- mock_execute.assert_called_once_with(request)
-
- @pytest.mark.asyncio
- async def test_execute_function_workspace_failure_stops_execution(self):
- """Test ExecuteFunction stops on workspace initialization failure."""
- request = FunctionRequest(
- function_name="test_func",
- function_code="def test_func(): return 'test'",
- args=[],
- kwargs={},
- )
-
- # Mock the workspace to have volume so initialization is triggered
- with patch.object(self.executor.workspace_manager, "has_runpod_volume", True):
- with patch.object(
- self.executor.workspace_manager, "initialize_workspace"
- ) as mock_init:
+ ) as mock_py_deps_async:
with patch.object(
self.executor.function_executor, "execute"
) as mock_execute:
- # Setup workspace failure - must return actual response-like object
- workspace_failure = Mock()
- workspace_failure.success = False
- workspace_failure.error = "Workspace init failed"
- workspace_failure.stdout = None # Must be string-like, not Mock
- mock_init.return_value = workspace_failure
+ # Mock async methods with proper FunctionResponse returns
+ from remote_execution import FunctionResponse
+
+ mock_sys_deps_async.return_value = FunctionResponse(
+ success=True, stdout="System deps installed"
+ )
+ mock_py_deps_async.return_value = FunctionResponse(
+ success=True, stdout="Python deps installed"
+ )
+ mock_execute.return_value = Mock(
+ success=True, result="encoded_result"
+ )
- response = await self.executor.ExecuteFunction(request)
+ await self.executor.ExecuteFunction(request)
- # Verify execution stopped and error returned
- assert response.success is False
- assert response.error and "Workspace init failed" in response.error
- mock_execute.assert_not_called()
+ # Verify all components were called in correct order
+ mock_sys_deps_async.assert_called_once_with(["curl"], True)
+ mock_py_deps_async.assert_called_once_with(["requests"], True)
+ mock_execute.assert_called_once_with(request)
@pytest.mark.asyncio
async def test_execute_function_dependency_failure_stops_execution(self):
@@ -190,34 +138,26 @@ async def test_execute_function_dependency_failure_stops_execution(self):
)
with patch.object(
- self.executor.workspace_manager, "initialize_workspace"
- ) as mock_init:
+ self.executor.dependency_installer,
+ "install_dependencies_async",
+ new_callable=AsyncMock,
+ ) as mock_py_deps_async:
with patch.object(
- self.executor.dependency_installer,
- "install_dependencies_async",
- new_callable=AsyncMock,
- ) as mock_py_deps_async:
- with patch.object(
- self.executor.function_executor, "execute"
- ) as mock_execute:
- # Setup successful workspace but failed dependencies
- mock_init.return_value = Mock(
- success=True, stdout="Workspace ready"
- )
-
- # Mock async method with FunctionResponse
- from remote_execution import FunctionResponse
+ self.executor.function_executor, "execute"
+ ) as mock_execute:
+ # Mock async method with FunctionResponse
+ from remote_execution import FunctionResponse
- mock_py_deps_async.return_value = FunctionResponse(
- success=False, error="Package not found"
- )
+ mock_py_deps_async.return_value = FunctionResponse(
+ success=False, error="Package not found"
+ )
- response = await self.executor.ExecuteFunction(request)
+ response = await self.executor.ExecuteFunction(request)
- # Verify execution stopped and error returned
- assert response.success is False
- assert response.error and "Package not found" in response.error
- mock_execute.assert_not_called()
+ # Verify execution stopped and error returned
+ assert response.success is False
+ assert response.error and "Package not found" in response.error
+ mock_execute.assert_not_called()
def test_component_access_methods(self):
"""Test that components can be accessed directly."""
@@ -229,14 +169,6 @@ def test_component_access_methods(self):
self.executor.dependency_installer.install_dependencies(["test"], True)
mock_install.assert_called_once_with(["test"], True)
- # Test workspace manager methods
- with patch.object(
- self.executor.workspace_manager, "initialize_workspace"
- ) as mock_init:
- mock_init.return_value = Mock(success=True)
- self.executor.workspace_manager.initialize_workspace(30)
- mock_init.assert_called_once_with(30) # default timeout
-
# Test function executor methods
request = FunctionRequest(
function_name="test",
@@ -250,17 +182,10 @@ def test_component_access_methods(self):
def test_component_attribute_exposure(self):
"""Test that component attributes are properly exposed."""
# Test that components are properly accessible
- assert hasattr(self.executor, "workspace_manager")
assert hasattr(self.executor, "dependency_installer")
assert hasattr(self.executor, "function_executor")
assert hasattr(self.executor, "class_executor")
- # Test workspace manager attributes through component
- assert hasattr(self.executor.workspace_manager, "has_runpod_volume")
- assert hasattr(self.executor.workspace_manager, "workspace_path")
- assert hasattr(self.executor.workspace_manager, "venv_path")
- assert hasattr(self.executor.workspace_manager, "cache_path")
-
# Test class executor attributes through component
assert hasattr(self.executor.class_executor, "class_instances")
assert hasattr(self.executor.class_executor, "instance_metadata")
diff --git a/tests/unit/test_workspace_manager.py b/tests/unit/test_workspace_manager.py
deleted file mode 100644
index 701ba70..0000000
--- a/tests/unit/test_workspace_manager.py
+++ /dev/null
@@ -1,428 +0,0 @@
-"""Tests for WorkspaceManager component."""
-
-import os
-import threading
-from unittest.mock import patch
-
-from workspace_manager import WorkspaceManager
-from remote_execution import FunctionResponse
-from constants import (
- RUNPOD_VOLUME_PATH,
- DEFAULT_WORKSPACE_PATH,
- VENV_DIR_NAME,
- UV_CACHE_DIR_NAME,
- HF_CACHE_DIR_NAME,
- RUNTIMES_DIR_NAME,
-)
-
-
-class TestEndpointIsolation:
- """Test endpoint-specific workspace isolation."""
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- def test_different_endpoints_get_different_workspaces(
- self, mock_exists, mock_makedirs
- ):
- """Test that different endpoint IDs create separate workspaces."""
- mock_exists.return_value = True
-
- # Test with endpoint-1
- with patch.dict("os.environ", {"RUNPOD_ENDPOINT_ID": "endpoint-1"}):
- manager1 = WorkspaceManager()
- expected_workspace1 = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/endpoint-1"
- assert manager1.workspace_path == expected_workspace1
- assert manager1.venv_path == f"{expected_workspace1}/{VENV_DIR_NAME}"
-
- # Test with endpoint-2
- with patch.dict("os.environ", {"RUNPOD_ENDPOINT_ID": "endpoint-2"}):
- manager2 = WorkspaceManager()
- expected_workspace2 = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/endpoint-2"
- assert manager2.workspace_path == expected_workspace2
- assert manager2.venv_path == f"{expected_workspace2}/{VENV_DIR_NAME}"
-
- # Workspaces should be different
- assert manager1.workspace_path != manager2.workspace_path
- assert manager1.venv_path != manager2.venv_path
-
- # But caches should be shared
- assert (
- manager1.cache_path
- == manager2.cache_path
- == f"{RUNPOD_VOLUME_PATH}/{UV_CACHE_DIR_NAME}"
- )
- assert (
- manager1.hf_cache_path
- == manager2.hf_cache_path
- == f"{RUNPOD_VOLUME_PATH}/{HF_CACHE_DIR_NAME}"
- )
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- def test_default_endpoint_id_when_not_set(self, mock_exists, mock_makedirs):
- """Test that 'default' is used when RUNPOD_ENDPOINT_ID is not set."""
- mock_exists.return_value = True
-
- with patch.dict("os.environ", {}, clear=True):
- manager = WorkspaceManager()
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- assert manager.workspace_path == expected_workspace
- assert manager.endpoint_id == "default"
-
-
-class TestVolumeDetection:
- """Test detection of RunPod volume availability."""
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- def test_detects_runpod_volume_exists(self, mock_exists, mock_makedirs):
- """Test that manager detects when /runpod-volume exists."""
- mock_exists.return_value = True
-
- manager = WorkspaceManager()
-
- assert manager.has_runpod_volume is True
- # Workspace is now endpoint-specific (using 'default' when RUNPOD_ENDPOINT_ID not set)
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- assert manager.workspace_path == expected_workspace
- assert manager.venv_path == f"{expected_workspace}/{VENV_DIR_NAME}"
- # Caches are shared at volume root
- assert manager.cache_path == f"{RUNPOD_VOLUME_PATH}/{UV_CACHE_DIR_NAME}"
- assert manager.hf_cache_path == f"{RUNPOD_VOLUME_PATH}/{HF_CACHE_DIR_NAME}"
- mock_exists.assert_called_with(RUNPOD_VOLUME_PATH)
-
- @patch("os.path.exists")
- def test_detects_runpod_volume_missing(self, mock_exists):
- """Test fallback behavior when no volume is present."""
- mock_exists.return_value = False
-
- manager = WorkspaceManager()
-
- assert manager.has_runpod_volume is False
- assert manager.workspace_path == DEFAULT_WORKSPACE_PATH
- assert manager.venv_path is None
- assert manager.cache_path is None
- assert manager.hf_cache_path is None
-
-
-class TestWorkspaceInitialization:
- """Test workspace initialization functionality."""
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- def test_workspace_initialization_creates_venv(self, mock_exists, mock_makedirs):
- """Test that workspace initialization creates virtual environment."""
- mock_exists.side_effect = lambda path: path == RUNPOD_VOLUME_PATH
-
- manager = WorkspaceManager()
-
- with patch.object(manager, "_create_virtual_environment") as mock_create:
- mock_create.return_value = FunctionResponse(
- success=True, stdout="venv created"
- )
- with (
- patch("os.makedirs"),
- patch("builtins.open"),
- patch("fcntl.flock"),
- patch("os.remove"),
- ):
- result = manager.initialize_workspace()
-
- assert result.success is True
- mock_create.assert_called_once()
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- def test_workspace_already_initialized_skips_creation(
- self, mock_exists, mock_validate, mock_makedirs
- ):
- """Test that existing workspace is not re-initialized."""
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_workspace,
- f"{expected_workspace}/{VENV_DIR_NAME}",
- ]
- mock_validate.return_value = FunctionResponse(success=True, stdout="Valid venv")
-
- manager = WorkspaceManager()
-
- result = manager.initialize_workspace()
-
- assert result.success is True
- assert "already initialized" in result.stdout
-
- @patch("os.path.exists")
- def test_no_volume_returns_success(self, mock_exists):
- """Test that no volume available returns success."""
- mock_exists.return_value = False
-
- manager = WorkspaceManager()
- result = manager.initialize_workspace()
-
- assert result.success is True
- assert "No volume available" in result.stdout
-
-
-class TestConcurrencySafety:
- """Test concurrent workspace initialization safety."""
-
- @patch("os.path.exists")
- @patch("os.makedirs")
- @patch("builtins.open")
- @patch("fcntl.flock")
- @patch("os.remove")
- def test_concurrent_workspace_initialization(
- self, mock_remove, mock_flock, mock_open, mock_makedirs, mock_exists
- ):
- """Test that concurrent initialization is handled safely."""
- mock_exists.side_effect = lambda path: path == RUNPOD_VOLUME_PATH
-
- results = []
-
- def init_workspace():
- manager = WorkspaceManager()
- with patch.object(manager, "_create_virtual_environment") as mock_create:
- mock_create.return_value = FunctionResponse(
- success=True, stdout="venv created"
- )
- result = manager.initialize_workspace()
- results.append(result)
-
- # Start multiple threads trying to initialize
- threads = [threading.Thread(target=init_workspace) for _ in range(3)]
- for thread in threads:
- thread.start()
- for thread in threads:
- thread.join()
-
- # All should succeed
- assert len([r for r in results if r.success]) >= 1
-
-
-class TestEnvironmentConfiguration:
- """Test environment variable configuration."""
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- def test_configure_volume_environment(self, mock_exists, mock_makedirs):
- """Test environment variables are set for volume usage."""
- mock_exists.return_value = True
-
- with patch.dict("os.environ", {}, clear=True):
- WorkspaceManager()
-
- # UV cache is shared at volume root
- assert (
- os.environ.get("UV_CACHE_DIR")
- == f"{RUNPOD_VOLUME_PATH}/{UV_CACHE_DIR_NAME}"
- )
- # HF cache is shared at volume root - HF manages subdirectories automatically
- assert (
- os.environ.get("HF_HOME") == f"{RUNPOD_VOLUME_PATH}/{HF_CACHE_DIR_NAME}"
- )
- # HF automatically creates and manages subdirectories, no need to set specific paths
- assert "TRANSFORMERS_CACHE" not in os.environ
- assert "HF_DATASETS_CACHE" not in os.environ
- assert "HUGGINGFACE_HUB_CACHE" not in os.environ
- # Virtual environment is endpoint-specific
- expected_venv = (
- f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default/{VENV_DIR_NAME}"
- )
- assert os.environ.get("VIRTUAL_ENV") == expected_venv
- assert f"{expected_venv}/bin" in os.environ.get("PATH", "")
-
- @patch("os.path.exists")
- def test_no_environment_changes_without_volume(self, mock_exists):
- """Test no environment changes when no volume present."""
- mock_exists.return_value = False
-
- with patch.dict("os.environ", {}, clear=True):
- WorkspaceManager()
-
- assert "UV_CACHE_DIR" not in os.environ
- assert "HF_HOME" not in os.environ
- assert "TRANSFORMERS_CACHE" not in os.environ
- assert "HF_DATASETS_CACHE" not in os.environ
- assert "HUGGINGFACE_HUB_CACHE" not in os.environ
- assert "VIRTUAL_ENV" not in os.environ
-
-
-class TestWorkspaceOperations:
- """Test workspace directory operations."""
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- @patch("os.getcwd")
- @patch("os.chdir")
- def test_change_to_workspace(
- self, mock_chdir, mock_getcwd, mock_exists, mock_makedirs
- ):
- """Test changing to workspace directory."""
- mock_exists.return_value = True
- mock_getcwd.return_value = "/original"
-
- manager = WorkspaceManager()
- original_cwd = manager.change_to_workspace()
-
- assert original_cwd == "/original"
- # Now changes to endpoint-specific workspace
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- mock_chdir.assert_called_once_with(expected_workspace)
-
- @patch("os.path.exists")
- def test_change_to_workspace_no_volume(self, mock_exists):
- """Test no directory change when no volume."""
- mock_exists.return_value = False
-
- manager = WorkspaceManager()
- original_cwd = manager.change_to_workspace()
-
- assert original_cwd is None
-
- @patch("os.makedirs")
- @patch("workspace_manager.WorkspaceManager._validate_virtual_environment")
- @patch("os.path.exists")
- @patch("glob.glob")
- def test_setup_python_path(
- self, mock_glob, mock_exists, mock_validate, mock_makedirs
- ):
- """Test Python path setup with virtual environment."""
- expected_workspace = f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/default"
- expected_venv = f"{expected_workspace}/{VENV_DIR_NAME}"
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- expected_venv,
- ]
- mock_glob.return_value = [f"{expected_venv}/lib/python3.12/site-packages"]
- mock_validate.return_value = FunctionResponse(success=True, stdout="Valid venv")
-
- manager = WorkspaceManager()
-
- import sys
-
- original_path = sys.path.copy()
- try:
- manager.setup_python_path()
- assert f"{expected_venv}/lib/python3.12/site-packages" in sys.path
- finally:
- sys.path = original_path
-
-
-class TestAppVenvSymlink:
- """Tests for /app/.venv symlink functionality."""
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- @patch("os.symlink")
- @patch("shutil.rmtree")
- @patch("os.path.isdir")
- @patch("os.path.islink")
- def test_create_app_venv_symlink_removes_existing_dir(
- self,
- mock_islink,
- mock_isdir,
- mock_rmtree,
- mock_symlink,
- mock_exists,
- mock_makedirs,
- ):
- """Test that existing /app/.venv directory is removed before creating symlink."""
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- "/app/.venv",
- ]
- mock_islink.return_value = False
- mock_isdir.return_value = True
-
- with patch.dict(os.environ, {"RUNPOD_ENDPOINT_ID": "test-endpoint"}):
- manager = WorkspaceManager()
- manager._create_app_venv_symlink()
-
- expected_venv = (
- f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/test-endpoint/{VENV_DIR_NAME}"
- )
- mock_rmtree.assert_called_once_with("/app/.venv")
- mock_symlink.assert_called_once_with(expected_venv, "/app/.venv")
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- @patch("os.symlink")
- @patch("os.remove")
- @patch("os.path.islink")
- def test_create_app_venv_symlink_removes_existing_symlink(
- self, mock_islink, mock_remove, mock_symlink, mock_exists, mock_makedirs
- ):
- """Test that existing /app/.venv symlink is removed before creating new one."""
- mock_exists.side_effect = lambda path: path in [
- RUNPOD_VOLUME_PATH,
- "/app/.venv",
- ]
- mock_islink.return_value = True
-
- with patch.dict(os.environ, {"RUNPOD_ENDPOINT_ID": "test-endpoint"}):
- manager = WorkspaceManager()
- manager._create_app_venv_symlink()
-
- expected_venv = (
- f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/test-endpoint/{VENV_DIR_NAME}"
- )
- mock_remove.assert_called_once_with("/app/.venv")
- mock_symlink.assert_called_once_with(expected_venv, "/app/.venv")
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- @patch("os.symlink")
- def test_create_app_venv_symlink_creates_new_symlink(
- self, mock_symlink, mock_exists, mock_makedirs
- ):
- """Test that symlink is created when /app/.venv doesn't exist."""
- mock_exists.side_effect = lambda path: path == RUNPOD_VOLUME_PATH
-
- with patch.dict(os.environ, {"RUNPOD_ENDPOINT_ID": "test-endpoint"}):
- manager = WorkspaceManager()
- manager._create_app_venv_symlink()
-
- expected_venv = (
- f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/test-endpoint/{VENV_DIR_NAME}"
- )
- mock_symlink.assert_called_once_with(expected_venv, "/app/.venv")
-
- @patch("os.makedirs")
- @patch("os.path.exists")
- @patch("os.path.islink")
- @patch("os.readlink")
- @patch("os.remove")
- def test_remove_app_venv_symlink_removes_matching_symlink(
- self, mock_remove, mock_readlink, mock_islink, mock_exists, mock_makedirs
- ):
- """Test that /app/.venv symlink is removed when it points to our venv."""
- mock_exists.return_value = True
- mock_islink.return_value = True
- expected_venv = (
- f"{RUNPOD_VOLUME_PATH}/{RUNTIMES_DIR_NAME}/test-endpoint/{VENV_DIR_NAME}"
- )
- mock_readlink.return_value = expected_venv
-
- with patch.dict(os.environ, {"RUNPOD_ENDPOINT_ID": "test-endpoint"}):
- manager = WorkspaceManager()
- manager._remove_app_venv_symlink()
-
- mock_remove.assert_called_once_with("/app/.venv")
-
- @patch("os.path.islink")
- @patch("os.readlink")
- @patch("os.remove")
- def test_remove_app_venv_symlink_skips_different_target(
- self, mock_remove, mock_readlink, mock_islink
- ):
- """Test that /app/.venv symlink is not removed when it points to a different venv."""
- mock_islink.return_value = True
- mock_readlink.return_value = "/different/venv/path"
-
- with patch.dict(os.environ, {"RUNPOD_ENDPOINT_ID": "test-endpoint"}):
- manager = WorkspaceManager()
- manager._remove_app_venv_symlink()
-
- mock_remove.assert_not_called()
diff --git a/uv.lock b/uv.lock
index 9ef9820..9dbdbab 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2610,7 +2610,7 @@ wheels = [
[[package]]
name = "worker-tetra"
-version = "0.5.0"
+version = "0.6.0"
source = { virtual = "." }
dependencies = [
{ name = "cloudpickle" },