Spaces:

factorstudios
/

NMFL

Runtime error

File size: 10,081 Bytes

import os
import json
import time
import threading
import shutil
from fastapi import FastAPI, HTTPException, BackgroundTasks, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import uvicorn
from typing import Dict, List
from pathlib import Path
from datetime import datetime
from fastapi.responses import FileResponse
import humanize

# Import storage and test module
from test_ai_integration_http import test_ai_integration_http
from http_storage import LocalStorage

# Initialize FastAPI app
app = FastAPI()

# Create templates directory if it doesn't exist
templates_dir = Path("templates")
templates_dir.mkdir(exist_ok=True)

# Create static directory for CSS
static_dir = Path("static")
static_dir.mkdir(exist_ok=True)

# Initialize templates and static files
templates = Jinja2Templates(directory="templates")
app.mount("/static", StaticFiles(directory="static"), name="static")

# Create CSS file
css_content = """

body {

    font-family: 'Arial', sans-serif;

    margin: 0;

    padding: 20px;

    background: #f5f5f5;

}

.container {

    max-width: 1200px;

    margin: 0 auto;

}

.stats-grid {

    display: grid;

    grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));

    gap: 20px;

    margin-bottom: 30px;

}

.stat-card {

    background: white;

    padding: 20px;

    border-radius: 10px;

    box-shadow: 0 2px 5px rgba(0,0,0,0.1);

}

.stat-title {

    font-size: 14px;

    color: #666;

    margin-bottom: 10px;

}

.stat-value {

    font-size: 24px;

    color: #2c3e50;

    font-weight: bold;

}

.files-section {

    background: white;

    padding: 20px;

    border-radius: 10px;

    box-shadow: 0 2px 5px rgba(0,0,0,0.1);

}

table {

    width: 100%;

    border-collapse: collapse;

}

th, td {

    padding: 12px;

    text-align: left;

    border-bottom: 1px solid #eee;

}

th {

    background: #f8f9fa;

    color: #666;

}

.file-icon {

    width: 20px;

    margin-right: 10px;

}

.header {

    margin-bottom: 30px;

}

.header h1 {

    color: #2c3e50;

    margin: 0;

}

.refresh-time {

    color: #666;

    font-size: 14px;

    margin-top: 5px;

}

"""

# Create CSS file
with open(static_dir / "style.css", "w") as f:
    f.write(css_content)

# Create HTML template
html_content = """

<!DOCTYPE html>

<html>

<head>

    <title>GPU Storage Dashboard</title>

    <link rel="stylesheet" href="{{ url_for('static', path='/style.css') }}">

    <meta http-equiv="refresh" content="30">

</head>

<body>

    <div class="container">

        <div class="header">

            <h1>GPU Storage Dashboard</h1>

            <div class="refresh-time">Last updated: {{ last_updated }}</div>

        </div>

        

        <div class="stats-grid">

            <div class="stat-card">

                <div class="stat-title">Total Storage Used</div>

                <div class="stat-value">{{ storage_stats.total_used }}</div>

            </div>

            <div class="stat-card">

                <div class="stat-title">Free Space</div>

                <div class="stat-value">{{ storage_stats.free_space }}</div>

            </div>

            <div class="stat-card">

                <div class="stat-title">Active Models</div>

                <div class="stat-value">{{ storage_stats.active_models }}</div>

            </div>

            <div class="stat-card">

                <div class="stat-title">Active Tensors</div>

                <div class="stat-value">{{ storage_stats.active_tensors }}</div>

            </div>

        </div>



        <div class="files-section">

            <h2>Models</h2>

            <table>

                <tr>

                    <th>Name</th>

                    <th>Size</th>

                    <th>Last Modified</th>

                </tr>

                {% for model in models %}

                <tr>

                    <td>{{ model.name }}</td>

                    <td>{{ model.size }}</td>

                    <td>{{ model.modified }}</td>

                </tr>

                {% endfor %}

            </table>

        </div>



        <div class="files-section">

            <h2>Tensors</h2>

            <table>

                <tr>

                    <th>Name</th>

                    <th>Size</th>

                    <th>Shape</th>

                    <th>Last Modified</th>

                </tr>

                {% for tensor in tensors %}

                <tr>

                    <td>{{ tensor.name }}</td>

                    <td>{{ tensor.size }}</td>

                    <td>{{ tensor.shape }}</td>

                    <td>{{ tensor.modified }}</td>

                </tr>

                {% endfor %}

            </table>

        </div>

    </div>

</body>

</html>

"""

