swarm-assistent/AssistentTrainingJobs.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.WebSockets;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Utils;
using SwarmUI.WebAPI;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>QLoRA training job runner with VRAM lock and progress streaming.</summary>
public partial class SwarmAssistentExtension
{
    static readonly TrainingJobManager TrainingJobManager = new();

    public async Task<JObject> AssistentStartTrainJob(Session session, JObject raw)
    {
        if (raw is null)
        {
            return new JObject { ["error"] = "body required" };
        }
        if (TrainingJobManager.IsRunning)
        {
            return new JObject { ["error"] = "Тренировка уже идёт" };
        }
        string hfBase = raw["base_model"]?.ToString()?.Trim();
        string outputName = raw["output_name"]?.ToString()?.Trim();
        if (string.IsNullOrWhiteSpace(hfBase) || string.IsNullOrWhiteSpace(outputName))
        {
            return new JObject { ["error"] = "base_model and output_name required" };
        }
        string hfDataset = null;
        JObject hfCheck = null;
        JObject hfMapping = raw["hf_mapping"] as JObject;
        if (!string.IsNullOrWhiteSpace(raw["hf_dataset"]?.ToString()))
        {
            hfDataset = NormalizeHfDatasetId(raw["hf_dataset"]?.ToString());
            if (hfDataset is null)
            {
                return new JObject { ["error"] = "invalid hf_dataset id" };
            }
            hfCheck = await CheckHfDatasetInternal(session, hfDataset, useCache: true);
            if (hfCheck["gate"]?.ToString() == "rejected")
            {
                return new JObject { ["error"] = hfCheck["reason"]?.ToString() ?? "hf dataset rejected" };
            }
            hfMapping = ResolveHfMapping(hfCheck, hfMapping);
            if (MappingRequired(hfCheck, hfMapping))
            {
                return new JObject { ["error"] = "Нужен маппинг колонок для HF набора", ["check"] = hfCheck };
            }
        }
        JObject runner = Config.LoadTrainingRunner();
        string kind = runner["kind"]?.ToString()?.Trim();
        if (string.IsNullOrWhiteSpace(kind))
        {
            kind = "builtin";
        }
        string baseUrl = NormalizeBaseUrl(raw["base_url"]?.ToString());
        string chatModel = raw["chat_model"]?.ToString()?.Trim();
        if (!string.IsNullOrWhiteSpace(chatModel))
        {
            await AssistentParkLlm(session, baseUrl, chatModel);
        }
        string datasetPath = null;
        int exportCount = 0;
        if (string.IsNullOrWhiteSpace(hfDataset))
        {
            JObject export = await AssistentExportDataset(session, "approved", "jsonl");
            datasetPath = export["path"]?.ToString();
            exportCount = export["count"]?.Value<int?>() ?? 0;
            if (string.IsNullOrWhiteSpace(datasetPath) || !File.Exists(datasetPath) || exportCount <= 0)
            {
                return new JObject { ["error"] = "Нет одобренных примеров для тренировки (или укажи hf_dataset)" };
            }
        }
        else
        {
            JObject export = await AssistentExportDataset(session, "approved", "jsonl");
            exportCount = export["count"]?.Value<int?>() ?? 0;
            if (exportCount > 0 && File.Exists(export["path"]?.ToString() ?? ""))
            {
                datasetPath = export["path"]?.ToString();
            }
        }
        long now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        string jobId = $"tj_{now}";
        string jobDir = Path.Combine(TrainingRoot(), "jobs", jobId);
        Directory.CreateDirectory(jobDir);
        string configPath = Path.Combine(jobDir, "config.json");
        string logPath = Path.Combine(jobDir, "log.txt");
        string adapterDir = Path.Combine(TrainingRoot(), "adapters", SanitizeAdapterName(outputName));
        Directory.CreateDirectory(adapterDir);
        JObject jobConfig = new()
        {
            ["base_model"] = hfBase,
            ["output_name"] = outputName,
            ["ollama_base"] = raw["ollama_base"]?.ToString()?.Trim(),
            ["gguf_base_path"] = raw["gguf_base_path"]?.ToString()?.Trim() ?? runner["gguf_base_path"]?.ToString()?.Trim(),
            ["dataset_path"] = datasetPath,
            ["hf_dataset"] = hfDataset,
            ["hf_mapping"] = hfMapping,
            ["hf_schema"] = hfCheck?["schema"],
            ["max_samples"] = raw["max_samples"] ?? 0,
            ["rank"] = raw["rank"] ?? 16,
            ["alpha"] = raw["alpha"] ?? 32,
            ["lr"] = raw["lr"] ?? 0.0002,
            ["epochs"] = raw["epochs"] ?? 3,
            ["seq_len"] = raw["seq_len"] ?? 2048,
            ["four_bit"] = raw["four_bit"] ?? true,
            ["batch_size"] = raw["batch_size"] ?? 1,
            ["gradient_accumulation_steps"] = raw["gradient_accumulation_steps"] ?? 4,
            ["adapter_dir"] = adapterDir,
            ["local_export_count"] = exportCount,
        };
        await File.WriteAllTextAsync(configPath, jobConfig.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
        Memory.SaveTrainJob(new JObject
        {
            ["id"] = jobId,
            ["kind"] = "qlora",
            ["status"] = "running",
            ["config"] = jobConfig,
            ["base_model"] = hfBase,
            ["output_name"] = outputName,
            ["log_path"] = logPath,
            ["created_at"] = now,
        });
        string cmdLine = BuildRunnerCommand(runner, configPath, logPath, jobDir);
        bool started = TrainingJobManager.Start(this, session, jobId, cmdLine, logPath, baseUrl, chatModel, GetHfToken(session));
        if (!started)
        {
            Memory.SaveTrainJob(new JObject { ["id"] = jobId, ["status"] = "failed", ["progress"] = new JObject { ["error"] = "process start failed" } });
            return new JObject { ["error"] = "Не удалось запустить процесс тренировки" };
        }
        return new JObject { ["success"] = true, ["job_id"] = jobId, ["log_path"] = logPath };
    }

    static string SanitizeAdapterName(string name)
    {
        if (string.IsNullOrWhiteSpace(name))
        {
            return "adapter";
        }
        char[] bad = Path.GetInvalidFileNameChars();
        StringBuilder sb = new();
        foreach (char c in name)
        {
            sb.Append(Array.IndexOf(bad, c) >= 0 ? '_' : c);
        }
        return sb.ToString().Trim();
    }

    string BuildRunnerCommand(JObject runner, string configPath, string logPath, string workDir)
    {
        string python = runner["python"]?.ToString()?.Trim();
        if (string.IsNullOrWhiteSpace(python))
        {
            python = "python";
        }
        string kind = runner["kind"]?.ToString()?.Trim() ?? "builtin";
        string custom = runner["cmd"]?.ToString()?.Trim();
        string scriptPath = Path.Combine(FilePath, "scripts", "train_qlora.py");
        if (string.Equals(kind, "custom", StringComparison.OrdinalIgnoreCase) && !string.IsNullOrWhiteSpace(custom))
        {
            return custom
                .Replace("{python}", python, StringComparison.OrdinalIgnoreCase)
                .Replace("{config}", configPath, StringComparison.OrdinalIgnoreCase)
                .Replace("{log}", logPath, StringComparison.OrdinalIgnoreCase)
                .Replace("{workdir}", workDir, StringComparison.OrdinalIgnoreCase);
        }
        return $"\"{python}\" \"{scriptPath}\" --config \"{configPath}\" --log \"{logPath}\"";
    }

    public async Task<JObject> AssistentCancelTrainJob(Session session, string id = null)
    {
        await Task.CompletedTask;
        TrainingJobManager.Cancel();
        string jobId = id ?? TrainingJobManager.CurrentJobId;
        if (!string.IsNullOrWhiteSpace(jobId))
        {
            Memory.SaveTrainJob(new JObject
            {
                ["id"] = jobId,
                ["status"] = "cancelled",
                ["finished_at"] = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(),
            });
        }
        return new JObject { ["success"] = true, ["cancelled"] = true };
    }

    public async Task<JObject> AssistentTrainWS(Session session, WebSocket ws, JObject raw)
    {
        await Task.CompletedTask;
        try
        {
            while (TrainingJobManager.IsRunning && ws.State == System.Net.WebSockets.WebSocketState.Open)
            {
                JObject progress = TrainingJobManager.GetProgress();
                string msg = progress.ToString(Newtonsoft.Json.Formatting.None);
                await ws.SendAsync(Encoding.UTF8.GetBytes(msg), System.Net.WebSockets.WebSocketMessageType.Text, true, CancellationToken.None);
                await Task.Delay(800);
            }
            JObject final = TrainingJobManager.GetProgress();
            final["done"] = true;
            await ws.SendAsync(Encoding.UTF8.GetBytes(final.ToString(Newtonsoft.Json.Formatting.None)), System.Net.WebSockets.WebSocketMessageType.Text, true, CancellationToken.None);
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentTrainWS: {ex.Message}");
        }
        return new JObject { ["success"] = true };
    }

    internal async Task FinishTrainJobAsync(
        string jobId,
        bool success,
        string logPath,
        Session session,
        string baseUrl,
        string chatModel,
        JObject jobConfig,
        JObject runner)
    {
        long now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        string adapterDir = jobConfig?["adapter_dir"]?.ToString() ?? "";
        string outputName = jobConfig?["output_name"]?.ToString() ?? "";
        JObject progress = TrainingJobManager.GetProgress();
        string finalStatus = success ? "completed" : "failed";
        if (success && Directory.Exists(adapterDir))
        {
            JObject reg = await RegisterAdapterPipeline(session, baseUrl, outputName, adapterDir, jobConfig, runner);
            progress["ollama"] = reg;
            if (reg["success"]?.Value<bool?>() != true && reg["skipped"]?.Value<bool?>() != true)
            {
                finalStatus = "completed_with_warnings";
            }
        }
        progress["status"] = finalStatus;
        Memory.SaveTrainJob(new JObject
        {
            ["id"] = jobId,
            ["status"] = finalStatus,
            ["finished_at"] = now,
            ["progress"] = progress,
        });
        if (!string.IsNullOrWhiteSpace(chatModel))
        {
            await AssistentWarmLlm(session, baseUrl, chatModel);
        }
        TrainingJobManager.ClearRunning();
    }

    async Task<JObject> RegisterAdapterPipeline(Session session, string baseUrl, string outputName, string adapterDir, JObject jobConfig, JObject runner)
    {
        string safetensors = Directory.GetFiles(adapterDir, "adapter_model.safetensors").FirstOrDefault();
        if (string.IsNullOrWhiteSpace(safetensors))
        {
            return new JObject { ["success"] = false, ["error"] = "adapter_model.safetensors not found" };
        }
        string ggufPath = Directory.GetFiles(adapterDir, "*.gguf").FirstOrDefault();
        string ggufScript = runner?["gguf_script"]?.ToString()?.Trim();
        string ggufBase = jobConfig?["gguf_base_path"]?.ToString()?.Trim() ?? runner?["gguf_base_path"]?.ToString()?.Trim();
        string python = runner?["python"]?.ToString()?.Trim();
        if (string.IsNullOrWhiteSpace(python))
        {
            python = "python";
        }
        if (string.IsNullOrWhiteSpace(ggufPath) && !string.IsNullOrWhiteSpace(ggufScript) && File.Exists(ggufScript))
        {
            if (string.IsNullOrWhiteSpace(ggufBase) || !File.Exists(ggufBase))
            {
                return new JObject
                {
                    ["success"] = false,
                    ["skipped"] = true,
                    ["error"] = "gguf_base_path не задан или файл не найден — адаптер сохранён как safetensors",
                    ["adapter_dir"] = adapterDir,
                };
            }
            ggufPath = Path.Combine(adapterDir, "adapter.gguf");
            string ggufCmd = runner?["gguf_cmd"]?.ToString()?.Trim();
            if (string.IsNullOrWhiteSpace(ggufCmd))
            {
                ggufCmd = "\"{python}\" \"{script}\" \"{base}\" \"{lora}\" \"{out}\"";
            }
            string cmd = ggufCmd
                .Replace("{python}", python, StringComparison.OrdinalIgnoreCase)
                .Replace("{script}", ggufScript, StringComparison.OrdinalIgnoreCase)
                .Replace("{base}", ggufBase, StringComparison.OrdinalIgnoreCase)
                .Replace("{lora}", adapterDir, StringComparison.OrdinalIgnoreCase)
                .Replace("{out}", ggufPath, StringComparison.OrdinalIgnoreCase);
            int code = await RunShellCommandAsync(cmd, adapterDir);
            if (code != 0 || !File.Exists(ggufPath))
            {
                return new JObject
                {
                    ["success"] = false,
                    ["error"] = $"GGUF convert failed exit={code}",
                    ["adapter_dir"] = adapterDir,
                };
            }
        }
        if (string.IsNullOrWhiteSpace(ggufPath) || !File.Exists(ggufPath))
        {
            return new JObject
            {
                ["success"] = false,
                ["skipped"] = true,
                ["note"] = "Настрой convert_lora_to_gguf.py и gguf_base_path для регистрации в Ollama",
                ["adapter_dir"] = adapterDir,
            };
        }
        string ollamaBase = jobConfig?["ollama_base"]?.ToString()?.Trim();
        if (string.IsNullOrWhiteSpace(ollamaBase))
        {
            return new JObject
            {
                ["success"] = false,
                ["skipped"] = true,
                ["error"] = "ollama_base не задан — укажи базовую Ollama-модель на форме QLoRA",
                ["adapter_dir"] = adapterDir,
                ["gguf"] = ggufPath,
            };
        }
        return await RegisterAdapterInOllama(baseUrl, outputName, ollamaBase, ggufPath);
    }

    static async Task<int> RunShellCommandAsync(string commandLine, string workDir)
    {
        try
        {
            ProcessStartInfo psi = new()
            {
                FileName = "cmd.exe",
                Arguments = $"/c {commandLine}",
                UseShellExecute = false,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
                CreateNoWindow = true,
                WorkingDirectory = workDir ?? Environment.CurrentDirectory,
            };
            using Process proc = Process.Start(psi);
            if (proc is null)
            {
                return -1;
            }
            await proc.WaitForExitAsync();
            return proc.ExitCode;
        }
        catch (Exception ex)
        {
            Logs.Debug($"RunShellCommand: {ex.Message}");
            return -1;
        }
    }

    async Task<JObject> RegisterAdapterInOllama(string baseUrl, string outputName, string ollamaBase, string adapterGguf)
    {
        StringBuilder mf = new();
        mf.AppendLine($"FROM {ollamaBase}");
        mf.AppendLine($"ADAPTER {adapterGguf.Replace("\\", "/")}");
        JObject payload = new()
        {
            ["name"] = outputName,
            ["modelfile"] = mf.ToString(),
            ["stream"] = false,
        };
        using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
        using HttpResponseMessage resp = await HttpClient.PostAsync($"{NormalizeBaseUrl(baseUrl)}/api/create", content);
        string body = await resp.Content.ReadAsStringAsync();
        if (!resp.IsSuccessStatusCode)
        {
            return new JObject
            {
                ["success"] = false,
                ["error"] = $"ollama create HTTP {(int)resp.StatusCode}: {Clip(body, 400)}",
                ["modelfile"] = mf.ToString(),
            };
        }
        return new JObject
        {
            ["success"] = true,
            ["name"] = outputName,
            ["ollama_base"] = ollamaBase,
            ["adapter"] = adapterGguf,
            ["response"] = body,
        };
    }
}

sealed class TrainingJobManager
{
    static readonly Regex LossRe = new(@"loss[:\s]+([0-9.]+)", RegexOptions.IgnoreCase | RegexOptions.Compiled);
    static readonly Regex StepRe = new(@"step\s+(\d+)\s*/\s*(\d+)", RegexOptions.IgnoreCase | RegexOptions.Compiled);

