cellmap_flow.finetune.job_manager.persistence

A run’s directory, as the job manager writes it and reads it back.

A job runs in <session>/runs/<model>_<YYYYmmdd_HHMMSS>/. Beside what the trainer writes there (its log, iterations and exports; finetune.run_outputs), the manager keeps its record of the job in metadata.json: - written at submit (submission_metadata, write_metadata), then

given the job’s lsf_job_id and status;

updated as the job’s status changes, with its inference_server_url, and when it ends, with its finetuned_model_name and model_yaml_path (update_metadata);
completed when it succeeds (record_completion).

Each write replaces the file whole (write_json_atomically). The trainer writes a restart’s settings into its params. The dashboard reads back the model_entry the trainer was given (recorded_model_entry).

A dashboard started later, perhaps newer than the one that submitted the jobs, finds a session’s jobs again from these records (rehydrate), so they are a format. It reads job_id, lsf_job_id, status, model_name, params, created_at and corrections_path, and records a job LSF says has ended as final, with a status_detail when LSF no longer knows the job.

finetune_export_kwargs and check_export say which export a finished run left: a LoRA run’s lora_adapter/, a full finetune’s full_finetune/model_state_dict.pt.

Attributes

`logger`
`METADATA_FILE`

Functions

`submission_metadata`(→ dict)	Build metadata persisted for a submitted finetuning job.
`write_metadata`(→ None)	Write a new run's metadata.json (write_json_atomically).
`update_metadata`(→ None)	Merge `fields` into the run's metadata.json, replacing it atomically
`recorded_model_entry`(→ Optional[dict])	The model entry the run's trainer was given (submit.model_entry), as its
`record_completion`(→ None)	Record in metadata.json that the job completed, with its model and last epoch.
`rehydrate`(...)	The jobs of a session still alive on the cluster, rebuilt from their records.
`finetune_export_kwargs`(→ dict)	Which artifact a finished run produced, as FinetuneModelConfig kwargs.
`check_export`(→ None)	Raise RuntimeError unless the run's export (finetune_export_kwargs) is all on disk.

Module Contents

cellmap_flow.finetune.job_manager.persistence.logger

cellmap_flow.finetune.job_manager.persistence.METADATA_FILE = 'metadata.json'

cellmap_flow.finetune.job_manager.persistence.submission_metadata(*, model_config, model_type: str, checkpoint_path: pathlib.Path | None, corrections_path: pathlib.Path, num_corrections: int, output_dir: pathlib.Path, lora_r: int, num_epochs: int, batch_size: int, learning_rate: float, loss_type: str, label_smoothing: float, distillation_lambda: float | None, distillation_scope: str, margin: float, balance_classes: bool, augment: bool, channels: List[str], input_voxel_size: List[int], output_voxel_size: List[int], output_type: str, queue: str, charge_group: str, command: str) → dict: Build metadata persisted for a submitted finetuning job.

cellmap_flow.finetune.job_manager.persistence.write_metadata(output_dir, metadata: dict) → None: Write a new run’s metadata.json (write_json_atomically).

cellmap_flow.finetune.job_manager.persistence.update_metadata(output_dir, **fields) → None

Merge fields into the run’s metadata.json, replacing it atomically (write_json_atomically).

A missing record is started with fields; one that cannot be read is left as it is, with a warning.

cellmap_flow.finetune.job_manager.persistence.recorded_model_entry(output_dir) → dict | None: The model entry the run’s trainer was given (submit.model_entry), as its metadata.json records it. None for a model that went by its own flags, and for a record that is missing or cannot be read.

cellmap_flow.finetune.job_manager.persistence.record_completion(finetune_job: cellmap_flow.finetune.job_manager.state.FinetuneJob) → None

Record in metadata.json that the job completed, with its model and last epoch.

Through update_metadata, so the record is replaced whole or not at all. A record that cannot be read is left as it is, with a warning, and the job is still COMPLETED: its export was found (check_export), and the record is only the dashboard’s note of the job, which a later dashboard skips when it cannot read it.

cellmap_flow.finetune.job_manager.persistence.rehydrate(session_path, known) → List[cellmap_flow.finetune.job_manager.state.FinetuneJob]

The jobs of a session still alive on the cluster, rebuilt from their records.

A run is rebuilt when its recorded status is not final, its job id is not in known (the jobs already followed), and bjobs still reports it as pending or running; its status is then LSF’s, and what only its log says (its server, its model, its epoch) is left for the monitor to read. Local runs (a PID, not an LSF job) are not reattached.

All the session’s candidates are asked about in one bjobs call, and whatever bjobs says has ended – including a job it no longer knows at all – is recorded as final, so it is not asked about again. This runs on every load of the finetune tab.

cellmap_flow.finetune.job_manager.persistence.finetune_export_kwargs(output_dir, params=None) → dict

Which artifact a finished run produced, as FinetuneModelConfig kwargs.

A LoRA run exports lora_adapter/; a full finetune (–lora-r 0) exports full_finetune/model_state_dict.pt. Decided by what is on disk first – the job’s own record of lora_r is the fallback for a run that has not written its export yet – so the dashboard never points a viewer at an adapter directory that a rank-0 run never made.

cellmap_flow.finetune.job_manager.persistence.check_export(output_dir, params) → None: Raise RuntimeError unless the run’s export (finetune_export_kwargs) is all on disk.