Finetuning refactor

[davidackerman]

This pull request introduces a new human-in-the-loop finetuning feature for CellMap-Flow models, and brings various UI and backend enhancements to support model finetuning and improve user experience. The most significant changes are the addition of a modular finetuning package (with LoRA-based training and a corrections dataset), integration of finetuning into the dashboard, and improvements to the CLI and UI for usability.

Finetuning functionality:

• Introduced the cellmap_flow.finetune package, which provides LoRA-based finetuning for CellMap-Flow models using user corrections as training data. This includes dataset utilities, model wrapping, and a trainer for lightweight adaptation. [1] [2]
• Added FinetuneJobManager to dashboard state for managing finetune jobs, and registered the finetuning blueprint in the dashboard application, enabling backend support for finetuning workflows. [1] [2] [3]

CLI improvements:

• Added a new cellmap_flow_viewer CLI (cellmap_flow/cli/viewer_cli.py) for launching a dataset viewer without requiring model configuration, simplifying dataset exploration.
• Improved logging configuration in the server CLI to ensure log level changes always take effect, and reserved short option names to prevent collisions. [1] [2] [3] [4]

Dashboard and UI enhancements:

• Added a new "Finetune" tab to the dashboard, with corresponding template and content inclusion, providing a user interface for launching and monitoring finetuning jobs. [1] [2]
• Improved dark mode styling for modals, forms, cards, and other UI elements to enhance readability and visual consistency in the dashboard.
  @stuarteberg

[Copilot]

Pull request overview

This PR adds a new human-in-the-loop finetuning workflow (LoRA-based) to CellMap-Flow, integrating annotation creation/sync (via MinIO), training job orchestration, and dashboard UI controls to iterate on model improvements.

Changes:

• Introduces cellmap_flow.finetune package (dataset utilities, LoRA wrapping, training loop, target transforms, job manager + CLI).
• Integrates finetuning into the dashboard (new Finetune tab, new routes, MinIO annotation workflows, job monitoring/log streaming).
• Improves CLI/server utilities and global dashboard state handling (viewer CLI, logging/state refactor, CLI option collision handling).

Reviewed changes

Copilot reviewed 32 out of 36 changed files in this pull request and generated 12 comments.

Show a summary per file

