feat: Workspace environment persisted in the network volume (#10)

* chore: make clean and make upgrade

* ci: make test|test-unit|test-integration|test-coverage|test-fast

* ci: make lint|lint-fix|format|format-check

* build: updated dependencies to latest

* chore: update tetra-rp

* ci: simplified CI workflow + manual deploy action

* test: sample unit and integration tests

Coverage fails when at below 50%. We can increase this later.

* chore: make lint-fix

* chore: make dev as default when running `make`

* chore: `make` shows help list of all available commands

* chore: forgot to include python 3.10 in the tests

* chore: ruff should exclude tetra-rp sub-module

* ci: test coverage fails at <35% (temp)

* test: auto-mark unit and integration tests appropriately

* ci: make typecheck with mypy as quality-check

* chore: reset uv.lock

* test: add comprehensive test foundation for volume workspace TDD

Establish testing infrastructure and protocol validation tests. Create shared fixtures and validates the FunctionRequest/FunctionResponse data models that will be extended for volume workspace functionality.

* test: volume detection and workspace initialization

Tests volume detection, virtual environment creation, file-based locking for concurrency, and timeout handling mechanisms.

* test: volume-aware execution context

Validates that functions execute in volume workspace, can access persistent packages, and fallback gracefully when volume is unavailable.

* test: integration tests for end-to-end workflows

Tests complete request workflows, concurrent access safety, mixed execution scenarios, and realistic error handling patterns.

* feat: implement core volume workspace detection and initialization

Adds volume detection logic, workspace initialization with file-based locking, virtual environment creation, and timeout handling to make some tests pass.

* feat: add differential dependency management for volume persistence

Implement smart dependency installation that only installs missing packages. Optimizes performance by leveraging persistent volume storage and avoiding redundant package installations.

* feat: implement volume-aware execution context and environment configuration

Enable functions to execute in volume workspace with access to persistent packages. Configures Python path, environment variables, and UV cache to utilize volume storage effectively.

* test: enhance existing test suites for volume workspace compatibility

Update existing tests to work with new volume workspace functionality. Ensures backward compatibility and validates that all existing functionality continues to work with the new volume-aware implementation.

* build: make update + cleanup

* chore: update tetra-rp submodule

* chore: updated CLAUDE.md

* chore: update CLAUDE.md

* chore: break out the large handler into modules also their tests

* refactor: constants for the magic values

* refactor: move Python files to src/ directory

- Move all Python modules to src/ for better organization
- Update Docker files to copy from src/ directory
- Update pyproject.toml with src/ in pythonpath
- Update Makefile to copy remote_execution.py to src/
- All tests pass with new structure

* feat: add virtual environment validation and repair

- Add _validate_virtual_environment() method to WorkspaceManager with symlink chain resolution using os.path.realpath()
- Add _remove_broken_virtual_environment() cleanup method
- Enhance initialize_workspace() with validation checks and automatic repair
- Add validation calls in setup_python_path() and dependency installer
- Update Docker files to work with src/ directory structure
- Update tests to mock new validation methods
- Fix pyproject.toml pythonpath configuration for tests

This resolves broken virtual environment symlinks when different endpoints
create venvs with different Python interpreter paths on shared volumes.

* feat: implement per-endpoint workspace isolation for shared volumes

- Add RUNPOD_ENDPOINT_ID environment variable support for endpoint isolation
- Workspace paths now: /runpod-volume/runtimes/{endpoint_id}
- Shared UV cache at volume root: /runpod-volume/.uv-cache
- Add RUNTIMES_DIR_NAME constant for endpoint workspace organization
- Update WorkspaceManager to create endpoint-specific workspace paths
- Add comprehensive endpoint isolation tests
- Update integration tests for new workspace structure
- Resolve merge conflicts from virtual environment validation features

* feat: configure Hugging Face cache to use persistent volume storage

- Add HF_CACHE_DIR_NAME constant for .hf-cache directory
- Implement _configure_huggingface_cache() method in WorkspaceManager
- Set HF environment variables (HF_HOME, TRANSFORMERS_CACHE, etc.) to use volume paths
- Update unit and integration tests to mock os.makedirs calls
- Fix "No space left on device" errors when downloading HF models

* fix: Test local handler execution should test all test_*.json

* feat: add comprehensive handler testing with make test-handler command

- Add make test-handler command that tests all test_*.json files locally
- Update CI to use make test-handler for consistency between local and CI testing
- Ensure local development environment matches CI testing exactly
- Remove code duplication between Makefile and CI configuration
- Support cross-platform testing (handles timeout command availability)
- Update CLAUDE.md documentation with new testing commands

* feat: enhance concurrency safety implementation for workspace initialization

- Add configurable timeout constants (WORKSPACE_INIT_TIMEOUT, WORKSPACE_LOCK_POLL_INTERVAL)
- Implement atomic lock file operations with proper file descriptor management
- Enhance lock file cleanup with comprehensive error handling
- Add workspace directory validation before lock acquisition
- Fix race condition in workspace functionality checks by making them atomic
- Add comprehensive timeout and edge case tests for concurrency scenarios
- Improve error messages and fallback behavior for various failure modes
- Maintain backward compatibility while significantly improving reliability

* Revert "feat: enhance concurrency safety implementation for workspace initialization"

This reverts commit a411bfe.

* refactor: introduce BaseExecutor base class for executor inheritance

- Add BaseExecutor abstract base class with common functionality
- Update FunctionExecutor and ClassExecutor to inherit from BaseExecutor
- Standardize execution environment setup via _setup_execution_environment
- Update ClassExecutor constructor to accept workspace_manager parameter
- Fix ClassExecutor tests to mock workspace_manager dependency

* refactor: replace print statements with structured logging

- Add logging support to DependencyInstaller and WorkspaceManager
- Replace print calls with appropriate log levels (info, warning, error)
- Improve debugging and monitoring capabilities

* chore: update submodule and dependencies

* fix: vLLM ignored the volume runtime

The fix will resolve the vLLM subprocess errors encountered while maintaining full compatibility with existing functionality. When deployed to RunPod with volumes, libraries like vLLM that hardcode /app/.venv paths will seamlessly use the volume's virtual environment.

* fix: resolve vLLM subprocess failures with volume virtual environments

- Add symlink from /app/.venv to volume venv to handle hardcoded paths
- Configure PYTHONPATH environment variable for subprocess compatibility
- Ensure libraries like vLLM can spawn subprocesses that find installed packages
- Add comprehensive test coverage for symlink functionality
- Maintain backward compatibility when no volume is present

Author: deanq

.dockerignore

@@ -17,3 +17,4 @@ build/
 *.egg-info/
 .eggs/
 *.egg
+*.pkl

.github/workflows/ci.yml

@@ -100,16 +100,12 @@ jobs:
 
       - name: Test local handler execution
         run: |
-          timeout 30s uv run python handler.py || exit_code=$?
+          timeout 30s make test-handler || exit_code=$?
           if [ ${exit_code:-0} -eq 124 ]; then
-            echo "Handler execution timed out as expected (running in serverless mode)"
-            exit 0
-          elif [ ${exit_code:-0} -eq 0 ]; then
-            echo "Handler executed successfully"
-            exit 0
+            echo "Handler testing timed out"
+            exit 1
           else
-            echo "Handler execution failed with exit code $exit_code"
-            exit $exit_code
+            exit ${exit_code:-0}
           fi
 
   release:

.gitignore

@@ -159,3 +159,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 remote_execution.py
+*.pkl

CLAUDE.md

@@ -4,18 +4,20 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 
 ## Project Overview
 
-This is `worker-tetra`, a RunPod Serverless worker template that provides dynamic GPU provisioning for ML workloads with transparent execution. The project consists of two main components:
+This is `worker-tetra`, a RunPod Serverless worker template that provides dynamic GPU provisioning for ML workloads with transparent execution and persistent workspace management. The project consists of two main components:
 
-1. **RunPod Worker Handler** (`handler.py`) - A serverless function that executes remote Python functions with dependency management
+1. **RunPod Worker Handler** (`handler.py`) - A serverless function that executes remote Python functions with dependency management and persistent volume workspace support
 2. **Tetra SDK** (`tetra-rp/` submodule) - Python library for distributed inference and serving of ML models
 
 ## Architecture
 
 ### Core Components
 
 - **`handler.py`**: Main RunPod serverless handler implementing `RemoteExecutor` class
-  - Executes arbitrary Python functions remotely
-  - Handles dynamic installation of Python and system dependencies
+  - Executes arbitrary Python functions remotely with persistent workspace support
+  - Handles dynamic installation of Python and system dependencies with differential updates
+  - Manages `/runpod-volume` workspace with virtual environment and shared package cache
+  - Implements concurrency-safe workspace initialization with file-based locking
   - Serializes/deserializes function arguments and results using cloudpickle
   - Captures stdout, stderr, and logs from remote execution
 
@@ -32,16 +34,29 @@ 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. **Dynamic Dependency Management**: Dependencies specified in decorators are installed at runtime
-3. **Serialization**: Uses cloudpickle + base64 encoding for function arguments and results
-4. **Resource Configuration**: `LiveServerless` objects define GPU requirements, scaling, and worker configuration
+2. **Persistent Workspace Management**: `/runpod-volume` provides persistent storage for packages and execution state
+3. **Dynamic Dependency Management**: Dependencies specified in decorators are installed at runtime with differential updates
+4. **Concurrency Safety**: File-based locking ensures safe workspace initialization across multiple workers
+5. **Serialization**: Uses cloudpickle + base64 encoding for function arguments and results
+6. **Resource Configuration**: `LiveServerless` objects define GPU requirements, scaling, and worker configuration
 
 ## Development Commands
 
 ### Setup and Dependencies
 ```bash
 make setup                    # Initialize project, sync dependencies, update submodules
-make dev                      # Install all development dependencies
+make dev                      # Install all development dependencies (includes pytest, ruff)
+uv sync                      # Sync production dependencies only
+uv sync --all-groups         # Sync all dependency groups (same as make dev)
+```
+
+### Code Quality
+```bash
+make lint                     # Check code with ruff linter
+make lint-fix                # Auto-fix linting issues
+make format                   # Format code with ruff
+make format-check            # Check if code is properly formatted
+make quality-check           # Run all quality checks (format, lint, test coverage)
 ```
 
 ### Docker Operations
@@ -53,19 +68,74 @@ make build-cpu               # Build CPU-only Docker image
 
 ### Local Testing  
 ```bash
-python handler.py            # Test handler locally with test_input.json
+# Test handler locally with test_input.json
+PYTHONPATH=src RUNPOD_TEST_INPUT="$(cat test_input.json)" uv run python src/handler.py
+
+# Test with other test files
+PYTHONPATH=src RUNPOD_TEST_INPUT="$(cat test_class_input.json)" uv run python src/handler.py
+PYTHONPATH=src RUNPOD_TEST_INPUT="$(cat test_hf_input.json)" uv run python src/handler.py
 ```
 
 ### Submodule Management
 ```bash
 git submodule update --remote --merge    # Update tetra-rp to latest
 ```
 
+## RunPod Volume Workspace
+
+The handler automatically detects and utilizes `/runpod-volume` for persistent workspace management when available:
+
+### Volume Features
+- **Automatic Detection**: Detects `/runpod-volume` presence on container startup
+- **Endpoint Isolation**: Each endpoint gets its own workspace at `/runpod-volume/runtimes/{endpoint_id}`
+- **Virtual Environment**: Creates and manages endpoint-specific `.venv` for persistent package installation
+- **Shared Package Cache**: Uses `/runpod-volume/.uv-cache` for efficient package caching across all endpoints
+- **Hugging Face Cache**: Configures HF model cache at `/runpod-volume/.hf-cache` to prevent storage issues
+- **Differential Installation**: Only installs missing packages, leveraging persistent storage
+- **Concurrency Safety**: File-based locking prevents race conditions during workspace initialization
+- **Graceful Fallback**: Works normally when no volume is present
+
+### Volume Structure
+```
+/runpod-volume/
+├── .uv-cache/               # Shared UV package cache (across all endpoints)
+├── .hf-cache/               # Shared Hugging Face model cache (across all endpoints)
+│   ├── transformers/        # Transformers model cache
+│   ├── datasets/            # HF datasets cache
+│   └── hub/                 # Hugging Face Hub cache
+├── runtimes/                # Per-endpoint runtime environments
+│   ├── endpoint-1/          # Workspace for endpoint-1
+│   │   ├── .venv/           # Endpoint-specific virtual environment
+│   │   ├── .initialization.lock  # Temporary workspace lock file
+│   │   └── <execution workspace>
+│   └── endpoint-2/          # Workspace for endpoint-2
+│       ├── .venv/           # Endpoint-specific virtual environment
+│       ├── .initialization.lock
+│       └── <execution workspace>
+```
+
+### Performance Benefits
+- **Faster Cold Starts**: Pre-installed packages and cached models reduce initialization time
+- **Reduced Network Usage**: Cached packages and models avoid redundant downloads
+- **Persistent State**: Function execution workspace survives across calls
+- **Endpoint Isolation**: Each endpoint maintains independent dependencies and state
+- **Optimized Resource Usage**: Shared caches across multiple endpoints while maintaining isolation
+- **ML Model Efficiency**: Large HF models cached on volume prevent "No space left on device" errors
+
 ## Configuration
 
 ### Environment Variables
 - `RUNPOD_API_KEY`: Required for RunPod Serverless integration
+- `RUNPOD_ENDPOINT_ID`: Used for workspace isolation (automatically set by RunPod)
 - `DEBIAN_FRONTEND=noninteractive`: Set during system package installation
+- `UV_CACHE_DIR`: Automatically set to `/runpod-volume/.uv-cache` when volume detected
+- `VIRTUAL_ENV`: Automatically set to `/runpod-volume/runtimes/{endpoint_id}/.venv` when available
+
+#### Hugging Face Cache Configuration (Auto-configured when volume available)
+- `HF_HOME`: Set to `/runpod-volume/.hf-cache` for main HF cache directory
+- `TRANSFORMERS_CACHE`: Set to `/runpod-volume/.hf-cache/transformers` for model cache
+- `HF_DATASETS_CACHE`: Set to `/runpod-volume/.hf-cache/datasets` for dataset cache
+- `HUGGINGFACE_HUB_CACHE`: Set to `/runpod-volume/.hf-cache/hub` for hub cache
 
 ### Resource Configuration
 Configure GPU resources using `LiveServerless` objects:
@@ -82,18 +152,40 @@ gpu_config = LiveServerless(
 
 ## Testing and Quality
 
-- No formal test suite currently exists
-- Testing is done via `test_input.json` with local handler execution
-- Uses `uv` for fast dependency management
-- Multi-stage Docker builds to minimize image size
+### Testing Commands
+```bash
+make test                     # Run all tests
+make test-unit               # Run unit tests only
+make test-integration        # Run integration tests only
+make test-coverage           # Run tests with coverage report
+make test-fast               # Run tests with fail-fast mode
+make test-handler            # Test handler locally with all test_*.json files (same as CI)
+
+# Test handler locally with specific test files
+PYTHONPATH=src RUNPOD_TEST_INPUT="$(cat test_input.json)" uv run python src/handler.py
+PYTHONPATH=src RUNPOD_TEST_INPUT="$(cat test_class_input.json)" uv run python src/handler.py
+PYTHONPATH=src RUNPOD_TEST_INPUT="$(cat test_hf_input.json)" uv run python src/handler.py
+```
+
+### Testing Framework
+- **pytest** with coverage reporting and async support
+- **Unit tests** (`tests/unit/`): Test individual components in isolation
+- **Integration tests** (`tests/integration/`): Test end-to-end workflows
+- **Coverage target**: 80% minimum, with HTML and XML reports
+- **Test fixtures**: Shared test data and mocks in `tests/conftest.py`
+- **CI Integration**: Tests run on all PRs and before releases/deployments
 
 ## Development Notes
 
 ### Dependency Management
 - Root project uses `uv` with `pyproject.toml`
 - Tetra SDK has separate `pyproject.toml` in `tetra-rp/`
 - System dependencies installed via `apt-get` in containerized environment
-- Python dependencies installed via `uv pip install` at runtime
+- Python dependencies installed via `uv pip install` at runtime with volume persistence
+- **Differential Installation**: Only installs packages missing from persistent volume
+- **Shared Cache**: UV cache in `/runpod-volume/.uv-cache` optimizes package downloads
+- **Virtual Environment**: Persistent `.venv` in volume survives across function calls
+- **ML Model Cache**: Hugging Face models cached in `/runpod-volume/.hf-cache` prevent storage issues
 
 ### Error Handling
 - All remote execution wrapped in try/catch with full traceback capture
@@ -103,17 +195,31 @@ gpu_config = LiveServerless(
 ### Security Considerations
 - Functions execute arbitrary Python code in sandboxed containers
 - System package installation requires root privileges in container
+- Volume workspace provides persistent storage but maintains container isolation
+- File-based locking prevents race conditions during concurrent workspace access
 - No secrets should be committed to repository
 - API keys passed via environment variables
 
 ## File Structure Highlights
 
 ```
-├── handler.py              # Main serverless function handler
+├── handler.py              # Main serverless function handler with volume support
 ├── remote_execution.py     # Protocol definitions
+├── PLAN.md                # TDD implementation plan for volume workspace
 ├── Dockerfile             # GPU container definition  
 ├── Dockerfile-cpu         # CPU container definition
-├── test_input.json        # Sample input for local testing
+├── test_input.json        # Basic function execution test
+├── test_class_input.json   # Class execution test
+├── test_hf_input.json      # HuggingFace model download test
+├── tests/                 # Comprehensive test suite
+│   ├── conftest.py        # Shared test fixtures
+│   ├── unit/              # Unit tests for individual components
+│   │   ├── test_runpod_volume_workspace.py  # Volume detection and initialization
+│   │   ├── test_volume_execution.py         # Volume-aware execution
+│   │   └── test_*.py      # Other unit tests
+│   └── integration/       # End-to-end integration tests
+│       ├── test_runpod_volume_integration.py # Volume workflow tests
+│       └── test_*.py      # Other integration tests
 ├── tetra-rp/              # Git submodule - Tetra SDK
 │   ├── src/tetra_rp/
 │   │   ├── client.py      # @remote decorator
@@ -130,9 +236,13 @@ gpu_config = LiveServerless(
 - Docker images are automatically built and pushed to Docker Hub (`runpod/tetra-rp`) on release
 
 ### GitHub Actions Workflows
-- **Release** (`.github/workflows/release-please.yml`): Manages releases and versioning
-- **Docker Images [Prod]** (`.github/workflows/docker-prod.yml`): Builds and pushes Docker images on release
-- **Docker Images [Dev]** (`.github/workflows/docker-dev.yml`): Builds and pushes `:dev` tagged images on main branch pushes
+- **CI/CD** (`.github/workflows/ci.yml`): Single workflow handling tests, linting, releases, and Docker builds
+  - Runs tests and linting on PRs and pushes to main
+  - **Local execution testing**: Automatically tests all `test_*.json` files in root directory to validate handler functionality
+  - Manages releases via `release-please` on main branch
+  - Builds and pushes `:dev` tagged images on main branch pushes
+  - Builds and pushes production images with semantic versioning on releases
+- **Deploy** (`.github/workflows/deploy.yml`): Manual deployment workflow for custom Docker tags and emergency deployments
 
 ### Required Secrets
 Configure these in GitHub repository settings:
@@ -141,5 +251,8 @@ Configure these in GitHub repository settings:
 
 ## Branch Information
 - Main branch: `main`
-- Current working branch: `dean/ae-518-cpu-live-serverless`
-- Submodule tracking: Updates pulled from remote automatically during setup
+- Submodule tracking: Updates pulled from remote automatically during setup
+
+## Development Best Practices
+
+- Always run `make quality-check` before committing changes

Dockerfile

@@ -10,7 +10,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
  && chmod +x /usr/local/bin/uv
 
 # Copy app code and install dependencies
-COPY README.md remote_execution.py handler.py pyproject.toml uv.lock test_input.json ./
+COPY README.md src/* pyproject.toml uv.lock test_*.json ./
 RUN uv sync
 
 

Dockerfile-cpu

@@ -11,7 +11,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
  && chmod +x /usr/local/bin/uv
 
 # Copy app files and install deps
-COPY README.md remote_execution.py handler.py pyproject.toml uv.lock ./
+COPY README.md src/* pyproject.toml uv.lock test_*.json ./
 RUN uv sync
 
 # Stage 2: Runtime stage

Makefile

@@ -32,7 +32,7 @@ upgrade: # Upgrade all dependencies
 setup: dev # Initialize project, sync deps, update submodules
 	git submodule init
 	git submodule update --remote --merge
-	cp tetra-rp/src/tetra_rp/protos/remote_execution.py .
+	cp tetra-rp/src/tetra_rp/protos/remote_execution.py src/
 
 build: setup # Build GPU Docker image (linux/amd64)
 	docker buildx build \
@@ -65,6 +65,30 @@ test-coverage: # Run tests with coverage report
 test-fast: # Run tests with fast-fail mode
 	uv run pytest tests/ -v -x --tb=short
 
+test-handler: # Test handler locally with all test_*.json files
+	@echo "Testing handler with all test_*.json files..."
+	@failed_tests=""; \
+	for test_file in test_*.json; do \
+		if [ ! -f "$$test_file" ]; then \
+			echo "No test_*.json files found"; \
+			exit 1; \
+		fi; \
+		echo "Testing with $$test_file..."; \
+		if env PYTHONPATH=src RUNPOD_TEST_INPUT="$$(cat "$$test_file")" uv run python src/handler.py >/dev/null 2>&1; then \
+			echo "✓ $$test_file: PASSED"; \
+		else \
+			exit_code=$$?; \
+			echo "✗ $$test_file: FAILED (exit code: $$exit_code)"; \
+			failed_tests="$$failed_tests $$test_file"; \
+		fi; \
+	done; \
+	if [ -z "$$failed_tests" ]; then \
+		echo "All tests passed!"; \
+	else \
+		echo "Failed tests:$$failed_tests"; \
+		exit 1; \
+	fi
+
 # Linting commands
 lint: # Check code with ruff
 	uv run ruff check .
@@ -78,5 +102,9 @@ format: # Format code with ruff
 format-check: # Check code formatting
 	uv run ruff format --check .
 
+# Type checking
+typecheck: # Check types with mypy
+	uv run mypy .
+
 # Quality gates (used in CI)
-quality-check: format-check lint test-coverage # Run all quality checks (format, lint, test)
+quality-check: format-check lint typecheck test-coverage