tetra-rp 0.12.0__py3-none-any.whl → 0.14.0__py3-none-any.whl

tetra_rp/cli/utils/skeleton.py

@@ -0,0 +1,347 @@
+"""Project skeleton creation utilities."""
+
+from pathlib import Path
+from typing import List
+
+# Template files content
+MAIN_PY_TEMPLATE = '''"""
+Flash Application - Flash Server
+
+This is the main entry point for your Flash application.
+It runs a FastAPI server that coordinates GPU workers.
+"""
+
+import asyncio
+from fastapi import FastAPI
+from pydantic import BaseModel
+from dotenv import load_dotenv
+
+from workers import ExampleWorker
+
+# Load environment variables
+load_dotenv()
+
+# Create FastAPI app
+app = FastAPI(title="Flash Application")
+
+
+class ProcessRequest(BaseModel):
+    """Request model for processing endpoint."""
+    data: str
+
+
+@app.get("/")
+def home():
+    """Health check endpoint."""
+    return {"status": "ok", "message": "Flash application running"}
+
+
+@app.get("/health")
+def health():
+    """Health check endpoint."""
+    return {"healthy": True}
+
+
+@app.post("/process")
+async def process(request: ProcessRequest):
+    """
+    Process data using GPU worker.
+
+    Example request:
+    {
+        "data": "test input"
+    }
+    """
+    # Instantiate worker
+    worker = ExampleWorker()
+
+    # Call worker's process method
+    result = await worker.process({"input": request.data})
+
+    return result
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="localhost", port=8888)
+'''
+
+WORKER_EXAMPLE_PY_TEMPLATE = '''"""
+Example GPU Worker
+
+This is an example of a GPU worker class that can be deployed
+to RunPod serverless endpoints.
+"""
+
+from tetra_rp import remote, LiveServerless
+
+
+# Configure GPU resource
+config = LiveServerless(
+    name="example_worker",
+    workersMax=3,
+)
+
+
+@remote(config)
+class ExampleWorker:
+    """Example GPU worker for processing tasks."""
+
+    def __init__(self):
+        """Initialize the worker."""
+        print("ExampleWorker initialized")
+
+    def process(self, input_data: dict) -> dict:
+        """
+        Process input data and return result.
+
+        Args:
+            input_data: Dictionary with input parameters
+
+        Returns:
+            Dictionary with processing results
+        """
+        # Your GPU processing logic here
+        result = {
+            "status": "success",
+            "input": input_data,
+            "output": f"Processed: {input_data}"
+        }
+
+        return result
+'''
+
+WORKERS_INIT_PY_TEMPLATE = '''"""GPU Workers package."""
+
+from .example_worker import ExampleWorker
+
+__all__ = ["ExampleWorker"]
+'''
+
+ENV_EXAMPLE_TEMPLATE = """# RunPod API Configuration
+RUNPOD_API_KEY=your_runpod_api_key_here
+
+# Development settings
+DEBUG=false
+LOG_LEVEL=INFO
+"""
+
+REQUIREMENTS_TXT_TEMPLATE = """# Core dependencies for Flash
+tetra-rp>=0.12.0
+fastapi>=0.104.0
+uvicorn[standard]>=0.24.0
+python-dotenv>=1.0.0
+pydantic>=2.0.0
+aiohttp>=3.9.0
+"""
+
+GITIGNORE_TEMPLATE = """# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+venv/
+ENV/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Environment
+.env
+.venv
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Tetra
+.tetra/
+
+# OS
+.DS_Store
+Thumbs.db
+"""
+
+FLASHIGNORE_TEMPLATE = """# Flash build ignores
+# Similar to .gitignore but specifically for flash build command
+
+# Version control
+.git/
+.gitignore
+
+# Python artifacts
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+
+# Virtual environments
+.env
+.venv/
+env/
+venv/
+ENV/
+
+# Build artifacts
+.build/
+.tetra/
+*.tar.gz
+*.egg-info/
+dist/
+build/
+
+# IDE files
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS files
+.DS_Store
+Thumbs.db
+
+# Test files
+tests/
+test_*.py
+*_test.py
+
+# Documentation (optional - comment out to include)
+# docs/
+# *.md
+
+# Logs
+*.log
+"""
+
+README_TEMPLATE = """# {{project_name}}
+
+Flash application with Flash Server and GPU workers.
+
+## Setup
+
+1. Activate the conda environment:
+```bash
+conda activate {{project_name}}
+```
+
+2. Configure your RunPod API key:
+```bash
+cp .env.example .env
+# Edit .env and add your RUNPOD_API_KEY
+```
+
+3. Run the development server:
+```bash
+flash run
+```
+
+## Project Structure
+
+```
+{{project_name}}/
+├── main.py              # Flash Server (FastAPI)
+├── workers/             # GPU workers
+│   ├── __init__.py
+│   └── example_worker.py
+├── .env.example         # Environment variables template
+├── requirements.txt     # Python dependencies
+└── README.md           # This file
+```
+
+## Development
+
+The Flash Server runs on `localhost:8888` and coordinates GPU workers.
+
+### Adding New Workers
+
+1. Create a new file in `workers/` directory
+2. Define a class with `@remote` decorator
+3. Import it in `workers/__init__.py`
+4. Use it in `main.py`
+
+Example:
+```python
+from tetra_rp import remote, LiveServerless
+
+config = LiveServerless(name="my_worker", workersMax=3)
+
+@remote(config)
+class MyWorker:
+    def process(self, data):
+        return {{"result": f"Processed: {{data}}"}}
+```
+
+## Deployment
+
+Deploy to production:
+```bash
+flash deploy send production
+```
+
+## Documentation
+
+- [Flash CLI Docs](./docs/)
+- [Tetra Documentation](https://docs.tetra.dev)
+"""
+
+
+def create_project_skeleton(project_dir: Path, force: bool = False) -> List[str]:
+    """
+    Create Flash project skeleton.
+
+    Args:
+        project_dir: Project directory path
+        force: Overwrite existing files
+
+    Returns:
+        List of created file paths
+    """
+    created_files = []
+
+    # Define project structure
+    files_to_create = {
+        "main.py": MAIN_PY_TEMPLATE,
+        "workers/__init__.py": WORKERS_INIT_PY_TEMPLATE,
+        "workers/example_worker.py": WORKER_EXAMPLE_PY_TEMPLATE,
+        ".env.example": ENV_EXAMPLE_TEMPLATE,
+        "requirements.txt": REQUIREMENTS_TXT_TEMPLATE,
+        ".gitignore": GITIGNORE_TEMPLATE,
+        ".flashignore": FLASHIGNORE_TEMPLATE,
+        "README.md": README_TEMPLATE.format(project_name=project_dir.name),
+    }
+
+    # Create files
+    for relative_path, content in files_to_create.items():
+        file_path = project_dir / relative_path
+
+        # Create parent directories if needed
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Skip existing files unless force is True
+        if file_path.exists() and not force:
+            continue
+
+        # Write file
+        file_path.write_text(content)
+        created_files.append(str(file_path.relative_to(project_dir)))
+
+    return created_files

