File size: 11,357 Bytes
63ed3a7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
import logging
import os
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from inference_server.models import list_supported_policies
from inference_server.session_manager import SessionManager

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Global session manager
session_manager = SessionManager()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Handle app startup and shutdown."""
    logger.info("🚀 Inference Server starting up...")
    yield
    logger.info("🔄 Inference Server shutting down...")
    await session_manager.cleanup_all_sessions()
    logger.info("✅ Inference Server shutdown complete")


# FastAPI app
app = FastAPI(
    title="Inference Server",
    description="Multi-Policy Model Inference Server for Real-time Robot Control",
    version="1.0.0",
    lifespan=lifespan,
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify actual origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Request/Response models
class CreateSessionRequest(BaseModel):
    session_id: str
    policy_path: str
    camera_names: list[str] = ["front"]  # Support multiple cameras
    arena_server_url: str = "http://localhost:8000"
    workspace_id: str | None = None  # Optional workspace ID
    policy_type: str = "act"  # Policy type: act, pi0, pi0fast, smolvla, diffusion
    language_instruction: str | None = None  # For vision-language policies


class CreateSessionResponse(BaseModel):
    workspace_id: str
    camera_room_ids: dict[str, str]  # {camera_name: room_id}
    joint_input_room_id: str
    joint_output_room_id: str


class SessionStatusResponse(BaseModel):
    session_id: str
    status: str
    policy_path: str
    policy_type: str
    camera_names: list[str]  # Multiple camera names
    workspace_id: str
    rooms: dict
    stats: dict
    inference_stats: dict | None = None
    error_message: str | None = None


# Health check
@app.get("/", tags=["Health"])
async def root():
    """Health check endpoint."""
    return {"message": "Inference Server is running", "status": "healthy"}


@app.get("/health", tags=["Health"])
async def health_check():
    """Detailed health check."""
    return {
        "status": "healthy",
        "active_sessions": len(session_manager.sessions),
        "session_ids": list(session_manager.sessions.keys()),
    }


@app.get("/policies", tags=["Policies"])
async def list_policies():
    """List supported policy types."""
    return {
        "supported_policies": list_supported_policies(),
        "description": "Available policy types for inference",
    }


# Session management endpoints
@app.post("/sessions", response_model=CreateSessionResponse, tags=["Sessions"])
async def create_session(request: CreateSessionRequest):
    """
    Create a new inference session.

    If workspace_id is provided, all rooms will be created in that workspace.
    If workspace_id is not provided, a new workspace will be generated automatically.
    All rooms for a session (cameras + joints) are always created in the same workspace.
    """
    try:
        room_ids = await session_manager.create_session(
            session_id=request.session_id,
            policy_path=request.policy_path,
            camera_names=request.camera_names,
            arena_server_url=request.arena_server_url,
            workspace_id=request.workspace_id,
            policy_type=request.policy_type,
            language_instruction=request.language_instruction,
        )
        return CreateSessionResponse(**room_ids)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.exception(f"Failed to create session {request.session_id}")
        raise HTTPException(status_code=500, detail=f"Failed to create session: {e!s}")


@app.get("/sessions", response_model=list[SessionStatusResponse], tags=["Sessions"])
async def list_sessions():
    """List all sessions."""
    sessions = await session_manager.list_sessions()
    return [SessionStatusResponse(**session) for session in sessions]


@app.get(
    "/sessions/{session_id}", response_model=SessionStatusResponse, tags=["Sessions"]
)
async def get_session_status(session_id: str):
    """Get status of a specific session."""
    try:
        status = await session_manager.get_session_status(session_id)
        return SessionStatusResponse(**status)
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")


@app.post("/sessions/{session_id}/start", tags=["Control"])
async def start_inference(session_id: str):
    """Start inference for a session."""
    try:
        await session_manager.start_inference(session_id)
        return {"message": f"Inference started for session {session_id}"}
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    except Exception as e:
        logger.exception(f"Failed to start inference for session {session_id}")
        raise HTTPException(status_code=500, detail=f"Failed to start inference: {e!s}")


@app.post("/sessions/{session_id}/stop", tags=["Control"])
async def stop_inference(session_id: str):
    """Stop inference for a session."""
    try:
        await session_manager.stop_inference(session_id)
        return {"message": f"Inference stopped for session {session_id}"}
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")


@app.post("/sessions/{session_id}/restart", tags=["Control"])
async def restart_inference(session_id: str):
    """Restart inference for a session."""
    try:
        await session_manager.restart_inference(session_id)
        return {"message": f"Inference restarted for session {session_id}"}
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    except Exception as e:
        logger.exception(f"Failed to restart inference for session {session_id}")
        raise HTTPException(
            status_code=500, detail=f"Failed to restart inference: {e!s}"
        )


@app.delete("/sessions/{session_id}", tags=["Sessions"])
async def delete_session(session_id: str):
    """Delete a session."""
    try:
        await session_manager.delete_session(session_id)
        return {"message": f"Session {session_id} deleted"}
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")


# Debug endpoints for enhanced monitoring
@app.get("/debug/system", tags=["Debug"])
async def get_system_info():
    """Get system information for debugging."""
    import psutil
    import torch

    try:
        # System info
        system_info = {
            "cpu_percent": psutil.cpu_percent(interval=1),
            "memory": {
                "total": psutil.virtual_memory().total,
                "available": psutil.virtual_memory().available,
                "percent": psutil.virtual_memory().percent,
            },
            "disk": {
                "total": psutil.disk_usage("/").total,
                "used": psutil.disk_usage("/").used,
                "percent": psutil.disk_usage("/").percent,
            },
        }

        # GPU info if available
        if torch.cuda.is_available():
            system_info["gpu"] = {
                "device_count": torch.cuda.device_count(),
                "current_device": torch.cuda.current_device(),
                "device_name": torch.cuda.get_device_name(),
                "memory_allocated": torch.cuda.memory_allocated(),
                "memory_cached": torch.cuda.memory_reserved(),
            }

        return system_info
    except Exception as e:
        return {"error": f"Failed to get system info: {e}"}


@app.get("/debug/logs", tags=["Debug"])
async def get_recent_logs():
    """Get recent log entries for debugging."""
    try:
        # This is a simple implementation - in production you might want to read from actual log files
        return {
            "message": "Log endpoint available",
            "note": "Implement actual log reading if needed",
            "active_sessions": len(session_manager.sessions),
        }
    except Exception as e:
        return {"error": f"Failed to get logs: {e}"}


@app.post("/debug/sessions/{session_id}/reset", tags=["Debug"])
async def debug_reset_session(session_id: str):
    """Reset a session's internal state for debugging."""
    try:
        if session_id not in session_manager.sessions:
            raise HTTPException(
                status_code=404, detail=f"Session {session_id} not found"
            )

        session = session_manager.sessions[session_id]

        # Reset inference engine if available
        if session.inference_engine:
            session.inference_engine.reset()

        # Clear action queue
        session.action_queue.clear()

        # Reset flags
        for camera_name in session.camera_names:
            session.images_updated[camera_name] = False
        session.joints_updated = False

        return {"message": f"Session {session_id} state reset successfully"}

    except Exception as e:
        logger.exception(f"Failed to reset session {session_id}")
        raise HTTPException(status_code=500, detail=f"Failed to reset session: {e}")


