Fix training lifecycle and dataset splitting bugs

This commit is contained in:
srvoyo-cell 2026-07-17 13:23:18 +04:00
parent 06deeb25d1
commit 3a5f502735
11 changed files with 661 additions and 132 deletions

73
.agents/PROJECT_ISSUES.md Normal file
View file

@ -0,0 +1,73 @@
# Исправленные проблемы проекта YOLO Train TUI
Дата исправления и повторной проверки: 2026-07-17
## Итог
Все 11 ранее зафиксированных дефектов исправлены и покрыты регрессионными
проверками.
| ID | Приоритет | Статус | Исправление |
|---|---|---|---|
| BUG-001 | Критический | Исправлено | `subprocess_runner.main()` возвращает код, а `SystemExit` создаётся только снаружи обрабатывающего блока |
| BUG-002 | Высокий | Исправлено | Перед каждым запуском `prepare_run()` сбрасывает состояние остановки |
| BUG-003 | Высокий | Исправлено | Родитель отправляет кооперативный сигнал; принудительный `kill()` используется только после таймаута |
| BUG-004 | Высокий | Исправлено | Запрос, сделанный до готовности subprocess, сохраняется и доставляется после маркера `READY` |
| BUG-005 | Высокий | Исправлено | Detection-style авторазбиение запрещено для `classify` в UI и конфигурации |
| BUG-006 | Средний | Исправлено | `.yaml`/`.yml` разбираются через `yaml.safe_load()`, поле `names` валидируется |
| BUG-007 | Средний | Исправлено | Датасет с одним изображением отклоняется с понятной ошибкой |
| BUG-008 | Средний | Исправлено | Изображения и метки ищутся рекурсивно с сохранением вложенных путей |
| BUG-009 | Средний | Исправлено | Каждый результат создаётся в уникальном `.yolo-tui/splits/<id>` без перезаписи пользовательского `split/` |
| BUG-010 | Низкий | Исправлено | Traceback выводится в журнал TUI; абсолютный путь другого пользователя удалён |
| BUG-011 | Низкий | Исправлено | Явно указанный отсутствующий или некорректный файл классов вызывает точную ошибку без fallback |
## Жизненный цикл обучения
- Дочерний процесс устанавливает обработчики остановки и только затем печатает
`__YOLO_TUI_READY__`.
- Если пользователь нажал «Остановить» раньше, родитель запоминает запрос и
отправляет его после получения маркера готовности.
- Дочерний `TrainingRunner` устанавливает `trainer.stop = True`; Ultralytics
останавливается между пакетами данных, затем выполняет штатную финализацию и
завершающие callbacks.
- Если процесс не завершился за 30 секунд, используется принудительный fallback.
- После завершения ссылка на subprocess и таймер очищаются; перед следующим
запуском флаг остановки сбрасывается.
## Работа с датасетами
- Текущий splitter предназначен для задач `detect`, `segment`, `pose` и `obb`
со структурой `images/` + `labels/`.
- Для `classify` требуется готовый каталог с `train`/`val` и подкаталогами
классов; несовместимый переключатель в UI отключён.
- Поддерживаются `classes.txt`, `.yaml` и `.yml`; YAML может хранить `names` как
список или словарь с последовательными ID от 0.
- Явный путь к классам считается обязательным и не заменяется автопоиском при
опечатке или ошибке формата.
- Split требует минимум два изображения и рекурсивно обрабатывает вложенные
каталоги.
- Файлы каждого запуска создаются эксклюзивно в отдельном управляемом каталоге.
## Проверка
Выполнены команды:
```text
uv run pytest -q
uv run python -m compileall -q src tests
git diff --check
```
Результат: `38 passed`; ошибок компиляции и форматирования diff нет.
Регрессионные тесты проверяют:
1. успешный и ошибочный коды `subprocess_runner.main()`;
2. сброс остановки между запусками;
3. кооперативный сигнал вместо немедленного `terminate()`;
4. доставку раннего запроса после готовности subprocess;
5. запрет авторазбиения для `classify`;
6. пользовательские YAML-файлы классов и ошибочный явный путь;
7. датасеты из одного и двух изображений;
8. вложенные изображения и метки;
9. сохранность пользовательского каталога `split/` и уникальность результатов.

View file

