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 (JobControl) 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. Stop Early reaches the trainer the same two ways, as a stop or a stop_signal.json, and ends the iteration’s training after the epoch it is in.

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

JobControl	Hands what the inference server's control endpoints are asked, a
TrainingSession	One finetune job, iteration after iteration, until it ends.

Module Contents

cellmap_flow.finetune.session_loop.logger

class cellmap_flow.finetune.session_loop.JobControl

Hands what the inference server’s control endpoints are asked, a restart or a stop, from the server’s thread to the job’s loop.

The server’s /__control__/restart and /__control__/stop, once they have checked the job’s token, call request_restart and request_stop with the request’s body. The wait for a restart (_wait_for_restart_signal) takes a restart with take_restart, and the trainer, before each epoch, a stop with take_stop (through TrainingSession._stop_requested). A request is kept until it is taken; a second one of the same kind replaces the first. A restart drops a stop that came before it: that stop was for the training before, which ended on its own before the stop was taken.

request_restart(payload: dict | None) → bool

take_restart() → dict | None

request_stop(payload: dict | None) → bool

take_stop() → 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

device

control

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.