aquin.run.Run Class Reference

Collaboration diagram for aquin.run.Run:

Public Member Functions
None	__init__ (self, str|None base_model=None, str|None run_name=None, dict[str, Any]|None config=None, *, str quant="none", str mode="external", str|Path|None run_dir=None)
None	log (self, int step, *, float loss, float|None learning_rate=None, float|None grad_norm=None, float|None momentum_norm=None, int|None epoch=None, int|None batch=None, int|None total_batches=None, float|None step_ms=None, **float extra)
None	observe (self, int step, str channel, float value)
None	signal (self, int step, str message, *, str severity="info")
None	checkpoint (self, Any model, int step)
None	finish (self, dict[str, Any]|None config=None)
str	run_id (self)

Protected Member Functions
None	_write_meta (self)
None	_write_metrics (self)
None	_write_obs (self)
None	_write_signals (self)
None	_write_config (self, dict[str, Any] config)

Protected Attributes
	_run_id = str(uuid.uuid4())
str	_run_name = run_name or f"run-{self._run_id[:8]}"
str	_base_model = base_model or os.environ.get("AQUIN_BASE_MODEL") or ""
	_started_at = time.time()
str	_run_dir = Path(run_dir) if run_dir is not None else Path.cwd() / "aquin_run"
list	_metrics = []
list	_obs = []
list	_signals = []
	_metrics_lock = threading.Lock()

Detailed Description

Records metrics and checkpoints locally during training.
Call aquin.init() to create one.

Writes under ``./aquin_run/`` (metrics, config, checkpoint).

Example
-------
run = aquin.init(
    base_model="meta-llama/Llama-3.2-1B-Instruct",
    run_name="my-run",
    config={"lr": 2e-4, "epochs": 3, "rank": 16, "lora_alpha": 32,
            "method": "qlora", "per_device_train_batch_size": 2,
            "gradient_accumulation_steps": 8, "dataset": "data.jsonl"},
)
for step, batch in enumerate(dataloader):
    loss = train_step(batch)
    run.log(step, loss=loss.item(), learning_rate=scheduler.get_last_lr()[0],
            grad_norm=grad_norm, epoch=epoch)
run.checkpoint(model, step=step)
run.finish()

Definition at line 30 of file run.py.

Constructor & Destructor Documentation

__init__()

None aquin.run.Run.__init__	(		self,
		str | None	base_model = None,
		str | None	run_name = None,
		dict[str, Any] | None	config = None,
		*	,
		str	quant = "none",
		str	mode = "external",
		str | Path | None	run_dir = None )

Definition at line 54 of file run.py.

Member Function Documentation

_write_config()

None aquin.run.Run._write_config ( self, dict[str, Any] config )

protected

Definition at line 197 of file run.py.

References _run_dir.

Referenced by finish().

_write_meta()

None aquin.run.Run._write_meta ( self )

protected

Definition at line 169 of file run.py.

References _base_model, _run_dir, _run_id, _run_name, and _started_at.

_write_metrics()

None aquin.run.Run._write_metrics ( self )

protected

Definition at line 179 of file run.py.

References _metrics, _metrics_lock, and _run_dir.

Referenced by checkpoint(), finish(), and log().

_write_obs()

None aquin.run.Run._write_obs ( self )

protected

Definition at line 185 of file run.py.

References _metrics_lock, _obs, and _run_dir.

Referenced by observe().

_write_signals()

None aquin.run.Run._write_signals ( self )

protected

Definition at line 191 of file run.py.

References _metrics_lock, _run_dir, and _signals.

Referenced by signal().

checkpoint()

None aquin.run.Run.checkpoint	(		self,
		Any	model,
		int	step )

Save the final model checkpoint. One per run :  replaces the previous save.

Definition at line 142 of file run.py.

References _run_dir, _write_metrics(), and signal().

finish()

None aquin.run.Run.finish	(		self,
		dict[str, Any] | None	config = None )

Finalise the run. Flushes metrics and optional config to disk.

Definition at line 156 of file run.py.

References _started_at, _write_config(), and _write_metrics().

log()

None aquin.run.Run.log	(		self,
		int	step,
		*	,
		float	loss,
		float | None	learning_rate = None,
		float | None	grad_norm = None,
		float | None	momentum_norm = None,
		int | None	epoch = None,
		int | None	batch = None,
		int | None	total_batches = None,
		float | None	step_ms = None,
		**float	extra )

Record metrics for one training step. Call every step inside your loop.

Definition at line 86 of file run.py.

References _metrics, _metrics_lock, and _write_metrics().

observe()

None aquin.run.Run.observe	(		self,
		int	step,
		str	channel,
		float	value )

Emit a custom scalar observation (eval metrics, etc.).

Definition at line 123 of file run.py.

References _metrics_lock, _obs, and _write_obs().

run_id()

str aquin.run.Run.run_id

(

self

)

Definition at line 166 of file run.py.

References _run_id.

signal()

None aquin.run.Run.signal	(		self,
		int	step,
		str	message,
		*	,
		str	severity = "info" )

Emit a non-scalar training signal (checkpoint saved, OOM, etc.).

Definition at line 130 of file run.py.

References _metrics_lock, _signals, and _write_signals().

Referenced by checkpoint().

Member Data Documentation

_base_model

str aquin.run.Run._base_model = base_model or os.environ.get("AQUIN_BASE_MODEL") or ""

protected

Definition at line 67 of file run.py.

Referenced by _write_meta().

_metrics

list aquin.run.Run._metrics = []

protected

Definition at line 74 of file run.py.

Referenced by _write_metrics(), and log().

_metrics_lock

aquin.run.Run._metrics_lock = threading.Lock()

protected

Definition at line 77 of file run.py.

Referenced by _write_metrics(), _write_obs(), _write_signals(), log(), observe(), and signal().

_obs

list aquin.run.Run._obs = []

protected

Definition at line 75 of file run.py.

Referenced by _write_obs(), and observe().

_run_dir

str aquin.run.Run._run_dir = Path(run_dir) if run_dir is not None else Path.cwd() / "aquin_run"

protected

Definition at line 70 of file run.py.

Referenced by _write_config(), _write_meta(), _write_metrics(), _write_obs(), _write_signals(), and checkpoint().

_run_id

aquin.run.Run._run_id = str(uuid.uuid4())

protected

Definition at line 65 of file run.py.

Referenced by _write_meta(), and run_id().

_run_name

str aquin.run.Run._run_name = run_name or f"run-{self._run_id[:8]}"

protected

Definition at line 66 of file run.py.

Referenced by _write_meta().

_signals

list aquin.run.Run._signals = []

protected

Definition at line 76 of file run.py.

Referenced by _write_signals(), and signal().

_started_at

aquin.run.Run._started_at = time.time()

protected

Definition at line 68 of file run.py.

Referenced by _write_meta(), and finish().

---

The documentation for this class was generated from the following file:

• aquin/run.py