7d11dfa688
Implementation of a mock scheduler for jobs/tasks that will evolve into a job manager for running training jobs over vast.ai instances.
17 lines
420 B
Python
17 lines
420 B
Python
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)
|