@ -329,9 +329,26 @@ class YoloTrainApp(App[None]):
@on(Switch.Changed, "#split-enabled") @on(Switch.Changed, "#split-enabled")
def toggle_split(self, event: Switch.Changed) -> None: def toggle_split(self, event: Switch.Changed) -> None:
if self.query_one("#task", Select).value == "classify" and event.value:
self.query_one("#split-enabled", Switch).value = False
self.notify(
"Для classify укажите готовый каталог с train/val по классам.",
severity="warning",
)
return
for control in self.query(".split-field Input"): for control in self.query(".split-field Input"):
control.disabled = not event.value control.disabled = not event.value
@on(Select.Changed, "#task")
def task_changed(self, event: Select.Changed) -> None:
split_switch = self.query_one("#split-enabled", Switch)
is_classify = event.value == "classify"
if is_classify and split_switch.value:
split_switch.value = False
split_switch.disabled = is_classify
for control in self.query(".split-field Input"):
control.disabled = is_classify or not split_switch.value
@on(Switch.Changed, "#mlflow-enabled") @on(Switch.Changed, "#mlflow-enabled")
def toggle_mlflow(self, event: Switch.Changed) -> None: def toggle_mlflow(self, event: Switch.Changed) -> None:
for widget_id in ("tracking-uri", "experiment-name", "mlflow-run-name"): for widget_id in ("tracking-uri", "experiment-name", "mlflow-run-name"):
@ -363,6 +380,7 @@ class YoloTrainApp(App[None]):
self.notify(str(exc), title="Проверьте параметры", severity="error") self.notify(str(exc), title="Проверьте параметры", severity="error")
return return
self.runner.prepare_run()
self._set_running(True) self._set_running(True)
progress = self.query_one("#progress", ProgressBar) progress = self.query_one("#progress", ProgressBar)
progress.update(total=config.epochs, progress=0) progress.update(total=config.epochs, progress=0)
@ -387,7 +405,7 @@ class YoloTrainApp(App[None]):
return return
self.runner.request_stop() self.runner.request_stop()
self.query_one("#status-title", Static).update("ОСТАНОВКА") self.query_one("#status-title", Static).update("ОСТАНОВКА")
self.query_one("#status-text", Static).update("Завершаю текущую эпоху") self.query_one("#status-text", Static).update("Корректно останавливаю обучение")
self.query_one("#stop-button", Button).disabled = True self.query_one("#stop-button", Button).disabled = True
self._write_log("[yellow]Запрошена остановка обучения.[/yellow]") self._write_log("[yellow]Запрошена остановка обучения.[/yellow]")
@ -415,7 +433,7 @@ class YoloTrainApp(App[None]):
text=True, text=True,
bufsize=1, bufsize=1,
) )
self.runner.set_subprocess(process) self.runner.set_subprocess(process, ready=False)
output_dir = None output_dir = None
@ -427,7 +445,9 @@ class YoloTrainApp(App[None]):
if not line_str: if not line_str:
continue continue
if line_str.startswith("__YOLO_TUI_EVENT__:"): if line_str == "__YOLO_TUI_READY__":
self.runner.mark_subprocess_ready()
elif line_str.startswith("__YOLO_TUI_EVENT__:"):
try: try:
event_data = json.loads(line_str[len("__YOLO_TUI_EVENT__:"):]) event_data = json.loads(line_str[len("__YOLO_TUI_EVENT__:"):])
event = TrainingEvent( event = TrainingEvent(
@ -446,6 +466,7 @@ class YoloTrainApp(App[None]):
process.wait() process.wait()
rc = process.returncode rc = process.returncode
self.runner.clear_subprocess()
if rc == 0: if rc == 0:
self.app.call_from_thread(self._training_finished, output_dir) self.app.call_from_thread(self._training_finished, output_dir)
@ -458,13 +479,21 @@ class YoloTrainApp(App[None]):
Exception("Процесс обучения завершился с ошибкой. Проверьте логи выше."), Exception("Процесс обучения завершился с ошибкой. Проверьте логи выше."),
) )
except BaseException as exc: except Exception as exc:
import traceback import traceback
try: if process is not None:
with open("/Users/vadim/.gemini/antigravity/brain/ab216120-71db-4a55-a2bf-2eff6a9caf26/error.log", "w", encoding="utf-8") as f: try:
traceback.print_exc(file=f) if process.poll() is None:
except Exception: process.kill()
pass process.wait(timeout=5)
except Exception:
pass
self.runner.clear_subprocess()
details = escape(traceback.format_exc())
self.app.call_from_thread(
self._write_log,
f"[red]{details}[/red]",
)
self.app.call_from_thread(self._training_failed, exc) self.app.call_from_thread(self._training_failed, exc)
finally: finally:
if temp_config_path and os.path.exists(temp_config_path): if temp_config_path and os.path.exists(temp_config_path):

View file

@ -157,6 +157,11 @@ class TrainingConfig:
raise ValueError("Укажите путь или имя модели.") raise ValueError("Укажите путь или имя модели.")
if self.task not in SUPPORTED_TASKS: if self.task not in SUPPORTED_TASKS:
raise ValueError(f"Неизвестный тип задачи: {self.task}.") raise ValueError(f"Неизвестный тип задачи: {self.task}.")
if self.task == "classify" and self.split.enabled:
raise ValueError(
"Автоматическое разделение доступно только для YOLO-датасетов "
"с папками images/labels и не поддерживает задачу classify."
)
if self.epochs < 1: if self.epochs < 1:
raise ValueError("Количество эпох должно быть не меньше 1.") raise ValueError("Количество эпох должно быть не меньше 1.")
if self.image_size < 32: if self.image_size < 32:

View file