@app.get("/debug/sessions/{session_id}/queue", tags=["Debug"])
async def get_session_queue_info(session_id: str):
    """Get detailed information about a session's action queue."""
    try:
        if session_id not in session_manager.sessions:
            raise HTTPException(
                status_code=404, detail=f"Session {session_id} not found"
            )

        session = session_manager.sessions[session_id]

        return {
            "session_id": session_id,
            "queue_length": len(session.action_queue),
            "queue_maxlen": session.action_queue.maxlen,
            "n_action_steps": session.n_action_steps,
            "control_frequency_hz": session.control_frequency_hz,
            "inference_frequency_hz": session.inference_frequency_hz,
            "last_queue_cleanup": session.last_queue_cleanup,
            "data_status": {
                "has_joint_data": session.latest_joint_positions is not None,
                "images_status": {
                    camera: camera in session.latest_images
                    for camera in session.camera_names
                },
                "images_updated": session.images_updated.copy(),
                "joints_updated": session.joints_updated,
            },
        }

    except Exception as e:
        logger.exception(f"Failed to get queue info for session {session_id}")
        raise HTTPException(status_code=500, detail=f"Failed to get queue info: {e}")


# Main entry point
if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("PORT", 8001))
    uvicorn.run(
        "inference_server.main:app",
        host="0.0.0.0",
        port=port,
        reload=False,
        log_level="info",
    )