from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class Job:
    script: str
    num_nodes: int = 2
    min_vram: int = 8
    docker_image: str = "pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime"
    rounds: int = 10
    local_steps: int = 100
    lr: float = 0.01
    batch_size: int = 64
    checkpoint_dir: str | None = None
    env: dict[str, str] = field(default_factory=dict)