@ -2,138 +2,217 @@ from __future__ import annotations
import random import random
from pathlib import Path from pathlib import Path
from typing import Any
from uuid import uuid4
import yaml import yaml
def _normalize_names(names: Any, source: Path) -> dict[int, str]:
if isinstance(names, list):
items = enumerate(names)
elif isinstance(names, dict):
normalized_items: list[tuple[int, Any]] = []
seen: set[int] = set()
for raw_key, value in names.items():
if isinstance(raw_key, bool):
raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.")
try:
class_id = int(raw_key)
except (TypeError, ValueError) as exc:
raise ValueError(
f"Некорректный ID класса в '{source}': {raw_key!r}."
) from exc
if class_id in seen:
raise ValueError(f"Повторяющийся ID класса {class_id} в '{source}'.")
seen.add(class_id)
normalized_items.append((class_id, value))
items = normalized_items
else:
raise ValueError(
f"Поле names в '{source}' должно быть списком или словарём."
)
result: dict[int, str] = {}
for class_id, raw_name in items:
if isinstance(raw_name, (dict, list)):
name = ""
else:
name = str(raw_name).strip()
if not name:
raise ValueError(f"Пустое имя класса {class_id} в '{source}'.")
result[class_id] = name
if not result:
raise ValueError(f"Список классов в '{source}' пуст.")
expected_ids = list(range(len(result)))
if sorted(result) != expected_ids:
raise ValueError(
f"ID классов в '{source}' должны идти подряд, начиная с 0."
)
return dict(sorted(result.items()))
def _parse_text_classes(path: Path) -> dict[int, str]:
with path.open("r", encoding="utf-8") as classes_file:
lines = [line.strip() for line in classes_file if line.strip()]
return _normalize_names(lines, path)
def _load_yaml(path: Path) -> Any:
try:
with path.open("r", encoding="utf-8") as yaml_file:
return yaml.safe_load(yaml_file)
except yaml.YAMLError as exc:
raise ValueError(f"Не удалось разобрать YAML классов '{path}': {exc}.") from exc
def _parse_yaml_classes(path: Path) -> dict[int, str]:
data = _load_yaml(path)
if not isinstance(data, dict) or "names" not in data:
raise ValueError(f"В YAML-файле '{path}' отсутствует поле names.")
return _normalize_names(data["names"], path)
def _parse_classes_file(path: Path) -> dict[int, str]: def _parse_classes_file(path: Path) -> dict[int, str]:
with open(path, "r", encoding="utf-8") as f: suffix = path.suffix.lower()
lines = [line.strip() for line in f if line.strip()] if suffix == ".txt":
return {i: name for i, name in enumerate(lines)} return _parse_text_classes(path)
if suffix in {".yaml", ".yml"}:
return _parse_yaml_classes(path)
raise ValueError(
f"Неподдерживаемый формат файла классов '{path}'. "
"Используйте .txt, .yaml или .yml."
)
def read_classes(dataset_dir: Path, custom_classes_path: str) -> dict[int, str]: def read_classes(dataset_dir: Path, custom_classes_path: str) -> dict[int, str]:
# 1. Custom path specified by user # An explicit path is authoritative: typos and malformed files must not fall back.
if custom_classes_path.strip(): if custom_classes_path.strip():
path = Path(custom_classes_path.strip()) path = Path(custom_classes_path.strip()).expanduser()
if path.exists(): if not path.exists():
return _parse_classes_file(path) raise ValueError(f"Указанный файл классов '{path}' не существует.")
if not path.is_file():
# 2. classes.txt in dataset_dir root raise ValueError(f"Путь к классам '{path}' не является файлом.")
path = dataset_dir / "classes.txt"
if path.exists():
return _parse_classes_file(path)
# 3. classes.txt inside dataset_dir/labels
path = dataset_dir / "labels" / "classes.txt"
if path.exists():
return _parse_classes_file(path)
# 4. Any .yaml file in dataset_dir (excluding split/dataset.yaml)
yaml_files = list(dataset_dir.glob("*.yaml"))
yaml_files = [f for f in yaml_files if f.name != "dataset.yaml"]
if yaml_files:
try: try:
with open(yaml_files[0], "r", encoding="utf-8") as f: return _parse_classes_file(path)
data = yaml.safe_load(f) except OSError as exc:
if isinstance(data, dict) and "names" in data: raise ValueError(f"Не удалось прочитать файл классов '{path}': {exc}.") from exc
names = data["names"]
if isinstance(names, dict):
return {int(k): str(v) for k, v in names.items()}
elif isinstance(names, list):
return {i: str(v) for i, v in enumerate(names)}
except Exception:
pass
# 5. Fallback: Scan label files to determine number of classes and use class_i names for path in (dataset_dir / "classes.txt", dataset_dir / "labels" / "classes.txt"):
class_ids = set() if path.is_file():
return _parse_text_classes(path)
yaml_files = sorted(
(*dataset_dir.glob("*.yaml"), *dataset_dir.glob("*.yml")),
key=lambda item: item.name,
)
for path in yaml_files:
try:
data = _load_yaml(path)
except (OSError, ValueError):
continue
if isinstance(data, dict) and "names" in data:
return _normalize_names(data["names"], path)
# Infer the range from every nested label file when no class list is available.
class_ids: set[int] = set()
labels_dir = dataset_dir / "labels" labels_dir = dataset_dir / "labels"
if labels_dir.exists(): if labels_dir.exists():
for txt_file in labels_dir.glob("*.txt"): for txt_file in labels_dir.rglob("*.txt"):
if txt_file.name == "classes.txt": if txt_file.name == "classes.txt":
continue continue
try: try:
with open(txt_file, "r", encoding="utf-8") as f: with txt_file.open("r", encoding="utf-8") as label_file:
for line in f: for line in label_file:
parts = line.strip().split() parts = line.strip().split()
if parts: if parts:
class_ids.add(int(parts[0])) class_ids.add(int(parts[0]))
except Exception: except (OSError, ValueError):
pass continue
if class_ids: if class_ids:
max_id = max(class_ids) max_id = max(class_ids)
return {i: f"class_{i}" for i in range(max_id + 1)} return {class_id: f"class_{class_id}" for class_id in range(max_id + 1)}
raise ValueError( raise ValueError(
"Не удалось найти список классов. Пожалуйста, создайте файл classes.txt " "Не удалось найти список классов. Создайте classes.txt или YAML с полем "
"в корневой папке датасета или укажите путь к нему." "names в корневой папке датасета либо укажите путь к нему."
) )
def _write_new(path: Path, content: str) -> None:
with path.open("x", encoding="utf-8") as output_file:
output_file.write(content)
def split_dataset( def split_dataset(
dataset_dir: str, train_ratio: float, classes_path: str dataset_dir: str, train_ratio: float, classes_path: str
) -> tuple[int, int, str]: ) -> tuple[int, int, str]:
base_dir = Path(dataset_dir.strip()).absolute() if not dataset_dir.strip():
raise ValueError("Укажите каталог датасета.")
if not 0.1 <= train_ratio <= 0.95:
raise ValueError("Доля обучающей выборки должна быть от 0.1 до 0.95.")
base_dir = Path(dataset_dir.strip()).expanduser().absolute()
images_dir = base_dir / "images" images_dir = base_dir / "images"
labels_dir = base_dir / "labels" labels_dir = base_dir / "labels"
if not base_dir.exists(): if not base_dir.is_dir():
raise ValueError(f"Каталог датасета '{base_dir}' не существует.") raise ValueError(f"Каталог датасета '{base_dir}' не существует.")
if not images_dir.exists(): if not images_dir.is_dir():
raise ValueError(f"Папка с изображениями '{images_dir}' не найдена.") raise ValueError(f"Папка с изображениями '{images_dir}' не найдена.")
if not labels_dir.exists(): if not labels_dir.is_dir():
raise ValueError(f"Папка с разметкой '{labels_dir}' не найдена.") raise ValueError(f"Папка с разметкой '{labels_dir}' не найдена.")
# Find images
valid_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} valid_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
image_files = [ image_files = sorted(
f for f in images_dir.iterdir() (
if f.is_file() and f.suffix.lower() in valid_extensions path
] for path in images_dir.rglob("*")
if path.is_file() and path.suffix.lower() in valid_extensions
),
key=lambda item: item.as_posix(),
)
if not image_files: if not image_files:
raise ValueError(f"В папке '{images_dir}' не найдено изображений.") raise ValueError(f"В папке '{images_dir}' не найдено изображений.")
if len(image_files) == 1:
raise ValueError(
"Для разделения нужно минимум 2 изображения; найдено: 1."
)
# Shuffle deterministically using a fixed seed
rng = random.Random(42) rng = random.Random(42)
rng.shuffle(image_files) rng.shuffle(image_files)
split_idx = int(len(image_files) * train_ratio) split_idx = int(len(image_files) * train_ratio)
if split_idx == 0: split_idx = max(1, min(split_idx, len(image_files) - 1))
split_idx = 1
if split_idx >= len(image_files):
split_idx = len(image_files) - 1
train_images = image_files[:split_idx] train_images = image_files[:split_idx]
val_images = image_files[split_idx:] val_images = image_files[split_idx:]
split_dir = base_dir / "split" # Resolve classes before creating output so invalid input leaves no partial split.
split_dir.mkdir(parents=True, exist_ok=True) classes = read_classes(base_dir, classes_path)
relative_split_dir = Path(".yolo-tui") / "splits" / uuid4().hex
split_dir = base_dir / relative_split_dir
split_dir.mkdir(parents=True, exist_ok=False)
train_txt_path = split_dir / "train.txt" train_txt_path = split_dir / "train.txt"
val_txt_path = split_dir / "val.txt" val_txt_path = split_dir / "val.txt"
with open(train_txt_path, "w", encoding="utf-8") as f:
for img in train_images:
f.write(f"{img}\n")
with open(val_txt_path, "w", encoding="utf-8") as f:
for img in val_images:
f.write(f"{img}\n")
# Load class mapping
classes = read_classes(base_dir, classes_path)
# Write data.yaml equivalent
dataset_yaml_path = split_dir / "dataset.yaml" dataset_yaml_path = split_dir / "dataset.yaml"
_write_new(train_txt_path, "".join(f"{image}\n" for image in train_images))
_write_new(val_txt_path, "".join(f"{image}\n" for image in val_images))
dataset_data = { dataset_data = {
"path": str(base_dir), "path": str(base_dir),
"train": f"split/{train_txt_path.name}", "train": (relative_split_dir / train_txt_path.name).as_posix(),
"val": f"split/{val_txt_path.name}", "val": (relative_split_dir / val_txt_path.name).as_posix(),
"names": classes, "names": classes,
} }
_write_new(
with open(dataset_yaml_path, "w", encoding="utf-8") as f: dataset_yaml_path,
yaml.safe_dump(dataset_data, f, allow_unicode=True, sort_keys=False) yaml.safe_dump(dataset_data, allow_unicode=True, sort_keys=False),
)
return len(train_images), len(val_images), str(dataset_yaml_path) return len(train_images), len(val_images), str(dataset_yaml_path)