    Process _process;
    readonly object _lock = new();
    JObject _progress = new() { ["status"] = "idle" };
    string _logPath;
    SwarmAssistentExtension _ext;
    Session _session;
    string _jobId;
    string _baseUrl;
    string _chatModel;
    long _lastProgressSaveMs;
    int _lastSavedStep = -1;

    public bool IsRunning { get; private set; }
    public string CurrentJobId => _jobId;

    public bool Start(SwarmAssistentExtension ext, Session session, string jobId, string commandLine, string logPath, string baseUrl, string chatModel, string hfToken)
    {
        lock (_lock)
        {
            if (IsRunning)
            {
                return false;
            }
            _ext = ext;
            _session = session;
            _jobId = jobId;
            _logPath = logPath;
            _baseUrl = baseUrl;
            _chatModel = chatModel;
            _lastProgressSaveMs = 0;
            _lastSavedStep = -1;
            _progress = new JObject { ["status"] = "running", ["step"] = 0, ["loss"] = null, ["log"] = "" };
            try
            {
                ProcessStartInfo psi = new()
                {
                    FileName = "cmd.exe",
                    Arguments = $"/c {commandLine}",
                    UseShellExecute = false,
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    CreateNoWindow = true,
                    WorkingDirectory = Path.GetDirectoryName(logPath) ?? Environment.CurrentDirectory,
                };
                if (!string.IsNullOrWhiteSpace(hfToken))
                {
                    psi.Environment["HF_TOKEN"] = hfToken;
                    psi.Environment["HUGGING_FACE_HUB_TOKEN"] = hfToken;
                }
                _process = new Process { StartInfo = psi, EnableRaisingEvents = true };
                _process.OutputDataReceived += (_, e) => AppendLog(e.Data);
                _process.ErrorDataReceived += (_, e) => AppendLog(e.Data);
                _process.Exited += async (_, _) => await OnExited();
                _process.Start();
                _process.BeginOutputReadLine();
                _process.BeginErrorReadLine();
                IsRunning = true;
                return true;
            }
            catch (Exception ex)
            {
                _progress["error"] = ex.Message;
                IsRunning = false;
                return false;
            }
        }
    }

