feat: Add download acceleration for dependencies & hugging face

CLAUDE.md

@@ -261,3 +261,6 @@ Configure these in GitHub repository settings:
 
 ### Docker Guidelines
 - Docker container should never refer to src/
+
+- Always run `make quality-check` before pronouncing you have finished your work
+- Always use `git mv` when moving existing files around

Dockerfile

@@ -4,13 +4,13 @@ WORKDIR /app
 
 # Install build tools and uv (only in builder stage)
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    git curl build-essential ca-certificates \
+    git curl build-essential ca-certificates aria2 \
  && curl -LsSf https://astral.sh/uv/install.sh | sh \
  && cp ~/.local/bin/uv /usr/local/bin/uv \
  && chmod +x /usr/local/bin/uv
 
 # Copy app code and install dependencies
-COPY README.md src/* pyproject.toml uv.lock test_*.json test-handler.sh ./
+COPY README.md src/* pyproject.toml uv.lock ./
 RUN uv sync
 
 
@@ -19,11 +19,12 @@ FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime
 
 WORKDIR /app
 
+# Install aria2 for download acceleration in runtime stage
+RUN apt-get update && apt-get install -y --no-install-recommends aria2 \
+ && rm -rf /var/lib/apt/lists/*
+
 # Copy app and uv binary from builder
 COPY --from=builder /app /app
 COPY --from=builder /usr/local/bin/uv /usr/local/bin/uv
 
-# Clean up any unnecessary system tools
-RUN rm -rf /var/lib/apt/lists/*
-
 CMD ["uv", "run", "handler.py"]

Dockerfile-cpu

@@ -5,13 +5,13 @@ WORKDIR /app
 
 # Install minimal OS deps and uv
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    curl ca-certificates git build-essential \
+    curl ca-certificates git build-essential aria2 \
  && curl -LsSf https://astral.sh/uv/install.sh | sh \
  && cp ~/.local/bin/uv /usr/local/bin/uv \
  && chmod +x /usr/local/bin/uv
 
 # Copy app files and install deps
-COPY README.md src/* pyproject.toml uv.lock test_*.json test-handler.sh ./
+COPY README.md src/* pyproject.toml uv.lock ./
 RUN uv sync
 
 # Stage 2: Runtime stage
@@ -21,7 +21,7 @@ WORKDIR /app
 
 # Install runtime dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    curl ca-certificates \
+    curl ca-certificates aria2 \
  && apt-get clean \
  && rm -rf /var/lib/apt/lists/*
 

Makefile

@@ -68,7 +68,7 @@ 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
-	./test-handler.sh
+	cd src && ./test-handler.sh
 
 # Smoke Tests (local on Mac OS)
 
@@ -97,7 +97,7 @@ format-check: # Check code formatting
 
 # Type checking
 typecheck: # Check types with mypy
-	uv run mypy .
+	uv run mypy src/
 
 # Quality gates (used in CI)
-quality-check: format-check lint typecheck test-coverage
+quality-check: format-check lint typecheck test-coverage test-handler

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.9,<3.13"
 dependencies = [
     "cloudpickle>=3.1.1",
     "pydantic>=2.11.4",
+    "requests>=2.25.0",
     "runpod",
 ]
 
@@ -18,6 +19,7 @@ dev = [
     "pytest-asyncio>=0.24.0",
     "ruff>=0.8.0",
     "mypy>=1.11.0",
+    "types-requests>=2.25.0",
 ]
 
 [tool.pytest.ini_options]
@@ -48,40 +50,35 @@ filterwarnings = [
     "ignore::pytest.PytestUnknownMarkWarning"
 ]
 
-[tool.ruff]
-# Exclude tetra-rp directory since it's a separate repository
-exclude = [
-    "tetra-rp/",
-]
-
 [tool.mypy]
-# Basic configuration
 python_version = "3.9"
-warn_return_any = true
-warn_unused_configs = true
-disallow_untyped_defs = false  # Start lenient, can be stricter later
-disallow_incomplete_defs = false
-check_untyped_defs = true
-
-# Import discovery
-mypy_path = "src"
+mypy_path = ["src"]
+explicit_package_bases = true
 namespace_packages = true
-
-# Error output
+check_untyped_defs = true
+disallow_any_generics = true
+disallow_untyped_defs = false
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_return_any = true
+strict_optional = true
 show_error_codes = true
 show_column_numbers = true
 pretty = true
-
-# Exclude directories
 exclude = [
     "tetra-rp/",
-    "tests/",  # Start by excluding tests, can add later
 ]
 
-# Per-module options
 [[tool.mypy.overrides]]
 module = [
-    "runpod.*",
-    "cloudpickle.*",
+    "cloudpickle",
+    "runpod",
+    "transformers",
 ]
 ignore_missing_imports = true
+
+[tool.ruff]
+# Exclude tetra-rp directory since it's a separate repository
+exclude = [
+    "tetra-rp/",
+]

src/__init__.py

@@ -0,0 +1 @@
+"""Worker Tetra package."""

src/class_executor.py

@@ -18,7 +18,7 @@ def __init__(self, workspace_manager):
         super().__init__(workspace_manager)
         # Instance registry for persistent class instances
         self.class_instances: Dict[str, Any] = {}
-        self.instance_metadata: Dict[str, Dict] = {}
+        self.instance_metadata: Dict[str, Dict[str, Any]] = {}
 
     def execute(self, request: FunctionRequest) -> FunctionResponse:
         """Execute class method - required by BaseExecutor interface."""

src/constants.py

@@ -20,3 +20,19 @@
 
 RUNTIMES_DIR_NAME = "runtimes"
 """Name of the runtimes directory containing per-endpoint workspaces."""
+
+# Download Acceleration Settings
+DEFAULT_DOWNLOAD_CONNECTIONS = 8
+"""Default number of parallel connections for accelerated downloads."""
+
+MIN_SIZE_FOR_ACCELERATION_MB = 10
+"""Minimum file size in MB to trigger download acceleration."""
+
+MAX_DOWNLOAD_CONNECTIONS = 16
+"""Maximum number of parallel connections for downloads."""
+
+DOWNLOAD_TIMEOUT_SECONDS = 600
+"""Default timeout for download operations in seconds."""
+
+DOWNLOAD_PROGRESS_UPDATE_INTERVAL = 1.0
+"""Interval in seconds for download progress updates."""

src/dependency_installer.py

@@ -5,6 +5,7 @@
 from typing import List, Dict
 
 from remote_execution import FunctionResponse
+from download_accelerator import DownloadAccelerator
 
 
 class DependencyInstaller:
@@ -13,6 +14,7 @@ class DependencyInstaller:
     def __init__(self, workspace_manager):
         self.workspace_manager = workspace_manager
         self.logger = logging.getLogger(__name__)
+        self.download_accelerator = DownloadAccelerator(workspace_manager)
 
     def install_system_dependencies(self, packages: List[str]) -> FunctionResponse:
         """
@@ -72,12 +74,16 @@ def install_system_dependencies(self, packages: List[str]) -> FunctionResponse:
                 error=f"Exception during system package installation: {e}",
             )
 
-    def install_dependencies(self, packages: List[str]) -> FunctionResponse:
+    def install_dependencies(
+        self, packages: List[str], accelerate_downloads: bool = True
+    ) -> FunctionResponse:
         """
         Install Python packages using uv with differential installation support.
+        Uses accelerated downloads for large packages when beneficial.
 
         Args:
             packages: List of package names or package specifications
+            accelerate_downloads: Whether to use accelerated downloads for large packages
         Returns:
             FunctionResponse: Object indicating success or failure with details
         """
@@ -117,6 +123,132 @@ def install_dependencies(self, packages: List[str]) -> FunctionResponse:
 
             packages = packages_to_install
 
+        # Check if we should use accelerated downloads for large packages
+        large_packages = self._identify_large_packages(packages)
+
+        if (
+            accelerate_downloads
+            and large_packages
+            and self.download_accelerator.aria2_downloader.aria2c_available
+        ):
+            self.logger.info(
+                f"Using accelerated downloads for large packages: {large_packages}"
+            )
+            return self._install_with_acceleration(packages, large_packages)
+        else:
+            return self._install_standard(packages)
+
+    def _identify_large_packages(self, packages: List[str]) -> List[str]:
+        """
+        Identify packages that are likely to be large and benefit from acceleration.
+
+        Args:
+            packages: List of package specifications
+
+        Returns:
+            List of package names that are likely large
+        """
+        # Known large packages that benefit from acceleration
+        large_package_patterns = [
+            "torch",
+            "pytorch",
+            "tensorflow",
+            "tf-nightly",
+            "transformers",
+            "diffusers",
+            "datasets",
+            "numpy",
+            "scipy",
+            "pandas",
+            "matplotlib",
+            "opencv",
+            "cv2",
+            "pillow",
+            "scikit-learn",
+            "huggingface-hub",
+            "safetensors",
+        ]
+
+        large_packages = []
+        for package in packages:
+            package_name = package.split("==")[0].split(">=")[0].split("<=")[0].lower()
+            if any(pattern in package_name for pattern in large_package_patterns):
+                large_packages.append(package)
+
+        return large_packages
+
+    def _install_with_acceleration(
+        self, packages: List[str], large_packages: List[str]
+    ) -> FunctionResponse:
+        """
+        Install packages with acceleration for large ones.
+
+        Args:
+            packages: All packages to install
+            large_packages: Packages that should use acceleration
+
+        Returns:
+            FunctionResponse with installation result
+        """
+        try:
+            # Prepare environment for virtual environment usage
+            env = os.environ.copy()
+            if (
+                self.workspace_manager.has_runpod_volume
+                and self.workspace_manager.venv_path
+            ):
+                env["VIRTUAL_ENV"] = self.workspace_manager.venv_path
+
+            # For now, we'll enhance UV's download behavior by setting optimal configurations
+            # UV internally uses efficient downloaders, but we can optimize the environment
+
+            # Set aria2c as a potential downloader for UV if it supports it
+            env["UV_CONCURRENT_DOWNLOADS"] = "8"  # Increase concurrent downloads
+
+            self.logger.info("Installing with optimized concurrent downloads")
+
+            # Use uv pip to install the packages with optimizations
+            command = ["uv", "pip", "install", "--no-cache-dir"] + packages
+            process = subprocess.Popen(
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                env=env,
+            )
+
+            stdout, stderr = process.communicate()
+            importlib.invalidate_caches()
+
+            if process.returncode != 0:
+                return FunctionResponse(
+                    success=False,
+                    error="Error installing packages with acceleration",
+                    stdout=stderr.decode(),
+                )
+            else:
+                self.logger.info(
+                    f"Successfully installed packages with acceleration: {packages}"
+                )
+                return FunctionResponse(
+                    success=True,
+                    stdout=f"Installed with acceleration: {stdout.decode()}",
+                )
+        except Exception as e:
+            self.logger.warning(
+                f"Accelerated installation failed, falling back to standard: {e}"
+            )
+            return self._install_standard(packages)
+
+    def _install_standard(self, packages: List[str]) -> FunctionResponse:
+        """
+        Install packages using standard UV method.
+
+        Args:
+            packages: Packages to install
+
+        Returns:
+            FunctionResponse with installation result
+        """
         try:
             # Prepare environment for virtual environment usage
             env = os.environ.copy()