View file

@ -1,6 +1,10 @@
import json import json
import os import os
import signal
import sys import sys
import traceback
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from pathlib import Path from pathlib import Path
# Force headless Matplotlib to avoid any thread/process GUI issues # Force headless Matplotlib to avoid any thread/process GUI issues
@ -9,41 +13,66 @@ os.environ["MPLBACKEND"] = "Agg"
from yolo_tui.config import TrainingConfig from yolo_tui.config import TrainingConfig
from yolo_tui.trainer import TrainingEvent, TrainingRunner from yolo_tui.trainer import TrainingEvent, TrainingRunner
def main():
if len(sys.argv) < 2:
print("Usage: python -m yolo_tui.subprocess_runner <config_json_path>", file=sys.stderr)
sys.exit(1)
config_path = sys.argv[1] @contextmanager
def _stop_signal_handlers(runner: TrainingRunner) -> Iterator[None]:
previous: dict[signal.Signals, signal.Handlers] = {}
def request_stop(_signum: int, _frame: object) -> None:
runner.request_stop()
for signum in (signal.SIGTERM, signal.SIGINT):
previous[signum] = signal.getsignal(signum)
signal.signal(signum, request_stop)
try: try:
with open(config_path, "r", encoding="utf-8") as f: yield
config_dict = json.load(f) finally:
config = TrainingConfig.from_dict(config_dict) for signum, handler in previous.items():
except Exception as e: signal.signal(signum, handler)
print(f"Error loading config: {e}", file=sys.stderr)
sys.exit(1)
def main(argv: Sequence[str] | None = None) -> int:
args = list(sys.argv[1:] if argv is None else argv)
if not args:
print(
"Usage: python -m yolo_tui.subprocess_runner <config_json_path>",
file=sys.stderr,
)
return 1
config_path = Path(args[0])
runner = TrainingRunner() runner = TrainingRunner()
runner.prepare_run()
def handle_event(event: TrainingEvent) -> None: with _stop_signal_handlers(runner):
event_dict = { # The parent waits for this marker before sending a cooperative signal.
"kind": event.kind, print("__YOLO_TUI_READY__", flush=True)
"message": event.message, try:
"epoch": event.epoch, with config_path.open("r", encoding="utf-8") as config_file:
"total_epochs": event.total_epochs, config_dict = json.load(config_file)
} config = TrainingConfig.from_dict(config_dict)
# Print structured JSON event so the parent process can parse it except Exception as exc:
print(f"__YOLO_TUI_EVENT__:{json.dumps(event_dict)}", flush=True) print(f"Error loading config: {exc}", file=sys.stderr)
return 1
try: def handle_event(event: TrainingEvent) -> None:
output_dir = runner.train(config, handle_event) event_dict = {
if output_dir: "kind": event.kind,
print(f"__YOLO_TUI_RESULT__:{output_dir}", flush=True) "message": event.message,
sys.exit(0) "epoch": event.epoch,
except BaseException as e: "total_epochs": event.total_epochs,
import traceback }
traceback.print_exc() # Print structured JSON event so the parent process can parse it
sys.exit(1) print(f"__YOLO_TUI_EVENT__:{json.dumps(event_dict)}", flush=True)
try:
output_dir = runner.train(config, handle_event)
if output_dir:
print(f"__YOLO_TUI_RESULT__:{output_dir}", flush=True)
return 0
except Exception:
traceback.print_exc()
return 1
if __name__ == "__main__": if __name__ == "__main__":
main() raise SystemExit(main())

