cellmap_flow.finetune.job_manager.manager

FinetuneJobManager: the finetune jobs a dashboard follows.

The finetune routes call it, as the session’s finetune_job_manager. It holds the jobs by job id, and their listeners, and puts the package’s modules together:

  • submit_finetuning_job: check the request and work out what to run (submit), make the run’s directory and its record (persistence), launch the trainer, and follow the job;

  • rehydrate_session: follow again the jobs an earlier dashboard left running (persistence.rehydrate);

  • monitor_job: what each job’s monitor thread runs (monitor);

  • restart_finetuning_job (restart) and cancel_job;

  • get_job_status, list_jobs, get_job_logs, get_job;

  • add_listener, remove_listener (listener).

Stopping a job early is between the finetune route and the trainer: the route writes the run’s stop_signal.json.

Attributes

logger

Classes

FinetuneJobManager

Orchestrate finetuning jobs from submission to completion.

Module Contents

cellmap_flow.finetune.job_manager.manager.logger
class cellmap_flow.finetune.job_manager.manager.FinetuneJobManager

Orchestrate finetuning jobs from submission to completion.

Manages the full lifecycle: 1. Validation and job submission to LSF 2. Background monitoring of training progress, told to the listeners 3. Post-training checks, and the job’s record in its metadata.json 4. Restarts and cancellation

jobs: Dict[str, cellmap_flow.finetune.job_manager.state.FinetuneJob]
add_listener(listener) → None

Tell listener about job events; see listener.FinetuneJobListener. Adding one already added does nothing.

remove_listener(listener) → None

Stop telling listener.

submit_finetuning_job(model_config, corrections_path: pathlib.Path, lora_r: int = 8, num_epochs: int = 10, batch_size: int = 8, learning_rate: float = 0.0001, output_base: pathlib.Path | None = None, queue: str | None = None, charge_group: str | None = None, walltime: str | None = None, checkpoint_path_override: pathlib.Path | None = None, auto_serve: bool = True, mask_unannotated: bool = False, loss_type: str = 'combined', label_smoothing: float = 0.0, distillation_lambda: float | None = None, distillation_scope: str = 'unlabeled', margin: float = 0.3, balance_classes: bool = False, augment: bool = False, output_type: str = 'binary', select_channel: int | None = None, offsets: str | None = None) → cellmap_flow.finetune.job_manager.state.FinetuneJob

Submit finetuning job to LSF cluster.

Parameters:

  • model_config – Model configuration object (FlyModelConfig, etc.)

  • corrections_path – Path to corrections.zarr directory

  • lora_r – LoRA rank (default: 8)

  • num_epochs – Number of training epochs (default: 10)

  • batch_size – Training batch size (default: 8)

  • learning_rate – Learning rate (default: 1e-4)

  • output_base – Base directory for outputs (default: output/finetuning)

  • queue – LSF queue name (default: the site’s, gpu_h100 at Janelia)

  • charge_group – LSF charge group (default: the site’s, cellmap at Janelia)

  • walltime – LSF run limit, “HH:MM” or minutes (default: the site’s, 08:00 at Janelia)

  • checkpoint_path_override – Optional path to override checkpoint detection (default: None)

  • auto_serve – Automatically start inference server after training (default: True)

Returns:

FinetuneJob object tracking the submitted job

Raises:

  • ValueError – If validation fails

  • RuntimeError – If job submission fails

rehydrate_session(session_path) → int

Pick up the jobs of a session that are still alive on the cluster.

Jobs live only in this process’s memory, so after a dashboard restart a running job disappeared: it could not be seen or cancelled, and with auto-serve it waited for a restart until walltime. Each job’s metadata.json now records its LSF job id and status; a run whose recorded status is not final and that bjobs still reports as pending or running is monitored again, which also tells the listeners again once the log shows its server. persistence.rehydrate says which runs those are, with one bjobs call per session, and records the others that have ended as final. Returns how many jobs were picked up.

monitor_job(finetune_job: cellmap_flow.finetune.job_manager.state.FinetuneJob)

Follow finetune_job until it ends (monitor.monitor_job); its monitor thread runs this.

cancel_job(job_id: str) → bool

Cancel a running job.

Parameters:

job_id – Job ID to cancel

Returns:

True if successfully cancelled, False otherwise

get_job_status(job_id: str) → Dict[str, Any] | None

Get detailed status of a specific job.

Parameters:

job_id – Job ID to query

Returns:

Dictionary with job status details, or None if not found

list_jobs() → List[Dict[str, Any]]

Get list of all jobs with their status.

Returns:

List of job status dictionaries

get_job_logs(job_id: str) → str | None

Get full log content for a job.

Parameters:

job_id – Job ID

Returns:

Log file content as string, or None if not found. A byte that is not UTF-8 is read as U+FFFD (see tailer.LogTailer).

get_job(job_id: str) → cellmap_flow.finetune.job_manager.state.FinetuneJob | None

Get a FinetuneJob object by ID.

Parameters:

job_id – Job ID to retrieve

Returns:

FinetuneJob object, or None if not found

restart_finetuning_job(job_id: str, updated_params: Dict[str, Any] | None = None) → cellmap_flow.finetune.job_manager.state.FinetuneJob

Restart training on the same GPU via control endpoint.

Primary path sends an HTTP restart request to the running inference server in the same process as the training loop. Falls back to file signal if control endpoint is unavailable.

Parameters:

  • job_id – ID of job to restart

  • updated_params – Dict of updated training parameters

Returns:

Same FinetuneJob object (updated in-place)

Raises:

ValueError – If job not found or not in a restartable state