tetra_rp/client.py

@@ -1,3 +1,4 @@
+import os
 import inspect
 import logging
 from functools import wraps
@@ -15,47 +16,69 @@ def remote(
     dependencies: Optional[List[str]] = None,
     system_dependencies: Optional[List[str]] = None,
     accelerate_downloads: bool = True,
-    hf_models_to_cache: Optional[List[str]] = None,
+    local: bool = False,
     **extra,
 ):
     """
     Decorator to enable dynamic resource provisioning and dependency management for serverless functions.
 
     This decorator allows a function to be executed in a remote serverless environment, with support for
-    dynamic resource provisioning and installation of required dependencies.
+    dynamic resource provisioning and installation of required dependencies. It can also bypass remote
+    execution entirely for local testing.
 
     Args:
         resource_config (ServerlessResource): Configuration object specifying the serverless resource
-            to be provisioned or used.
+            to be provisioned or used. Not used when local=True.
         dependencies (List[str], optional): A list of pip package names to be installed in the remote
-            environment before executing the function. Defaults to None.
+            environment before executing the function. Not used when local=True. Defaults to None.
         system_dependencies (List[str], optional): A list of system packages to be installed in the remote
-            environment before executing the function. Defaults to None.
+            environment before executing the function. Not used when local=True. Defaults to None.
         accelerate_downloads (bool, optional): Enable download acceleration for dependencies and models.
-            Defaults to True.
-        hf_models_to_cache (List[str], optional): List of HuggingFace model IDs to pre-cache using
-            download acceleration. Defaults to None.
+            Only applies to remote execution. Defaults to True.
+        local (bool, optional): Execute function/class locally instead of provisioning remote servers.
+            Returns the unwrapped function/class for direct local execution. Users must ensure all required
+            dependencies are already installed in their local environment. Defaults to False.
         extra (dict, optional): Additional parameters for the execution of the resource. Defaults to an empty dict.
 
     Returns:
-        Callable: A decorator that wraps the target function, enabling remote execution with the
-        specified resource configuration and dependencies.
+        Callable: A decorator that wraps the target function, enabling remote execution with the specified
+        resource configuration and dependencies, or returns the unwrapped function/class for local execution.
 
     Example:
     ```python
+        # Remote execution (production)
         @remote(
             resource_config=my_resource_config,
             dependencies=["numpy", "pandas"],
             accelerate_downloads=True,
-            hf_models_to_cache=["gpt2", "bert-base-uncased"]
         )
         async def my_function(data):
             # Function logic here
             pass
+
+        # Local execution (testing/development)
+        # Note: Ensure numpy and pandas are installed locally first
+        @remote(
+            resource_config=my_resource_config,
+            dependencies=["numpy", "pandas"],  # Only used for remote execution
+            local=True,
+        )
+        async def my_test_function(data):
+            # Runs locally - dependencies must be pre-installed
+            pass
     ```
     """
 
     def decorator(func_or_class):
