improve wandb integration

app.py

@@ -17,6 +17,7 @@ import socket
 
 import gradio as gr
 import pandas as pd
+import wandb
 from autotrain.project import AutoTrainProject
 from autotrain.params import (
     LLMTrainingParams,
@@ -198,10 +199,25 @@ def run_training_background(run_id: str, params: Any, backend: str):
         # Create AutoTrain project - this will handle W&B internally
         project = AutoTrainProject(params=params, backend=backend, process=True)
 
-        # Generate approximate W&B URL
+        # Actually run the training - this blocks until completion
+        print(f"Executing training job for run {run_id}...")
+        result = project.create()
+
+        print(f"Training completed successfully for run {run_id}")
+        print(f"Result: {result}")
+
+        # Get the actual W&B run URL after training starts
         wandb_url = f"https://wandb.ai/{WANDB_PROJECT}"
+        try:
+            if wandb.run is not None:
+                wandb_url = wandb.run.url
+                print(f"Got actual W&B URL: {wandb_url}")
+            else:
+                print("No active W&B run found, using default URL")
+        except Exception as e:
+            print(f"Could not get W&B URL: {e}")
 
-        # Update with W&B URL
+        # Update with actual W&B URL
         runs = load_runs()
         for run in runs:
             if run["run_id"] == run_id:
@@ -209,13 +225,6 @@ def run_training_background(run_id: str, params: Any, backend: str):
                 break
         save_runs(runs)
 
-        # Actually run the training - this blocks until completion
-        print(f"Executing training job for run {run_id}...")
-        result = project.create()
-
-        print(f"Training completed successfully for run {run_id}")
-        print(f"Result: {result}")
-
         # Update status to completed
         runs = load_runs()
         for run in runs:
@@ -604,6 +613,25 @@ def get_system_status(random_string: str = "") -> str:
             "spaces-l4x4",
         ]
 
+        # Extract nested expressions to avoid f-string nesting
+        tasks_list = chr(10).join(f"  • {task}" for task in available_tasks)
+        backends_list = chr(10).join(
+            f"  • {backend}" for backend in available_backends[:10]
+        )
+        backends_more = (
+            f"  ... and {len(available_backends) - 10} more"
+            if len(available_backends) > 10
+            else ""
+        )
+        wandb_api_status = (
+            "✅ Configured" if os.environ.get("WANDB_API_KEY") else "❌ Missing"
+        )
+        wandb_metrics_status = (
+            "✅ Enabled"
+            if os.environ.get("WANDB_API_KEY")
+            else "❌ System metrics only"
+        )
+
         status_text = f"""🚀 AutoTrain Gradio MCP Server - System Status
 
 **Server Status:** Running
@@ -613,15 +641,11 @@ def get_system_status(random_string: str = "") -> str:
 **Failed Runs:** {failed_runs}
 
 **Available Tasks:** {len(available_tasks)}
-{chr(10).join(f"  • {task}" for task in available_tasks)}
+{tasks_list}
 
 **Available Backends:** {len(available_backends)}
-{chr(10).join(f"  • {backend}" for backend in available_backends[:10])}
-{
-            f"  ... and {len(available_backends) - 10} more"
-            if len(available_backends) > 10
-            else ""
-        }
+{backends_list}
+{backends_more}
 
 💡 **Access Points:**
 • Gradio UI: http://SPACE_URL
@@ -630,12 +654,8 @@ def get_system_status(random_string: str = "") -> str:
 
 🛠️ **W&B Integration:**
 • Project: {WANDB_PROJECT}
-• API Key: {"✅ Configured" if os.environ.get("WANDB_API_KEY") else "❌ Missing"}
-• Training Metrics: {
-            "✅ Enabled"
-            if os.environ.get("WANDB_API_KEY")
-            else "❌ System metrics only"
-        }
+• API Key: {wandb_api_status}
+• Training Metrics: {wandb_metrics_status}
 • Set WANDB_API_KEY for complete training metrics logging"""
 
         return status_text
@@ -742,7 +762,7 @@ with gr.Blocks(
     """,
 ) as app:
     gr.Markdown("""
-    # 🚀 AutoTrain Gradio MCP Server
+    # 🚀 AutoTrain MCP Server
     
     Get your AI models to train your AI models! 
     
@@ -856,6 +876,17 @@ with gr.Blocks(
 
         # MCP Info Tab
         with gr.Tab("🔗 MCP Integration"):
+            # Extract nested expressions to avoid f-string nesting in Gradio markdown
+            total_runs = len(load_runs())
+            wandb_auth = (
+                "✅ Configured"
+                if os.environ.get("WANDB_API_KEY")
+                else "❌ Missing WANDB_API_KEY"
+            )
+            hub_auth = (
+                "✅ Configured" if os.environ.get("HF_TOKEN") else "❌ Missing HF_TOKEN"
+            )
+
             gr.Markdown(f"""
             ## MCP Server Information
             
@@ -905,24 +936,39 @@ with gr.Blocks(
             - `push_to_hub="true"` - Push to Hub using project name as repo
             - `hub_repo_id="my-org/my-model"` - Push to custom repository
             
-            ### Claude Desktop Configuration:
+            ### Connection to the MCP Server
             
-            ```json
-            {{
-              "mcpServers": {{
-                "autotrain": {{
-                  "url": "http://localhost:7860/gradio_api/mcp/sse"
-                }}
-              }}
-            }}
+            Connect to it like this:
+    
+            ```javascript
+            {"mcpServers": {"autotrain": {"url": "http://SPACE_URL/gradio_api/mcp/sse",
+                        "headers": {"Authorization": "Bearer <YOUR-HUGGING-FACE-TOKEN>"}
+                    }
+                }
+            }
+            ```
+            
+            Or like this for Claude Desktop:
+            
+            ```javascript
+            {"mcpServers": {"autotrain": {"command": "npx",
+                        "args": [
+                            "mcp-remote",
+                            "http://SPACE_URL/gradio_api/mcp/sse",
+                            "--header",
+                            "Authorization: Bearer <YOUR-HUGGING-FACE-TOKEN>"
+                        ]
+                    }
+                }
+            }
             ```
             
             ### Current Stats:
             
-            Total Runs: {len(load_runs())}
+            Total Runs: {total_runs}
             W&B Project: {WANDB_PROJECT}
-            W&B Auth: {"✅ Configured" if os.environ.get("WANDB_API_KEY") else "❌ Missing WANDB_API_KEY"}
-            Hub Auth: {"✅ Configured" if os.environ.get("HF_TOKEN") else "❌ Missing HF_TOKEN"}
+            W&B Auth: {wandb_auth}
+            Hub Auth: {hub_auth}
             """)
 
         # MCP Tools Tab