{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# AI Studio Worker - Google Colab\n",
        "\n",
        "Ce notebook prépare une machine GPU Google Colab pour qu’elle se connecte à ton **AI Studio Hugging Face Space**.\n",
        "\n",
        "Il fait :\n",
        "\n",
        "1. cloner le projet AI Studio,\n",
        "2. installer les dépendances Python,\n",
        "3. créer un tunnel ngrok,\n",
        "4. lancer le worker GPU,\n",
        "5. afficher les logs en direct."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Configuration\n",
        "\n",
        "Remplace :\n",
        "\n",
        "- `SERVER_URL` par l’URL de ton Space Hugging Face,\n",
        "- `WORKER_TOKEN` par le token worker généré dans l’interface AI Studio."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "SERVER_URL = \"https://ton-space.hf.space\"\n",
        "WORKER_TOKEN = \"TON_WORKER_TOKEN\"\n",
        "WORKER_PORT = 8765\n",
        "\n",
        "# Optionnel : si tu as un token ngrok, tu peux le mettre ici.\n",
        "# Sinon, ngrok peut fonctionner avec une URL temporaire.\n",
        "NGROK_AUTH_TOKEN = \"\"\n",
        "\n",
        "REPO_URL = \"https://github.com/NathMen12/AI-Studio.git\"\n",
        "PROJECT_DIR = \"/content/AI-Studio\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. Cloner le projet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "if not os.path.exists(PROJECT_DIR):\n",
        "    !git clone $REPO_URL $PROJECT_DIR\n",
        "\n",
        "%cd $PROJECT_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 3. Installer les dépendances GPU\n",
        "\n",
        "Cette cellule installe les dépendances du worker et de l’entraînement Hugging Face LoRA/PEFT."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install -q -r worker/requirements.txt pyngrok"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Démarrer ngrok\n",
        "\n",
        "L’URL affichée ici est celle à enregistrer dans l’interface AI Studio si elle n’y est pas déjà."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from pyngrok import ngrok\n",
        "\n",
        "if NGROK_AUTH_TOKEN:\n",
        "    ngrok.set_auth_token(NGROK_AUTH_TOKEN)\n",
        "\n",
        "public_url = ngrok.connect(WORKER_PORT, \"http\").public_url\n",
        "\n",
        "with open(\"/content/ai_studio_ngrok_url.txt\", \"w\", encoding=\"utf-8\") as file:\n",
        "    file.write(public_url)\n",
        "\n",
        "print(\"URL ngrok du worker :\")\n",
        "print(public_url)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 5. Lancer le worker GPU\n",
        "\n",
        "Le worker va se connecter automatiquement au serveur AI Studio avec :\n",
        "\n",
        "- ton URL de Space,\n",
        "- ton worker token,\n",
        "- l’URL ngrok."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import subprocess\n",
        "import threading\n",
        "\n",
        "command = [\n",
        "    \"python\",\n",
        "    \"worker.py\",\n",
        "    \"--server-url\", SERVER_URL,\n",
        "    \"--user-token\", WORKER_TOKEN,\n",
        "    \"--worker-name\", \"Colab GPU\",\n",
        "    \"--ngrok-url\", public_url,\n",
        "    \"--port\", str(WORKER_PORT)\n",
        "]\n",
        "\n",
        "print(\"Commande lancée :\")\n",
        "print(\" \".join(command))\n",
        "\n",
        "worker_process = subprocess.Popen(\n",
        "    command,\n",
        "    stdout=subprocess.PIPE,\n",
        "    stderr=subprocess.STDOUT,\n",
        "    text=True,\n",
        "    bufsize=1\n",
        ")\n",
        "\n",
        "def read_logs():\n",
        "    for line in worker_process.stdout:\n",
        "        print(line, end=\"\")\n",
        "\n",
        "log_thread = threading.Thread(target=read_logs, daemon=True)\n",
        "log_thread.start()\n",
        "\n",
        "print(\"Worker lancé. Retourne dans l’interface AI Studio pour vérifier que la machine est online.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Arrêter le worker et ngrok\n",
        "\n",
        "À exécuter quand tu as terminé."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "try:\n",
        "    worker_process.terminate()\n",
        "    worker_process.wait(timeout=10)\n",
        "except Exception:\n",
        "    pass\n",
        "\n",
        "try:\n",
        "    ngrok.disconnect()\n",
        "    ngrok.kill()\n",
        "except Exception:\n",
        "    pass\n",
        "\n",
        "print(\"Worker et ngrok arrêtés.\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}