Add external keep-alive script for Tranception Space and update settings

.claude/settings.local.json

@@ -10,7 +10,8 @@
       "Bash(git rm:*)",
       "Bash(python test:*)",
       "Bash(rm:*)",
-      "Bash(chmod:*)"
+      "Bash(chmod:*)",
+      "Bash(cp:*)"
     ],
     "deny": []
   }

app.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Tranception Design App - Hugging Face Spaces Version
+Tranception Design App - Hugging Face Spaces Version (Zero GPU Fixed)
 """
 import os
 import sys
@@ -13,129 +13,31 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 import gradio as gr
 from huggingface_hub import hf_hub_download
-import zipfile
 import shutil
 import uuid
 import gc
 import time
-import signal
-import threading
 import datetime
+import threading
 
-# Check if Zero GPU should be disabled via environment variable
-DISABLE_ZERO_GPU = os.environ.get('DISABLE_ZERO_GPU', 'false').lower() == 'true'
-
-# Keep-alive settings
-KEEP_ALIVE_INTERVAL = 300  # 5 minutes
-last_activity_time = datetime.datetime.now()
-keep_alive_thread = None
-
-# Auto-refresh component to keep connection alive
-AUTO_REFRESH_INTERVAL = 240  # 4 minutes
-
-# Create a mock spaces module for fallback
-class MockSpaces:
-    """Mock spaces module for when Zero GPU is not available"""
-    def GPU(self, *args, **kwargs):
-        """Mock GPU decorator that just returns the function as-is"""
-        def decorator(func):
-            return func
-        return decorator
-
-# Try to import spaces for Zero GPU support
-if DISABLE_ZERO_GPU:
+# Simplified Zero GPU handling
+try:
+    import spaces
+    SPACES_AVAILABLE = True
+    print("Zero GPU support detected")
+except ImportError:
     SPACES_AVAILABLE = False
-    spaces = MockSpaces()
-    print("Zero GPU disabled via environment variable")
-else:
-    try:
-        import spaces as real_spaces
-        # Test if spaces is working properly
-        test_decorator = real_spaces.GPU()
-        spaces = real_spaces
-        SPACES_AVAILABLE = True
-        print("Zero GPU support detected and available")
-    except ImportError:
-        SPACES_AVAILABLE = False
-        spaces = MockSpaces()
-        print("Warning: spaces module not available. Running without Zero GPU support.")
-    except Exception as e:
-        SPACES_AVAILABLE = False
-        spaces = MockSpaces()
-        print(f"Warning: Error with spaces module: {e}. Running without Zero GPU support.")
-
-# Flag to track if we should avoid Zero GPU due to initialization errors
-USE_ZERO_GPU = SPACES_AVAILABLE and not DISABLE_ZERO_GPU
+    print("Running without Zero GPU support")
 
-# Global flag to track initialization state
-INIT_FAILED = False
-
-def handle_init_error(signum, frame):
-    """Handle initialization errors gracefully"""
-    global INIT_FAILED
-    INIT_FAILED = True
-    print("Handling initialization error...")
-    sys.exit(1)
-
-# Set up signal handler for graceful shutdown
-signal.signal(signal.SIGTERM, handle_init_error)
+# Keep-alive state
+last_activity = datetime.datetime.now()
+activity_lock = threading.Lock()
 
 def update_activity():
-    """Update last activity time"""
-    global last_activity_time
-    last_activity_time = datetime.datetime.now()
-
-def keep_alive_worker():
-    """Background thread to keep the Space alive"""
-    while True:
-        try:
-            time.sleep(KEEP_ALIVE_INTERVAL)
-            current_time = datetime.datetime.now()
-            time_since_activity = (current_time - last_activity_time).total_seconds()
-            
-            print(f"Keep-alive check: Last activity {time_since_activity:.0f} seconds ago")
-            
-            # Update activity timestamp
-            if time_since_activity > KEEP_ALIVE_INTERVAL:
-                print("Updating activity timestamp...")
-                update_activity()
-                
-                # Create a dummy inference to keep Zero GPU warm
-                try:
-                    if USE_ZERO_GPU and hasattr(score_and_create_matrix_all_singles, '__wrapped__'):
-                        print("Triggering keep-alive inference...")
-                        # Use the smallest possible input
-                        dummy_result = score_and_create_matrix_all_singles(
-                            sequence="MSKGE",
-                            mutation_range_start=1,
-                            mutation_range_end=2,
-                            model_type="Small",
-                            scoring_mirror=False,
-                            batch_size_inference=1
-                        )
-                        print("Keep-alive inference completed")
-                        # Clean up results
-                        if dummy_result and len(dummy_result) > 2:
-                            for file in dummy_result[2]:  # CSV files
-                                try:
-                                    os.remove(file)
-                                except:
-                                    pass
-                except Exception as e:
-                    print(f"Keep-alive inference error (non-critical): {e}")
-        except Exception as e:
-            print(f"Keep-alive thread error: {e}")
-            time.sleep(60)  # Wait a bit before retrying
-
-def warm_up_zero_gpu():
-    """Warm up Zero GPU after idle period"""
-    if not USE_ZERO_GPU:
-        return False
-    
-    print("Warming up Zero GPU...")
-    # Note: Cannot reliably warm up Zero GPU outside of decorated functions
-    # This is a limitation of the Zero GPU system
-    return False
+    """Update last activity timestamp"""
+    global last_activity
+    with activity_lock:
+        last_activity = datetime.datetime.now()
 
 # Add current directory to path
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
@@ -315,7 +217,6 @@ def check_valid_mutant(sequence,mutant,AA_vocab=AA_vocab):
 def cleanup_old_files(max_age_minutes=30):
     """Clean up old inference files"""
     import glob
-    import time
     current_time = time.time()
     patterns = ["fitness_scoring_substitution_matrix_*.png", 
                 "fitness_scoring_substitution_matrix_*.csv",
@@ -344,7 +245,7 @@ def get_mutated_protein(sequence,mutant):
   return ''.join(mutated_sequence)
 
 def score_and_create_matrix_all_singles_impl(sequence,mutation_range_start=None,mutation_range_end=None,model_type="Large",scoring_mirror=False,batch_size_inference=20,max_number_positions_per_heatmap=50,num_workers=0,AA_vocab=AA_vocab):
-  # Update activity time
+  # Update activity
   update_activity()
   
   # Clean up old files periodically
@@ -384,10 +285,6 @@ def score_and_create_matrix_all_singles_impl(sequence,mutation_range_start=None,
   # Device selection - Zero GPU will provide CUDA when decorated with @spaces.GPU
   print(f"GPU Available: {torch.cuda.is_available()}")
   
-  # Try to ensure GPU is available when using Zero GPU
-  if USE_ZERO_GPU and not torch.cuda.is_available():
-    print("Zero GPU enabled but CUDA not available - this is expected before GPU allocation")
-  
   if torch.cuda.is_available():
     device = torch.device("cuda")
     model = model.to(device)
@@ -401,9 +298,6 @@ def score_and_create_matrix_all_singles_impl(sequence,mutation_range_start=None,
     device = torch.device("cpu")
     model = model.to(device)
     print("Inference will take place on CPU")
-    if USE_ZERO_GPU:
-      print("WARNING: Zero GPU is enabled but CUDA is not available!")
-      print("The Space may need to be restarted from the Hugging Face interface.")
     # Reduce batch size for CPU inference
     batch_size_inference = min(batch_size_inference, 10)
     
@@ -460,8 +354,8 @@ def score_and_create_matrix_all_singles_impl(sequence,mutation_range_start=None,
       if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-# Apply Zero GPU decorator - will use real decorator if available, mock otherwise
-if USE_ZERO_GPU:
+# Apply Zero GPU decorator if available
+if SPACES_AVAILABLE:
     score_and_create_matrix_all_singles = spaces.GPU(duration=300)(score_and_create_matrix_all_singles_impl)
 else:
     score_and_create_matrix_all_singles = score_and_create_matrix_all_singles_impl
@@ -475,19 +369,6 @@ def clear_inputs(protein_sequence_input,mutation_range_start,mutation_range_end)
   mutation_range_end = None
   return protein_sequence_input,mutation_range_start,mutation_range_end
 
-# Health check endpoint
-def health_check():
-    """Simple health check that returns current status"""
-    update_activity()
-    status = {
-        "status": "healthy",
-        "zero_gpu": USE_ZERO_GPU,
-        "cuda_available": torch.cuda.is_available(),
-        "last_activity": last_activity_time.isoformat(),
-        "timestamp": datetime.datetime.now().isoformat()
-    }
-    return status
-
 # Create Gradio app
 tranception_design = gr.Blocks()
 
@@ -496,21 +377,29 @@ with tranception_design:
     gr.Markdown("## 🧬 BASIS-China iGEM Team 2025 - Protein Engineering Platform")
     gr.Markdown("### Welcome to BASIS-China's implementation of Tranception on Hugging Face Spaces!")
     gr.Markdown("We are the BASIS-China iGEM team, and we're excited to present our deployment of the Tranception model for protein fitness prediction. This tool enables in silico directed evolution to iteratively improve protein fitness through single amino acid substitutions. At each step, Tranception computes log likelihood ratios for all possible mutations compared to the starting sequence, generating fitness heatmaps and recommendations to guide protein engineering.")
-    gr.Markdown("**Technical Details**: This deployment leverages Hugging Face's Zero GPU infrastructure, which dynamically allocates H200 GPU resources when available. This allows for efficient inference while managing computational resources effectively. The system includes automatic keep-alive mechanisms to maintain GPU availability.")
+    gr.Markdown("**Technical Details**: This deployment leverages Hugging Face's Zero GPU infrastructure, which dynamically allocates H200 GPU resources when available. This allows for efficient inference while managing computational resources effectively.")
+    
+    # Hidden keep-alive component
+    with gr.Row(visible=False):
+        keep_alive_component = gr.Number(value=0, visible=False)
+        
+        def keep_alive_update():
+            update_activity()
+            return time.time()
+        
+        # Update every 2 minutes to keep websocket alive
+        keep_alive_timer = gr.Timer(value=120)
+        keep_alive_timer.tick(keep_alive_update, outputs=[keep_alive_component])
     
     # Status indicator
     with gr.Row():
         with gr.Column(scale=1):
             def get_gpu_status():
-                time_since = (datetime.datetime.now() - last_activity_time).total_seconds()
-                if USE_ZERO_GPU:
-                    if torch.cuda.is_available():
-                        gpu_name = torch.cuda.get_device_name(0)
-                        return f"🔥 Zero GPU Active: {gpu_name} | Last activity: {int(time_since)}s ago"
-                    else:
-                        return f"⚠️ Zero GPU: Ready | Last activity: {int(time_since)}s ago"
-                else:
-                    return "💻 Running on CPU"
+                with activity_lock:
+                    time_since = (datetime.datetime.now() - last_activity).total_seconds()
+                
+                status = "🔥 Zero GPU" if SPACES_AVAILABLE else "💻 CPU Mode"
+                return f"{status} | Last activity: {int(time_since)}s ago"
             
             gpu_status = gr.Textbox(
                 label="Compute Status", 
@@ -520,19 +409,6 @@ with tranception_design:
                 elem_id="gpu_status"
             )
     
-    # Hidden components for keep-alive
-    with gr.Row(visible=False):
-        # Auto-refresh component to maintain WebSocket connection
-        keep_alive_refresh = gr.Number(value=0, visible=False)
-        
-        def increment_counter():
-            update_activity()
-            return gr.update(value=time.time())
-        
-        # This will trigger every 4 minutes to keep the connection alive
-        keep_alive_timer = gr.Timer(value=AUTO_REFRESH_INTERVAL)
-        keep_alive_timer.tick(increment_counter, outputs=[keep_alive_refresh])
-    
     with gr.Tabs():
         with gr.TabItem("Input"):
             with gr.Row():
@@ -622,66 +498,9 @@ with tranception_design:
     gr.Markdown("Links: <a href='https://proceedings.mlr.press/v162/notin22a.html' target='_blank'>Paper</a>  <a href='https://github.com/OATML-Markslab/Tranception' target='_blank'>Code</a>  <a href='https://sites.google.com/view/proteingym/substitutions' target='_blank'>ProteinGym</a>  <a href='https://igem.org/teams/5247' target='_blank'>BASIS-China iGEM Team</a>")
 
 if __name__ == "__main__":
-    # Start keep-alive thread
-    if USE_ZERO_GPU:
-        print("Starting keep-alive thread for Zero GPU...")
-        keep_alive_thread = threading.Thread(target=keep_alive_worker, daemon=True)
-        keep_alive_thread.start()
-        
-        # Schedule periodic dummy inferences to keep alive
-        print("Keep-alive system activated - will perform dummy inferences every 5 minutes")
-    
-    # Configure queue for better resource management
-    tranception_design.queue(
-        max_size=10,  # Limit queue size
-        status_update_rate="auto",  # Show status updates
-        api_open=False  # Disable API to prevent external requests
-    )
-    
-    # Launch with appropriate settings for HF Spaces
-    # Wrap launch in try-except to handle Zero GPU initialization errors gracefully
-    launch_retries = 0
-    max_launch_retries = 3
-    
-    while launch_retries < max_launch_retries:
-        try:
-            # Add a small delay on retries to allow system to stabilize
-            if launch_retries > 0:
-                print(f"Retry attempt {launch_retries}/{max_launch_retries}...")
-                time.sleep(5)
-            
-            tranception_design.launch(
-                max_threads=2,  # Limit concurrent threads
-                show_error=True,
-                server_name="0.0.0.0",
-                server_port=7860,
-                quiet=False,  # Show all logs
-                prevent_thread_lock=True,  # Prevent thread locking issues
-                share=False,  # Don't create public link
-                inbrowser=False  # Don't open browser
-            )
-            break  # If successful, exit the retry loop
-            
-        except RuntimeError as e:
-            error_msg = str(e)
-            if "ZeroGPU" in error_msg or "Unknown" in error_msg:
-                print(f"Zero GPU initialization error: {e}")
-                launch_retries += 1
-                
-                if launch_retries < max_launch_retries:
-                    print("Retrying with Zero GPU after warm-up...")
-                    # Wait longer before retry
-                    time.sleep(10)
-                else:
-                    print("Max retries reached. The Space may need to be restarted.")
-                    print("Note: Zero GPU containers can crash after idle periods.")
-                    print("Consider restarting the Space from the Hugging Face interface.")
-                    sys.exit(1)
-            else:
-                # Non-Zero GPU error, re-raise
-                raise
-        except Exception as e:
-            print(f"Unexpected error during launch: {e}")
-            launch_retries += 1
-            if launch_retries >= max_launch_retries:
-                raise
+    # Simple launch without queue to avoid Zero GPU conflicts
+    tranception_design.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )

keep_alive_external.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""
+External Keep-Alive for Tranception Space
+This script runs on your local machine to keep the Space active
+"""
+import requests
+import time
+from datetime import datetime
+import sys
+
+# Space configuration
+SPACE_EMBED_URL = "https://moraxcheng-transeption-igem-basischina-2025.hf.space"
+GRADIO_API_URL = f"{SPACE_EMBED_URL}/run/predict"
+
+# Keep-alive settings
+PING_INTERVAL = 240  # 4 minutes
+WAKE_UP_RETRIES = 3
+
+def simple_ping():
+    """Simple HTTP GET to keep connection alive"""
+    try:
+        response = requests.get(SPACE_EMBED_URL, timeout=30)
+        return response.status_code == 200
+    except:
+        return False
+
+def gradio_keep_alive():
+    """Send minimal Gradio API request"""
+    try:
+        # Minimal prediction request
+        payload = {
+            "data": [
+                "MSKGE",  # 5 amino acids
+                1,        # start
+                3,        # end  
+                "Small",  # smallest model
+                False,    # no mirror
+                1         # batch size 1
+            ]
+        }
+        
+        # First, initiate the prediction
+        response = requests.post(
+            GRADIO_API_URL,
+            json=payload,
+            timeout=60
+        )
+        
+        if response.status_code == 200:
+            return True
+    except Exception as e:
+        print(f"API request error: {e}")
+    
+    return False
+
+def keep_alive_loop():
+    """Main keep-alive loop"""
+    print("="*60)
+    print("Tranception Space Keep-Alive")
+    print(f"Space: {SPACE_EMBED_URL}")
+    print(f"Ping interval: {PING_INTERVAL} seconds")
+    print("Press Ctrl+C to stop")
+    print("="*60)
+    print()
+    
+    consecutive_failures = 0
+    
+    while True:
+        try:
+            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            
+            # Try simple ping first
+            print(f"[{timestamp}] Pinging Space...", end=" ", flush=True)
+            
+            if simple_ping():
+                print("✓ Online")
+                consecutive_failures = 0
+            else:
+                print("✗ No response")
+                consecutive_failures += 1
+                
+                # Try Gradio API as fallback
+                print(f"[{timestamp}] Trying Gradio API...", end=" ", flush=True)
+                if gradio_keep_alive():
+                    print("✓ Success")
+                    consecutive_failures = 0
+                else:
+                    print("✗ Failed")
+            
+            if consecutive_failures >= 3:
+                print("\n⚠️  WARNING: Space appears to be down!")
+                print("You may need to manually restart it at:")
+                print(f"https://huggingface.co/spaces/MoraxCheng/Transeption_iGEM_BASISCHINA_2025/settings\n")
+            
+            # Wait for next ping
+            print(f"Next ping in {PING_INTERVAL} seconds...\n")
+            time.sleep(PING_INTERVAL)
+            
+        except KeyboardInterrupt:
+            print("\n\nKeep-alive stopped by user")
+            break
+        except Exception as e:
+            print(f"\nUnexpected error: {e}")
+            print("Continuing...\n")
+            time.sleep(60)
+
+if __name__ == "__main__":
+    # Test connection
+    print("Testing connection...")
+    if simple_ping():
+        print("✓ Space is online\n")
+    else:
+        print("⚠ Space appears to be offline")
+        print("Starting keep-alive anyway...\n")
+    
+    # Start keep-alive
+    keep_alive_loop()