View file

@ -1,11 +1,12 @@
from __future__ import annotations from __future__ import annotations
import os import os
import signal
from collections.abc import Callable, Iterator from collections.abc import Callable, Iterator
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from threading import Event, Lock from threading import Event, RLock, Timer
from typing import Any from typing import Any
from .config import MlflowConfig, TrainingConfig from .config import MlflowConfig, TrainingConfig
@ -51,39 +52,108 @@ def mlflow_environment(config: MlflowConfig) -> Iterator[None]:
class TrainingRunner: class TrainingRunner:
"""Owns a single YOLO training run and exposes cooperative cancellation.""" """Owns a single YOLO training run and exposes cooperative cancellation."""
FORCE_STOP_TIMEOUT_SECONDS = 30.0
def __init__(self) -> None: def __init__(self) -> None:
self._model: Any | None = None self._model: Any | None = None
self._state_lock = Lock() self._state_lock = RLock()
self._stop_requested = Event() self._stop_requested = Event()
self._subprocess: Any | None = None self._subprocess: Any | None = None
self._subprocess_ready = False
self._force_stop_timer: Timer | None = None
def set_subprocess(self, process: Any) -> None: def prepare_run(self) -> None:
"""Reset cancellation state before starting a new training run."""
with self._state_lock:
timer = self._force_stop_timer
self._force_stop_timer = None
self._subprocess_ready = False
self._stop_requested.clear()
if timer is not None:
timer.cancel()
def set_subprocess(self, process: Any, *, ready: bool = False) -> None:
"""Register the child process without losing an earlier stop request."""
with self._state_lock: with self._state_lock:
self._subprocess = process self._subprocess = process
self._subprocess_ready = ready
stop_requested = self._stop_requested.is_set()
if stop_requested:
if ready:
self._send_cooperative_stop(process)
self._schedule_force_stop(process)
def mark_subprocess_ready(self) -> None:
"""Mark the child signal handler as ready and deliver any pending stop."""
with self._state_lock:
process = self._subprocess
self._subprocess_ready = process is not None
stop_requested = self._stop_requested.is_set()
if process is not None and stop_requested:
self._send_cooperative_stop(process)
self._schedule_force_stop(process)
def clear_subprocess(self) -> None: def clear_subprocess(self) -> None:
with self._state_lock: with self._state_lock:
self._subprocess = None self._subprocess = None
self._subprocess_ready = False
timer = self._force_stop_timer
self._force_stop_timer = None
if timer is not None:
timer.cancel()
def request_stop(self) -> None: def request_stop(self) -> None:
self._stop_requested.set() self._stop_requested.set()
with self._state_lock: with self._state_lock:
if self._subprocess is not None: process = self._subprocess
try: process_ready = self._subprocess_ready
self._subprocess.terminate()
except Exception:
pass
trainer = getattr(self._model, "trainer", None) trainer = getattr(self._model, "trainer", None)
if trainer is not None: if process is not None:
trainer.stop = True if process_ready:
self._send_cooperative_stop(process)
self._schedule_force_stop(process)
if trainer is not None:
trainer.stop = True
@property @property
def stop_requested(self) -> bool: def stop_requested(self) -> bool:
return self._stop_requested.is_set() return self._stop_requested.is_set()
@staticmethod
def _send_cooperative_stop(process: Any) -> None:
try:
process.send_signal(signal.SIGTERM)
except (AttributeError, OSError, ProcessLookupError):
try:
process.terminate()
except (AttributeError, OSError, ProcessLookupError):
pass
def _schedule_force_stop(self, process: Any) -> None:
with self._state_lock:
if self._subprocess is not process or self._force_stop_timer is not None:
return
timer = Timer(
self.FORCE_STOP_TIMEOUT_SECONDS,
self._force_stop,
args=(process,),
)
timer.daemon = True
self._force_stop_timer = timer
timer.start()
def _force_stop(self, process: Any) -> None:
with self._state_lock:
if self._subprocess is not process:
return
try:
if process.poll() is None:
process.kill()
except (AttributeError, OSError, ProcessLookupError):
pass
def train(self, config: TrainingConfig, on_event: EventHandler) -> Path | None: def train(self, config: TrainingConfig, on_event: EventHandler) -> Path | None:
config.validate() config.validate()
self._stop_requested.clear()
train_args = config.train_kwargs() train_args = config.train_kwargs()

