Feature/inference scaled#59
Open
Panzax wants to merge 377 commits into
Open
Conversation
Introduces a canonical checkpoint_meta.json sidecar written alongside every saved checkpoint. Stores model class name, experiment name, W&B run info, epoch/iter/loss, and a hash of the Hydra config so checkpoints are self-describing and reproducible. CheckpointManager.save() now requires a metadata dict built via build_metadata(); load() returns the metadata instead of the raw DeepSpeed client_state dict, decoupling callers from DS internals.
Renames the ambiguous `input_format` parameter to `data_format` across all zarr IO functions for clarity — the format describes the data layout on disk, not just input. Also overhauls save_masks to accept an explicit annotation_name instead of deriving it from a task string, removes the hardcoded TASK_MASK_CHANNEL_NAMES mapping, and makes save_zarr_annotations overwrite-upsert (creates if missing) rather than raising on a missing annotation.
Renames the `pin_to_numa_node` keyword argument to `pin_numa_node` throughout the buffer management stack for naming consistency. Also refactors buffer metrics from cumulative counters to per-call lists so callers can compute arbitrary statistics (median, p90, etc.) and reset cleanly between logging windows. Adds buffer_tensors field to outputs_metadata so only large spatial tensors use pinned SHM pools; small sparse tensors (boxes, labels) pass inline as numpy through the Ray object store.
Replaces the old sliding-window/hypercube inference plumbing with a cleaner tile-mode-only flow. The worker now accepts structured outputs_metadata with buffer_tensors/save_tensors/visualize_tensors separation, acquires SHM buffer slots before the D2H transfer loop, and falls back to inline numpy for tensors not in buffer_tensors. Adds _tree_to_cpu_numpy for recursive CUDA→CPU conversion, a _attach_save_worker_metainfo helper that stamps model_name/task/ channel_names onto metainfo, and a get_metrics/clear_metrics pair that replaces the old aggregated get_step_metrics.
Replaces the flat save_tensors_dtypes+input_format approach with a structured save_tensors_metadata dict that carries dtype, data_format, annotation_type, and annotation name per tensor. SaveWorker save_mode is clarified to 'create'/'overwrite' (drops 'append'). Metrics switch from cumulative counters to per-call lists matching buffer/viz workers. Shard/chunk spatial shapes are now required constructor arguments.
Adds module-level helpers (_infer_batch_size, _metainfo_scalar_at, _get_base_sample_name_for_index) so all visualization handlers can produce per-element filenames from batched metainfo arrays. Updates _handle_semantic_map, _handle_save_predictions, and _handle_feature_viz to unpack tensors and iterate per batch element. Removes deprecated save_bbox_overlay import and stubs _handle_bbox_overlay with NotImplementedError. Metrics switch to per-call lists.
Fixes unpack_batched_tensors to squeeze the batch axis from numpy chunks (matching torch.unbind behaviour). Extends _ensure_numpy_tzyxc to accept 3-D ZYX input. Updates visualize_semantic_labels to read pred_masks/pred_classes/masks_labelmap/class_labels keys from the new model output dict and adds a merged integer-label column. Extends reduce_queries_to_semantic_map to return average_probability alongside the semantic map. Fixes save_semantic_predictions to support optional targets and per-element names.
Mask2Former.predict() now calls reduce_queries_to_semantic_map to produce pred_masks (soft per-class maps), masks_labelmap (foreground integer labels), pred_classes, and class_labels in a unified output dict. Adds validate_outputs to verify tensor shapes against output_metadata.tensor_info at inference time. Removes num_classes from Mask2FormerHead (moved to meta-arch). Adds MaskDINO.validate_outputs and updates its predict() to return a batched dict instead of a list. MaskDINO collapse_instance_masks folds per-instance binary masks into a single label map.
Adds an `idx` parameter to every preprocessor forward() signature and propagates it into the metainfo dict. This lets downstream workers trace which dataset index produced each sample, aiding debugging and ordering guarantees. Also fixes a variable name collision in UpsamplePreprocessor where the loop variable `idx` shadowed the parameter.
Occupancy/CDF threshold filters now short-circuit and return the dataframe unchanged when the threshold is None or ≤0, preventing silent empty-dataframe bugs when aggregation produces NaN values. _expand_tiles_timepoints_df is fully implemented: it duplicates rows for each timepoint offset, sets time_size=1 per row, and handles missing time_start columns. Fixes oversampling when CDF pass is skipped by setting max_hypercubes_before_cdf=max_hypercubes. Adds logging and directory creation to get_rois_dataframe.
Adds ray_assigned_gpu_to_torch_ordinal() to translate a physical GPU id from ray.get_gpu_ids() to the CUDA_VISIBLE_DEVICES-relative ordinal that PyTorch expects. FinetuneCollatorActor and CollatorActor now use this helper for NUMA affinity binding instead of raw int(gpu_ids[0]). batch_size_actual is stored as an int64 ndarray (not a Python int) so Ray map_batches can serialise it; collators unwrap it with np.asarray. LoaderActor.save_mode defaults to None instead of 'create'. Adds docstring to get_context_spec.
Replaces DeepSpeed-initialised Inferencer with a lighter weight setup: model is moved to CUDA directly (no DS wrapper), checkpoint metadata drives model_name_slug, and config paths are reorganised under cfg.inference.inferencer_worker / save_worker / viz_worker. Buffer budget reads from cfg.datasets.buffers instead of cfg.inference.memory. Adds put_scalar_batch to EventRecorder and InferenceMetricsHook collects per-step lists from all workers.
Splits the flat inference YAML configs into nested inferencer_worker / save_worker / viz_worker sub-configs, aligning with the new Inferencer setup. Moves buffers budget from a standalone buffers.yaml into pretrain_dataset_ray.yaml. Updates the semantic segmentation test experiment config with new dataset paths, inference settings, and additional default config groups. Updates the instance segmentation inference config similarly.
Updates the metrics structure in HostMemoryBuffer to store metrics as lists instead of cumulative counters. This change allows for better tracking of metrics over time, including `try_get_free_drops` and `in_use_current`. Adjusts the `get_metrics` and `clear_metrics` methods to accommodate the new structure. Additionally, removes unnecessary clear_metrics methods from other classes in the inference module to streamline metric management.
Updated the parameter name from `drop_last` to `last_batch_policy` in `get_dataset_ray`, `get_dataloader_ray`, and `SampleIndexPlanner`. Adjusted logic to handle "drop" and "pad".
…etric Introduced 'channel_mapping' to the metainfo in the InferencerWorker, ensuring it is included during inference. Updated SaveWorker and VizWorker to track queue time metrics (time between when .remote is called and when the method actually gets executed), covering a gap in previous timing metrics. Adjusted related methods to accommodate these changes.
…servatory/cell_observatory_platform into feature/inference-scaled
Validation steps share the frozen training iteration counter, so step-scoped metrics produced sparse points on the training iter axis joined by misleading wandb interpolation lines. Switches val timing and loss logging to epoch scope to avoid this. Extracts the reduce logic into a shared `reduce_values` helper and adds `reduce_epoch_metric` on EventRecorder so hooks (BestMetricSaver, EarlyStopHook) compute the same reduction as PeriodicWriter — pooling raw per-rank values then reducing once, rather than reducing per-rank reductions. This removes the direct `gather_and_reduce` calls from hooks entirely. Prevents double-sourcing of evaluator metrics that were already logged per-step by `log_loss_dict`, which would skew epoch reductions. Updates tests to match the new epoch scoping and clear the recorder between simulated epochs as PeriodicWriter normally would.
Raises BestMetricSaver priority to HIGH so it runs before BestCheckpointer and PeriodicWriter (both MEDIUM), preventing stale/inf metric reporting and reading from an already-cleared buffer. Adds an early check in epoch-based validation to reject evaluators that declare a predict_method, since in-loop validation only runs the forward pass (loss + outputs), not model predictions. This avoids confusing failures deep in process() and guides users toward either a loss-based evaluator or job_type=test.
…y/cell_observatory_platform into feature/inference-scaled
Consolidates the duplicated value reduction logic into a single canonical helper in the utils/context module. Previously, loggers defined reduce_values locally while metrics classes inlined their own subset of the same logic, risking inconsistent behavior. Now all call sites (loggers, metrics, evaluators) share one implementation, ensuring plotted, hook-selected, and evaluated values are computed identically.
- Changed experiment name for clarity. - Increased batch size from 8 to 16 and updated worker nodes from 1 to 2. - Set max_rows to null for dataset configuration. - Updated epochs from 2 to 100 in the scheduler settings.
- Updated experiment name to reflect changes in configuration. - Increased batch size from 16 to 24 and updated worker nodes from 2 to 3 for improved resource allocation.
…configuration from 1.33e-4 to 5e-4
…y/cell_observatory_platform into feature/inference-scaled
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.