File	Description
pyproject.toml	Adds finetune optional deps and a new cellmap_flow_viewer entrypoint.
cellmap_flow/utils/load_py.py	Adjusts AST safety analysis behavior for config scripts.
cellmap_flow/utils/ds.py	Fixes dataset path splitting by using rsplit.
cellmap_flow/utils/bsub_utils.py	Changes local job execution behavior (now via shell).
cellmap_flow/server.py	Adds restart control endpoint for iterative finetuning runs.
cellmap_flow/models/models_config.py	Adds support for loading a full model.pt checkpoint.
cellmap_flow/globals.py	Moves dashboard state into the Flow singleton and adds log buffering handler.
cellmap_flow/dashboard/state.py	Re-exports dashboard state from globals for backward compatibility.
cellmap_flow/dashboard/app.py	Registers finetuning blueprint and shifts to globals-based state.
cellmap_flow/dashboard/routes/*	Adds finetune routes and updates existing routes to use globals-based state.
cellmap_flow/dashboard/templates/_dashboard.html	Adds a new Finetune tab to the UI.
cellmap_flow/dashboard/static/css/dark.css	Improves dark-mode styling for modals/forms/cards.
docs/finetuning.md	Adds a dashboard-driven finetuning guide with screenshots.
cellmap_flow/finetune/*	New finetuning implementation (LoRA wrapper/trainer, CLI, job manager, templates, target transforms, dataset).
cellmap_flow/cli/viewer_cli.py	Adds a simple dataset viewer CLI without needing model config.
cellmap_flow/cli/server_cli.py	Makes logging config changes take effect reliably and avoids short-option collisions.
.gitignore	Adds project-specific ignores for outputs and zarr artifacts.

[Copilot]

This test module lives under cellmap_flow/finetune/, so it will be included as part of the installed package. If these are intended to be pytest tests, consider moving them under tests/ (and naming accordingly) to avoid shipping test code in the runtime distribution.

[Copilot]

Duplicate import: from cellmap_flow.globals import g appears twice. Drop one to avoid confusion and keep imports clean.

Suggested change

from cellmap_flow.globals import g

[Copilot]

run_locally() now executes the command with shell=True, which allows shell injection if any part of the command string contains user-controlled data (e.g., paths or parameters coming from API requests). Prefer passing an argv list with shell=False (e.g., subprocess.Popen(shlex.split(command)) or building the args list explicitly), and only use shell=True if you fully control/escape all inputs.

[Copilot]

Loading model.pt via torch.load(..., weights_only=False) will unpickle arbitrary code and is unsafe if the checkpoint path can come from users/config files. Consider requiring an explicit “trusted checkpoint” flag, restricting this path to a vetted allowlist, or switching to a safe format (TorchScript / safetensors / weights_only=True state_dict) to avoid RCE.

[Copilot]

Default output_base is set to output/finetuning, but output_dir appends another /finetuning/runs/..., producing output/finetuning/finetuning/runs/.... If the intent is output/finetuning/runs/..., either set output_base = Path("output") or drop the extra "finetuning" segment when building output_dir.

Suggested change

	output_dir = output_base / "finetuning" / "runs" / run_dir_name
	output_dir = output_base / "runs" / run_dir_name

[Copilot]

torch.cuda.empty_cache() is called unconditionally. In CPU-only environments (or builds without CUDA), this can raise an error (e.g., “Torch not compiled with CUDA enabled”). Guard these calls with if torch.cuda.is_available(): ... (and similarly for other CUDA-only cleanup).

Suggested change

	torch.cuda.empty_cache()
	if torch.cuda.is_available():
	torch.cuda.empty_cache()

[Copilot]

torch.cuda.empty_cache() is called unconditionally when preparing for retraining. This can fail on CPU-only environments; guard with torch.cuda.is_available() (or skip entirely if training is CPU).

Suggested change

	torch.cuda.empty_cache()
	if torch.cuda.is_available():
	torch.cuda.empty_cache()

[Copilot]

This trainer uses torch.cuda.amp.autocast/GradScaler but use_mixed_precision defaults to True even when device resolves to CPU. On CPU-only runs, CUDA AMP utilities can error or behave unexpectedly. Consider automatically disabling mixed precision unless self.device.type == "cuda", or switching to torch.amp.autocast(device_type=...) for broader support.

[Copilot]

The new /__control__/restart endpoint can trigger training restarts without any authentication/authorization checks. Since the server binds to 0.0.0.0, this endpoint may be reachable by other users on the network/cluster. Consider restricting it to localhost, requiring a shared secret/token, or disabling it by default unless explicitly enabled.

[Copilot]

analyze_script() no longer flags attribute calls for disallowed functions. This makes it easier to bypass the check via builtins.eval(...) / __builtins__.eval(...) (attribute access) while still executing dangerous builtins. Consider restoring attribute-call detection but limit it to known unsafe roots (e.g., builtins, __builtins__) so model.eval() remains allowed.

Suggested change

	# Note: We intentionally do NOT flag method calls like `model.eval()` here
	# Method calls on objects (e.g., model.eval()) are safe - only direct calls
	# to dangerous builtin functions (e.g., eval()) are a security risk
	# If function is accessed as an attribute of builtins/__builtins__,
	# e.g., builtins.eval() or __builtins__.eval(), also disallow it.
	elif isinstance(node.func, ast.Attribute):
	attr_name = node.func.attr
	base = node.func.value
	if (
	attr_name in DISALLOWED_FUNCTIONS
	and isinstance(base, ast.Name)
	and base.id in {"builtins", "__builtins__"}
	):
	issues.append(
	"Disallowed function call detected via attribute access: "
	f"{base.id}.{attr_name}"
	)
	# Note: Other method calls like model.eval() remain allowed; we only
	# treat attribute calls on known unsafe roots (builtins/__builtins__)
	# as security risks here.