    void AppendLog(string line)
    {
        if (string.IsNullOrWhiteSpace(line))
        {
            return;
        }
        lock (_lock)
        {
            try
            {
                File.AppendAllText(_logPath, line + Environment.NewLine);
            }
            catch
            {
                // ignore
            }
            string prev = _progress["log"]?.ToString() ?? "";
            string combined = prev + line + "\n";
            if (combined.Length > 12000)
            {
                combined = combined[^12000..];
            }
            _progress["log"] = combined;
            Match lossM = LossRe.Match(line);
            if (lossM.Success)
            {
                _progress["loss"] = lossM.Groups[1].Value;
            }
            Match stepM = StepRe.Match(line);
            if (stepM.Success)
            {
                int step = int.Parse(stepM.Groups[1].Value);
                int total = int.Parse(stepM.Groups[2].Value);
                _progress["step"] = step;
                _progress["total_steps"] = total;
                _progress["percent"] = total > 0 ? (int)(100.0 * step / total) : 0;
            }
            MaybeSaveProgress(stepM.Success ? int.Parse(stepM.Groups[1].Value) : -1);
        }
    }

    void MaybeSaveProgress(int step)
    {
        long now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        bool stepChanged = step >= 0 && step != _lastSavedStep;
        if (!stepChanged && now - _lastProgressSaveMs < 2500)
        {
            return;
        }
        _lastProgressSaveMs = now;
        if (step >= 0)
        {
            _lastSavedStep = step;
        }
        try
        {
            _ext?.Memory?.SaveTrainJob(new JObject
            {
                ["id"] = _jobId,
                ["status"] = "running",
                ["progress"] = _progress,
            });
        }
        catch
        {
            // ignore
        }
    }