View file

@ -98,3 +98,21 @@ def test_read_config_ignores_disabled_split() -> None:
assert config.split.train_ratio == DatasetSplitConfig(enabled=False).train_ratio assert config.split.train_ratio == DatasetSplitConfig(enabled=False).train_ratio
asyncio.run(exercise()) asyncio.run(exercise())
def test_classify_disables_detection_style_split() -> None:
async def exercise() -> None:
app = YoloTrainApp()
async with app.run_test(size=(140, 45)) as pilot:
app.query_one("#split-enabled", Switch).value = True
await pilot.pause()
app.query_one("#task", Select).value = "classify"
await pilot.pause()
assert app.query_one("#split-enabled", Switch).value is False
assert app.query_one("#split-enabled", Switch).disabled is True
assert app.query_one("#split-ratio", Input).disabled is True
assert app.query_one("#split-classes", Input).disabled is True
asyncio.run(exercise())

View file

@ -96,3 +96,15 @@ def test_dataset_split_ratio_is_validated(ratio: float) -> None:
config = DatasetSplitConfig(enabled=True, train_ratio=ratio) config = DatasetSplitConfig(enabled=True, train_ratio=ratio)
with pytest.raises(ValueError, match="Доля обучающей выборки"): with pytest.raises(ValueError, match="Доля обучающей выборки"):
config.validate() config.validate()
def test_classification_rejects_detection_style_auto_split() -> None:
config = TrainingConfig(
dataset="classification-dataset",
model="model.pt",
task="classify",
split=DatasetSplitConfig(enabled=True),
)
with pytest.raises(ValueError, match="classify"):
config.validate()

