cellmap_flow.finetune.session_loop

The finetune job’s loop: train, export, serve, wait for a restart, reset, again.

TrainingSession.run is the job. Each iteration builds its data, target and trainer, trains, exports into its own iterations/NNN_<ts>/ (see run_outputs), and announces the result on stdout (see markers). A served job (–auto-serve) then starts, or keeps, its inference server, which shares the model object, and waits for a restart: from the server’s control endpoint (RestartController) or a restart_signal.json in the output directory. A restart may change training settings only (cli.RESTARTABLE_ARGS); training then starts again from the model the job started from.

The job manager follows the job through those stdout lines, and a job outlives a dashboard upgrade, so their text and order only change compatibly.

Attributes

logger

Classes

RestartController

Hands a restart from the inference server's thread to the training loop.

TrainingSession

One finetune job, iteration after iteration, until it ends.

Module Contents

cellmap_flow.finetune.session_loop.logger
class cellmap_flow.finetune.session_loop.RestartController

Hands a restart from the inference server’s thread to the training loop.

The server’s /__control__/restart endpoint, once it has checked the job’s token, calls request_restart with the request’s body; the loop’s wait (_wait_for_restart_signal) polls get_if_triggered. A request that comes while the job is still training is kept until the wait; a second one replaces the first.

request_restart(payload: dict | None) → bool
get_if_triggered() → dict | None
class cellmap_flow.finetune.session_loop.TrainingSession(args, model_config: cellmap_flow.models.models_config.ModelConfig, model, strategy)

One finetune job, iteration after iteration, until it ends.

model is the model to train, already prepared by strategy (the request decides the strategy, once, in finetune_cli.main; from then on the model does, through strategy_for). It is built once and reused across restarts: the inference server shares the object, so a restart retrains what is being served.

args
model_config
model
restart_controller
server_started = False
inference_server = None
iteration = 0
teacher_model = None
initial_state
tb_position = (0, 0)
applied_restarts
pending_reset = False
model_holds = 'the starting weights'
run() → int

Train (and serve, and retrain on each restart) until the job ends; its exit code.