+        if os.getenv("RUNPOD_POD_ID") or os.getenv("RUNPOD_ENDPOINT_ID"):
+            # Worker mode when running on RunPod platform
+            return func_or_class
+
+        # Local execution mode - execute without provisioning remote servers
+        if local:
+            return func_or_class
+
+        # Remote execution mode
         if inspect.isclass(func_or_class):
             # Handle class decoration
             return create_remote_class(
@@ -64,11 +87,10 @@ def remote(
                 dependencies,
                 system_dependencies,
                 accelerate_downloads,
-                hf_models_to_cache,
                 extra,
             )
         else:
-            # Handle function decoration (unchanged)
+            # Handle function decoration
             @wraps(func_or_class)
             async def wrapper(*args, **kwargs):
                 resource_manager = ResourceManager()
@@ -82,7 +104,6 @@ def remote(
                     dependencies,
                     system_dependencies,
                     accelerate_downloads,
-                    hf_models_to_cache,
                     *args,
                     **kwargs,
                 )

tetra_rp/config.py

@@ -0,0 +1,29 @@
+"""Configuration management for tetra-rp CLI."""
+
+from pathlib import Path
+from typing import NamedTuple
+
+
+class TetraPaths(NamedTuple):
+    """Paths for tetra-rp configuration and data."""
+
+    tetra_dir: Path
+    config_file: Path
+    deployments_file: Path
+
+    def ensure_tetra_dir(self) -> None:
+        """Ensure the .tetra directory exists."""
+        self.tetra_dir.mkdir(exist_ok=True)
+
+
+def get_paths() -> TetraPaths:
+    """Get standardized paths for tetra-rp configuration."""
+    tetra_dir = Path.cwd() / ".tetra"
+    config_file = tetra_dir / "config.json"
+    deployments_file = tetra_dir / "deployments.json"
+
+    return TetraPaths(
+        tetra_dir=tetra_dir,
+        config_file=config_file,
+        deployments_file=deployments_file,
+    )

tetra_rp/core/resources/cpu.py

@@ -16,6 +16,9 @@ class CpuInstanceType(str, Enum):
     - cpu5g: Not available
     """
 
+    ANY = "any"
+    """Any CPU"""
+
     # 3rd Generation General Purpose (RAM multiplier: 4.0)
 
     CPU3G_1_4 = "cpu3g-1-4"
@@ -58,6 +61,11 @@ class CpuInstanceType(str, Enum):
     CPU5C_8_16 = "cpu5c-8-16"
     """8 vCPU, 16GB RAM, max 120GB container disk"""
 
+    @classmethod
+    def all(cls) -> List["CpuInstanceType"]:
+        """Returns all CPU Instance Types."""
+        return [c for c in cls if c != cls.ANY]
+
 
 def calculate_max_disk_size(instance_type: CpuInstanceType) -> int:
     """
@@ -104,6 +112,7 @@ def calculate_max_disk_size(instance_type: CpuInstanceType) -> int:
 CPU_INSTANCE_DISK_LIMITS = {
     instance_type: calculate_max_disk_size(instance_type)
     for instance_type in CpuInstanceType
+    if instance_type != CpuInstanceType.ANY
 }
 
 

tetra_rp/core/resources/serverless_cpu.py

@@ -6,7 +6,7 @@ This module contains all CPU-related serverless functionality, separate from GPU
 
 from typing import List, Optional
 
-from pydantic import field_serializer, model_validator
+from pydantic import field_serializer, model_validator, field_validator
 
 from .cpu import (
     CpuInstanceType,
@@ -105,7 +105,7 @@ class CpuServerlessEndpoint(CpuEndpointMixin, ServerlessEndpoint):
     Represents a CPU-only serverless endpoint distinct from a live serverless.
     """
 
-    instanceIds: Optional[List[CpuInstanceType]] = [CpuInstanceType.CPU3G_2_8]
+    instanceIds: Optional[List[CpuInstanceType]] = [CpuInstanceType.ANY]
 
     def _create_new_template(self) -> PodTemplate:
         """Create a new PodTemplate with CPU-appropriate disk sizing."""
@@ -133,6 +133,14 @@ class CpuServerlessEndpoint(CpuEndpointMixin, ServerlessEndpoint):
         # Apply CPU-specific disk sizing
         self._apply_cpu_disk_sizing(self.template)
 
+    @field_validator("instanceIds")
+    @classmethod
+    def validate_cpus(cls, value: List[CpuInstanceType]) -> List[CpuInstanceType]:
+        """Expand ANY to all GPU groups"""
+        if value == [CpuInstanceType.ANY]:
+            return CpuInstanceType.all()
+        return value
+
     @model_validator(mode="after")
     def set_serverless_template(self):
         # Sync CPU-specific fields first

tetra_rp/execute_class.py

@@ -203,7 +203,6 @@ def create_remote_class(
     dependencies: Optional[List[str]],
     system_dependencies: Optional[List[str]],
     accelerate_downloads: bool,
-    hf_models_to_cache: Optional[List[str]],
     extra: dict,
 ):
     """
@@ -222,7 +221,6 @@ def create_remote_class(
             self._dependencies = dependencies or []
             self._system_dependencies = system_dependencies or []
             self._accelerate_downloads = accelerate_downloads
-            self._hf_models_to_cache = hf_models_to_cache
             self._extra = extra
             self._constructor_args = args
             self._constructor_kwargs = kwargs
@@ -307,7 +305,6 @@ def create_remote_class(
                     dependencies=self._dependencies,
                     system_dependencies=self._system_dependencies,
                     accelerate_downloads=self._accelerate_downloads,
-                    hf_models_to_cache=self._hf_models_to_cache,
                     instance_id=self._instance_id,
                     create_new_instance=not hasattr(
                         self, "_stub"

tetra_rp/protos/remote_execution.py

@@ -81,10 +81,6 @@ class FunctionRequest(BaseModel):
         default=True,
         description="Enable download acceleration for dependencies and models",
     )
-    hf_models_to_cache: Optional[List[str]] = Field(
-        default=None,
-        description="List of HuggingFace model IDs to pre-cache using acceleration",
-    )
 
     @model_validator(mode="after")
     def validate_execution_requirements(self) -> "FunctionRequest":

tetra_rp/stubs/live_serverless.py

@@ -68,7 +68,6 @@ class LiveServerlessStub(RemoteExecutorStub):
         dependencies,
         system_dependencies,
         accelerate_downloads,
-        hf_models_to_cache,
         *args,
         **kwargs,
     ):
@@ -79,7 +78,6 @@ class LiveServerlessStub(RemoteExecutorStub):
             "dependencies": dependencies,
             "system_dependencies": system_dependencies,
             "accelerate_downloads": accelerate_downloads,
-            "hf_models_to_cache": hf_models_to_cache,
         }
 
         # Thread-safe cache access
@@ -110,7 +108,7 @@ class LiveServerlessStub(RemoteExecutorStub):
 
         if response.stdout:
             for line in response.stdout.splitlines():
-                log.info(f"Remote | {line}")
+                print(line)
 
         if response.success:
             if response.result is None:

tetra_rp/stubs/registry.py

@@ -31,7 +31,6 @@ def _create_live_serverless_stub(resource, **extra):
         dependencies,
         system_dependencies,
         accelerate_downloads,
-        hf_models_to_cache,
         *args,
         **kwargs,
     ) -> dict:
@@ -43,7 +42,6 @@ def _create_live_serverless_stub(resource, **extra):
             dependencies,
             system_dependencies,
             accelerate_downloads,
-            hf_models_to_cache,
             *args,
             **kwargs,
         )