# Create HTML template file
with open(templates_dir / "dashboard.html", "w") as f:
    f.write(html_content)

# Initialize storage
storage = LocalStorage()

@app.get("/", response_class=HTMLResponse)
async def get_dashboard(request: Request):
    """Serve the storage dashboard"""
    stats = {}
    
    # Get storage statistics
    total_size = 0
    for path in [storage.models_path, storage.vram_path, storage.cache_path, storage.state_path]:
        if path.exists():
            total_size += sum(f.stat().st_size for f in path.rglob('*') if f.is_file())
    
    # Get disk usage
    disk_usage = shutil.disk_usage(str(storage.base_path))
    
    stats['storage_stats'] = {
        'total_used': humanize.naturalsize(total_size),
        'free_space': humanize.naturalsize(disk_usage.free),
        'active_models': len(storage.model_registry),
        'active_tensors': len(storage.tensor_registry)
    }
    
    # Get model files
    models = []
    if storage.models_path.exists():
        for model_dir in storage.models_path.iterdir():
            if model_dir.is_dir():
                config_file = model_dir.joinpath('config.json')
                if config_file.exists():
                    size = sum(f.stat().st_size for f in model_dir.rglob('*') if f.is_file())
                    models.append({
                        'name': model_dir.name,
                        'size': humanize.naturalsize(size),
                        'modified': datetime.fromtimestamp(model_dir.stat().st_mtime).strftime('%Y-%m-%d %H:%M:%S')
                    })
    
    # Get tensor files
    tensors = []
    if storage.vram_path.exists():
        for tensor_file in storage.vram_path.glob('*.npy'):
            if tensor_file.exists():
                meta_file = tensor_file.with_name(tensor_file.stem + '_meta.json')
                shape = "Unknown"
                if meta_file.exists():
                    with open(meta_file) as f:
                        meta = json.load(f)
                        shape = str(meta.get('shape', 'Unknown'))
                
                tensors.append({
                    'name': tensor_file.stem,
                    'size': humanize.naturalsize(tensor_file.stat().st_size),
                    'shape': shape,
                    'modified': datetime.fromtimestamp(tensor_file.stat().st_mtime).strftime('%Y-%m-%d %H:%M:%S')
                })
    
    return templates.TemplateResponse("dashboard.html", {
        "request": request,
        "storage_stats": stats['storage_stats'],
        "models": models,
        "tensors": tensors,
        "last_updated": datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    })

# Run the server
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)


# FastAPI App Definition
app = FastAPI()

# Add CORS middleware to allow cross-origin requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],
)

# Global variable to track if processing is running
processing_thread = None

def log_message(message):
    """Add a log message with timestamp"""
    timestamp = datetime.now().strftime("%H:%M:%S")
    log_entry = f"[{timestamp}] {message}"
    processing_status["logs"].append(log_entry)
    
    # Keep only the last 100 logs
    if len(processing_status["logs"]) > 100:
        processing_status["logs"] = processing_status["logs"][-100:]
    
    print(log_entry)

@app.on_event("startup")
async def startup_event():
    global processing_thread
    if not (processing_thread and processing_thread.is_alive()):
        processing_thread = threading.Thread(target=test_ai_integration_http)
        processing_thread.daemon = True
        processing_thread.start()

from fastapi.staticfiles import StaticFiles

# app.mount("/static", StaticFiles(directory="static"), name="static")

# Serve your main HTML file
@app.get("/")
async def root():
    return ()

from fastapi.encoders import jsonable_encoder

def get_disk_usage(path: str) -> Dict[str, float]:
    """Get disk usage statistics in GB"""
    statvfs = os.statvfs(path)
    total = statvfs.f_frsize * statvfs.f_blocks / (1024**3)
    free = statvfs.f_frsize * statvfs.f_bavail / (1024**3)
    used = total - free
    return {"total": total, "free": free, "used": used}

class SafeJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        try:
            if isinstance(obj, float):
                if obj != obj:  # Check for NaN
                    return None
                if obj == float('inf') or obj == float('-inf'):
                    return None
            return super().default(obj)
        except:
            return None



if __name__ == "__main__":
    get_disk_usage()
    uvicorn.run(
        app, 
        host="0.0.0.0", 
        port=8000, 
        log_level="info",
        reload=False  # Set to False for production
    )