View file

@ -1,7 +1,7 @@
from __future__ import annotations from __future__ import annotations
import os
from pathlib import Path from pathlib import Path
import pytest import pytest
import yaml import yaml
@ -16,6 +16,32 @@ def test_read_classes_custom_path(tmp_path: Path) -> None:
assert classes == {0: "classA", 1: "classB"} assert classes == {0: "classA", 1: "classB"}
@pytest.mark.parametrize(
("names", "expected"),
[
(["cat", "dog"], {0: "cat", 1: "dog"}),
({0: "cat", 1: "dog"}, {0: "cat", 1: "dog"}),
],
)
def test_read_classes_custom_yaml(
tmp_path: Path, names: object, expected: dict[int, str]
) -> None:
custom_file = tmp_path / "custom_classes.yaml"
custom_file.write_text(
yaml.safe_dump({"names": names}, allow_unicode=True),
encoding="utf-8",
)
assert read_classes(tmp_path, str(custom_file)) == expected
def test_missing_custom_classes_path_does_not_fall_back(tmp_path: Path) -> None:
(tmp_path / "classes.txt").write_text("fallback\n", encoding="utf-8")
with pytest.raises(ValueError, match="не существует"):
read_classes(tmp_path, str(tmp_path / "typo.yaml"))
def test_read_classes_root_classes_txt(tmp_path: Path) -> None: def test_read_classes_root_classes_txt(tmp_path: Path) -> None:
classes_file = tmp_path / "classes.txt" classes_file = tmp_path / "classes.txt"
classes_file.write_text("class0\nclass1\nclass2\n", encoding="utf-8") classes_file.write_text("class0\nclass1\nclass2\n", encoding="utf-8")
@ -84,13 +110,15 @@ def test_split_dataset_flow(tmp_path: Path) -> None:
with open(yaml_path, "r", encoding="utf-8") as f: with open(yaml_path, "r", encoding="utf-8") as f:
data = yaml.safe_load(f) data = yaml.safe_load(f)
assert data["path"] == str(tmp_path) assert data["path"] == str(tmp_path)
assert data["train"] == "split/train.txt" relative_output_dir = Path(yaml_path).parent.relative_to(tmp_path)
assert data["val"] == "split/val.txt" assert data["train"] == (relative_output_dir / "train.txt").as_posix()
assert data["val"] == (relative_output_dir / "val.txt").as_posix()
assert data["names"] == {0: "dummy_class"} assert data["names"] == {0: "dummy_class"}
# Verify lists content # Verify lists content
train_list = (tmp_path / "split" / "train.txt").read_text(encoding="utf-8").strip().split("\n") output_dir = Path(yaml_path).parent
val_list = (tmp_path / "split" / "val.txt").read_text(encoding="utf-8").strip().split("\n") train_list = (output_dir / "train.txt").read_text(encoding="utf-8").strip().split("\n")
val_list = (output_dir / "val.txt").read_text(encoding="utf-8").strip().split("\n")
assert len(train_list) == 3 assert len(train_list) == 3
assert len(val_list) == 1 assert len(val_list) == 1
@ -98,3 +126,69 @@ def test_split_dataset_flow(tmp_path: Path) -> None:
# Verify paths are absolute # Verify paths are absolute
assert Path(train_list[0]).is_absolute() assert Path(train_list[0]).is_absolute()
assert Path(val_list[0]).is_absolute() assert Path(val_list[0]).is_absolute()
def test_split_dataset_rejects_single_image_without_writing_output(
tmp_path: Path,
) -> None:
images_dir = tmp_path / "images"
labels_dir = tmp_path / "labels"
images_dir.mkdir()
labels_dir.mkdir()
(images_dir / "only.jpg").write_bytes(b"")
(labels_dir / "only.txt").write_text("0 0.5 0.5 1 1\n", encoding="utf-8")
(tmp_path / "classes.txt").write_text("item\n", encoding="utf-8")
with pytest.raises(ValueError, match="минимум 2"):
split_dataset(str(tmp_path), 0.8, "")
assert not (tmp_path / ".yolo-tui").exists()
def test_split_dataset_finds_nested_images_and_labels(tmp_path: Path) -> None:
images_dir = tmp_path / "images" / "day"
labels_dir = tmp_path / "labels" / "day"
images_dir.mkdir(parents=True)
labels_dir.mkdir(parents=True)
for index in range(2):
(images_dir / f"nested-{index}.jpg").write_bytes(b"")
(labels_dir / f"nested-{index}.txt").write_text(
"2 0.5 0.5 0.2 0.2\n",
encoding="utf-8",
)
train_count, val_count, yaml_path = split_dataset(str(tmp_path), 0.8, "")
assert (train_count, val_count) == (1, 1)
data = yaml.safe_load(Path(yaml_path).read_text(encoding="utf-8"))
assert data["names"] == {0: "class_0", 1: "class_1", 2: "class_2"}
listed_images = "".join(
(Path(yaml_path).parent / filename).read_text(encoding="utf-8")
for filename in ("train.txt", "val.txt")
)
assert "images/day/nested-0.jpg" in listed_images
assert "images/day/nested-1.jpg" in listed_images
def test_split_dataset_preserves_existing_split_files(tmp_path: Path) -> None:
images_dir = tmp_path / "images"
labels_dir = tmp_path / "labels"
images_dir.mkdir()
labels_dir.mkdir()
for index in range(2):
(images_dir / f"image-{index}.jpg").write_bytes(b"")
(labels_dir / f"image-{index}.txt").write_text(
"0 0.5 0.5 0.2 0.2\n",
encoding="utf-8",
)
(tmp_path / "classes.txt").write_text("item\n", encoding="utf-8")
user_split = tmp_path / "split"
user_split.mkdir()
(user_split / "train.txt").write_text("user data\n", encoding="utf-8")
first_yaml = Path(split_dataset(str(tmp_path), 0.5, "")[2])
second_yaml = Path(split_dataset(str(tmp_path), 0.5, "")[2])
assert (user_split / "train.txt").read_text(encoding="utf-8") == "user data\n"
assert first_yaml.parent != second_yaml.parent
assert user_split not in first_yaml.parents

