Spaces:

Fred808
/

tserv

Paused

tserv / tensor_server.py

Update tensor_server.py

819c548 verified about 1 year ago

16.9 kB

import os
import json
import torch
import psutil
import asyncio
from datetime import datetime
from typing import Dict, List, Optional
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

# ===== Config =====
class Settings:
    # Server configuration
    HOST = "0.0.0.0"  # Listen on all interfaces
    PORT = 8001
    SERVER_ID = os.getenv("SERVER_ID", "tensor1")  # Unique ID for this tensor server
    
    # The IP or hostname where this tensor server is accessible
    PUBLIC_URL = os.getenv("PUBLIC_URL", f"http://192.168.1.101:8001")
    
    # URLs for other services (should be actual IP addresses or hostnames)
    CONTROLLER_URL = os.getenv("CONTROLLER_URL", "http://192.168.1.100:8000")
    AGGREGATOR_URL = os.getenv("AGGREGATOR_URL", "http://192.168.1.104:8002")
    
    # Model settings
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    MAX_BATCH_SIZE = 32
    METRICS_UPDATE_INTERVAL = 5  # seconds
    MODEL_DIR = "model_chunks"
    
    @classmethod
    def from_env(cls):
        """Load settings from environment variables"""
        cls.HOST = os.getenv("TENSOR_HOST", cls.HOST)
        cls.PORT = int(os.getenv("TENSOR_PORT", cls.PORT))
        cls.SERVER_ID = os.getenv("SERVER_ID", cls.SERVER_ID)
        cls.CONTROLLER_URL = os.getenv("CONTROLLER_URL", cls.CONTROLLER_URL)
        cls.AGGREGATOR_URL = os.getenv("AGGREGATOR_URL", cls.AGGREGATOR_URL)
        return cls

# ===== Models =====
class ModelChunk(BaseModel):
    """Represents a received model chunk configuration"""
    chunk_id: int
    files: List[str]
    config: Dict
    
class InferenceRequest(BaseModel):
    """Represents an inference request"""
    inputs: List[List[float]]
    batch_size: Optional[int] = None
    top_k: Optional[int] = 5

class MetricsData(BaseModel):
    """Server metrics data"""
    cpu_usage: float
    memory_usage: float
    gpu_usage: Optional[float]
    active_requests: int
    total_requests: int
    average_response_time: float
    last_error: Optional[str]
    error_count: int

# ===== FastAPI App =====
app = FastAPI(
    title="Tensor Server",
    description="Handles model chunk computations",
    version="1.0.0"
)

# ===== State =====
class ServerState:
    def __init__(self):
        self.loaded_chunks: Dict[int, torch.nn.Module] = {}
        self.active_requests: int = 0
        self.total_requests: int = 0
        self.request_times: List[float] = []
        self.error_count: int = 0
        self.last_error: Optional[str] = None
        self.is_computing: bool = False
        
state = ServerState()

# ===== Metrics Collection =====
async def collect_metrics() -> MetricsData:
    """Collect current server metrics"""
    # CPU and memory metrics
    cpu_usage = psutil.cpu_percent()
    memory = psutil.virtual_memory()
    memory_usage = memory.percent
    
    # GPU metrics if available
    gpu_usage = None
    if torch.cuda.is_available():
        try:
            gpu_usage = torch.cuda.memory_allocated() / torch.cuda.max_memory_allocated() * 100
        except:
            pass
    
    # Calculate average response time
    avg_response_time = sum(state.request_times) / len(state.request_times) if state.request_times else 0
    
    return MetricsData(
        cpu_usage=cpu_usage,
        memory_usage=memory_usage,
        gpu_usage=gpu_usage,
        active_requests=state.active_requests,
        total_requests=state.total_requests,
        average_response_time=avg_response_time,
        last_error=state.last_error,
        error_count=state.error_count
    )

