tetra-rp 0.5.5__py3-none-any.whl

tetra_rp/__init__.py

@@ -0,0 +1,37 @@
+# Load .env vars from file
+# before everything else
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+from .logger import setup_logging  # noqa: E402
+
+setup_logging()
+
+from .client import remote  # noqa: E402
+from .core.resources import (  # noqa: E402
+    CpuServerlessEndpoint,
+    CpuInstanceType,
+    CudaVersion,
+    GpuGroup,
+    LiveServerless,
+    PodTemplate,
+    ResourceManager,
+    ServerlessEndpoint,
+    runpod,
+)
+
+
+__all__ = [
+    "remote",
+    "CpuServerlessEndpoint",
+    "CpuInstanceType",
+    "CudaVersion",
+    "GpuGroup",
+    "LiveServerless",
+    "PodTemplate",
+    "ResourceManager",
+    "ServerlessEndpoint",
+    "runpod",
+]

tetra_rp/client.py

@@ -0,0 +1,59 @@
+import logging
+from functools import wraps
+from typing import List
+from .core.resources import ServerlessResource, ResourceManager
+from .stubs import stub_resource
+
+
+log = logging.getLogger(__name__)
+
+
+def remote(
+    resource_config: ServerlessResource,
+    dependencies: List[str] = None,
+    system_dependencies: List[str] = None,
+    **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.
+
+        resource_config (ServerlessResource): Configuration object specifying the serverless resource
+            to be provisioned or used.
+        dependencies (List[str], optional): A list of pip package names to be installed in the remote
+            environment before executing the function. Defaults to None.
+        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.
+
+    Example:
+    ```python
+        @remote(
+            resource_config=my_resource_config,
+            dependencies=["numpy", "pandas"],
+            sync=True  # Optional, to run synchronously
+        )
+        async def my_function(data):
+            # Function logic here
+            pass
+    ```
+    """
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            resource_manager = ResourceManager()
+            remote_resource = await resource_manager.get_or_deploy_resource(
+                resource_config
+            )
+
+            stub = stub_resource(remote_resource, **extra)
+            return await stub(func, dependencies, system_dependencies, *args, **kwargs)
+
+        return wrapper
+
+    return decorator

tetra_rp/core/api/__init__.py

@@ -0,0 +1,5 @@
+from .runpod import RunpodGraphQLClient
+
+__all__ = [
+    "RunpodGraphQLClient",
+]

tetra_rp/core/api/runpod.py

@@ -0,0 +1,212 @@
+"""
+Direct GraphQL communication with Runpod API.
+Bypasses the outdated runpod-python SDK limitations.
+"""
+
+import os
+import json
+import aiohttp
+from typing import Dict, Any, Optional
+import logging
+
+log = logging.getLogger(__name__)
+
+RUNPOD_API_BASE_URL = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io")
+
+
+class RunpodGraphQLClient:
+    """
+    Runpod GraphQL client for Runpod API.
+    Communicates directly with Runpod's GraphQL endpoint without SDK limitations.
+    """
+
+    GRAPHQL_URL = f"{RUNPOD_API_BASE_URL}/graphql"
+
+    def __init__(self, api_key: Optional[str] = None):
+        self.api_key = api_key or os.getenv("RUNPOD_API_KEY")
+        if not self.api_key:
+            raise ValueError("Runpod API key is required")
+
+        self.session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        """Get or create an aiohttp session."""
+        if self.session is None or self.session.closed:
+            timeout = aiohttp.ClientTimeout(total=300)  # 5 minute timeout
+            self.session = aiohttp.ClientSession(
+                timeout=timeout,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                },
+            )
+        return self.session
+
+    async def _execute_graphql(
+        self, query: str, variables: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Execute a GraphQL query/mutation."""
+        session = await self._get_session()
+
+        payload = {"query": query, "variables": variables or {}}
+
+        log.debug(f"GraphQL Query: {query}")
+        log.debug(f"GraphQL Variables: {json.dumps(variables, indent=2)}")
+
+        try:
+            async with session.post(self.GRAPHQL_URL, json=payload) as response:
+                response_data = await response.json()
+
+                log.debug(f"GraphQL Response Status: {response.status}")
+                log.debug(f"GraphQL Response: {json.dumps(response_data, indent=2)}")
+
+                if response.status >= 400:
+                    raise Exception(
+                        f"GraphQL request failed: {response.status} - {response_data}"
+                    )
+
+                if "errors" in response_data:
+                    errors = response_data["errors"]
+                    error_msg = "; ".join([e.get("message", str(e)) for e in errors])
+                    raise Exception(f"GraphQL errors: {error_msg}")
+
+                return response_data.get("data", {})
+
+        except aiohttp.ClientError as e:
+            log.error(f"HTTP client error: {e}")
+            raise Exception(f"HTTP request failed: {e}")
+
+    async def create_endpoint(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Create a serverless endpoint using direct GraphQL mutation.
+        Supports both GPU and CPU endpoints with full field support.
+        """
+        # GraphQL mutation for saveEndpoint (based on actual schema)
+        mutation = """
+        mutation saveEndpoint($input: EndpointInput!) {
+            saveEndpoint(input: $input) {
+                aiKey
+                gpuIds
+                id
+                idleTimeout
+                locations
+                name
+                networkVolumeId
+                scalerType
+                scalerValue
+                templateId
+                type
+                userId
+                version
+                workersMax
+                workersMin
+                workersStandby
+                workersPFBTarget
+                gpuCount
+                allowedCudaVersions
+                executionTimeoutMs
+                instanceIds
+                activeBuildid
+                idePodId
+            }
+        }
+        """
+
+        variables = {"input": input_data}
+
+        log.debug(
+            f"Creating endpoint with GraphQL: {input_data.get('name', 'unnamed')}"
+        )
+
+        result = await self._execute_graphql(mutation, variables)
+
+        if "saveEndpoint" not in result:
+            raise Exception("Unexpected GraphQL response structure")
+
+        endpoint_data = result["saveEndpoint"]
+        log.info(
+            f"Created endpoint: {endpoint_data.get('id', 'unknown')} - {endpoint_data.get('name', 'unnamed')}"
+        )
+
+        return endpoint_data
+
+    async def get_cpu_types(self) -> Dict[str, Any]:
+        """Get available CPU types."""
+        query = """
+        query getCpuTypes {
+            cpuTypes {
+                id
+                displayName
+                manufacturer
+                cores
+                threadsPerCore
+                groupId
+            }
+        }
+        """
+
+        result = await self._execute_graphql(query)
+        return result.get("cpuTypes", [])
+
+    async def get_gpu_types(
+        self, gpu_filter: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Get available GPU types."""
+        query = """
+        query getGpuTypes($input: GpuTypeFilter) {
+            gpuTypes(input: $input) {
+                id
+                displayName
+                manufacturer
+                memoryInGb
+                cudaCores
+                secureCloud
+                communityCloud
+                securePrice
+                communityPrice
+                communitySpotPrice
+                secureSpotPrice
+                maxGpuCount
+                maxGpuCountCommunityCloud
+                maxGpuCountSecureCloud
+                minPodGpuCount
+                nodeGroupGpuSizes
+                throughput
+            }
+        }
+        """
+
+        variables = {"input": gpu_filter} if gpu_filter else {}
+        result = await self._execute_graphql(query, variables)
+        return result.get("gpuTypes", [])
+
+    async def get_endpoint(self, endpoint_id: str) -> Dict[str, Any]:
+        """Get endpoint details."""
+        # Note: The schema doesn't show a specific endpoint query
+        # This would need to be implemented if such query exists
+        raise NotImplementedError("Get endpoint query not available in current schema")
+
+    async def delete_endpoint(self, endpoint_id: str) -> Dict[str, Any]:
+        """Delete a serverless endpoint."""
+        mutation = """
+        mutation deleteEndpoint($id: String!) {
+            deleteEndpoint(id: $id)
+        }
+        """
+
+        variables = {"id": endpoint_id}
+        log.info(f"Deleting endpoint: {endpoint_id}")
+
+        result = await self._execute_graphql(mutation, variables)
+        return {"success": result.get("deleteEndpoint") is not None}
+
+    async def close(self):
+        """Close the HTTP session."""
+        if self.session and not self.session.closed:
+            await self.session.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()

tetra_rp/core/pool/cluster_manager.py

@@ -0,0 +1,177 @@
+import time
+from worker import Worker
+from job import Job
+
+from dataclass import WorkerStatus, JobStatus
+
+import logging
+import inspect
+
+
+def setup_logging(level=logging.INFO, fmt=None):
+    if fmt is None:
+        fmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    logging.basicConfig(level=level, format=fmt)
+
+
+def get_logger(name=None):
+    """
+    Returns a logger. If no name is provided, it infers the caller's module name.
+    """
+    if name is None:
+        # Get the caller's module name.
+        frame = inspect.stack()[1]
+        module = inspect.getmodule(frame[0])
+        name = module.__name__ if module else "__main__"
+    return logging.getLogger(name)
+
+
+logger = get_logger(__name__)
+
+
+class ClusterManager:
+    """
+    Manages workers and Jobs currently in Memory:
+    - Runpod for provisioning
+    - Real remote execution
+    - Data base for the
+    """
+
+    def __init__(self):
+        self.workers = {}  # Worker ID -> Worker
+        self.jobs = {}  # Job ID -> Job
+
+    # ----------------- Worker Management -----------------
+    # ------------------------------------------------------
+    def add_worker(self, resource_config: dict):
+        """
+        Add a new worker to the cluster
+        """
+        # here will go the logic to create a worker and add it to the cluster: RUNPOD LOGIC will be added here.
+        worker = Worker(resource_config)
+        self.workers[worker.worker_id] = worker
+
+        logger.info(f"Added worker {worker.worker_id} to the cluster")
+        return worker.worker_id
+
+    def remove_worker(self, worker_id):
+        """
+        Remove a worker from the cluster
+        """
+        worker = self.workers.get(worker_id)
+        if not worker:
+            logger.error(f"Worker {worker_id} not found")
+            return False
+        if worker.status == WorkerStatus.RUNNING:
+            logger.error(f"Worker {worker_id} is still running")
+            return False
+        del self.workers[worker_id]
+        logger.info(f"Removed worker {worker_id} from the cluster")
+        return True
+
+    def list_workers(self):
+        """
+        List all workers in the cluster
+        """
+        return list(self.workers.values())
+
+    # ----------------- Job Management -----------------
+    # ---------------------------------------------------
+
+    def submit_job(self, resource_config: dict):
+        """
+        Submit a new job to the cluster (Queueud). Then attempt to scheduel it.
+        """
+        job = Job(resource_config)
+        self.jobs[job.job_id] = job
+        logger.info(f"Submitted job {job.job_id} to the cluster")
+        # attempt to schedule the job
+        self.schedule_job(job)
+        return job.job_id
+
+    def schedule_job(self, job: Job):
+        """
+        find a suitable worker for the job. It none, Job remains queued.
+        If we want to a auto provision we can actually add a logic here to add a worker if none is available.
+        """
+        if job.status != JobStatus.QUEUED:
+            logger.error(f"Job {job.job_id} is not pending")
+            return False
+
+        # Find worker candidate
+        candidate = self.find_idle_worker(job.resource_config)
+        if candidate:
+            self.assign_job_to_worker(job, candidate)
+        else:
+            logger.info(f"No worker available for job {job.job_id}")
+            # we cn either provision new worker from here and then scehediule the job from here.
+
+    def find_idle_worker(self, resource_config: dict):
+        """
+        Find an idle worker that can run the job
+        """
+        for w in self.workers.values():
+            if w.status == WorkerStatus.IDLE:
+                # check the resource config
+                if w.resource_config == resource_config:
+                    continue
+                return w
+        return None
+
+    def assign_job_to_worker(self, job: Job, worker: Worker):
+        """
+        Mark the job as running and the worker as Running and 'execute' the job.
+        In a real system, we would send a remote command to the worker (eg: gRPC) to execute the job.
+        """
+        job.worker_id = worker.worker_id
+        job.status = JobStatus.RUNNING
+        worker.status = WorkerStatus.RUNNING
+        worker.current_job_id = job.job_id
+        logger.info(f"Assigned job {job.job_id} to worker {worker.worker_id}")
+        self._execute_job(job, worker)
+
+    def _execute_job(self, job: Job, worker: Worker):
+        """
+        Simulate the remote execution. right now, we jsut sleep for 1s.
+        In production, what we we can do is:
+        - Open a gRPC connection to the worker
+        - pass the job details
+        - wait for the compeltion call back
+        """
+        try:
+            logger.info(f"Executing job {job.job_id} on worker {worker.worker_id}")
+            time.sleep(
+                1
+            )  # Here we can add the actual execution logic, currently it mimics the execution.
+
+            # mark the job as completed
+            job.status = JobStatus.COMPLETED
+            job.result = "Job completed successfully"
+            logger.info(f"[Cluster Manager] Job {job.job_id} completed successfully")
+        except Exception as e:
+            job.status = JobStatus.FAILED
+            job.result = f"Job failed: {str(e)}"
+            logger.error(f"[Cluster Manager] Job {job.job_id} failed: {str(e)}")
+        finally:
+            worker.status = WorkerStatus.IDLE
+            worker.current_job_id = None
+
+    def get_job_status(self, job_id):
+        """
+        Get the job details
+        """
+        job = self.jobs.get(job_id)
+        if not job:
+            logger.error(f"Job {job_id} not found")
+            return None
+        return job
+
+    # this function has retry logic but it's currently fuzzy, we might have to change it.
+
+    def retry_queued_jobs(self):
+        """
+        Retry all queued jobs
+        """
+        for job in self.jobs.values():
+            if job.status == JobStatus.QUEUED:
+                self.schedule_job(job)

tetra_rp/core/pool/dataclass.py

@@ -0,0 +1,18 @@
+from enum import Enum
+
+
+class WorkerStatus(Enum):
+    """Enum representing the status of a worker"""
+
+    IDLE = "idle"
+    RUNNING = "running"
+    OFFLINE = "offline"
+
+
+class JobStatus(Enum):
+    """Enum representing the status of a job"""
+
+    QUEUED = "queued"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"

tetra_rp/core/pool/ex.py

@@ -0,0 +1,38 @@
+from cluster_manager import ClusterManager
+
+
+if __name__ == "__main__":
+    cm = ClusterManager()
+
+    # 1) Submit a job with no existing workers (use resource_config dict)
+    job_id = cm.submit_job(
+        resource_config={"gpu": "H100", "memory": 16, "network_volume": 50}
+    )
+    print(
+        "Job status:", cm.get_job_status(job_id)
+    )  # should be QUEUED, no suitable worker
+
+    # 2) Add a worker that doesn't match the GPU
+    w1 = cm.add_worker(
+        resource_config={"gpu": "H100", "memory": 16, "network_volume": 50}
+    )
+    # Re-try scheduling
+    cm.retry_queued_jobs()
+    print("Job status (still queued):", cm.get_job_status(job_id))
+
+    # 3) Add a matching worker
+    w2 = cm.add_worker(
+        resource_config={"gpu": "H100", "memory": 16, "network_volume": 50}
+    )
+    # Re-try scheduling
+    cm.retry_queued_jobs()
+    print("Job status (should complete):", cm.get_job_status(job_id))
+
+    # 4) Submit another job that requires less resources
+    job_id2 = cm.submit_job(resource_config={"memory": 8, "network_volume": 10})
+    # Should be assigned to w1 if it's idle
+    print("Job2 final status:", cm.get_job_status(job_id2))
+
+    # 5) Show final state of workers
+    for worker in cm.list_workers():
+        print("Worker:", worker)

tetra_rp/core/pool/job.py

@@ -0,0 +1,22 @@
+import uuid
+from dataclass import JobStatus
+
+
+class Job:
+    """Represents a 'job' in the system
+
+    In a real system, this might contain the function to run,
+    arguments, and reference to data or code.
+    """
+
+    def __init__(self, resource_config: dict):
+        self.job_id = str(uuid.uuid4())[:8]
+        self.resource_config = resource_config
+        self.status = JobStatus.QUEUED
+
+        self.worker_id = None
+        self.result = None
+        self.error = None
+
+    def __repr__(self):
+        return f"Job(job_id={self.job_id}, status={self.status})"

tetra_rp/core/pool/worker.py

@@ -0,0 +1,19 @@
+import uuid
+from dataclass import WorkerStatus
+
+
+class Worker:
+    """Represents a single worker in the pool
+
+    For Now we store ressources in memory
+    """
+
+    def __init__(self, resource_config: dict):
+        self.worker_id = str(uuid.uuid4())[:8]
+        self.resource_config = resource_config
+        self.status = WorkerStatus.IDLE
+
+        self.current_job_id = None
+
+    def __repr__(self):
+        return f"Worker(worker_id={self.worker_id}, status={self.status})"

tetra_rp/core/resources/__init__.py

@@ -0,0 +1,33 @@
+from .base import BaseResource, DeployableResource
+from .cloud import runpod
+from .cpu import CpuInstanceType
+from .gpu import GpuGroup, GpuType, GpuTypeDetail
+from .resource_manager import ResourceManager
+from .live_serverless import LiveServerless
+from .serverless import (
+    CpuServerlessEndpoint,
+    ServerlessResource,
+    ServerlessEndpoint,
+    JobOutput,
+    CudaVersion,
+)
+from .template import PodTemplate
+
+
+__all__ = [
+    "runpod",
+    "BaseResource",
+    "CpuInstanceType",
+    "CpuServerlessEndpoint",
+    "CudaVersion",
+    "DeployableResource",
+    "GpuGroup",
+    "GpuType",
+    "GpuTypeDetail",
+    "JobOutput",
+    "LiveServerless",
+    "ResourceManager",
+    "ServerlessResource",
+    "ServerlessEndpoint",
+    "PodTemplate",
+]

tetra_rp/core/resources/base.py

@@ -0,0 +1,47 @@
+import hashlib
+from abc import ABC, abstractmethod
+from typing import Optional
+from pydantic import BaseModel, ConfigDict
+
+
+class BaseResource(BaseModel):
+    """Base class for all resources."""
+
+    model_config = ConfigDict(
+        validate_by_name=True,
+        validate_default=True,
+        serialize_by_alias=True,
+    )
+
+    id: Optional[str] = None
+
+    @property
+    def resource_id(self) -> str:
+        """Unique resource ID based on configuration."""
+        resource_type = self.__class__.__name__
+        config_str = self.model_dump_json(exclude_none=True)
+        hash_obj = hashlib.md5(f"{resource_type}:{config_str}".encode())
+        return f"{resource_type}_{hash_obj.hexdigest()}"
+
+
+class DeployableResource(BaseResource, ABC):
+    """Base class for deployable resources."""
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}"
+
+    @property
+    @abstractmethod
+    def url(self) -> str:
+        """Public URL of the resource."""
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    @abstractmethod
+    def is_deployed(self) -> bool:
+        """Check the resource if it's still valid or available."""
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    @abstractmethod
+    async def deploy(self) -> "DeployableResource":
+        """Deploy the resource."""
+        raise NotImplementedError("Subclasses should implement this method.")

tetra_rp/core/resources/cloud.py

@@ -0,0 +1,4 @@
+import os
+import runpod
+
+runpod.api_key = os.getenv("RUNPOD_API_KEY")