core/task_manager.py

import time
from datetime import datetime
from typing import Optional, Dict, List, Any

from config import MODEL_PRESETS, DATA_SOURCES, config
from core.state import (
    tasks, completed_task_ids, results,
    training_active, training_target, training_total_tasks,
    training_completed_tasks, task_counter,
    result_save_counter, RESULT_SAVE_BATCH,
    add_task, update_task, get_task,
    get_pending_tasks, get_assigned_tasks, get_completed_tasks,
    workers,
)
from core.database import save_results_to_hf

def create_tasks(target_size: str, language: str = "python", data_source: str = None) -> Dict[str, Any]:
    global task_counter, training_active, training_target, training_total_tasks, training_completed_tasks

    target_size_lower = target_size.lower()
    matched_target = None
    for key in MODEL_PRESETS.keys():
        if key.lower() == target_size_lower:
            matched_target = key
            break
    if matched_target is None:
        raise ValueError(f"Unknown target size: {target_size}. Available: {', '.join(list(MODEL_PRESETS.keys()))}")

    target_size = matched_target
    language = language or "python"
    data_source = data_source or config.get("data_source") or DATA_SOURCES[0]
    if data_source not in DATA_SOURCES:
        raise ValueError(f"Unknown data source: {data_source}")

    config.set("target_size", target_size)
    config.set("data_source", data_source)
    config.set("language", language)

    preset = MODEL_PRESETS[target_size]
    total_steps = preset.get("tasks", 1000)
    steps_per_task = config.get("steps_per_task", 100)
    num_tasks = max(1, total_steps // steps_per_task)

    tasks.clear()
    completed_task_ids.clear()
    global results
    results = []
    task_counter = 0

    for i in range(num_tasks):
        task_counter += 1
        task_id = f"task-{str(task_counter).zfill(6)}"
        tasks[task_id] = {
            "task_id": task_id,
            "type": "training",
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "steps": steps_per_task,
            "language": language,
            "target_size": target_size,
            "data_source": data_source,
            "assigned_to": None,
            "assigned_at": None,
            "completed_at": None,
            "result": None,
            "loss": None,
        }

    training_active = True
    training_target = target_size
    training_total_tasks = num_tasks
    training_completed_tasks = 0

    print(f"🎯 Created training tasks: {target_size}, total steps {total_steps}, tasks {num_tasks}, steps per task {steps_per_task}")
    return {
        "target_size": target_size,
        "total_steps": total_steps,
        "num_tasks": num_tasks,
        "steps_per_task": steps_per_task,
        "language": language,
        "data_source": data_source,
    }

def get_next_task_for_worker(worker_id: str) -> Optional[Dict[str, Any]]:
    if not training_active:
        return None
    if worker_id not in workers:
        return None

    capabilities = workers[worker_id].get("capabilities", {})
    languages = capabilities.get("languages", [config.get("language")])

    for task_id, task in tasks.items():
        if task.get("status") == "pending":
            task_lang = task.get("language", config.get("language"))
            if task_lang in languages:
                task["status"] = "assigned"
                task["assigned_to"] = worker_id
                task["assigned_at"] = datetime.now().isoformat()
                print(f"📤 Assigned task {task_id} to {worker_id}")
                return task
    return None

def complete_task(task_id: str, worker_id: str, loss: float, steps: int, result_data: Dict) -> bool:
    global training_completed_tasks, training_active, result_save_counter

    if task_id not in tasks:
        print(f"❌ Task {task_id} not found")
        return False
    task = tasks[task_id]
    if task.get("status") != "assigned":
        print(f"❌ Task {task_id} status not assigned")
        return False

    task["status"] = "completed"
    task["completed_at"] = datetime.now().isoformat()
    task["result"] = result_data
    task["loss"] = loss
    task["steps"] = steps
    completed_task_ids.add(task_id)

    if worker_id in workers:
        workers[worker_id]["completed_tasks"] = workers[worker_id].get("completed_tasks", 0) + 1
        workers[worker_id]["current_task"] = None
        workers[worker_id]["status"] = "idle"

    result_entry = {
        "task_id": task_id,
        "worker_id": worker_id,
        "loss": loss,
        "steps": steps,
        "data_source": task.get("data_source", "unknown"),
        "target_size": task.get("target_size", "unknown"),
        "completed_at": datetime.now().isoformat(),
    }
    results.append(result_entry)
    result_save_counter += 1
    if result_save_counter >= RESULT_SAVE_BATCH:
        save_results_to_hf(results)
        result_save_counter = 0

    training_completed_tasks += 1
    print(f"✅ Task {task_id} completed, loss={loss:.4f} (progress: {training_completed_tasks}/{training_total_tasks})")

    if training_completed_tasks >= training_total_tasks:
        training_active = False
        save_results_to_hf(results)
        print(f"🎉 All tasks completed! Total {training_total_tasks} tasks")

    return True

def get_task_status() -> Dict[str, Any]:
    pending = sum(1 for t in tasks.values() if t.get("status") == "pending")
    assigned = sum(1 for t in tasks.values() if t.get("status") == "assigned")
    completed = sum(1 for t in tasks.values() if t.get("status") == "completed")
    progress = 0.0
    if training_total_tasks > 0:
        progress = (training_completed_tasks / training_total_tasks * 100)
    return {
        "total": len(tasks),
        "pending": pending,
        "assigned": assigned,
        "completed": completed,
        "training_active": training_active,
        "target": training_target,
        "completed_tasks": training_completed_tasks,
        "total_tasks": training_total_tasks,
        "progress": progress,
        "total_completed": len(completed_task_ids),
        "data_source": config.get("data_source"),
        "language": config.get("language"),
    }

def delete_completed_tasks() -> int:
    to_delete = [tid for tid, t in tasks.items() if t.get("status") == "completed"]
    for tid in to_delete:
        del tasks[tid]
        if tid in completed_task_ids:
            completed_task_ids.remove(tid)
    global results
    results = [r for r in results if r.get("task_id") not in to_delete]
    print(f"🗑️ Deleted {len(to_delete)} completed tasks")
    return len(to_delete)

def clear_all_tasks() -> int:
    count = len(tasks)
    tasks.clear()
    completed_task_ids.clear()
    global results, training_completed_tasks, training_active
    results = []
    training_completed_tasks = 0
    training_active = False
    return count