async def update_metrics_loop():
    """Background task to update metrics periodically"""
    while True:
        try:
            metrics = await collect_metrics()
            # Store metrics for health checks
            state.current_metrics = metrics
        except Exception as e:
            print(f"[ERROR] Failed to update metrics: {str(e)}")
        await asyncio.sleep(Settings.METRICS_UPDATE_INTERVAL)

# ===== Helper Functions =====
def load_chunk(chunk: ModelChunk) -> torch.nn.Module:
    """Load a model chunk into memory"""
    try:
        # Create chunk directory if it doesn't exist
        os.makedirs(Settings.MODEL_DIR, exist_ok=True)
        
        # Get chunk configuration
        chunk_config = chunk.config
        if "original_file" not in chunk_config:
            raise ValueError("Missing original_file in chunk configuration")
            
        # Save chunk data to file
        chunk_file = os.path.join(Settings.MODEL_DIR, chunk.files[0])
        if not os.path.exists(chunk_file):
            # We'll need to receive the actual chunk data in a separate request
            raise ValueError(f"Chunk file not found: {chunk_file}")
            
        # For raw binary chunks, we'll create a simple buffer module
        class ChunkBuffer(nn.Module):
            """
            A single Florence-2 caption chunk that receives pre-encoded image embeddings
            and produces partial vocabulary logits.
            """
        
            def __init__(self, chunk_path: str, config: dict):
                super().__init__()
        
                # Get dimensions from config
                input_dim = config.get("input_dim", 1024)            # Florence-2 embedding dim
                output_dim = config.get("output_dim", 1000)          # size of vocab shard
                dropout = config.get("dropout", 0.1)
        
                # Optional: chunk_path can point to pretrained weights
                self.chunk_path = chunk_path
        
                # Main projection layer: embedding → partial vocab logits
                self.linear = nn.Linear(input_dim, output_dim)
        
                # Optional normalization + dropout (stabilizes training or inference variance)
                self.norm = nn.LayerNorm(input_dim)
                self.dropout = nn.Dropout(dropout)
        
                # Initialize weights (small variance, stable logits)
                nn.init.xavier_uniform_(self.linear.weight)
                nn.init.zeros_(self.linear.bias)
        
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                """
                Args:
                    x: Florence-2 image embedding tensor, shape [batch, 1024]
                Returns:
                    logits for this vocab shard, shape [batch, output_dim]
                """
                # Normalize + dropout
                x = self.norm(x)
                x = self.dropout(x)
        
                # Linear projection to vocab slice
                logits = self.linear(x)
        
                # (Optional) softmax for probabilities, but usually the main model handles this
                # probs = F.softmax(logits, dim=-1)
                return logits
        
                        
        # Create and return the chunk buffer
        chunk_model = ChunkBuffer(chunk_file, chunk_config)
        # Ensure the chunk_model.config is the up-to-date config (including any assigned offsets)
        chunk_model.config = chunk_config
        print(f"[INFO] Loaded chunk {chunk.chunk_id} ({chunk_config.get('size_bytes', 0)} bytes) from {chunk.files[0]}")

        return chunk_model
                
    except Exception as e:
        raise Exception(f"Failed to load chunk: {str(e)}")

async def process_tensor(chunk_id: int, inputs: torch.Tensor) -> torch.Tensor:
    """Process input tensor through the specified chunk"""
    if chunk_id not in state.loaded_chunks:
        raise HTTPException(status_code=400, detail=f"Chunk {chunk_id} not loaded")
        
    chunk_model = state.loaded_chunks[chunk_id]
    with torch.no_grad():
        outputs = chunk_model(inputs)
    return outputs

# ===== API Endpoints =====
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    metrics = await collect_metrics()
    return {
        "status": "healthy",
        "device": Settings.DEVICE,
        "loaded_chunks": list(state.loaded_chunks.keys()),
        "metrics": metrics.dict()
    }

@app.get("/metrics")
async def get_metrics():
    """Get current server metrics"""
    return await collect_metrics()

from fastapi import File, UploadFile

