Add Model Registry & Storage

[alicup29]

Model Registry & Storage

This PR implements a comprehensive model registry and the beginning to a checkpoint recovery system for tracking trained models on workers, enabling model discovery, and supporting training resumption after connection loss.

Overview

The model registry solves two key problems:

1. Model Identification: Models are now identified by deterministic hash-based IDs, making it easy to reference specific trained models
2. Checkpoint Recovery: Training can resume from checkpoints if connections drop during training

Key Features

1. Model Registry System

• Hash-based Model IDs: 8-character SHA256 hashes generated from config + dataset + run name
• Persistent Storage: JSON/YAML registry file at models/.registry/manifest.json
• Model Metadata Tracking:
  • Status (training, completed, interrupted, failed)
  • Checkpoint paths
  • Training metrics
  • Timestamps and run information
  • GPU model and dataset info

2. Remote Registry Queries

• Client-to-Worker Queries: Clients can query remote worker registries via WebRTC
• Worker Discovery: Automatic worker discovery in rooms using discover_peers
• Dual-Mode CLI: Commands work both locally (on worker) and remotely (via network)

3. Enhanced CLI Commands

List Models:

# Local (on worker machine)
sleap-rtc list-models
sleap-rtc list-models --status completed --model-type centroid

# Remote (from client)
sleap-rtc list-models --room-id ROOM --token TOKEN
sleap-rtc list-models --session-string SESSION

Model Info:

# Local
sleap-rtc model-info a3f5e8c9

# Remote  
sleap-rtc model-info a3f5e8c9 --room-id ROOM --token TOKEN

Model ID-Based Inference:

# Use model ID instead of paths
sleap-rtc client-track \
  --model a3f5e8c9 \
  --data_path video.slp \
  --room-id ROOM --token TOKEN

4. Checkpoint Recovery

• Models marked as "interrupted" on connection loss
• Training automatically resumes from last checkpoint when same config is submitted again
• Uses PyTorch Lightning's native checkpoint recovery (resume_ckpt_path)

Implementation Details

New Files

• sleap_rtc/worker/model_registry.py (432 lines) - Core registry implementation
• sleap_rtc/client/registry_query.py (295 lines) - Client-side remote query support
• tests/worker/test_model_registry.py (650+ lines) - Comprehensive test suite

Modified Files

• sleap_rtc/worker/worker_class.py - Registry integration, training lifecycle tracking
• sleap_rtc/cli.py - New registry commands with remote query support

Registry Message Protocol

• REGISTRY_QUERY_LIST::<filters_json> - Query for model list
• REGISTRY_QUERY_INFO::<model_id> - Query for specific model info
• REGISTRY_RESPONSE_LIST::<response_json> - List response
• REGISTRY_RESPONSE_INFO::<response_json> - Info response
• REGISTRY_RESPONSE_ERROR::<error_json> - Error response

Use Cases Enabled

1. Client Discovers and Uses Remote Models

# 1. Query available models
sleap-rtc list-models --room-id ROOM --token TOKEN

# 2. Get model details
sleap-rtc model-info a3f5e8c9 --room-id ROOM --token TOKEN

# 3. Run inference using model ID
sleap-rtc client-track --model a3f5e8c9 --data video.slp --room-id ROOM --token TOKEN

2. Worker Provides Model Options

Workers can respond to registry queries, allowing clients to discover what models are available before running inference.

3. Training Resumption After Failure

If connection drops during training:

1. Worker marks model as "interrupted" in registry
2. Client reconnects and submits same training package
3. Worker detects interrupted job and resumes from checkpoint

Testing

Test Coverage

• 22/24 tests passing in test_model_registry.py
• 2 skipped (YAML-specific, PyYAML not required)

Manual Testing Performed

• ✅ Local registry queries
• ✅ Remote registry queries via WebRTC
• ✅ Model completion status tracking
• ✅ Worker discovery and auto-selection
• ✅ Model ID-based inference
• ✅ Registry file persistence and corruption recovery

Bug Fixes Included

Issue 1: Model Completion Status

Problem: Models stuck in "training" status after completion
Cause: Loop variable scope issue - self.current_model_id overwritten before mark_completed() called
Fix: Use local variable current_training_model_id within loop scope

Issue 2: Client Registry Access

Problem: Clients could access local registry on same machine
Cause: No validation that registry directory should exist
Fix: Check registry exists before local access, provide helpful error messages

Issue 3: Worker Discovery Message

Problem: Used discover_workers message type not supported by signaling server
Cause: Wrong message type
Fix: Changed to discover_peers with role filtering

Issue 4: Async/Await in Anonymous Signin

Problem: Synchronous requests.post() called in async function
Cause: Blocking I/O in async context
Fix: Use asyncio.run_in_executor() to run in thread pool

Breaking Changes

None - this is entirely new functionality.

Migration Notes

• Registry is created automatically on first training job
• Existing model directories are not affected
• New models will use hash-based directory names: {model_type}_{model_id}/

Future Work

• Parse training logs to extract actual validation loss metrics
• Track actual epoch numbers during training for better resumption
• Support for model deletion/archival in registry
• Registry synchronization across multiple workers
• Web UI for registry browsing

Related Issues

Addresses the model identification and checkpoint recovery concerns discussed in initial requirements.

---

🤖 Generated with Claude Code

Co-Authored-By: Claude [email protected]