Update main.py

main.py

@@ -164,6 +164,172 @@ def trigger_kaggle_test():
         
     return "\n".join(logs)
 
+def create_worker_notebook(worker_dir, config_name):
+    notebook_content = {
+        "cells": [
+            {
+                "cell_type": "code",
+                "execution_count": None,
+                "metadata": {},
+                "outputs": [],
+                "source": [
+                    "!pip install openai huggingface_hub yt-dlp opencv-python-headless\n"
+                ]
+            },
+            {
+                "cell_type": "code",
+                "execution_count": None,
+                "metadata": {},
+                "outputs": [],
+                "source": [
+                    "import json\n",
+                    "import os\n",
+                    "import subprocess\n",
+                    "import sys\n",
+                    "\n",
+                    f"with open('{config_name}') as f:\n",
+                    "    config = json.load(f)\n",
+                    "\n",
+                    "links = config['links']\n",
+                    "api_key = config['api_key']\n",
+                    "os.environ['FEATHERLESS_API_KEY'] = api_key\n",
+                    "\n",
+                    "print(f'Starting worker with {len(links)} links using API Key: {api_key[:5]}...')\n",
+                    "\n",
+                    "for link in links:\n",
+                    "    print(f'\\n[+] Processing: {link}')\n",
+                    "    # 1. Download Video\n",
+                    "    dl_cmd = [sys.executable, 'worker/youtube_worker/downloader.py', link, '--resolution', 'HD', '-o', 'test_output']\n",
+                    "    res = subprocess.run(dl_cmd, capture_output=True, text=True)\n",
+                    "    if res.returncode != 0:\n",
+                    "        print(f'Download failed:\\n{res.stderr}')\n",
+                    "        continue\n",
+                    "    \n",
+                    "    # Find the downloaded video path from stdout (last line)\n",
+                    "    lines = res.stdout.strip().split('\\n')\n",
+                    "    video_path = lines[-1]\n",
+                    "    if not os.path.exists(video_path):\n",
+                    "        print(f'Could not find downloaded video at: {video_path}')\n",
+                    "        continue\n",
+                    "        \n",
+                    "    print(f'[+] Video downloaded to: {video_path}')\n",
+                    "    \n",
+                    "    # 2. Curate Model (Video Worker)\n",
+                    "    prompt = 'Analyze this video comprehensively. Focus on visuals, objects, and text.'\n",
+                    "    vw_cmd = [sys.executable, 'worker/models_worker/video_worker.py', '--video-url', video_path, '--prompt', prompt, '--api-key', api_key]\n",
+                    "    print('[+] Running model curation...')\n",
+                    "    subprocess.run(vw_cmd)\n",
+                    "    \n",
+                    "    # 3. Optional: Upload to HF\n",
+                    "    # hf_cmd = [sys.executable, 'worker/hf_uploader.py', 'test_output']\n",
+                    "    # subprocess.run(hf_cmd)\n",
+                    "\n",
+                    "print('\\nAll tasks completed!')\n"
+                ]
+            }
+        ],
+        "metadata": {
+            "kernelspec": {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "python3"
+            },
+            "language_info": {
+                "name": "python"
+            }
+        },
+        "nbformat": 4,
+        "nbformat_minor": 4
+    }
+    
+    with open(worker_dir / "notebook.ipynb", "w") as f:
+        json.dump(notebook_content, f, indent=2)
+
+def trigger_kaggle_workers(db_path, links_per_worker, max_workers):
+    api_keys, links = fetch_data(db_path, limit_links=int(links_per_worker) * int(max_workers))
+    if not api_keys:
+        return "Error: No API keys found in DB."
+    if not links:
+        return "Error: No pending links found in DB."
+        
+    username = get_kaggle_username()
+    workers_dir = Path("kaggle_workers_out")
+    if workers_dir.exists():
+        shutil.rmtree(workers_dir)
+    workers_dir.mkdir(parents=True)
+    
+    logs = []
+    logs.append(f"Loaded {len(api_keys)} API keys and {len(links)} links.")
+    
+    link_chunks = [links[i:i + int(links_per_worker)] for i in range(0, len(links), int(links_per_worker))]
+    logs.append(f"Divided into {len(link_chunks)} worker chunks.")
+    
+    for idx, chunk in enumerate(link_chunks):
+        worker_id = f"worker-{uuid.uuid4().hex[:8]}"
+        worker_title = f"yt-worker-{idx+1}-{uuid.uuid4().hex[:4]}"
+        
+        api_key = api_keys[idx % len(api_keys)]
+        worker_path = workers_dir / worker_id
+        worker_path.mkdir(parents=True)
+        
+        # Copy worker if exists locally
+        if Path("worker").exists():
+            shutil.copytree("worker", worker_path / "worker")
+        else:
+            # We copy from parent directory if this is running locally, otherwise it's an error on HF unless packed
+            parent_worker = Path("../hf_space_worker/worker")
+            if parent_worker.exists():
+                shutil.copytree(parent_worker, worker_path / "worker")
+            else:
+                logs.append("Warning: 'worker' directory not found. Kaggle notebook may fail.")
+                
+        if Path("cookies.txt").exists():
+            shutil.copy("cookies.txt", worker_path / "cookies.txt")
+        elif Path("../cookies.txt").exists():
+            shutil.copy("../cookies.txt", worker_path / "cookies.txt")
+            
+        config_name = "worker_config.json"
+        config_data = {
+            "links": chunk,
+            "api_key": api_key,
+        }
+        with open(worker_path / config_name, "w") as f:
+            json.dump(config_data, f, indent=2)
+            
+        create_worker_notebook(worker_path, config_name)
+        
+        metadata = {
+            "id": f"{username}/{worker_title}",
+            "title": worker_title,
+            "code_file": "notebook.ipynb",
+            "language": "python",
+            "kernel_type": "notebook",
+            "is_private": True,
+            "enable_gpu": False,
+            "enable_internet": True,
+            "dataset_sources": [],
+            "competition_sources": [],
+            "kernel_sources": []
+        }
+        with open(worker_path / "kernel-metadata.json", "w") as f:
+            json.dump(metadata, f, indent=2)
+            
+        logs.append(f"Prepared Worker {idx+1}: {worker_title} with {len(chunk)} links.")
+        
+        command = ["kaggle", "kernels", "push", "-p", str(worker_path)]
+        logs.append(f"Pushing {worker_title} to Kaggle...")
+        try:
+            res = subprocess.run(command, capture_output=True, text=True)
+            if res.returncode == 0:
+                logs.append("Push successful.")
+            else:
+                logs.append("Push failed.")
+                logs.append(res.stderr)
+        except Exception as e:
+            logs.append(f"Error: {e}")
+            
+    return "\n".join(logs)
+
 with gr.Blocks() as app:
     gr.Markdown("# YouTube Video Orchestrator Host")
     gr.Markdown("Deploy this to Hugging Face Spaces or run locally. It fetches pending links from your database and delegates them to FastAPI workers or Kaggle.")