@app.post("/load_chunk")
async def load_model_chunk(chunk: ModelChunk):
    """Register a chunk configuration"""
    try:
        # Create model directory if it doesn't exist
        os.makedirs(Settings.MODEL_DIR, exist_ok=True)
        
        # Store the chunk metadata
        chunk_file = os.path.join(Settings.MODEL_DIR, chunk.files[0])
        state.chunk_configs = getattr(state, 'chunk_configs', {})

        # Ensure a vocab_offset is present; if not, assign a non-overlapping offset
        cfg = chunk.config or {}
        if 'vocab_offset' not in cfg:
            # Compute next available offset from existing registered chunks
            max_end = 0
            for existing in state.chunk_configs.values():
                try:
                    e_cfg = existing.config if hasattr(existing, 'config') else existing
                    e_offset = int(e_cfg.get('vocab_offset', 0))
                    e_shard = int(e_cfg.get('shard_dim', e_cfg.get('size', 1) or 1))
                    max_end = max(max_end, e_offset + e_shard)
                except Exception:
                    continue

            # If this chunk declares a shard_dim, use it; otherwise default to 1
            shard_dim = int(cfg.get('shard_dim', cfg.get('size', 1) or 1))
            cfg['vocab_offset'] = max_end
            cfg['shard_dim'] = cfg.get('shard_dim', shard_dim)

        # Store back the possibly-updated config
        chunk.config = cfg
        state.chunk_configs[chunk.chunk_id] = chunk
        
        print(f"[INFO] Registered chunk {chunk.chunk_id} configuration")
        print(f"[INFO] Waiting for chunk data: {chunk.files[0]}")
        
        return {
            "status": "configured",
            "chunk_id": chunk.chunk_id,
            "ready_for_data": True
        }
        
    except Exception as e:
        state.error_count += 1
        state.last_error = str(e)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/upload_chunk_data/{chunk_id}")
async def upload_chunk_data(chunk_id: int, file: UploadFile = File(...)):
    """Receive the actual chunk data"""
    try:
        if chunk_id not in getattr(state, 'chunk_configs', {}):
            raise HTTPException(status_code=400, detail="Chunk configuration not registered")
            
        chunk = state.chunk_configs[chunk_id]
        chunk_file = os.path.join(Settings.MODEL_DIR, chunk.files[0])
        
        # Save the uploaded file
        with open(chunk_file, 'wb') as f:
            content = await file.read()
            f.write(content)
            
        # Now load the chunk
        chunk_model = load_chunk(chunk)
        # Ensure the loaded module has the registered config (including vocab_offset)
        try:
            registered = getattr(state, 'chunk_configs', {}).get(chunk_id)
            if registered is not None:
                # registered is a ModelChunk; merge config into module
                reg_cfg = registered.config or {}
                if hasattr(chunk_model, 'config'):
                    chunk_model.config.update(reg_cfg)
                else:
                    chunk_model.config = reg_cfg
                # expose vocab_offset on module
                chunk_model.vocab_offset = int(reg_cfg.get('vocab_offset', 0))
        except Exception:
            pass
        state.loaded_chunks[chunk_id] = chunk_model
        
        file_size = os.path.getsize(chunk_file)
        print(f"[INFO] Received and loaded chunk {chunk_id} data ({file_size} bytes)")
        
        return {
            "status": "loaded",
            "chunk_id": chunk_id,
            "size_bytes": file_size,
            "file": chunk.files[0]
        }
        
    except Exception as e:
        state.error_count += 1
        state.last_error = str(e)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/compute/{chunk_id}")
