@fastino-ai/pioneer-cli 0.1.0 → 0.2.0

package/src/tools/training.ts

@@ -0,0 +1,443 @@
+/**
+ * Model training orchestration tool
+ * Combines Modal for compute and W&B for tracking
+ */
+
+import type { Tool, ToolResult } from "../agent/types.js";
+import { spawn } from "child_process";
+import * as fs from "fs";
+import * as path from "path";
+import * as os from "os";
+
+export interface TrainingConfig {
+  modalTokenId?: string;
+  modalTokenSecret?: string;
+  wandbApiKey?: string;
+  wandbEntity?: string;
+  defaultProject?: string;
+  timeout?: number;
+}
+
+export function createTrainingTool(config: TrainingConfig = {}): Tool {
+  return {
+    name: "train_model",
+    description: "Train a machine learning model using Modal for compute and W&B for tracking. Supports fine-tuning and training from scratch.",
+    parameters: [
+      {
+        name: "training_code",
+        type: "string",
+        description: "Python training code. Should include model definition, data loading, training loop, and logging to W&B.",
+        required: true,
+      },
+      {
+        name: "model_type",
+        type: "string",
+        description: "Type of model: 'transformer', 'cnn', 'custom'",
+        required: true,
+      },
+      {
+        name: "dataset_path",
+        type: "string",
+        description: "Path to training dataset (local or remote URL)",
+        required: true,
+      },
+      {
+        name: "hyperparameters",
+        type: "object",
+        description: "Training hyperparameters as JSON (e.g., {\"learning_rate\": 0.001, \"epochs\": 10})",
+        required: false,
+      },
+      {
+        name: "gpu",
+        type: "string",
+        description: "GPU type for Modal: 'T4', 'A10G', 'A100' (default: T4)",
+        required: false,
+      },
+      {
+        name: "run_name",
+        type: "string",
+        description: "Name for the W&B run",
+        required: true,
+      },
+    ],
+
+    async execute(args: Record<string, unknown>): Promise<ToolResult> {
+      const trainingCode = args.training_code as string;
+      const modelType = args.model_type as string;
+      const datasetPath = args.dataset_path as string;
+      const hyperparameters = (args.hyperparameters as Record<string, unknown>) || {};
+      const gpu = (args.gpu as string) || "T4";
+      const runName = args.run_name as string;
+
+      if (!trainingCode || !modelType || !datasetPath || !runName) {
+        return {
+          toolCallId: "",
+          output: "",
+          success: false,
+          error: "training_code, model_type, dataset_path, and run_name are required",
+        };
+      }
+
+      const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "training-"));
+      const appFile = path.join(tempDir, "training_app.py");
+
+      // Generate Modal training app
+      const modalCode = `
+import modal
+import wandb
+import json
+
+app = modal.App("training-${runName}")
+
+# Define the training image with dependencies
+image = modal.Image.debian_slim(python_version="3.11").pip_install([
+    "torch",
+    "transformers",
+    "datasets",
+    "wandb",
+    "numpy",
+    "scikit-learn",
+])
+
+@app.function(
+    image=image,
+    gpu="${gpu}",
+    timeout=3600,
+    secrets=[
+        modal.Secret.from_name("wandb-secret", required=False),
+    ],
+)
+def train():
+    import os
+    
+    # Initialize W&B
+    wandb.init(
+        project="${config.defaultProject || "pioneer-training"}",
+        name="${runName}",
+        config=${JSON.stringify(hyperparameters)},
+    )
+    
+    # User's training code
+${trainingCode.split('\n').map(line => '    ' + line).join('\n')}
+    
+    wandb.finish()
+    return {"status": "completed", "run_name": "${runName}"}
+
+@app.local_entrypoint()
+def main():
+    result = train.remote()
+    print(json.dumps(result))
+`;
+
+      fs.writeFileSync(appFile, modalCode);
+
+      try {
+        const timeout = config.timeout || 3600000; // 1 hour default
+
+        return new Promise((resolve) => {
+          let stdout = "";
+          let stderr = "";
+          let timedOut = false;
+
+          const env: NodeJS.ProcessEnv = { ...process.env };
+          if (config.modalTokenId) env.MODAL_TOKEN_ID = config.modalTokenId;
+          if (config.modalTokenSecret) env.MODAL_TOKEN_SECRET = config.modalTokenSecret;
+          if (config.wandbApiKey) env.WANDB_API_KEY = config.wandbApiKey;
+
+          const proc = spawn("modal", ["run", appFile], {
+            stdio: ["pipe", "pipe", "pipe"],
+            env,
+          });
+
+          const timer = setTimeout(() => {
+            timedOut = true;
+            proc.kill("SIGTERM");
+          }, timeout);
+
+          proc.stdout.on("data", (data: Buffer) => {
+            stdout += data.toString();
+          });
+
+          proc.stderr.on("data", (data: Buffer) => {
+            stderr += data.toString();
+          });
+
+          proc.on("close", (code) => {
+            clearTimeout(timer);
+
+            let output = stdout;
+            if (stderr) {
+              output += (output ? "\n\nLogs:\n" : "") + stderr;
+            }
+            if (timedOut) {
+              output += "\n\n(Training timed out)";
+            }
+
+            resolve({
+              toolCallId: "",
+              output: output || "(no output)",
+              success: !timedOut && code === 0,
+              error: timedOut
+                ? "Training timed out"
+                : code !== 0
+                ? `Exit code: ${code}`
+                : undefined,
+            });
+          });
+
+          proc.on("error", (err) => {
+            clearTimeout(timer);
+            resolve({
+              toolCallId: "",
+              output: "",
+              success: false,
+              error: err.message,
+            });
+          });
+        });
+      } finally {
+        try {
+          fs.rmSync(tempDir, { recursive: true, force: true });
+        } catch {
+          // Ignore
+        }
+      }
+    },
+  };
+}
+
+export function createFineTuningTool(config: TrainingConfig = {}): Tool {
+  return {
+    name: "finetune_model",
+    description: "Fine-tune a pre-trained model (GPT, BERT, etc.) on custom data using Modal for compute.",
+    parameters: [
+      {
+        name: "base_model",
+        type: "string",
+        description: "Base model to fine-tune (e.g., 'gpt2', 'bert-base-uncased', 'meta-llama/Llama-2-7b')",
+        required: true,
+      },
+      {
+        name: "dataset_path",
+        type: "string",
+        description: "Path to training dataset (JSONL format with 'text' or 'prompt'/'completion' fields)",
+        required: true,
+      },
+      {
+        name: "output_dir",
+        type: "string",
+        description: "Directory to save the fine-tuned model",
+        required: true,
+      },
+      {
+        name: "hyperparameters",
+        type: "object",
+        description: "Training hyperparameters (learning_rate, epochs, batch_size, etc.)",
+        required: false,
+      },
+      {
+        name: "run_name",
+        type: "string",
+        description: "Name for the training run",
+        required: true,
+      },
+      {
+        name: "use_lora",
+        type: "boolean",
+        description: "Use LoRA for efficient fine-tuning (default: true for large models)",
+        required: false,
+      },
+    ],
+
+    async execute(args: Record<string, unknown>): Promise<ToolResult> {
+      const baseModel = args.base_model as string;
+      const datasetPath = args.dataset_path as string;
+      const outputDir = args.output_dir as string;
+      const hyperparameters = (args.hyperparameters as Record<string, unknown>) || {};
+      const runName = args.run_name as string;
+      const useLora = args.use_lora !== false;
+
+      if (!baseModel || !datasetPath || !outputDir || !runName) {
+        return {
+          toolCallId: "",
+          output: "",
+          success: false,
+          error: "base_model, dataset_path, output_dir, and run_name are required",
+        };
+      }
+
+      const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "finetune-"));
+      const appFile = path.join(tempDir, "finetune_app.py");
+
+      const hp = {
+        learning_rate: 2e-5,
+        num_epochs: 3,
+        batch_size: 8,
+        warmup_steps: 100,
+        ...hyperparameters,
+      };
+
+      const modalCode = `
+import modal
+
+app = modal.App("finetune-${runName}")
+
+image = modal.Image.debian_slim(python_version="3.11").pip_install([
+    "torch",
+    "transformers",
+    "datasets",
+    "accelerate",
+    "peft",
+    "bitsandbytes",
+    "wandb",
+    "trl",
+])
+
+@app.function(
+    image=image,
+    gpu="A10G",
+    timeout=7200,
+    secrets=[modal.Secret.from_name("wandb-secret", required=False)],
+)
+def finetune():
+    import torch
+    import wandb
+    from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
+    from datasets import load_dataset
+    from peft import LoraConfig, get_peft_model
+    from trl import SFTTrainer
+    
+    wandb.init(project="${config.defaultProject || "pioneer-finetune"}", name="${runName}")
+    
+    # Load model and tokenizer
+    model = AutoModelForCausalLM.from_pretrained(
+        "${baseModel}",
+        torch_dtype=torch.float16,
+        device_map="auto",
+    )
+    tokenizer = AutoTokenizer.from_pretrained("${baseModel}")
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    ${useLora ? `
+    # Apply LoRA
+    lora_config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=["q_proj", "v_proj"],
+        lora_dropout=0.05,
+        bias="none",
+    )
+    model = get_peft_model(model, lora_config)
+    ` : ""}
+    
+    # Load dataset
+    dataset = load_dataset("json", data_files="${datasetPath}")["train"]
+    
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir="${outputDir}",
+        num_train_epochs=${hp.num_epochs},
+        per_device_train_batch_size=${hp.batch_size},
+        learning_rate=${hp.learning_rate},
+        warmup_steps=${hp.warmup_steps},
+        logging_steps=10,
+        save_steps=500,
+        report_to="wandb",
+    )
+    
+    # Train
+    trainer = SFTTrainer(
+        model=model,
+        train_dataset=dataset,
+        args=training_args,
+        tokenizer=tokenizer,
+    )
+    trainer.train()
+    
+    # Save model
+    trainer.save_model("${outputDir}")
+    tokenizer.save_pretrained("${outputDir}")
+    
+    wandb.finish()
+    return {"status": "completed", "output_dir": "${outputDir}"}
+
+@app.local_entrypoint()
+def main():
+    import json
+    result = finetune.remote()
+    print(json.dumps(result))
+`;
+
+      fs.writeFileSync(appFile, modalCode);
+
+      try {
+        const timeout = config.timeout || 7200000; // 2 hours
+
+        return new Promise((resolve) => {
+          let stdout = "";
+          let stderr = "";
+          let timedOut = false;
+
+          const env: NodeJS.ProcessEnv = { ...process.env };
+          if (config.modalTokenId) env.MODAL_TOKEN_ID = config.modalTokenId;
+          if (config.modalTokenSecret) env.MODAL_TOKEN_SECRET = config.modalTokenSecret;
+          if (config.wandbApiKey) env.WANDB_API_KEY = config.wandbApiKey;
+
+          const proc = spawn("modal", ["run", appFile], {
+            stdio: ["pipe", "pipe", "pipe"],
+            env,
+          });
+
+          const timer = setTimeout(() => {
+            timedOut = true;
+            proc.kill("SIGTERM");
+          }, timeout);
+
+          proc.stdout.on("data", (data: Buffer) => {
+            stdout += data.toString();
+          });
+
+          proc.stderr.on("data", (data: Buffer) => {
+            stderr += data.toString();
+          });
+
+          proc.on("close", (code) => {
+            clearTimeout(timer);
+
+            let output = stdout;
+            if (stderr) output += (output ? "\n\n" : "") + stderr;
+            if (timedOut) output += "\n\n(Fine-tuning timed out)";
+
+            resolve({
+              toolCallId: "",
+              output: output || "(no output)",
+              success: !timedOut && code === 0,
+              error: timedOut ? "Timed out" : code !== 0 ? `Exit code: ${code}` : undefined,
+            });
+          });
+
+          proc.on("error", (err) => {
+            clearTimeout(timer);
+            resolve({
+              toolCallId: "",
+              output: "",
+              success: false,
+              error: err.message,
+            });
+          });
+        });
+      } finally {
+        try {
+          fs.rmSync(tempDir, { recursive: true, force: true });
+        } catch {
+          // Ignore
+        }
+      }
+    },
+  };
+}
+
+export function createTrainingTools(config: TrainingConfig = {}): Tool[] {
+  return [createTrainingTool(config), createFineTuningTool(config)];
+}
+