View file

@ -0,0 +1,63 @@
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
from yolo_tui import subprocess_runner
class FakeRunner:
def __init__(self, *, error: Exception | None = None) -> None:
self.error = error
self.prepared = False
self.stop_requested = False
def prepare_run(self) -> None:
self.prepared = True
def request_stop(self) -> None:
self.stop_requested = True
def train(self, config: Any, on_event: Any) -> Path:
if self.error is not None:
raise self.error
return Path("/tmp/successful-run")
def _write_config(tmp_path: Path) -> Path:
config_path = tmp_path / "config.json"
config_path.write_text(
json.dumps({"dataset": "dataset.yaml", "model": "model.pt"}),
encoding="utf-8",
)
return config_path
def test_main_returns_zero_after_successful_training(
monkeypatch: Any, tmp_path: Path, capsys: Any
) -> None:
runner = FakeRunner()
monkeypatch.setattr(subprocess_runner, "TrainingRunner", lambda: runner)
return_code = subprocess_runner.main([str(_write_config(tmp_path))])
output = capsys.readouterr()
assert return_code == 0
assert runner.prepared is True
assert "__YOLO_TUI_READY__" in output.out
assert "__YOLO_TUI_RESULT__:/tmp/successful-run" in output.out
assert "Traceback" not in output.err
def test_main_returns_one_when_training_raises(
monkeypatch: Any, tmp_path: Path, capsys: Any
) -> None:
runner = FakeRunner(error=RuntimeError("training failed"))
monkeypatch.setattr(subprocess_runner, "TrainingRunner", lambda: runner)
return_code = subprocess_runner.main([str(_write_config(tmp_path))])
output = capsys.readouterr()
assert return_code == 1
assert "RuntimeError: training failed" in output.err

View file

@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import sys import sys
import signal
from pathlib import Path from pathlib import Path
from types import ModuleType, SimpleNamespace from types import ModuleType, SimpleNamespace
from typing import Any from typing import Any
@ -9,6 +10,25 @@ from yolo_tui.config import MlflowConfig, TrainingConfig
from yolo_tui.trainer import TrainingEvent, TrainingRunner from yolo_tui.trainer import TrainingEvent, TrainingRunner
class FakeProcess:
def __init__(self) -> None:
self.signals: list[int] = []
self.terminate_calls = 0
self.kill_calls = 0
def send_signal(self, signum: int) -> None:
self.signals.append(signum)
def terminate(self) -> None:
self.terminate_calls += 1
def kill(self) -> None:
self.kill_calls += 1
def poll(self) -> None:
return None
def test_runner_wires_yolo_callbacks_and_returns_output( def test_runner_wires_yolo_callbacks_and_returns_output(
monkeypatch: Any, tmp_path: Path monkeypatch: Any, tmp_path: Path
) -> None: ) -> None:
@ -66,3 +86,40 @@ def test_runner_wires_yolo_callbacks_and_returns_output(
assert train_arguments[0]["data"] == "dataset.yaml" assert train_arguments[0]["data"] == "dataset.yaml"
assert train_arguments[0]["verbose"] is False assert train_arguments[0]["verbose"] is False
assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"] assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]
def test_prepare_run_clears_previous_stop_request() -> None:
runner = TrainingRunner()
runner.request_stop()
assert runner.stop_requested is True
runner.prepare_run()
assert runner.stop_requested is False
def test_early_stop_is_delivered_after_subprocess_ready() -> None:
runner = TrainingRunner()
process = FakeProcess()
runner.request_stop()
runner.set_subprocess(process, ready=False)
assert process.signals == []
runner.mark_subprocess_ready()
assert process.signals == [signal.SIGTERM]
assert process.terminate_calls == 0
runner.clear_subprocess()
def test_ready_subprocess_receives_cooperative_signal_not_terminate() -> None:
runner = TrainingRunner()
process = FakeProcess()
runner.set_subprocess(process, ready=True)
runner.request_stop()
assert process.signals == [signal.SIGTERM]
assert process.terminate_calls == 0
runner.clear_subprocess()