async def compute(chunk_id: int, request: InferenceRequest):
    """Perform computation on inputs using specified chunk"""
    try:
        start_time = datetime.now()
        state.active_requests += 1
        state.total_requests += 1
        
        # Convert inputs to tensor
        inputs = torch.tensor(request.inputs, dtype=torch.float32, device=Settings.DEVICE)
        
        # Split into batches if needed
        batch_size = request.batch_size or Settings.MAX_BATCH_SIZE
        if len(inputs) > batch_size:
            batches = torch.split(inputs, batch_size)
            outputs = []
            for batch in batches:
                batch_output = await process_tensor(chunk_id, batch)
                outputs.append(batch_output)
            output_tensor = torch.cat(outputs, dim=0)
        else:
            output_tensor = await process_tensor(chunk_id, inputs)
        
        # Convert output to numpy for diagnostics
        try:
            shard_np = output_tensor.cpu().numpy()
        except Exception:
            shard_np = None

        chunk_details = {}
        try:
            # Normalize to 2D: (seq_len, shard_dim)
            if shard_np is None:
                raise ValueError("Unable to convert output tensor to numpy")

            seq_len = shard_np.shape[0]
            shard_2d = shard_np.reshape(seq_len, -1)

            k = int(request.top_k or 5)
            k = min(k, shard_2d.shape[1]) if shard_2d.shape[1] > 0 else 0

            # compute local top-k per position
            if k > 0:
                topk_idx = np.argpartition(-shard_2d, k-1, axis=1)[:, :k]
                topk_vals = np.take_along_axis(shard_2d, topk_idx, axis=1)
            else:
                topk_idx = np.zeros((seq_len, 0), dtype=int)
                topk_vals = np.zeros((seq_len, 0), dtype=float)

            # determine vocab_offset from the loaded chunk config if available
            cfg = None
            try:
                chunk_model = state.loaded_chunks.get(chunk_id)
                cfg = getattr(chunk_model, 'config', None) or getattr(state, 'chunk_configs', {}).get(chunk_id, {}).config if chunk_id in getattr(state, 'chunk_configs', {}) else None
            except Exception:
                cfg = None

            vocab_offset = 0
            if isinstance(cfg, dict):
                vocab_offset = int(cfg.get('vocab_offset', 0))
            elif cfg is not None and hasattr(cfg, 'get'):
                vocab_offset = int(cfg.get('vocab_offset', 0))

            per_position_topk = []
            for pos_idx in range(seq_len):
                toks = []
                for jj in range(topk_idx.shape[1]):
                    local_idx = int(topk_idx[pos_idx, jj])
                    token_id = int(vocab_offset + local_idx)
                    score = float(topk_vals[pos_idx, jj])
                    toks.append([token_id, score])
                per_position_topk.append(toks)

            chunk_details[chunk_id] = {
                'logits_shard': shard_2d.tolist(),
                'topk': per_position_topk,
                'vocab_offset': vocab_offset,
                'shard_dim': shard_2d.shape[1]
            }
        except Exception as e:
            # If diagnostics fail, include error info but keep main outputs
            chunk_details = {chunk_id: {'error': str(e)}}

        # Convert output to list for backward compatibility
        output_list = output_tensor.cpu().numpy().tolist()

        # Update metrics
        end_time = datetime.now()
        processing_time = (end_time - start_time).total_seconds()
        state.request_times.append(processing_time)
        # Keep only last 100 request times
        state.request_times = state.request_times[-100:]

        return {
            "outputs": output_list,
            "processing_time": processing_time,
            "chunk_details": chunk_details
        }
        
    except Exception as e:
        state.error_count += 1
        state.last_error = str(e)
        raise HTTPException(status_code=500, detail=str(e))
        
    finally:
        state.active_requests -= 1

@app.on_event("startup")
async def startup_event():
    """Start background tasks"""
    asyncio.create_task(update_metrics_loop())