    async Task OnExited()
    {
        bool ok = false;
        JObject jobConfig = new();
        lock (_lock)
        {
            ok = _process?.ExitCode == 0;
            IsRunning = false;
            _progress["status"] = ok ? "completed" : "failed";
            _progress["exit_code"] = _process?.ExitCode;
        }
        if (_ext != null)
        {
            JObject job = _ext.Memory.GetTrainJob(_jobId);
            try
            {
                string cfgRaw = job?["config_json"]?.ToString();
                if (!string.IsNullOrWhiteSpace(cfgRaw))
                {
                    jobConfig = JObject.Parse(cfgRaw);
                }
            }
            catch
            {
                // ignore
            }
            JObject runner = _ext.Config.LoadTrainingRunner();
            await _ext.FinishTrainJobAsync(_jobId, ok, _logPath, _session, _baseUrl, _chatModel, jobConfig, runner);
        }
    }

    public void Cancel()
    {
        lock (_lock)
        {
            try
            {
                if (_process != null && !_process.HasExited)
                {
                    _process.Kill(entireProcessTree: true);
                }
            }
            catch
            {
                // ignore
            }
            IsRunning = false;
            _progress["status"] = "cancelled";
        }
    }

    public JObject GetProgress()
    {
        lock (_lock)
        {
            return (JObject)_progress.DeepClone();
        }
    }

    public void ClearRunning()
    {
        lock (_lock)
        {
            IsRunning = false;
        }
    }
}