cellmap_flow.finetune.finetune_job_manager

Job manager for orchestrating finetuning jobs on LSF cluster.

This module provides: - FinetuneJob: Track metadata and status of a single finetuning job - FinetuneJobManager: Orchestrate job lifecycle from submission to completion

Attributes

logger

Classes

JobStatus	Status of a finetuning job.
FinetuneJob	Track a finetuning job with metadata, status, and training progress.
FinetuneJobManager	Orchestrate finetuning jobs from submission to completion.

Module Contents

cellmap_flow.finetune.finetune_job_manager.logger

class cellmap_flow.finetune.finetune_job_manager.JobStatus(*args, **kwds)

Status of a finetuning job.

PENDING = 'PENDING'

RUNNING = 'RUNNING'

COMPLETED = 'COMPLETED'

FAILED = 'FAILED'

CANCELLED = 'CANCELLED'

class cellmap_flow.finetune.finetune_job_manager.FinetuneJob

Track a finetuning job with metadata, status, and training progress.

Manages lifecycle from submission through completion, including inference server state and restart chain linkage.

job_id: str

lsf_job: cellmap_flow.utils.bsub_utils.LSFJob | None

model_name: str

output_dir: pathlib.Path

params: Dict[str, Any]

status: JobStatus

created_at: datetime.datetime

log_file: pathlib.Path

finetuned_model_name: str | None = None

model_script_path: pathlib.Path | None = None

model_yaml_path: pathlib.Path | None = None

current_epoch: int = 0

total_epochs: int = 10

latest_loss: float | None = None

inference_server_url: str | None = None

inference_server_ready: bool = False

previous_job_id: str | None = None

next_job_id: str | None = None

to_dict() → Dict[str, Any]

Convert to dictionary for JSON serialization.

class cellmap_flow.finetune.finetune_job_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 3. Post-training model registration 4. Job cancellation and cleanup

jobs: Dict[str, FinetuneJob]

logger

submit_finetuning_job(model_config, corrections_path: pathlib.Path, lora_r: int = 8, num_epochs: int = 10, batch_size: int = 2, learning_rate: float = 0.0001, output_base: pathlib.Path | None = None, queue: str = 'gpu_h100', charge_group: str = 'cellmap', 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 = 0.0, distillation_scope: str = 'unlabeled', margin: float = 0.3, balance_classes: bool = False, output_type: str = 'binary', select_channel: int | None = None, offsets: str | None = None) → 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: 2)

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

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

• queue – LSF queue name (default: gpu_h100)

• charge_group – LSF charge group (default: cellmap)

• 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

monitor_job(finetune_job: FinetuneJob)

Background thread for job monitoring.

Polls LSF status and tails log file to track training progress. Triggers completion when job finishes.

Parameters:

finetune_job – The FinetuneJob to monitor

complete_job(finetune_job: FinetuneJob)

Post-training actions after job completes successfully.

1. Verify adapter files exist

2. Generate model script and YAML

3. Register in g.models_config

4. Update job status and metadata

Parameters:

finetune_job – The completed job

Raises:

RuntimeError – If adapter files missing or registration fails

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

get_job(job_id: str) → 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) → 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