# ===== Main Execution =====
if __name__ == "__main__":
    port = int(os.getenv("PORT", 8001))  # Default to 8001 to avoid conflict with controller
    print(f"[INFO] Starting tensor server on port {port}")
    print(f"[INFO] Using device: {Settings.DEVICE}")
    print(f"[INFO] API Documentation available at http://localhost:{port}/docs")
    
    uvicorn.run(
        "tensor_server:app",
        host="0.0.0.0",
        port=port,
        reload=False
    )

	import os
	import json
	import torch
	import psutil
	import asyncio
	from datetime import datetime
	from typing import Dict, List, Optional
	from fastapi import FastAPI, HTTPException
	from pydantic import BaseModel
	import uvicorn
	import numpy as np
	import torch.nn as nn
	import torch.nn.functional as F

	# ===== Config =====
	class Settings:
	# Server configuration
	HOST = "0.0.0.0" # Listen on all interfaces
	PORT = 8001
	SERVER_ID = os.getenv("SERVER_ID", "tensor1") # Unique ID for this tensor server

	# The IP or hostname where this tensor server is accessible
	PUBLIC_URL = os.getenv("PUBLIC_URL", f"http://192.168.1.101:8001")

	# URLs for other services (should be actual IP addresses or hostnames)
	CONTROLLER_URL = os.getenv("CONTROLLER_URL", "http://192.168.1.100:8000")
	AGGREGATOR_URL = os.getenv("AGGREGATOR_URL", "http://192.168.1.104:8002")

	# Model settings
	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
	MAX_BATCH_SIZE = 32
	METRICS_UPDATE_INTERVAL = 5 # seconds
	MODEL_DIR = "model_chunks"

	@classmethod
	def from_env(cls):
	"""Load settings from environment variables"""
	cls.HOST = os.getenv("TENSOR_HOST", cls.HOST)
	cls.PORT = int(os.getenv("TENSOR_PORT", cls.PORT))
	cls.SERVER_ID = os.getenv("SERVER_ID", cls.SERVER_ID)
	cls.CONTROLLER_URL = os.getenv("CONTROLLER_URL", cls.CONTROLLER_URL)
	cls.AGGREGATOR_URL = os.getenv("AGGREGATOR_URL", cls.AGGREGATOR_URL)
	return cls

	# ===== Models =====
	class ModelChunk(BaseModel):
	"""Represents a received model chunk configuration"""
	chunk_id: int
	files: List[str]
	config: Dict

	class InferenceRequest(BaseModel):
	"""Represents an inference request"""
	inputs: List[List[float]]
	batch_size: Optional[int] = None
	top_k: Optional[int] = 5

	class MetricsData(BaseModel):
	"""Server metrics data"""
	cpu_usage: float
	memory_usage: float
	gpu_usage: Optional[float]
	active_requests: int
	total_requests: int
	average_response_time: float
	last_error: Optional[str]
	error_count: int

	# ===== FastAPI App =====
	app = FastAPI(
	title="Tensor Server",
	description="Handles model chunk computations",
	version="1.0.0"
	)

	# ===== State =====
	class ServerState:
	def __init__(self):
	self.loaded_chunks: Dict[int, torch.nn.Module] = {}
	self.active_requests: int = 0
	self.total_requests: int = 0
	self.request_times: List[float] = []
	self.error_count: int = 0
	self.last_error: Optional[str] = None
	self.is_computing: bool = False

	state = ServerState()

	# ===== Metrics Collection =====
	async def collect_metrics() -> MetricsData:
	"""Collect current server metrics"""
	# CPU and memory metrics
	cpu_usage = psutil.cpu_percent()
	memory = psutil.virtual_memory()
	memory_usage = memory.percent

	# GPU metrics if available
	gpu_usage = None
	if torch.cuda.is_available():
	try:
	gpu_usage = torch.cuda.memory_allocated() / torch.cuda.max_memory_allocated() * 100
	except:
	pass

	# Calculate average response time
	avg_response_time = sum(state.request_times) / len(state.request_times) if state.request_times else 0

	return MetricsData(
	cpu_usage=cpu_usage,
	memory_usage=memory_usage,
	gpu_usage=gpu_usage,
	active_requests=state.active_requests,
	total_requests=state.total_requests,
	average_response_time=avg_response_time,
	last_error=state.last_error,
	error_count=state.error_count
	)

	async def update_metrics_loop():
	"""Background task to update metrics periodically"""
	while True:
	try:
	metrics = await collect_metrics()
	# Store metrics for health checks
	state.current_metrics = metrics
	except Exception as e:
	print(f"[ERROR] Failed to update metrics: {str(e)}")
	await asyncio.sleep(Settings.METRICS_UPDATE_INTERVAL)

	# ===== Helper Functions =====
	def load_chunk(chunk: ModelChunk) -> torch.nn.Module:
	"""Load a model chunk into memory"""
	try:
	# Create chunk directory if it doesn't exist
	os.makedirs(Settings.MODEL_DIR, exist_ok=True)

	# Get chunk configuration
	chunk_config = chunk.config
	if "original_file" not in chunk_config:
	raise ValueError("Missing original_file in chunk configuration")

	# Save chunk data to file
	chunk_file = os.path.join(Settings.MODEL_DIR, chunk.files[0])
	if not os.path.exists(chunk_file):
	# We'll need to receive the actual chunk data in a separate request
	raise ValueError(f"Chunk file not found: {chunk_file}")

	# For raw binary chunks, we'll create a simple buffer module
	class ChunkBuffer(nn.Module):
	"""
	A single Florence-2 caption chunk that receives pre-encoded image embeddings
	and produces partial vocabulary logits.
	"""

	def __init__(self, chunk_path: str, config: dict):
	super().__init__()

	# Get dimensions from config
	input_dim = config.get("input_dim", 1024) # Florence-2 embedding dim
	output_dim = config.get("output_dim", 1000) # size of vocab shard
	dropout = config.get("dropout", 0.1)

	# Optional: chunk_path can point to pretrained weights
	self.chunk_path = chunk_path

	# Main projection layer: embedding → partial vocab logits
	self.linear = nn.Linear(input_dim, output_dim)

	# Optional normalization + dropout (stabilizes training or inference variance)
	self.norm = nn.LayerNorm(input_dim)
	self.dropout = nn.Dropout(dropout)

	# Initialize weights (small variance, stable logits)
	nn.init.xavier_uniform_(self.linear.weight)
	nn.init.zeros_(self.linear.bias)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	"""
	Args:
	x: Florence-2 image embedding tensor, shape [batch, 1024]
	Returns:
	logits for this vocab shard, shape [batch, output_dim]
	"""
	# Normalize + dropout
	x = self.norm(x)
	x = self.dropout(x)

	# Linear projection to vocab slice
	logits = self.linear(x)

	# (Optional) softmax for probabilities, but usually the main model handles this
	# probs = F.softmax(logits, dim=-1)
	return logits


	# Create and return the chunk buffer
	chunk_model = ChunkBuffer(chunk_file, chunk_config)
	# Ensure the chunk_model.config is the up-to-date config (including any assigned offsets)
	chunk_model.config = chunk_config
	print(f"[INFO] Loaded chunk {chunk.chunk_id} ({chunk_config.get('size_bytes', 0)} bytes) from {chunk.files[0]}")

	return chunk_model

	except Exception as e:
	raise Exception(f"Failed to load chunk: {str(e)}")

	async def process_tensor(chunk_id: int, inputs: torch.Tensor) -> torch.Tensor:
	"""Process input tensor through the specified chunk"""
	if chunk_id not in state.loaded_chunks:
	raise HTTPException(status_code=400, detail=f"Chunk {chunk_id} not loaded")

	chunk_model = state.loaded_chunks[chunk_id]
	with torch.no_grad():
	outputs = chunk_model(inputs)
	return outputs

	# ===== API Endpoints =====
	@app.get("/health")
	async def health_check():
	"""Health check endpoint"""
	metrics = await collect_metrics()
	return {
	"status": "healthy",
	"device": Settings.DEVICE,
	"loaded_chunks": list(state.loaded_chunks.keys()),
	"metrics": metrics.dict()
	}

	@app.get("/metrics")
	async def get_metrics():
	"""Get current server metrics"""
	return await collect_metrics()

	from fastapi import File, UploadFile

	@app.post("/load_chunk")
	async def load_model_chunk(chunk: ModelChunk):
	"""Register a chunk configuration"""
	try:
	# Create model directory if it doesn't exist
	os.makedirs(Settings.MODEL_DIR, exist_ok=True)

	# Store the chunk metadata
	chunk_file = os.path.join(Settings.MODEL_DIR, chunk.files[0])
	state.chunk_configs = getattr(state, 'chunk_configs', {})

	# Ensure a vocab_offset is present; if not, assign a non-overlapping offset
	cfg = chunk.config or {}
	if 'vocab_offset' not in cfg:
	# Compute next available offset from existing registered chunks
	max_end = 0
	for existing in state.chunk_configs.values():
	try:
	e_cfg = existing.config if hasattr(existing, 'config') else existing
	e_offset = int(e_cfg.get('vocab_offset', 0))
	e_shard = int(e_cfg.get('shard_dim', e_cfg.get('size', 1) or 1))
	max_end = max(max_end, e_offset + e_shard)
	except Exception:
	continue

	# If this chunk declares a shard_dim, use it; otherwise default to 1
	shard_dim = int(cfg.get('shard_dim', cfg.get('size', 1) or 1))
	cfg['vocab_offset'] = max_end
	cfg['shard_dim'] = cfg.get('shard_dim', shard_dim)

	# Store back the possibly-updated config
	chunk.config = cfg
	state.chunk_configs[chunk.chunk_id] = chunk

	print(f"[INFO] Registered chunk {chunk.chunk_id} configuration")
	print(f"[INFO] Waiting for chunk data: {chunk.files[0]}")

	return {
	"status": "configured",
	"chunk_id": chunk.chunk_id,
	"ready_for_data": True
	}

	except Exception as e:
	state.error_count += 1
	state.last_error = str(e)
	raise HTTPException(status_code=500, detail=str(e))

	@app.post("/upload_chunk_data/{chunk_id}")
	async def upload_chunk_data(chunk_id: int, file: UploadFile = File(...)):
	"""Receive the actual chunk data"""
	try:
	if chunk_id not in getattr(state, 'chunk_configs', {}):
	raise HTTPException(status_code=400, detail="Chunk configuration not registered")

	chunk = state.chunk_configs[chunk_id]
	chunk_file = os.path.join(Settings.MODEL_DIR, chunk.files[0])

	# Save the uploaded file
	with open(chunk_file, 'wb') as f:
	content = await file.read()
	f.write(content)

	# Now load the chunk
	chunk_model = load_chunk(chunk)
	# Ensure the loaded module has the registered config (including vocab_offset)
	try:
	registered = getattr(state, 'chunk_configs', {}).get(chunk_id)
	if registered is not None:
	# registered is a ModelChunk; merge config into module
	reg_cfg = registered.config or {}
	if hasattr(chunk_model, 'config'):
	chunk_model.config.update(reg_cfg)
	else:
	chunk_model.config = reg_cfg
	# expose vocab_offset on module
	chunk_model.vocab_offset = int(reg_cfg.get('vocab_offset', 0))
	except Exception:
	pass
	state.loaded_chunks[chunk_id] = chunk_model

	file_size = os.path.getsize(chunk_file)
	print(f"[INFO] Received and loaded chunk {chunk_id} data ({file_size} bytes)")

	return {
	"status": "loaded",
	"chunk_id": chunk_id,
	"size_bytes": file_size,
	"file": chunk.files[0]
	}

	except Exception as e:
	state.error_count += 1
	state.last_error = str(e)
	raise HTTPException(status_code=500, detail=str(e))

	@app.post("/compute/{chunk_id}")
	async def compute(chunk_id: int, request: InferenceRequest):
	"""Perform computation on inputs using specified chunk"""
	try:
	start_time = datetime.now()
	state.active_requests += 1
	state.total_requests += 1

	# Convert inputs to tensor
	inputs = torch.tensor(request.inputs, dtype=torch.float32, device=Settings.DEVICE)

	# Split into batches if needed
	batch_size = request.batch_size or Settings.MAX_BATCH_SIZE
	if len(inputs) > batch_size:
	batches = torch.split(inputs, batch_size)
	outputs = []
	for batch in batches:
	batch_output = await process_tensor(chunk_id, batch)
	outputs.append(batch_output)
	output_tensor = torch.cat(outputs, dim=0)
	else:
	output_tensor = await process_tensor(chunk_id, inputs)

	# Convert output to numpy for diagnostics
	try:
	shard_np = output_tensor.cpu().numpy()
	except Exception:
	shard_np = None

	chunk_details = {}
	try:
	# Normalize to 2D: (seq_len, shard_dim)
	if shard_np is None:
	raise ValueError("Unable to convert output tensor to numpy")

	seq_len = shard_np.shape[0]
	shard_2d = shard_np.reshape(seq_len, -1)

	k = int(request.top_k or 5)
	k = min(k, shard_2d.shape[1]) if shard_2d.shape[1] > 0 else 0

	# compute local top-k per position
	if k > 0:
	topk_idx = np.argpartition(-shard_2d, k-1, axis=1)[:, :k]
	topk_vals = np.take_along_axis(shard_2d, topk_idx, axis=1)
	else:
	topk_idx = np.zeros((seq_len, 0), dtype=int)
	topk_vals = np.zeros((seq_len, 0), dtype=float)

	# determine vocab_offset from the loaded chunk config if available
	cfg = None
	try:
	chunk_model = state.loaded_chunks.get(chunk_id)
	cfg = getattr(chunk_model, 'config', None) or getattr(state, 'chunk_configs', {}).get(chunk_id, {}).config if chunk_id in getattr(state, 'chunk_configs', {}) else None
	except Exception:
	cfg = None

	vocab_offset = 0
	if isinstance(cfg, dict):
	vocab_offset = int(cfg.get('vocab_offset', 0))
	elif cfg is not None and hasattr(cfg, 'get'):
	vocab_offset = int(cfg.get('vocab_offset', 0))

	per_position_topk = []
	for pos_idx in range(seq_len):
	toks = []
	for jj in range(topk_idx.shape[1]):
	local_idx = int(topk_idx[pos_idx, jj])
	token_id = int(vocab_offset + local_idx)
	score = float(topk_vals[pos_idx, jj])
	toks.append([token_id, score])
	per_position_topk.append(toks)

	chunk_details[chunk_id] = {
	'logits_shard': shard_2d.tolist(),
	'topk': per_position_topk,
	'vocab_offset': vocab_offset,
	'shard_dim': shard_2d.shape[1]
	}
	except Exception as e:
	# If diagnostics fail, include error info but keep main outputs
	chunk_details = {chunk_id: {'error': str(e)}}

	# Convert output to list for backward compatibility
	output_list = output_tensor.cpu().numpy().tolist()

	# Update metrics
	end_time = datetime.now()
	processing_time = (end_time - start_time).total_seconds()
	state.request_times.append(processing_time)
	# Keep only last 100 request times
	state.request_times = state.request_times[-100:]

	return {
	"outputs": output_list,
	"processing_time": processing_time,
	"chunk_details": chunk_details
	}

	except Exception as e:
	state.error_count += 1
	state.last_error = str(e)
	raise HTTPException(status_code=500, detail=str(e))

	finally:
	state.active_requests -= 1

	@app.on_event("startup")
	async def startup_event():
	"""Start background tasks"""
	asyncio.create_task(update_metrics_loop())

	# ===== Main Execution =====
	if __name__ == "__main__":
	port = int(os.getenv("PORT", 8001)) # Default to 8001 to avoid conflict with controller
	print(f"[INFO] Starting tensor server on port {port}")
	print(f"[INFO] Using device: {Settings.DEVICE}")
	print(f"[INFO] API Documentation available at http://localhost:{port}/docs")

	uvicorn.run(
	"tensor_server:app",
	host="0.0.0.0",
	port=port,
	reload=False
	)