@@ -78,7 +76,6 @@ def _(resource, **extra):
         dependencies,
         system_dependencies,
         accelerate_downloads,
-        hf_models_to_cache,
         *args,
         **kwargs,
     ) -> dict:
@@ -103,7 +100,6 @@ def _(resource, **extra):
         dependencies,
         system_dependencies,
         accelerate_downloads,
-        hf_models_to_cache,
         *args,
         **kwargs,
     ) -> dict:

{tetra_rp-0.12.0.dist-info → tetra_rp-0.14.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.12.0
+Version: 0.14.0
 Summary: A Python library for distributed inference and serving of machine learning models
 Author-email: Marut Pandya <pandyamarut@gmail.com>, Patrick Rachford <prachford@icloud.com>, Dean Quinanola <dean.quinanola@runpod.io>
 License: MIT
@@ -14,6 +14,10 @@ Requires-Dist: cloudpickle>=3.1.1
 Requires-Dist: runpod
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: pydantic>=2.0.0
+Requires-Dist: rich>=14.0.0
+Requires-Dist: typer>=0.12.0
+Requires-Dist: questionary>=2.0.0
+Requires-Dist: pathspec>=0.11.0
 
 # Tetra: Serverless computing for AI workloads
 
@@ -359,6 +363,7 @@ if __name__ == "__main__":
 ```python
 import asyncio
 from tetra_rp import remote, LiveServerless, GpuGroup, PodTemplate
+import base64
 
 # Advanced GPU configuration with consolidated template overrides
 sd_config = LiveServerless(