@@ -182,10 +348,18 @@ with gr.Blocks() as app:
 
         with gr.Tab("Kaggle Workers"):
             gr.Markdown("Trigger Kaggle notebook workers. Note: You must mount `kaggle.json` to `~/.kaggle/kaggle.json` in the container for this to work.")
-            test_run_btn = gr.Button("Run Simple Kaggle Test Worker", variant="secondary")
-            kaggle_log = gr.Textbox(label="Kaggle Logs", lines=10)
             
-            test_run_btn.click(fn=trigger_kaggle_test, inputs=[], outputs=kaggle_log)
+            with gr.Row():
+                k_db_input = gr.Textbox(label="DB Path", value="/data/db/accounts.db")
+                k_links_per_worker = gr.Number(label="Links per worker", value=2)
+                k_max_workers = gr.Number(label="Max Workers", value=5)
+                
+            k_run_btn = gr.Button("Run Kaggle Orchestration", variant="primary")
+            k_test_run_btn = gr.Button("Run Simple Kaggle Test Worker", variant="secondary")
+            kaggle_log = gr.Textbox(label="Kaggle Logs", lines=15)
+            
+            k_run_btn.click(fn=trigger_kaggle_workers, inputs=[k_db_input, k_links_per_worker, k_max_workers], outputs=kaggle_log)
+            k_test_run_btn.click(fn=trigger_kaggle_test, inputs=[], outputs=kaggle_log)
 
 if __name__ == "__main__":
     app.launch(server_name="0.0.0.0", server_port=7860)