diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..42b2a94 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,47 @@ +name: build + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + lint: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install .[dev] + - name: Lint check + run: make lint + + test-examples: + needs: [lint] + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + # To re-enable live GCP submission, restore the auth step above and add a + # GCP_SA_KEY repo secret with batch.jobs.create + batch.jobs.get permissions. + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install . + - name: Test examples + run: make test-examples diff --git a/.gitignore b/.gitignore index 48c565d..1e9b036 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,12 @@ +my_configs/ + +# nextflow +.nextflow* +.nextflow/ +local_outputs/ +work/ +nextflow*.html + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index fa5c296..0000000 --- a/MANIFEST.in +++ /dev/null @@ -1,2 +0,0 @@ -include README.md -include requirements/ \ No newline at end of file diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..a4572f2 --- /dev/null +++ b/Makefile @@ -0,0 +1,26 @@ +.PHONY: install lint format FORCE + +install: FORCE + pip install -e .[dev] + +lint: FORCE + ruff check . + ruff format --check . + +format: FORCE + ruff check --fix . + ruff format . + +test-examples: FORCE + cellarium-workflow submit-batch-component \ + --tool onepass_mean_var_std \ + --subcommand fit \ + --config cellarium/workflows/example/onepass_train_config.yaml \ + --project dsp-cellarium \ + --location us-central1 \ + --machine-type n1-standard-4 \ + --accelerator-type nvidia-tesla-t4 \ + --accelerator-count 1 \ + --dry-run + +FORCE: diff --git a/README.md b/README.md index 4827c2e..4b17f87 100644 --- a/README.md +++ b/README.md @@ -1,21 +1,227 @@ -# Cellarium Workflows Example -This code contains helper functions and example to run scripts in Vertex AI platform (powered by Kubeflow) +# Cellarium Workflows + +A CLI tool for submitting [cellarium-ml](https://github.com/cellarium-ai/cellarium-ml) training and inference jobs to different compute backends. + +Example: start an scVI run on the entire 20250811 CZI cellxgene census on specified Google Compute Engine hardware, all from your command line: -## Quick start -* [Install gcloud CLI](https://cloud.google.com/sdk/docs/install) -* [Authenticate gcloud CLI util](https://cloud.google.com/docs/authentication/gcloud) -* Install project requirements like: ```bash -$ pip isntall -r requirements/base.txt +cellarium-workflow submit-batch-component \ + --tool scvi \ + --subcommand fit \ + --config folder/scvi_config.yaml \ + --machine-type n1-standard-16 \ + --accelerator-type nvidia-tesla-t4 \ + --accelerator-count 1 \ + --extract-bucket gs://bucket/path/to/extract_files ``` -## Example -Go to example dir +In this example, `gs://bucket/path/to/extract_files` is the bucket created by extracting data using `Cellarium Nexus`. If data or configs or output directories in the config file are in Google Cloud Storage, file localization will be handled automatically. + +## Installation + ```bash -$ cd cellarium/workflows/example +pip install -e . ``` -Submit an example pipeline: + +This installs the `cellarium-workflow` command. + +## Backends + +The best tested method is `submit-batch-component`, and we recommend it for most use cases. + +| Sub-command | Where it runs | +|---|---| +| `local-single-component` | Local machine | +| `submit-vertex-component` | Vertex AI Pipelines (single step) | +| `submit-vertex-pipeline` | Vertex AI Pipelines (sequential multi-step) | +| `submit-batch-component` | Google Cloud Batch (single job) | +| `submit-batch-pipeline` | Google Cloud Batch (sequential multi-step) | + +## Usage + +### Local + +Run a job on your local machine — useful for development and smoke-testing. + +```bash +cellarium-workflow local-single-component \ + --tool onepass_mean_var_std \ + --subcommand fit \ + --config /path/to/config.yaml +``` + +### Vertex AI — Single Component + +Submit a single job to Vertex AI Pipelines. ```bash -$ python submit_example_pipeline.py --project_id dsp-cell-annotation-service --location us-central1 --display_name test --component_1_config gs://test-bucket/test-config-1.yaml --component_2_config gs://test-bucket/test-config-2.yaml -``` \ No newline at end of file +cellarium-workflow submit-vertex-component \ + --tool onepass_mean_var_std \ + --subcommand fit \ + --config gs://bucket/config.yaml \ + --project my-project \ + --machine-type n1-standard-8 \ + --accelerator-type NVIDIA_TESLA_T4 \ + --accelerator-count 1 +``` + +### Vertex AI — Pipeline + +Submit a multi-step sequential pipeline to Vertex AI. + +```bash +cellarium-workflow submit-vertex-pipeline \ + --pipeline-config pipeline_config.yaml \ + --project my-project +``` + +### Google Cloud Batch — Single Component + +Submit a single job to Cloud Batch. Supports local SSD data staging, custom networking, and optional log capture to GCS. + +```bash +cellarium-workflow submit-batch-component \ + --tool scvi \ + --subcommand fit \ + --config gs://bucket/config.yaml \ + --project my-project \ + --machine-type n1-standard-4 \ + --accelerator-type nvidia-tesla-t4 \ + --accelerator-count 1 +``` + +### Google Cloud Batch — Pipeline + +Submit a multi-step pipeline as individual Cloud Batch jobs. Use `--submit-sequentially` to run them one at a time. + +```bash +cellarium-workflow submit-batch-pipeline \ + --config pipeline_config.yaml \ + --project my-project \ + --submit-sequentially true +``` + +## Configuration + +### Training Config + +A standard [PyTorch Lightning CLI](https://lightning.ai/docs/pytorch/stable/cli/lightning_cli.html) YAML. The key fields used by this tool are: + +```yaml +trainer: + default_root_dir: gs://my-bucket/outputs/ +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://bucket/path/extract_{0..9446}.h5ad + shard_size: 10000 + last_shard_size: 3147 +``` + +The `filenames` field uses brace-expansion to reference sharded `.h5ad` files. You can pass `--extract-bucket gs://bucket/path/` to have the tool auto-populate `filenames`, `shard_size`, and `last_shard_size` from GCS at submission time instead of hardcoding them. + +### Pipeline Config + +A YAML file listing the steps to run in order: + +```yaml +my_pipeline_name: + - tool: onepass_mean_var_std + subcommand: fit + config: gs://bucket/configs/onepass_config.yaml + machine_type: n1-standard-4 + accelerator_type: nvidia-tesla-t4 + accelerator_count: 1 + + - tool: scvi + subcommand: fit + config: gs://bucket/configs/scvi_config.yaml + machine_type: n1-standard-16 + accelerator_type: nvidia-tesla-v100 + accelerator_count: 1 + max_run_duration: 7200s +``` + +Each step can specify its own machine type and GPU configuration. Fields not set on a step inherit the pipeline defaults. + +## Authentication + +```bash +gcloud auth application-default login +``` + +## Google Batch setup + +For a new Google project which has never used Batch before, you will need to set up a few things. We have distilled this into a helper script, included here. Fill in `PROJECT_ID` with the appropriate value for your Google project. + +```bash +#!/bin/bash + +# ======================================================================= +# Set your Google Cloud Project ID here +# ======================================================================= +PROJECT_ID="your-project-id-here" + +echo "Starting Google Batch setup for project: $PROJECT_ID..." + +# 1. Enable the Google Batch API (and Compute Engine API, which is a prerequisite) +echo "Enabling necessary APIs..." +gcloud services enable batch.googleapis.com compute.googleapis.com \ + --project="$PROJECT_ID" + +# 2. Create a "default" VPC network +# (An auto-mode network automatically creates subnets in all regions, including us-central1) +echo "Creating 'default' VPC network..." +gcloud compute networks create default \ + --subnet-mode=auto \ + --project="$PROJECT_ID" || true + +# 3. Turn on Private Google Access for the us-central1 subnet +echo "Enabling Private Google Access for us-central1..." +gcloud compute networks subnets update default \ + --region=us-central1 \ + --enable-private-ip-google-access \ + --project="$PROJECT_ID" + +# 4. Create an egress firewall rule for the Batch Agent +# This ensures the VM can reach out to Google APIs on port 443 even if standard egress is blocked +echo "Creating egress firewall rule for TCP 443..." +gcloud compute firewall-rules create allow-batch-agent-egress \ + --network=default \ + --direction=EGRESS \ + --action=ALLOW \ + --destination-ranges=0.0.0.0/0 \ + --rules=tcp:443 \ + --description="Allows Google Batch agent to communicate with Google APIs" \ + --project="$PROJECT_ID" || true + +# 5. Find the Compute Engine default service account +echo "Locating the Compute Engine default service account..." +PROJECT_NUMBER=$(gcloud projects describe "$PROJECT_ID" --format="value(projectNumber)") +COMPUTE_SA="${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" +echo "Found Service Account: $COMPUTE_SA" + +# 6. Grant the required IAM permissions to the service account +echo "Applying IAM roles..." + +# Batch Agent Reporter (Critical for the VM to report status) +gcloud projects add-iam-policy-binding "$PROJECT_ID" \ + --member="serviceAccount:$COMPUTE_SA" \ + --role="roles/batch.agentReporter" \ + --condition=None > /dev/null + +# Logs Writer (Critical for stdout/stderr logs from the container) +gcloud projects add-iam-policy-binding "$PROJECT_ID" \ + --member="serviceAccount:$COMPUTE_SA" \ + --role="roles/logging.logWriter" \ + --condition=None > /dev/null + +# Storage Object Admin (Critical for mounting GCS buckets) +gcloud projects add-iam-policy-binding "$PROJECT_ID" \ + --member="serviceAccount:$COMPUTE_SA" \ + --role="roles/storage.objectAdmin" \ + --condition=None > /dev/null + +echo "Setup complete! The project $PROJECT_ID is ready for Google Batch." +``` diff --git a/cellarium/workflows/README.md b/cellarium/workflows/README.md new file mode 100644 index 0000000..88e5dab --- /dev/null +++ b/cellarium/workflows/README.md @@ -0,0 +1,128 @@ +# Cellarium Workflows - Shared Components + +This directory contains code for submitting cellarium-ml tasks to Vertex AI Pipelines. + +## Structure + +``` +cellarium/workflows/ +├── scripts/ # Individual Python scripts with full IDE support +│ ├── __init__.py +│ ├── git_install.py # Git installation logic +│ ├── pytorch_setup.py # PyTorch environment setup +│ ├── data_download.py # GCS data download with brace expansion +│ └── cellarium_cli.py # Cellarium CLI execution +├── shared_components.py # Module that converts scripts to strings and provides train_op +├── submit_single_component.py # Single component pipeline submission (Vertex AI) +├── submit_pipeline.py # Multi-component pipeline submission (Vertex AI) +├── local_single_component.py # Single component local execution (no Vertex AI) +└── test_shared_components.py # Test script for validation +``` + +## Usage + +### Local Execution (for testing) +```bash +# Run locally for testing and development, or interactive runs on a VM +(cellarium)$ python local_single_component.py --tool scvi --subcommand fit --config /path/to/config.yaml +``` + +### Vertex AI Single Component +```bash +# Submit a single component to Vertex AI +(vertex)$ python submit_single_component.py --tool scvi --subcommand fit --config gs://path/to/config.yaml +``` + +### Vertex AI Multi-Component Pipeline +```bash +# Submit a multi-component pipeline to Vertex AI +(vertex)$ python submit_pipeline.py --pipeline-config pipeline_config.yaml +``` + +## How It Works + +### Core Architecture + +1. **Scripts**: Individual Python files in `scripts/` contain the actual logic with full IDE support +2. **Shared Components**: The `shared_components.py` module provides: + - `_get_train_op_text()`: Ground truth train_op implementation as text + - `create_train_op_function()`: Creates a function for local execution using `exec()` on the ground truth + - `get_train_op_code()`: Returns the ground truth text for Vertex AI components + - Individual code generators for each script +3. **Execution Modes**: + - **Local**: Uses `create_train_op_function()` which executes the ground truth text + - **Vertex AI**: Uses `get_train_op_code()` to get the same ground truth text for containerized execution + +### Train Op Function + +The core training logic is defined once as a ground-truth text template in `_get_train_op_text()` that: +- Installs the specified git SHA of cellarium-ml +- Optionally downloads data from GCS with parallel processing +- Sets up PyTorch environment for optimal performance +- Executes the cellarium CLI + +This ground truth is used by: +- `create_train_op_function()` - Wraps the text in a function for local execution +- `get_train_op_code()` - Returns the text directly for Vertex AI component execution +- Both `submit_single_component.py` and `submit_pipeline.py` - Use the text in containerized components +- `local_single_component.py` - Uses the function wrapper for direct execution + +This ensures zero redundancy - there's exactly one definition of the training logic. + +### Flexible Execution +- **Local Development**: Test your workflows locally before submitting to Vertex AI +- **Single Components**: Submit individual training jobs +- **Multi-Component Pipelines**: Chain multiple training steps with dependencies + +## Testing + +Run the comprehensive test suite: +```bash +python test_shared_components.py +``` + +This tests: +- Code generation functions +- Train op function creation +- Import validation for all workflow files +- Utility functions + +## Development + +### Adding New Shared Functionality + +1. Create a new script in `scripts/` with proper imports and logic +2. Add a corresponding function in `shared_components.py` +3. Import and use in your workflow files +4. Update tests as needed + +### Using the Train Op Function + +For local execution: +```python +from shared_components import create_train_op_function + +train_op = create_train_op_function(copy_data_to_local_disk=True) +train_op(tool="scvi", subcommand="fit", config="config.yaml", git_sha="") +``` + +For Vertex AI components: +```python +from shared_components import get_train_op_code + +@dsl.component(...) +def train_op(...): + exec(get_train_op_code(copy_data_to_local_disk=True)) +``` + +### Variable Context + +The scripts can reference variables that will be available in the execution context: +- `tool`, `subcommand`, `config`, `git_sha` - always available +- `copy_data_to_local_disk` - available in functions that use it + +## Example Workflow + +1. **Develop locally**: Use `local_single_component.py` to test your configuration +2. **Submit single job**: Use `submit_single_component.py` for one-off training +3. **Scale to pipeline**: Use `submit_pipeline.py` for multi-step workflows diff --git a/cellarium/workflows/cli.py b/cellarium/workflows/cli.py new file mode 100644 index 0000000..faa3bfa --- /dev/null +++ b/cellarium/workflows/cli.py @@ -0,0 +1,26 @@ +"""Main CLI entry point for cellarium-workflows.""" + +import click + +from .local_single_component import run_local_single_component +from .submit_batch_component import submit_batch_component +from .submit_batch_pipeline import submit_batch_pipeline +from .submit_vertex_component import submit_single_component_pipeline +from .submit_vertex_pipeline import submit_sequential_pipeline + + +@click.group() +def cli(): + """Cellarium workflow submission tools.""" + pass + + +cli.add_command(submit_batch_component, name="submit-batch-component") +cli.add_command(submit_batch_pipeline, name="submit-batch-pipeline") +cli.add_command(submit_single_component_pipeline, name="submit-vertex-component") +cli.add_command(submit_sequential_pipeline, name="submit-vertex-pipeline") +cli.add_command(run_local_single_component, name="local-single-component") + + +if __name__ == "__main__": + cli() diff --git a/cellarium/workflows/example/components.py b/cellarium/workflows/example/components.py deleted file mode 100644 index 89289ea..0000000 --- a/cellarium/workflows/example/components.py +++ /dev/null @@ -1,19 +0,0 @@ -from kfp import dsl - - -@dsl.component(base_image="python:3.10-slim") -def example_component_1(config: str): - """ - Test Example component - """ - print(config) - print("Example component is being executed....") - - -@dsl.component(base_image="python:3.10-slim") -def example_component_2(config: str): - """ - Test Example component - """ - print(config) - print("Second example component is being executed...") diff --git a/cellarium/workflows/example/gpu_monitor.py b/cellarium/workflows/example/gpu_monitor.py new file mode 100644 index 0000000..2e7602e --- /dev/null +++ b/cellarium/workflows/example/gpu_monitor.py @@ -0,0 +1,86 @@ +""" +Custom GPU monitoring callback for comprehensive GPU stats including utilization. +""" + +from typing import Any +import lightning.pytorch as pl +from lightning.pytorch.callbacks import Callback +from lightning.pytorch.accelerators.cuda import get_nvidia_gpu_stats +from lightning.pytorch.utilities.types import STEP_OUTPUT + + +class GPUUtilizationMonitor(Callback): + """ + A callback that monitors comprehensive GPU stats including utilization, memory usage, + temperature, and fan speed using nvidia-smi. + + This provides more detailed GPU information than the default DeviceStatsMonitor + which only logs PyTorch CUDA memory stats. + """ + + def __init__(self, log_every_n_steps: int = 50): + """ + Args: + log_every_n_steps: How frequently to log GPU stats (in training steps) + """ + super().__init__() + self.log_every_n_steps = log_every_n_steps + self._step_count = 0 + + def _log_gpu_stats(self, trainer: "pl.Trainer", stage: str) -> None: + """Log comprehensive GPU stats if nvidia-smi is available.""" + if not trainer._logger_connector.should_update_logs: + return + + device = trainer.strategy.root_device + if device.type != "cuda": + return + + try: + # Get comprehensive GPU stats from nvidia-smi + gpu_stats = get_nvidia_gpu_stats(device) + + # Prefix the metrics with our callback name and stage + prefixed_stats = {} + for key, value in gpu_stats.items(): + prefixed_stats[f"GPUMonitor.{stage}.{key}"] = value + + # Log to all loggers + for logger in trainer.loggers: + logger.log_metrics(prefixed_stats, step=trainer.global_step) + + except FileNotFoundError: + # nvidia-smi not available, fall back to basic torch stats + if hasattr(trainer.accelerator, "get_device_stats"): + basic_stats = trainer.accelerator.get_device_stats(device) + prefixed_stats = { + f"GPUMonitor.{stage}.{k}": v for k, v in basic_stats.items() + } + for logger in trainer.loggers: + logger.log_metrics(prefixed_stats, step=trainer.global_step) + + def on_train_batch_end( + self, + trainer: "pl.Trainer", + pl_module: "pl.LightningModule", + outputs: STEP_OUTPUT, + batch: Any, + batch_idx: int, + ) -> None: + """Log GPU stats every N training steps.""" + self._step_count += 1 + if self._step_count % self.log_every_n_steps == 0: + self._log_gpu_stats(trainer, "train") + + def on_validation_batch_end( + self, + trainer: "pl.Trainer", + pl_module: "pl.LightningModule", + outputs: STEP_OUTPUT, + batch: Any, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + """Log GPU stats during validation.""" + if batch_idx == 0: # Log only on first validation batch to avoid spam + self._log_gpu_stats(trainer, "val") diff --git a/cellarium/workflows/example/ipca_config.yaml b/cellarium/workflows/example/ipca_config.yaml new file mode 100644 index 0000000..150da75 --- /dev/null +++ b/cellarium/workflows/example/ipca_config.yaml @@ -0,0 +1,130 @@ +# lightning.pytorch==2.2.1 +seed_everything: true +trainer: + accelerator: auto + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + init_args: + accelerator: null + parallel_devices: null + cluster_environment: null + checkpoint_io: null + precision_plugin: null + ddp_comm_state: null + ddp_comm_hook: null + ddp_comm_wrapper: null + model_averaging_period: null + process_group_backend: null + timeout: 0:30:00 + start_method: popen + dict_kwargs: + broadcast_buffers: false + devices: auto + num_nodes: 1 + precision: null + logger: null + callbacks: null + fast_dev_run: false + max_epochs: 1 + min_epochs: null + max_steps: -1 + min_steps: null + max_time: null + limit_train_batches: null + limit_val_batches: null + limit_test_batches: null + limit_predict_batches: null + overfit_batches: 0.0 + val_check_interval: null + check_val_every_n_epoch: 1 + num_sanity_val_steps: null + log_every_n_steps: null + enable_checkpointing: null + enable_progress_bar: null + enable_model_summary: null + accumulate_grad_batches: 1 + gradient_clip_val: null + gradient_clip_algorithm: null + deterministic: null + benchmark: null + inference_mode: true + use_distributed_sampler: true + profiler: null + detect_anomaly: false + barebones: false + plugins: null + sync_batchnorm: false + reload_dataloaders_every_n_epochs: 0 + default_root_dir: ipca +model: + transforms: + - class_path: cellarium.ml.transforms.NormalizeTotal + init_args: + target_count: 10_000 + - cellarium.ml.transforms.Log1p + - class_path: cellarium.ml.transforms.ZScore + init_args: + mean_g: + !CheckpointLoader + file_path: gs://broad-bican-cellarium-file-system/curriculum/20240611_full/trained_models/onepass_mean_var_std/epoch=0-step=2425.ckpt + attr: model.mean_g + convert_fn: null + std_g: + !CheckpointLoader + file_path: gs://broad-bican-cellarium-file-system/curriculum/20240611_full/trained_models/onepass_mean_var_std/epoch=0-step=2425.ckpt + attr: model.std_g + convert_fn: null + var_names_g: + !CheckpointLoader + file_path: gs://broad-bican-cellarium-file-system/curriculum/20240611_full/trained_models/onepass_mean_var_std/epoch=0-step=2425.ckpt + attr: model.var_names_g + convert_fn: null + model: + class_path: cellarium.ml.models.IncrementalPCA + init_args: + n_components: 50 + svd_lowrank_niter: 2 + perform_mean_correction: true + optim_fn: null + optim_kwargs: null + scheduler_fn: null + scheduler_kwargs: null + is_initialized: false +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: /data/extract_{0..248}.h5ad + limits: null + shard_size: 10000 + last_shard_size: 3147 + max_cache_size: 2 + cache_size_strictly_enforced: true + label: null + keys: null + index_unique: null + convert: null + indices_strict: true + obs_columns_to_validate: + - total_mrna_umis + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var_names + total_mrna_umis_n: + attr: obs + key: total_mrna_umis + batch_size: 2048 + iteration_strategy: cache_efficient + shuffle: false + shuffle_seed: 0 + drop_last_indices: false + drop_incomplete_batch: false + worker_seed: null + test_mode: false + num_workers: 4 + prefetch_factor: null + persistent_workers: false +ckpt_path: null \ No newline at end of file diff --git a/cellarium/workflows/example/ipca_pipeline_config.yaml b/cellarium/workflows/example/ipca_pipeline_config.yaml new file mode 100644 index 0000000..7d9eebc --- /dev/null +++ b/cellarium/workflows/example/ipca_pipeline_config.yaml @@ -0,0 +1,8 @@ +ipca: + - tool: incremental_pca + subcommand: fit + config: gs://broad-bican-cellarium-file-system/curriculum/20240611_full/configs/20241202_ipca_config.yaml + machine_type: n1-standard-16 + accelerator_type: NVIDIA_TESLA_T4 + accelerator_count: 4 + git_sha: main diff --git a/cellarium/workflows/example/onepass_train_config.yaml b/cellarium/workflows/example/onepass_train_config.yaml new file mode 100644 index 0000000..8f704f8 --- /dev/null +++ b/cellarium/workflows/example/onepass_train_config.yaml @@ -0,0 +1,110 @@ +# lightning.pytorch==2.2.1 +seed_everything: true +trainer: + accelerator: auto + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + init_args: + accelerator: null + parallel_devices: null + cluster_environment: null + checkpoint_io: null + precision_plugin: null + ddp_comm_state: null + ddp_comm_hook: null + ddp_comm_wrapper: null + model_averaging_period: null + process_group_backend: null + timeout: 0:30:00 + start_method: popen + dict_kwargs: + broadcast_buffers: false + devices: auto + num_nodes: 1 + precision: null + callbacks: null + fast_dev_run: false + max_epochs: 1 + min_epochs: null + max_steps: -1 + min_steps: null + max_time: null + limit_train_batches: null + limit_val_batches: null + limit_test_batches: null + limit_predict_batches: null + overfit_batches: 0.0 + val_check_interval: null + check_val_every_n_epoch: 1 + num_sanity_val_steps: null + log_every_n_steps: null + enable_checkpointing: null + enable_progress_bar: null + enable_model_summary: null + accumulate_grad_batches: 1 + gradient_clip_val: null + gradient_clip_algorithm: null + deterministic: null + benchmark: null + inference_mode: true + use_distributed_sampler: true + profiler: null + detect_anomaly: false + barebones: false + plugins: null + sync_batchnorm: false + reload_dataloaders_every_n_epochs: 0 + default_root_dir: /gcs/cellarium-dev-central/workflows/onepass_train_tmp +model: + transforms: + - class_path: cellarium.ml.transforms.NormalizeTotal + init_args: + target_count: 10_000 + - cellarium.ml.transforms.Log1p + model: + class_path: cellarium.ml.models.OnePassMeanVarStd + init_args: + algorithm: shifted_data + optim_fn: null + optim_kwargs: null + scheduler_fn: null + scheduler_kwargs: null + is_initialized: false +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_homosapiens20251108_primary_umi_count_gte_300_donor_dataset/extract_files/extract_{000000..000010}.h5ad + limits: null + shard_size: 10_000 + # last_shard_size: 6428 + max_cache_size: 2 + cache_size_strictly_enforced: true + label: null + keys: null + index_unique: null + convert: null + indices_strict: true + obs_columns_to_validate: + - raw_sum + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var_names + total_mrna_umis_n: + attr: obs + key: raw_sum + batch_size: 5000 + iteration_strategy: cache_efficient + shuffle: false + shuffle_seed: 0 + drop_last_indices: false + drop_incomplete_batch: false + worker_seed: null + test_mode: false + num_workers: 4 + prefetch_factor: null + persistent_workers: true +ckpt_path: null diff --git a/cellarium/workflows/example/onepass_train_smoketest.sh b/cellarium/workflows/example/onepass_train_smoketest.sh new file mode 100755 index 0000000..780ee27 --- /dev/null +++ b/cellarium/workflows/example/onepass_train_smoketest.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# copy the local yaml file to the bucket +gsutil cp onepass_train_smoketest_config.yaml gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/onepass_train_smoketest_config.yaml + +# submit a pipeline +python ../submit_single_component.py \ + --tool onepass_mean_var_std \ + --subcommand fit \ + --config gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/onepass_train_smoketest_config.yaml \ + --accelerator-count 0 diff --git a/cellarium/workflows/example/onepass_train_smoketest_config.yaml b/cellarium/workflows/example/onepass_train_smoketest_config.yaml new file mode 100644 index 0000000..d9455ee --- /dev/null +++ b/cellarium/workflows/example/onepass_train_smoketest_config.yaml @@ -0,0 +1,110 @@ +# lightning.pytorch==2.2.1 +seed_everything: true +trainer: + accelerator: auto + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + init_args: + accelerator: null + parallel_devices: null + cluster_environment: null + checkpoint_io: null + precision_plugin: null + ddp_comm_state: null + ddp_comm_hook: null + ddp_comm_wrapper: null + model_averaging_period: null + process_group_backend: null + timeout: 0:30:00 + start_method: popen + dict_kwargs: + broadcast_buffers: false + devices: auto + num_nodes: 1 + precision: null + callbacks: null + fast_dev_run: false + max_epochs: 1 + min_epochs: null + max_steps: -1 + min_steps: null + max_time: null + limit_train_batches: null + limit_val_batches: null + limit_test_batches: null + limit_predict_batches: null + overfit_batches: 0.0 + val_check_interval: null + check_val_every_n_epoch: 1 + num_sanity_val_steps: null + log_every_n_steps: null + enable_checkpointing: null + enable_progress_bar: null + enable_model_summary: null + accumulate_grad_batches: 1 + gradient_clip_val: null + gradient_clip_algorithm: null + deterministic: null + benchmark: null + inference_mode: true + use_distributed_sampler: true + profiler: null + detect_anomaly: false + barebones: false + plugins: null + sync_batchnorm: false + reload_dataloaders_every_n_epochs: 0 + default_root_dir: /gcs/cellarium-human-primary-data/curriculum/human_all_primary_20241108/trained_models/tmp/ +model: + transforms: + - class_path: cellarium.ml.transforms.NormalizeTotal + init_args: + target_count: 10_000 + - cellarium.ml.transforms.Log1p + model: + class_path: cellarium.ml.models.OnePassMeanVarStd + init_args: + algorithm: shifted_data + optim_fn: null + optim_kwargs: null + scheduler_fn: null + scheduler_kwargs: null + is_initialized: false +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/extract_files/extract_{0..1}.h5ad # 4440}.h5ad + limits: null + shard_size: 10_000 + # last_shard_size: 6428 + max_cache_size: 2 + cache_size_strictly_enforced: true + label: null + keys: null + index_unique: null + convert: null + indices_strict: true + obs_columns_to_validate: + - total_mrna_umis + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var_names + total_mrna_umis_n: + attr: obs + key: total_mrna_umis + batch_size: 5000 + iteration_strategy: cache_efficient + shuffle: false + shuffle_seed: 0 + drop_last_indices: false + drop_incomplete_batch: false + worker_seed: null + test_mode: false + num_workers: 4 + prefetch_factor: null + persistent_workers: true +ckpt_path: null diff --git a/cellarium/workflows/example/pipeline_config.yaml b/cellarium/workflows/example/pipeline_config.yaml new file mode 100644 index 0000000..8ebdcbb --- /dev/null +++ b/cellarium/workflows/example/pipeline_config.yaml @@ -0,0 +1,11 @@ +test_pipelines: + - tool: onepass_mean_var_std + subcommand: fit + config: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/20241114_onepass_train_config_test.yaml + machine_type: n1-standard-4 + accelerator_count: 0 + - tool: onepass_mean_var_std + subcommand: fit + config: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/20241114_onepass_train_config_test.yaml + machine_type: n1-standard-4 + accelerator_count: 0 \ No newline at end of file diff --git a/cellarium/workflows/example/pipelines.py b/cellarium/workflows/example/pipelines.py deleted file mode 100644 index 0a3fda5..0000000 --- a/cellarium/workflows/example/pipelines.py +++ /dev/null @@ -1,31 +0,0 @@ -from kfp import dsl -from cellarium.workflows import kfp_helpers -from cellarium.workflows.example import components - - -@dsl.pipeline() -def example_pipeline(component_1_config: str, component_2_config: str): - """ - KFP pipeline to run PCA train pipeline. - - """ - component_job_1 = kfp_helpers.create_job( - component_func=components.example_component_1, - display_name="Example Job 1", - config=component_1_config - ) - - component_job_2 = kfp_helpers.create_job( - component_func=components.example_component_2, - display_name="Example Job 2", - config=component_2_config - ) - - task_1 = component_job_1() - task_2 = component_job_2() - - task_2.after(task_1) - - -if __name__ == '__main__': - example_pipeline() diff --git a/cellarium/workflows/example/scvi_pipeline_config.yaml b/cellarium/workflows/example/scvi_pipeline_config.yaml new file mode 100644 index 0000000..1a83489 --- /dev/null +++ b/cellarium/workflows/example/scvi_pipeline_config.yaml @@ -0,0 +1,8 @@ +scvi_cpu_vanilla_with_full_latent_batch: + - tool: scvi + subcommand: fit + config: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/20241127_scvi_train_config.yaml + machine_type: n1-standard-16 + # accelerator_type: NVIDIA_TESLA_T4 + accelerator_count: 0 + git_sha: c14705370d2a7a805286fa3dd0e4795c10e6cefd diff --git a/cellarium/workflows/example/scvi_train_config.yaml b/cellarium/workflows/example/scvi_train_config.yaml new file mode 100644 index 0000000..14032dd --- /dev/null +++ b/cellarium/workflows/example/scvi_train_config.yaml @@ -0,0 +1,182 @@ +# lightning.pytorch==2.4.0 +seed_everything: true +trainer: + accelerator: auto + strategy: auto + # class_path: lightning.pytorch.strategies.DDPStrategy + # dict_kwargs: + # broadcast_buffers: false + # find_unused_parameters: true + devices: auto + num_nodes: 1 + precision: 32 + logger: null + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + log_momentum: false + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + every_n_epochs: 1 + # every_n_train_steps: 2000 + # save_top_k: 1 + # - class_path: lightning.pytorch.callbacks.DeviceStatsMonitor + # init_args: + # cpu_stats: true + # - class_path: lightning.pytorch.callbacks.EarlyStopping + # init_args: + # monitor: train_loss + # min_delta: 0.0 + # patience: 45 + # mode: min + fast_dev_run: false + max_epochs: 3 + min_epochs: null + max_steps: -1 + min_steps: null + max_time: null + limit_train_batches: null + limit_val_batches: null + limit_test_batches: null + limit_predict_batches: null + overfit_batches: 0.0 + val_check_interval: null + check_val_every_n_epoch: null + num_sanity_val_steps: null + log_every_n_steps: 1 + enable_checkpointing: true + enable_progress_bar: true + enable_model_summary: null + accumulate_grad_batches: 1 + gradient_clip_val: 0.5 + gradient_clip_algorithm: norm + deterministic: null + benchmark: null + inference_mode: false + use_distributed_sampler: true + profiler: null + detect_anomaly: false + barebones: false + plugins: null + sync_batchnorm: false + reload_dataloaders_every_n_epochs: 0 + default_root_dir: /gcs/cellarium-human-primary-data/curriculum/human_all_primary_20241108/trained_models/20241122_scvi_vanilla_with_full_latent_batch/ +model: + cpu_transforms: + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/scvi_hvg.txt + loader_fn: pandas.read_csv + attr: "ensembl_id" + convert_fn: pandas.Series.to_list + + model: + class_path: cellarium.ml.models.SingleCellVariationalInference + init_args: + + # these are pulled from data, leave null + var_names_g: null + n_batch: null + + # these can be changed + batch_embedded: true + batch_representation_sampled: true + n_latent_batch: null + batch_kl_weight: 0.01 + encoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 128 + batch_to_bias_hidden_layers: [] # one linear transformation: what scvi-tools does + final_layer: + class_path: torch.nn.Linear + init_args: {} + decoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 128 + batch_to_bias_hidden_layers: [] + # hidden_layers: + # - class_path: torch.nn.Linear + # init_args: + # out_features: 10 + # - class_path: torch.nn.Linear + # init_args: + # out_features: 20 + final_layer: + class_path: torch.nn.Linear # cellarium.ml.models.scvi.LinearWithBatch + init_args: {} + # batch_to_bias_hidden_layers: [] + final_additive_bias: false + n_latent: 50 + + # probably leave these alone + n_continuous_cov: 0 + n_cats_per_cov: null + dropout_rate: 0.1 + dispersion: gene + log_variational: true + gene_likelihood: nb + latent_distribution: normal + use_batch_norm: both + use_layer_norm: none + use_size_factor_key: false + use_observed_lib_size: true + optim_fn: torch.optim.AdamW + optim_kwargs: + lr: 1e-4 + scheduler_fn: null + scheduler_kwargs: null + is_initialized: False +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/extract_files/extract_{0..4440}.h5ad + limits: null + shard_size: 10_000 + last_shard_size: 6428 + max_cache_size: 2 + cache_size_strictly_enforced: true + label: null + keys: null + index_unique: null + convert: null + indices_strict: true + obs_columns_to_validate: [] + # - donor_dataset_concat + # - suspension_type + # - assay + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var_names + batch_index_n: + attr: obs + key: donor_dataset_concat + convert_fn: cellarium.ml.utilities.data.categories_to_codes + categorical_covariate_index_nd: + attr: obs + key: + - suspension_type + - assay + convert_fn: cellarium.ml.utilities.data.categories_to_codes + batch_size: 512 + iteration_strategy: cache_efficient + shuffle: false + shuffle_seed: 0 + drop_last_indices: false + drop_incomplete_batch: false + worker_seed: null + test_mode: false + num_workers: 4 + prefetch_factor: null + persistent_workers: true +ckpt_path: null diff --git a/cellarium/workflows/example/scvi_train_config_small_test.yaml b/cellarium/workflows/example/scvi_train_config_small_test.yaml new file mode 100644 index 0000000..fa39868 --- /dev/null +++ b/cellarium/workflows/example/scvi_train_config_small_test.yaml @@ -0,0 +1,197 @@ +# lightning.pytorch==2.4.0 +seed_everything: true +trainer: + accelerator: auto + strategy: auto + # class_path: lightning.pytorch.strategies.DDPStrategy + # dict_kwargs: + # broadcast_buffers: false + # find_unused_parameters: true + devices: 1 + num_nodes: 1 + precision: 32 + logger: null + callbacks: + # - class_path: lightning.pytorch.callbacks.LearningRateMonitor + # init_args: + # logging_interval: step + # log_momentum: false + - class_path: lightning.pytorch.callbacks.TQDMProgressBar + init_args: + refresh_rate: 10 + leave: true + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + save_last: true # Always save the last checkpoint + save_on_train_epoch_end: true # Save at end of each epoch + every_n_epochs: 1 + # every_n_train_steps: 2000 + save_top_k: -1 + # - class_path: lightning.pytorch.callbacks.DeviceStatsMonitor + # init_args: + # cpu_stats: false # Disable CPU stats since we're using GPU + # Note: This only logs PyTorch CUDA memory stats, not GPU utilization. + # - class_path: cellarium.workflows.example.system_monitor.SystemMonitor + # init_args: + # log_every_n_steps: 25 # Log key system metrics every 25 training steps + # - class_path: lightning.pytorch.callbacks.EarlyStopping + # init_args: + # monitor: train_loss + # min_delta: 0.0 + # patience: 45 + # mode: min + fast_dev_run: false + max_epochs: 5 + min_epochs: null + max_steps: -1 + min_steps: null + max_time: null + limit_train_batches: null + limit_val_batches: null + limit_test_batches: null + limit_predict_batches: null + overfit_batches: 0.0 + val_check_interval: null + check_val_every_n_epoch: null + num_sanity_val_steps: null + log_every_n_steps: 1 + enable_checkpointing: true + enable_progress_bar: true + enable_model_summary: null + accumulate_grad_batches: 1 + gradient_clip_val: 50 + gradient_clip_algorithm: norm + deterministic: null + benchmark: null + inference_mode: false + use_distributed_sampler: true + profiler: null + detect_anomaly: false + barebones: false + plugins: null + sync_batchnorm: false + reload_dataloaders_every_n_epochs: 0 + default_root_dir: /gcs/cellarium-human-primary-data/curriculum/human_all_primary_20241108/trained_models/20250801_scvi_small_test/ +model: + cpu_transforms: + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/scvi_hvg.txt + loader_fn: pandas.read_csv + attr: "ensembl_id" + convert_fn: pandas.Series.to_list + + model: + class_path: cellarium.ml.models.SingleCellVariationalInference + init_args: + + # these are pulled from data, leave null + var_names_g: null + n_batch: null + + # these can be changed + batch_embedded: false + batch_representation_sampled: false + n_latent_batch: null + batch_kl_weight_max: 0.0 + kl_warmup_epochs: 5 + kl_annealing_start: 0.0 + z_kl_weight_max: 1.0 + encoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 128 + batch_to_bias_hidden_layers: [] # one linear transformation: what scvi-tools does + final_layer: + class_path: torch.nn.Linear + init_args: {} + decoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 128 + batch_to_bias_hidden_layers: [] + # hidden_layers: + # - class_path: torch.nn.Linear + # init_args: + # out_features: 10 + # - class_path: torch.nn.Linear + # init_args: + # out_features: 20 + final_layer: + class_path: torch.nn.Linear # cellarium.ml.models.scvi.LinearWithBatch + init_args: {} + # batch_to_bias_hidden_layers: [] + final_additive_bias: false + n_latent: 50 + + # probably leave these alone + n_continuous_cov: 0 + n_cats_per_cov: null + dropout_rate: 0.1 + dispersion: gene + log_variational: true + gene_likelihood: nb + latent_distribution: normal + use_batch_norm: both + use_layer_norm: none + use_size_factor_key: false + use_observed_lib_size: true + optim_fn: torch.optim.Adam + optim_kwargs: + lr: 1e-3 + weight_decay: 1e-6 + eps: 0.01 + scheduler_fn: null + scheduler_kwargs: null + is_initialized: false +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/extract_files/extract_{0..9}.h5ad + limits: null + shard_size: 10_000 + last_shard_size: 10_000 + max_cache_size: 8 # Increased from 4 - cache more files in memory + cache_size_strictly_enforced: false # Allow adaptive caching + label: null + keys: null + index_unique: null + convert: null + indices_strict: true + obs_columns_to_validate: [] + # - donor_dataset_concat + # - suspension_type + # - assay + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var_names + batch_index_n: + attr: obs + key: donor_dataset_concat + convert_fn: cellarium.ml.utilities.data.categories_to_codes + categorical_covariate_index_nd: + attr: obs + key: + - suspension_type + - assay + convert_fn: cellarium.ml.utilities.data.categories_to_codes + batch_size: 500 + iteration_strategy: cache_efficient + shuffle: true + shuffle_seed: 0 + drop_last_indices: false + drop_incomplete_batch: false + worker_seed: null + test_mode: false + num_workers: 16 # Increased from 12 - more I/O parallelism + prefetch_factor: 4 # Increased from 3 - more prefetching per worker + persistent_workers: true +ckpt_path: null diff --git a/cellarium/workflows/example/submit_example_pipeline.py b/cellarium/workflows/example/submit_example_pipeline.py deleted file mode 100644 index fce4e90..0000000 --- a/cellarium/workflows/example/submit_example_pipeline.py +++ /dev/null @@ -1,25 +0,0 @@ -import click - -from cellarium.workflows.example import pipelines -from cellarium.workflows import kfp_helpers - - -@click.command() -@click.option("--project_id") -@click.option("--location") -@click.option("--display_name") -@click.option("--component_1_config") -@click.option("--component_2_config") -def submit_example(project_id: str, location: str, display_name: str, component_1_config: str, component_2_config: str): - kfp_helpers.submit_pipeline( - pipeline_func=pipelines.example_pipeline, - project_id=project_id, - location=location, - pipeline_display_name=display_name, - pipeline_kwargs={"component_1_config": component_1_config, "component_2_config": component_2_config} - ) - print("Submitted pipeline!") - - -if __name__ == "__main__": - submit_example() diff --git a/cellarium/workflows/example/system_monitor.py b/cellarium/workflows/example/system_monitor.py new file mode 100644 index 0000000..a5fac91 --- /dev/null +++ b/cellarium/workflows/example/system_monitor.py @@ -0,0 +1,120 @@ +""" +Lightweight system monitoring callback for key metrics. +""" + +from typing import Any, Dict +import lightning.pytorch as pl +from lightning.pytorch.callbacks import Callback +from lightning.pytorch.utilities.types import STEP_OUTPUT + +try: + import psutil + + PSUTIL_AVAILABLE = True +except ImportError: + PSUTIL_AVAILABLE = False + +try: + from lightning.pytorch.accelerators.cuda import get_nvidia_gpu_stats + + NVIDIA_SMI_AVAILABLE = True +except ImportError: + NVIDIA_SMI_AVAILABLE = False + + +class SystemMonitor(Callback): + """ + Lightweight callback that logs only the essential system metrics: + - CPU utilization percentage + - Memory utilization percentage + - GPU utilization percentage (if available) + - GPU memory usage (if available) + """ + + def __init__(self, log_every_n_steps: int = 50): + """ + Args: + log_every_n_steps: How frequently to log system stats (in training steps) + """ + super().__init__() + self.log_every_n_steps = log_every_n_steps + self._step_count = 0 + + def _get_system_stats(self, trainer: "pl.Trainer") -> Dict[str, float]: + """Get essential system stats.""" + stats = {} + + # CPU and Memory stats + if PSUTIL_AVAILABLE: + stats["cpu_percent"] = psutil.cpu_percent() + stats["memory_percent"] = psutil.virtual_memory().percent + + # GPU stats + device = trainer.strategy.root_device + if device.type == "cuda": + try: + if NVIDIA_SMI_AVAILABLE: + # Get comprehensive GPU stats from nvidia-smi + gpu_stats = get_nvidia_gpu_stats(device) + # Extract just the key metrics we care about + stats["gpu_utilization_percent"] = gpu_stats.get( + "utilization.gpu (%)", 0.0 + ) + stats["gpu_memory_used_mb"] = gpu_stats.get("memory.used (MB)", 0.0) + stats["gpu_memory_free_mb"] = gpu_stats.get("memory.free (MB)", 0.0) + stats["gpu_memory_utilization_percent"] = gpu_stats.get( + "utilization.memory (%)", 0.0 + ) + else: + # Fallback to basic torch memory stats + import torch + + if torch.cuda.is_available(): + mem_allocated = ( + torch.cuda.memory_allocated(device) / 1024**2 + ) # MB + mem_reserved = ( + torch.cuda.memory_reserved(device) / 1024**2 + ) # MB + stats["gpu_memory_allocated_mb"] = mem_allocated + stats["gpu_memory_reserved_mb"] = mem_reserved + except Exception: + # Silently skip GPU stats if there's any error + pass + + return stats + + def _log_stats(self, trainer: "pl.Trainer", stage: str) -> None: + """Log system stats.""" + if not trainer._logger_connector.should_update_logs: + return + + stats = self._get_system_stats(trainer) + if not stats: + return + + # Prefix metrics with stage + prefixed_stats = {f"system_{stage}_{k}": v for k, v in stats.items()} + + # Log to all loggers + for logger in trainer.loggers: + logger.log_metrics(prefixed_stats, step=trainer.global_step) + + def on_train_batch_end( + self, + trainer: "pl.Trainer", + pl_module: "pl.LightningModule", + outputs: STEP_OUTPUT, + batch: Any, + batch_idx: int, + ) -> None: + """Log system stats every N training steps.""" + self._step_count += 1 + if self._step_count % self.log_every_n_steps == 0: + self._log_stats(trainer, "train") + + def on_validation_epoch_start( + self, trainer: "pl.Trainer", pl_module: "pl.LightningModule" + ) -> None: + """Log system stats at start of validation.""" + self._log_stats(trainer, "val") diff --git a/cellarium/workflows/kfp_helpers.py b/cellarium/workflows/kfp_helpers.py deleted file mode 100644 index d18db3c..0000000 --- a/cellarium/workflows/kfp_helpers.py +++ /dev/null @@ -1,79 +0,0 @@ -import tempfile -import typing as t -import os - -from kfp import compiler -from kfp.components import BaseComponent -from google.cloud import aiplatform -from google_cloud_pipeline_components.v1.custom_job import create_custom_training_job_from_component - - -def create_job( - component_func: t.Callable[..., t.Any], - config: str, - display_name: str = "", - replica_count: int = 1, - machine_type: str = "n1-standard-4", - accelerator_type: str = "", - accelerator_count: int = 1, - boot_disk_size_gb: int = 100 -) -> t.Callable[[], t.Any]: - """ - Create a custom training Google Vertex AI job for running a custom training component. - - :param component_func: Custom training component. - :param config: Config file path on GCS. - :param display_name: Display name of component in Vertex AI - :param replica_count: The count of instances in the cluster. - :param machine_type: The type of the machine to run the CustomJob. - :param accelerator_type: The type of accelerator(s) that may be attached to the machine per `accelerator_count`. - :param accelerator_count: The number of accelerators to attach to the machine. - :param boot_disk_size_gb: Size in GB of the boot disk - - :return: Callable custom training job. - """ - job = create_custom_training_job_from_component( - component_func, - display_name=display_name, - replica_count=replica_count, - machine_type=machine_type, - accelerator_type=accelerator_type, - accelerator_count=accelerator_count, - boot_disk_size_gb=boot_disk_size_gb, - ) - - return lambda: job(config=config) - - -def submit_pipeline( - pipeline_func: t.Union[BaseComponent, t.Callable], - pipeline_display_name: str, - pipeline_kwargs: t.Dict[str, t.Any], - project_id: str, - location: str, -) -> None: - """ - Create and run a pipeline on Vertex AI Pipelines. Use a temporary file to compile the pipeline config, - then run the pipeline job and delete the temporary file. - - :param pipeline_func: Pipeline function, must be a function wrapped :func:`kfp.dsl.pipeline` decorator. - :param pipeline_display_name: A name displayed in the Vertex AI Pipelines UI. - :param pipeline_kwargs: Keyword arguments to pass to the pipeline function. - :param project_id: Google Cloud Project ID - :param location: Datacenter location of Google Cloud Platform to run the pipeline job. - """ - temp_file = tempfile.NamedTemporaryFile(suffix=".yaml") - os.environ["GRPC_DNS_RESOLVER"] = "native" - - aiplatform.init(project=project_id, location=location) - - compiler.Compiler().compile(pipeline_func=pipeline_func, package_path=temp_file.name) - - job = aiplatform.PipelineJob( - display_name=pipeline_display_name, - template_path=temp_file.name, - parameter_values=pipeline_kwargs, - ) - - job.submit() - temp_file.close() diff --git a/cellarium/workflows/local_single_component.py b/cellarium/workflows/local_single_component.py new file mode 100644 index 0000000..84c92ca --- /dev/null +++ b/cellarium/workflows/local_single_component.py @@ -0,0 +1,78 @@ +"""Run a single component cellarium-ml job locally. Useful for testing.""" + +import click +from .shared_components import create_train_op_function, prepare_config_with_overrides + + +@click.command(short_help="Run a single-component cellarium-ml job locally.") +@click.option( + "--tool", + required=True, + help="Tool to run, e.g. 'onepass_mean_var_std'.", +) +@click.option( + "--subcommand", + required=True, + type=click.Choice(["fit", "predict"]), + help="Subcommand to run, either 'fit' or 'predict'.", +) +@click.option( + "--config", + required=True, + help="Local or GCS path to the training config YAML file.", +) +@click.option( + "--copy-data-to-local-disk", + default=True, + type=bool, + help="True copies GCS data to local disk fully (once) before training. False is ephemeral.", +) +@click.option( + "--git-sha", + default="", + type=str, + help="Cellarium-ML git SHA to install (if provided).", +) +@click.option( + "--extract-bucket", + default=None, + help="GCS URI prefix containing extract_*.h5ad files, e.g. gs://my-bucket/my-prefix.", +) +def run_local_single_component( + tool: str, + subcommand: str, + config: str, + copy_data_to_local_disk: bool, + git_sha: str, + extract_bucket=None, +): + """ + Run a single component cellarium-ml job locally without Vertex AI. + + This is useful for testing and development before submitting to Vertex AI. + """ + print(f"Running {tool} {subcommand} locally...") + print(f"Config: {config}") + print(f"Git SHA: {git_sha}") + print(f"Copy data to local disk: {copy_data_to_local_disk}") + + config = prepare_config_with_overrides(config, extract_bucket) + + # Create and run the train operation locally + train_op = create_train_op_function(copy_data_to_local_disk=copy_data_to_local_disk) + + try: + train_op( + tool=tool, + subcommand=subcommand, + config=config, + git_sha=git_sha, + ) + print("Local execution completed successfully!") + except Exception as e: + print(f"Local execution failed: {e}") + raise + + +if __name__ == "__main__": + run_local_single_component() diff --git a/cellarium/workflows/scripts/__init__.py b/cellarium/workflows/scripts/__init__.py new file mode 100644 index 0000000..8b38d75 --- /dev/null +++ b/cellarium/workflows/scripts/__init__.py @@ -0,0 +1 @@ +"""Scripts for kubeflow components.""" diff --git a/cellarium/workflows/scripts/batch_setup.sh b/cellarium/workflows/scripts/batch_setup.sh new file mode 100644 index 0000000..6a88866 --- /dev/null +++ b/cellarium/workflows/scripts/batch_setup.sh @@ -0,0 +1,99 @@ +#!/bin/bash +# Google Cloud Batch setup script for Cellarium ML training jobs +# This script installs Python packages and sets up the environment +# Note: GPU drivers are automatically installed by Google Cloud Batch when GPUs are allocated + +set -euxo pipefail + +# Set environment variables to prevent interactive prompts +export DEBIAN_FRONTEND=noninteractive +export TZ=UTC +export NEEDRESTART_MODE=a # Automatic restart services without prompting +export PYTHONFAULTHANDLER=1 # better logging for python crashes + +echo "🚀 Starting Google Cloud Batch setup..." + +# Comprehensive GPU diagnostics +echo "🔍 GPU Diagnostics..." +echo "Environment variables:" +echo " CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-not set}" +echo " NVIDIA_VISIBLE_DEVICES: ${NVIDIA_VISIBLE_DEVICES:-not set}" + +# echo "Checking for GPU hardware..." +# if lspci | grep -i nvidia; then +# echo "✅ NVIDIA hardware detected" +# else +# echo "❌ No NVIDIA hardware found in lspci" +# fi + +echo "Checking for nvidia-smi..." +if command -v nvidia-smi &> /dev/null; then + echo "✅ nvidia-smi command available" + echo "Running nvidia-smi..." + nvidia-smi || echo "❌ nvidia-smi failed to run" +else + echo "❌ nvidia-smi command not found" + echo "Checking if GPU drivers are installed..." + ls /usr/bin/nvidia-* 2>/dev/null || echo "No nvidia binaries found in /usr/bin/" + ls /dev/nvidia* 2>/dev/null || echo "No nvidia device files found in /dev/" +fi + +echo "Checking Docker GPU runtime..." +if docker info 2>/dev/null | grep -i nvidia; then + echo "✅ Docker NVIDIA runtime detected" +else + echo "⚠️ Docker NVIDIA runtime not detected" +fi + +# # increase file i/o read-ahead +# LOCAL_SSD_MOUNT="/mnt/disks/local-ssd" +# LOCAL_SSD_DEVICE=$(findmnt -n -o SOURCE "$LOCAL_SSD_MOUNT" 2>/dev/null || df -P "$LOCAL_SSD_MOUNT" 2>/dev/null | tail -1 | awk '{print $1}') +# if [ -n "$LOCAL_SSD_DEVICE" ]; then +# echo "⚡ Increasing file I/O read-ahead for $LOCAL_SSD_DEVICE (mounted at $LOCAL_SSD_MOUNT)..." +# blockdev --setra 8192 "$LOCAL_SSD_DEVICE" +# else +# echo "⚠️ Could not determine block device for $LOCAL_SSD_MOUNT, skipping read-ahead tuning" +# fi + +# Install required Python packages +echo "📦 Installing Python packages..." +if [ ! -z "$TRAIN_OP_REQUIREMENTS" ]; then + echo "Installing: $TRAIN_OP_REQUIREMENTS" + pip install -q $TRAIN_OP_REQUIREMENTS + echo "✅ Python packages installed successfully" +else + echo "⚠️ No TRAIN_OP_REQUIREMENTS specified" +fi + +# PyTorch GPU test +echo "🧪 Testing PyTorch GPU availability..." +python3 -c " +import torch +print(f'PyTorch version: {torch.__version__}') +print(f'CUDA available: {torch.cuda.is_available()}') +print(f'CUDA version: {torch.version.cuda}') +if torch.cuda.is_available(): + print(f'GPU count: {torch.cuda.device_count()}') + for i in range(torch.cuda.device_count()): + print(f'GPU {i}: {torch.cuda.get_device_name(i)}') + print(f'GPU {i} memory: {torch.cuda.get_device_properties(i).total_memory / 1e9:.1f} GB') +else: + print('❌ PyTorch cannot access GPU') + print('Possible causes:') + print(' 1. Container not started with --gpus flag') + print(' 2. CUDA version mismatch') + print(' 3. Missing nvidia-container-toolkit') +" + +# Verify environment variables are set +echo "🔍 Verifying environment variables..." +required_vars=("TOOL" "SUBCOMMAND" "CONFIG" "GIT_SHA" "COPY_DATA_TO_LOCAL_DISK") +for var in "${required_vars[@]}"; do + if [ -z "${!var}" ]; then + echo "⚠️ Warning: $var is not set" + else + echo "✅ $var=${!var}" + fi +done + +echo "🎉 Batch setup completed successfully!" diff --git a/cellarium/workflows/scripts/cellarium_cli.py b/cellarium/workflows/scripts/cellarium_cli.py new file mode 100644 index 0000000..1771a51 --- /dev/null +++ b/cellarium/workflows/scripts/cellarium_cli.py @@ -0,0 +1,267 @@ +"""Cellarium CLI execution code for kubeflow components.""" + +from cellarium.ml.cli import main as cellarium_ml_cli +import gcsfs +import yaml +import re +import os +import sys +import contextlib +from datetime import datetime + + +def detect_gcs_paths_in_config(config_path: str): + """Parse config file and detect GCS output paths.""" + gcs_paths = {} + + try: + # Load the config file + if config_path.startswith("gs://"): + fs = gcsfs.GCSFileSystem() + with fs.open(config_path, "r") as f: + config_content = f.read() + else: + with open(config_path, "r") as f: + config_content = f.read() + + # Find all /gcs/ paths in the config + gcs_pattern = r"/gcs/([^/\s]+)(/[^\s]*)?" + matches = re.findall(gcs_pattern, config_content) + + for bucket, path in matches: + full_gcs_path = f"/gcs/{bucket}{path}" + gcs_url = f"gs://{bucket}{path}" + gcs_paths[full_gcs_path] = gcs_url + + return gcs_paths + except Exception as e: + print(f"Warning: Could not parse config for GCS paths: {e}") + return {} + + +def setup_gcs_output_handling(config_path: str) -> str: + """Main function to set up GCS output handling.""" + print("Setting up GCS output handling...") + + # Always create the job_logs directory for Batch logging + # (even if CELLARIUM_CAPTURE_LOGS is false, Batch might write logs here) + job_logs_dir = "/tmp/gcs_output/job_logs" + os.makedirs(job_logs_dir, exist_ok=True) + + # Detect GCS paths in config + gcs_paths = detect_gcs_paths_in_config(config_path) + + if not gcs_paths: + print("No GCS output paths detected in config") + return config_path + + print(f"Detected {len(gcs_paths)} GCS output paths:") + for gcs_path, gcs_url in gcs_paths.items(): + print(f" {gcs_path} -> {gcs_url}") + + # Set up local directories and update config + path_mapping = {} + + # Load the config + if config_path.startswith("gs://"): + fs = gcsfs.GCSFileSystem() + with fs.open(config_path, "r") as f: + config_content = f.read() + else: + with open(config_path, "r") as f: + config_content = f.read() + + # Replace GCS paths with local paths + updated_content = config_content + for gcs_path, gcs_url in gcs_paths.items(): + # Create local directory + local_path = f"/tmp/gcs_output{gcs_path[4:]}" # Remove /gcs prefix + os.makedirs(local_path, exist_ok=True) + path_mapping[gcs_path] = local_path + print(f"Mapped {gcs_path} -> {local_path} (will sync to {gcs_url})") + + # Update config content + updated_content = updated_content.replace(gcs_path, local_path) + + # Write updated config to local file + local_config_path = "/tmp/config_with_local_paths.yaml" + with open(local_config_path, "w") as f: + f.write(updated_content) + + # Store paths for later syncing + with open("/tmp/gcs_sync_info.yaml", "w") as f: + yaml.dump({"gcs_paths": gcs_paths, "path_mapping": path_mapping}, f) + + print(f"Updated config saved to {local_config_path}") + return local_config_path + + +def setup_log_capture(): + """Set up stdout/stderr capture to files for later GCS sync.""" + # Create logs directory + log_dir = "/tmp/gcs_output/job_logs" + os.makedirs(log_dir, exist_ok=True) + + # Generate timestamped log files + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + stdout_file = f"{log_dir}/stdout_{timestamp}.log" + stderr_file = f"{log_dir}/stderr_{timestamp}.log" + + return stdout_file, stderr_file + + +@contextlib.contextmanager +def capture_logs_to_files(stdout_file, stderr_file): + """Context manager to capture stdout/stderr to files while still showing output.""" + # Open log files + stdout_log = open(stdout_file, "w") + stderr_log = open(stderr_file, "w") + + class TeeWriter: + def __init__(self, original, log_file): + self.original = original + self.log_file = log_file + + def write(self, text): + self.original.write(text) + self.log_file.write(text) + self.original.flush() + self.log_file.flush() + + def flush(self): + self.original.flush() + self.log_file.flush() + + # Save original stdout/stderr + original_stdout = sys.stdout + original_stderr = sys.stderr + + try: + # Replace with tee writers + sys.stdout = TeeWriter(original_stdout, stdout_log) + sys.stderr = TeeWriter(original_stderr, stderr_log) + yield stdout_file, stderr_file + finally: + # Restore original stdout/stderr + sys.stdout = original_stdout + sys.stderr = original_stderr + stdout_log.close() + stderr_log.close() + + +# Staging directory for bare-filename task outputs (e.g. output_path: results.csv). +# Placed on the local SSD when available so the post-container host-VM runnable can +# sweep it up with `gsutil -m cp -r`. Falls back to /tmp when no SSD is mounted +# (in which case post-run GCS upload is skipped — handled by the caller). +_SSD_MOUNT = "/mnt/disks/local-ssd" +TASK_OUTPUT_DIR = ( + f"{_SSD_MOUNT}/task_outputs" + if os.path.ismount(_SSD_MOUNT) + else "/tmp/cellarium_task_outputs" +) + + +def setup_task_output_dir() -> str: + """Create the staging directory for bare-filename task outputs and set CWD to it. + + Any relative output paths written by the model (e.g. ``output_path: results.csv``) + will land here, making them easy to sweep up and upload to GCS afterwards. + Lightning's ``default_root_dir`` is always absolute so checkpoint paths are unaffected. + """ + os.makedirs(TASK_OUTPUT_DIR, exist_ok=True) + os.chdir(TASK_OUTPUT_DIR) + print(f" Task output staging dir: {TASK_OUTPUT_DIR} (CWD changed)") + if not os.path.ismount(_SSD_MOUNT): + print( + f" WARNING: local SSD not mounted at {_SSD_MOUNT}; " + "task outputs will NOT be uploaded by the post-run gsutil step." + ) + return TASK_OUTPUT_DIR + + +def finalize_gcs_output_sync(): + """Final step: sync all outputs back to GCS.""" + try: + print("Finalizing GCS output sync...") + + # Load sync info + if not os.path.exists("/tmp/gcs_sync_info.yaml"): + print("No GCS sync info found") + return + + with open("/tmp/gcs_sync_info.yaml", "r") as f: + sync_info = yaml.safe_load(f) + + gcs_paths = sync_info.get("gcs_paths", {}) + path_mapping = sync_info.get("path_mapping", {}) + + if not gcs_paths: + print("No GCS paths to sync") + return + + fs = gcsfs.GCSFileSystem() + + for gcs_path, gcs_url in gcs_paths.items(): + local_path = path_mapping.get(gcs_path) + if not local_path or not os.path.exists(local_path): + print(f"Skipping {gcs_path}: local path {local_path} not found") + continue + + try: + print(f"Syncing {local_path} to {gcs_url}") + + # Upload all files in the local directory + for root, dirs, files in os.walk(local_path): + for file in files: + local_file = os.path.join(root, file) + + # Calculate relative path and GCS destination + rel_path = os.path.relpath(local_file, local_path) + gcs_file = f"{gcs_url.rstrip('/')}/{rel_path}" + + # Ensure parent directory exists in GCS + gcs_dir = "/".join(gcs_file.split("/")[:-1]) + fs.makedirs(gcs_dir, exist_ok=True) + + # Upload file + fs.put(local_file, gcs_file) + print(f" Uploaded {rel_path}") + + print(f" Successfully synced {local_path} to {gcs_url}") + + except Exception as e: + print(f" Failed to sync {local_path} to {gcs_url}: {e}") + + except Exception as e: + print(f"Warning: Could not finalize GCS sync: {e}") + + +# Set up GCS output handling before training +updated_config = setup_gcs_output_handling(config) # noqa: F821 + +# Change CWD to the staging directory so that relative output paths (e.g. +# output_path: results.csv) land in a known location for the post-run gsutil upload. +setup_task_output_dir() + +# Check if we should capture logs to files (for GCS sync) +capture_logs = os.environ.get("CELLARIUM_CAPTURE_LOGS", "").lower() == "true" # noqa: F821 + +try: + if capture_logs: + # Set up log capture + stdout_file, stderr_file = setup_log_capture() + print(f" Capturing logs to {stdout_file} and {stderr_file}") + + # Run with log capture + with capture_logs_to_files(stdout_file, stderr_file): + cellarium_ml_cli(args=[tool, subcommand, "--config", updated_config]) # noqa: F821 + + print(" Logs captured and will be synced to GCS") + else: + # Run normally without log capture + cellarium_ml_cli(args=[tool, subcommand, "--config", updated_config]) # noqa: F821 +finally: + # Sync /gcs/-style path outputs (legacy mechanism) + finalize_gcs_output_sync() + # Task outputs in TASK_OUTPUT_DIR are uploaded by the post-container host-VM + # runnable using `gsutil -m cp -r` (bulk, parallel). Nothing to do here. diff --git a/cellarium/workflows/scripts/data_download.py b/cellarium/workflows/scripts/data_download.py new file mode 100644 index 0000000..767d47e --- /dev/null +++ b/cellarium/workflows/scripts/data_download.py @@ -0,0 +1,105 @@ +"""Data download and processing code for kubeflow components.""" + +import os +import glob +import gcsfs +import subprocess +from ruamel.yaml import YAML + +# Get config from environment variable +config = os.environ.get("CONFIG") +if not config: + raise RuntimeError("CONFIG environment variable not set") + +# 0. localize the config file and set up local data directory +fs = gcsfs.GCSFileSystem() + +# Create a dedicated directory for training data on local disk +# Prefer Local SSD if available for high-performance I/O +if os.path.exists("/mnt/disks/local-ssd"): + LOCAL_DATA_DIR = "/mnt/disks/local-ssd/training_data" + print(" Using Local SSD for high-performance data storage") + print("df -h /mnt/disks/local-ssd:") + subprocess.run(["df", "-h", "/mnt/disks/local-ssd"]) +else: + LOCAL_DATA_DIR = "/tmp/training_data" + print(" Using boot disk for data storage") + +os.makedirs(LOCAL_DATA_DIR, exist_ok=True) +print(f" Created local data directory: {LOCAL_DATA_DIR}") + + +# Handle config file localization +if config.startswith("gs://"): + print(f" Downloading config from GCS: {config}") + # Use Local SSD for config if available, otherwise use /tmp + if os.path.exists("/mnt/disks/local-ssd"): + config_local_path = "/mnt/disks/local-ssd/downloaded_config.yaml" + else: + config_local_path = "/tmp/downloaded_config.yaml" + + with fs.open(config, "r") as fsrc: + with open(config_local_path, "w") as fdst: + fdst.write(fsrc.read()) + print(f" Config downloaded to: {config_local_path}") +else: + print(f" Using local config: {config}") + config_local_path = config + +# 1. find data reference +yaml = YAML() +yaml.preserve_quotes = True + +with open(config_local_path, "r") as f: + config_data = yaml.load(f) +try: + original_data_reference = config_data["data"]["dadc"]["init_args"]["filenames"] +except KeyError: + raise RuntimeError( + f"Could not find dataset in {config_local_path} when attempting to access data.dadc.init_args.filenames\n\n" + f"{os.system('cat ' + config_local_path)}" + ) + +# 2. download data to local disk +print(f"Copying data from GCS {original_data_reference} to local disk {LOCAL_DATA_DIR}") + +sentinel_path = os.path.join(LOCAL_DATA_DIR, ".download_complete") +if os.path.exists(sentinel_path): + print(f" Sentinel found at {sentinel_path} — data already downloaded, skipping.") +else: + print(f" Downloading data via gsutil: {original_data_reference} -> {LOCAL_DATA_DIR}/") + subprocess.run( + f"gsutil -m cp {original_data_reference} {LOCAL_DATA_DIR}/", + shell=True, + executable="/bin/bash", + check=True, + ) + open(sentinel_path, "w").close() + +print("Listing local .h5ad files (at most 10):") +h5ad_files = glob.glob(os.path.join(LOCAL_DATA_DIR, "*.h5ad")) +print("\n".join(h5ad_files[:10])) + +# 3. rewrite the config file to point to the local data +if isinstance(original_data_reference, str): + local_data_reference = os.path.join( + LOCAL_DATA_DIR, os.path.basename(original_data_reference) + ) +else: + local_data_reference = [ + os.path.join(LOCAL_DATA_DIR, os.path.basename(f)) + for f in original_data_reference + ] +config_data["data"]["dadc"]["init_args"]["filenames"] = local_data_reference +with open(config_local_path, "w") as f: + yaml.dump(config_data, f) +print(f"Re-writing config file {config_local_path} to point to local data:") +print(f" Data files now point to: {local_data_reference}") +print(" Config file contents:\n") +with open(config_local_path, "r") as f: + print(f.read()) + +# Update the CONFIG environment variable to point to the local config file +# This ensures that cellarium_cli.py will use the updated config with local data paths +os.environ["CONFIG"] = config_local_path +print(f" Updated CONFIG environment variable to: {config_local_path}") diff --git a/cellarium/workflows/scripts/git_install.py b/cellarium/workflows/scripts/git_install.py new file mode 100644 index 0000000..e8e8a1b --- /dev/null +++ b/cellarium/workflows/scripts/git_install.py @@ -0,0 +1,17 @@ +"""Git installation code for kubeflow components.""" + +import os + +# re-install cellarium-ml if a git sha is provided +# if git_sha != "": # noqa: F821 +# os.system("apt-get update") +# os.system("apt-get install -y git") +# cmd = f"pip install -U -q git+https://github.com/cellarium-ai/cellarium-ml.git@{git_sha}" # noqa: F821 +# print(cmd) +# os.system(cmd) + +# faster +if git_sha != "": # noqa: F821 + cmd = f"pip install -q https://github.com/cellarium-ai/cellarium-ml/archive/{git_sha}.tar.gz" # noqa: F821 + print(cmd) + os.system(cmd) diff --git a/cellarium/workflows/scripts/pytorch_setup.py b/cellarium/workflows/scripts/pytorch_setup.py new file mode 100644 index 0000000..a9a1d57 --- /dev/null +++ b/cellarium/workflows/scripts/pytorch_setup.py @@ -0,0 +1,19 @@ +"""PyTorch environment setup code for kubeflow components.""" + +import psutil +import torch +import os + +# set env variables to allow pytorch to use all CPUs +num_physical_cores = psutil.cpu_count(logical=False) +os.environ["OMP_NUM_THREADS"] = str(num_physical_cores) +os.environ["MKL_NUM_THREADS"] = str(num_physical_cores) +os.environ["OPENBLAS_NUM_THREADS"] = str(num_physical_cores) # Only if using OpenBLAS +os.environ["NUMEXPR_NUM_THREADS"] = str(num_physical_cores) # Not critical for PyTorch + +# handle multi-node training +if os.environ.get("RANK") is not None: + os.environ["NODE_RANK"] = os.environ.get("RANK") + +# set number of threads for torch +torch.set_num_threads(num_physical_cores) diff --git a/cellarium/workflows/shared_components.py b/cellarium/workflows/shared_components.py new file mode 100644 index 0000000..53c54d1 --- /dev/null +++ b/cellarium/workflows/shared_components.py @@ -0,0 +1,1134 @@ +"""Shared component code for cellarium workflows.""" + +import gcsfs +from pathlib import Path +from typing import Optional + + +local_output_path = "/mnt/disks/local-ssd/run_outputs" + + +def _assert_gcs_bucket_exists(gcs_path: str) -> None: + """Raise an error if the GCS bucket in gcs_path does not exist or is not accessible.""" + # Extract bucket name from gs://bucket/... or gs://bucket + without_scheme = gcs_path[5:] # strip "gs://" + bucket_name = without_scheme.split("/")[0] + fs = gcsfs.GCSFileSystem() + if not fs.exists(f"gs://{bucket_name}"): + raise FileNotFoundError( + f" Output GCS bucket does not exist or is not accessible: gs://{bucket_name}\n" + f" (from config path: {gcs_path})" + ) + print(f" Output GCS bucket exists: gs://{bucket_name}") + + +def extract_data_gcs_bucket_from_config(config_path: str) -> str: + """ + Extract the data GCS bucket name from the config file's data.dadc.init_args.filenames. + + Args: + config_path: Path to the config YAML file (local or gs://) + + Returns: + The GCS bucket name (e.g. 'my-bucket') if found, otherwise an empty string + """ + import re + + try: + if config_path.startswith("gs://"): + fs = gcsfs.GCSFileSystem() + with fs.open(config_path, "r") as f: + content = f.read() + else: + with open(config_path, "r") as f: + content = f.read() + + # Find filenames: value under data.dadc.init_args + match = re.search(r"^\s*filenames:\s*([^\s\n]+)", content, re.MULTILINE) + if not match: + print(" No filenames field found in config data section") + return "" + + filenames = match.group(1).strip() + if not filenames.startswith("gs://"): + print(f" Data filenames path is not a GCS path: {filenames}") + return "" + + # Extract bucket name — handles brace-expansion paths like + # gs://bucket/path/extract_{0..10}.h5ad + bucket = filenames[5:].split("/")[0] + print(f" Detected data GCS bucket from config: {bucket}") + return bucket + + except Exception as e: + print(f" Warning: Could not parse config for data bucket: {e}") + return "" + + +def extract_data_gcs_glob_from_config(config_path: str) -> str: + """ + Extract a gcloud-compatible glob URI for the data files in the config. + + Brace-expansion patterns like gs://bucket/path/extract_{000000..009446}.h5ad + are converted to gs://bucket/path/extract_*.h5ad. + + Returns the glob URI string, or an empty string if not found. + """ + import re + + try: + if config_path.startswith("gs://"): + fs = gcsfs.GCSFileSystem() + with fs.open(config_path, "r") as f: + content = f.read() + else: + with open(config_path, "r") as f: + content = f.read() + + match = re.search(r"^\s*filenames:\s*([^\s\n]+)", content, re.MULTILINE) + if not match: + return "" + + filenames = match.group(1).strip() + if not filenames.startswith("gs://"): + return "" + + print(f" Data GCS filenames: {filenames}") + return filenames + + except Exception as e: + print(f" Warning: Could not parse config for data GCS glob: {e}") + return "" + + +def assert_gcs_bucket_accessible_as_service_account( + project: str, bucket_name: str +) -> None: + """ + Verify that the Compute Engine default service account can access a GCS bucket. + + Impersonates the default Compute Engine SA + ({project_number}-compute@developer.gserviceaccount.com) and checks bucket + existence/access. If your user account lacks + roles/iam.serviceAccountTokenCreator on that SA, falls back to checking with + your current credentials and prints a warning. + + Args: + project: Google Cloud project ID + bucket_name: GCS bucket name (without gs:// prefix) + + Raises: + PermissionError: If the bucket is not accessible + """ + import requests + from google.auth import default + from google.auth.transport.requests import Request + + # --- get project number to build the default SA email --- + credentials, _ = default() + credentials.refresh(Request()) + token = credentials.token + resp = requests.get( + f"https://cloudresourcemanager.googleapis.com/v1/projects/{project}", + headers={"Authorization": f"Bearer {token}"}, + timeout=10, + ) + resp.raise_for_status() + project_number = resp.json()["projectNumber"] + sa_email = f"{project_number}-compute@developer.gserviceaccount.com" + print(f" Checking data bucket access as Compute Engine SA: {sa_email}") + + # --- try to impersonate the SA and check the bucket --- + try: + from google.auth import impersonated_credentials + + impersonated = impersonated_credentials.Credentials( + source_credentials=credentials, + target_principal=sa_email, + target_scopes=["https://www.googleapis.com/auth/devstorage.read_only"], + lifetime=60, + ) + # Force a token fetch so impersonation errors surface here rather than lazily + impersonated.refresh(Request()) + fs = gcsfs.GCSFileSystem(token=impersonated) + accessible = fs.exists(f"gs://{bucket_name}") + check_label = f"SA {sa_email}" + except Exception as impersonation_err: + print( + f" WARNING: Could not impersonate SA {sa_email}: {impersonation_err}\n" + " (Your account may need roles/iam.serviceAccountTokenCreator on that SA.)\n" + " Falling back to checking with your current user credentials — " + "this may not reflect actual SA permissions." + ) + fs = gcsfs.GCSFileSystem() + accessible = fs.exists(f"gs://{bucket_name}") + check_label = "current user" + + if not accessible: + raise PermissionError( + f" Data bucket 'gs://{bucket_name}' is not accessible to {check_label}.\n" + " The Batch job will fail when it tries to read training data.\n" + " Fix: grant Storage Object Viewer on gs://{bucket_name} to {sa_email}." + ) + print(f" Data bucket 'gs://{bucket_name}' is accessible to {check_label}.") + + +def assert_data_first_file_exists(config_path: str) -> None: + """ + Check that the first data file referenced in the config's filenames field exists in GCS. + + Resolves brace-expansion patterns (e.g. ``extract_{0..10}.h5ad`` or + ``extract_{000000..000010}.h5ad``) to the first concrete filename and verifies + it exists, giving a fast-fail check for path typos. + + Args: + config_path: Path to the config YAML file (local or gs://) + + Raises: + FileNotFoundError: If the first data file does not exist in GCS + """ + import re + + try: + if config_path.startswith("gs://"): + fs = gcsfs.GCSFileSystem() + with fs.open(config_path, "r") as f: + content = f.read() + else: + with open(config_path, "r") as f: + content = f.read() + + match = re.search(r"^\s*filenames:\s*([^\s\n]+)", content, re.MULTILINE) + if not match: + print(" No filenames field found — skipping file existence check") + return + + filenames = match.group(1).strip() + if not filenames.startswith("gs://"): + print(" Filenames is not a GCS path — skipping file existence check") + return + + # Resolve {START..END} brace expansion to START, preserving leading zeros + first_file = re.sub(r"\{(\d+)\.\.(\d+)\}", r"\1", filenames) + + print(f" Checking first data file exists: {first_file}") + fs = gcsfs.GCSFileSystem() + if not fs.exists(first_file): + raise FileNotFoundError( + f" First data file does not exist: {first_file}\n" + f" (from filenames pattern: {filenames})\n" + " Check for typos in the filenames path in your config." + ) + print(f" First data file exists: {first_file}") + + except FileNotFoundError: + raise + except Exception as e: + print(f" Warning: Could not check data file existence: {e}") + + +def extract_ckpt_path_from_config(config_path: str) -> str: + """ + Extract the ckpt_path value from a Lightning CLI config YAML. + + Args: + config_path: Path to the config YAML file (local or gs://) + + Returns: + The raw ckpt_path string if present and not null, otherwise an empty string. + """ + import re + + try: + if config_path.startswith("gs://"): + fs = gcsfs.GCSFileSystem() + with fs.open(config_path, "r") as f: + content = f.read() + else: + with open(config_path, "r") as f: + content = f.read() + + match = re.search(r"^ckpt_path:\s*([^\s\n]+)", content, re.MULTILINE) + if not match: + return "" + + value = match.group(1).strip() + if value.lower() in ("null", "~", ""): + return "" + + return value + + except Exception as e: + print(f" Warning: Could not parse config for ckpt_path: {e}") + return "" + + +def assert_ckpt_path_exists(ckpt_path: str) -> None: + """ + Verify that the checkpoint file referenced by ckpt_path exists in GCS. + + Args: + ckpt_path: GCS path to the checkpoint file (gs://...) + + Raises: + FileNotFoundError: If the checkpoint file does not exist + """ + if not ckpt_path or ckpt_path.lower() in ("null", "~"): + return + + if not ckpt_path.startswith("gs://"): + print(f" ckpt_path is not a GCS path, skipping existence check: {ckpt_path}") + return + + print(f" Checking ckpt_path exists: {ckpt_path}") + fs = gcsfs.GCSFileSystem() + if not fs.exists(ckpt_path): + raise FileNotFoundError( + f" Checkpoint file does not exist: {ckpt_path}\n" + " Check for typos in ckpt_path in your config." + ) + print(f" Checkpoint file exists: {ckpt_path}") + + +def extract_output_gcs_bucket_from_config(config_path: str) -> str: + """ + Extract the output GCS bucket path from the config file's trainer.default_root_dir. + + Args: + config_path: Path to the config YAML file (local or gs://) + + Returns: + The GCS bucket path if found, otherwise an empty string + """ + try: + # Load the config file content as text + if config_path.startswith("gs://"): + fs = gcsfs.GCSFileSystem() + with fs.open(config_path, "r") as f: + content = f.read() + else: + with open(config_path, "r") as f: + content = f.read() + + # Use regex to find default_root_dir value + import re + + # Look for default_root_dir: followed by the path + pattern = r"default_root_dir:\s*([^\s\n]+)" + match = re.search(pattern, content) + + if not match: + print(" No default_root_dir found in config") + return "" + + default_root_dir = match.group(1).strip() + + # Check if it's a GCS path + if default_root_dir.startswith("gs://"): + print(f" Detected GCS output path: {default_root_dir}") + return default_root_dir + elif default_root_dir.startswith("/gcs/"): + # Convert /gcs/bucket/path format to gs://bucket/path + gcs_path = default_root_dir[5:] # Remove '/gcs/' prefix + if "/" in gcs_path: + bucket, path = gcs_path.split("/", 1) + gcs_url = f"gs://{bucket}/{path}" + else: + gcs_url = f"gs://{gcs_path}" + print(f" Detected GCS output path: {default_root_dir} -> {gcs_url}") + return gcs_url + else: + print(f" Local output path detected: {default_root_dir}") + return "" + + except Exception as e: + print(f" Warning: Could not parse config for output path: {e}") + return "" + + +def prepare_config_with_overrides( + config_path: str, + extract_bucket: Optional[str] = None, +) -> str: + """ + Prepare a config file for job submission, optionally patching dataset fields. + + If ``extract_bucket`` is provided (e.g. ``gs://my-bucket/my-prefix``), the + function lists all ``extract_*.h5ad`` files at that GCS prefix, reads the + first and last files to determine ``shard_size`` and ``last_shard_size``, + constructs a brace-expansion ``filenames`` pattern, patches those three + fields in a copy of the config YAML, and returns the path to the patched + temporary file. When ``extract_bucket`` is ``None`` the original + ``config_path`` is returned unchanged. + + Args: + config_path: Local or ``gs://`` path to a Lightning CLI config YAML. + extract_bucket: GCS URI prefix containing ``extract_*.h5ad`` shards, + e.g. ``gs://my-bucket/my-prefix``. + + Returns: + Path to a (possibly modified) config YAML file. + """ + if extract_bucket is None: + return config_path + + import re as _re + import tempfile + + import h5py + from ruamel.yaml import YAML + + fs = gcsfs.GCSFileSystem() + prefix = extract_bucket.rstrip("/") + matched = fs.glob(f"{prefix}/extract_*.h5ad") + if not matched: + raise FileNotFoundError(f"No extract_*.h5ad files found under {prefix}") + + def _shard_digits(p: str) -> str: + """Return the raw digit string from an extract filename (preserves leading zeros).""" + m = _re.search(r"extract_(\d+)\.h5ad$", p) + return m.group(1) if m else "" + + def _shard_index(p: str) -> int: + d = _shard_digits(p) + return int(d) if d else -1 + + matched = sorted(matched, key=_shard_index) + first_digits = _shard_digits(matched[0]) + last_digits = _shard_digits(matched[-1]) + # Preserve leading-zero padding: width is taken from the first filename + width = len(first_digits) + last_idx = int(last_digits) + start_fmt = first_digits # e.g. "0" or "000000" + end_fmt = str(last_idx).zfill(width) # e.g. "9446" or "009446" + filenames = f"{prefix}/extract_{{{start_fmt}..{end_fmt}}}.h5ad" + + def _obs_count(gcs_path: str) -> int: + with fs.open(gcs_path, "rb") as raw: + with h5py.File(raw, "r") as h5f: + obs = h5f["obs"] + idx_col = obs.attrs.get("_index", None) + if idx_col is not None: + if isinstance(idx_col, bytes): + idx_col = idx_col.decode() + if idx_col in obs: + return len(obs[idx_col]) + x = h5f["X"] + if isinstance(x, h5py.Group) and "indptr" in x: + return len(x["indptr"]) - 1 + return x.shape[0] + + shard_size = _obs_count(matched[0]) + last_shard_size = _obs_count(matched[-1]) + + print(f" Found {len(matched)} shards under {prefix}") + print(f" shard_size (from first file): {shard_size}") + print(f" last_shard_size (from last file): {last_shard_size}") + print(f" filenames: {filenames}") + + yaml = YAML() + yaml.preserve_quotes = True + if config_path.startswith("gs://"): + with fs.open(config_path, "r") as f: + doc = yaml.load(f) + else: + with open(config_path, "r") as f: + doc = yaml.load(f) + + try: + dadc_args = doc["data"]["dadc"]["init_args"] + except KeyError as exc: + raise KeyError( + f"Config {config_path!r} is missing expected key {exc}. " + "Expected structure: data.dadc.init_args" + ) from exc + + dadc_args["filenames"] = filenames + dadc_args["shard_size"] = shard_size + dadc_args["last_shard_size"] = last_shard_size + + tmp = tempfile.NamedTemporaryFile( + mode="w", suffix=".yaml", delete=False, prefix="cellarium_config_" + ) + yaml.dump(doc, tmp) + tmp.close() + + print(f" Patched config written to: {tmp.name}") + return tmp.name + + +def get_machine_type_resources(machine_type: str) -> tuple[int, int]: + """ + Get CPU (in milliCPU) and memory (in MiB) for a given machine type. + + Args: + machine_type: Machine type string (e.g., 'n1-standard-4') + + Returns: + Tuple of (cpu_milli, memory_mib) + """ + # Common machine type mappings + # Format: machine_type -> (cpu_milli, memory_mib) + machine_type_specs = { + # N1 Standard series + "n1-standard-1": (1000, 3840), # 1 vCPU, 3.75 GB + "n1-standard-2": (2000, 7680), # 2 vCPU, 7.5 GB + "n1-standard-4": (4000, 15360), # 4 vCPU, 15 GB + "n1-standard-8": (8000, 30720), # 8 vCPU, 30 GB + "n1-standard-16": (16000, 61440), # 16 vCPU, 60 GB + "n1-standard-32": (32000, 122880), # 32 vCPU, 120 GB + "n1-standard-64": (64000, 245760), # 64 vCPU, 240 GB + "n1-standard-96": (96000, 368640), # 96 vCPU, 360 GB + # N1 High-memory series + "n1-highmem-1": (1000, 6656), # 1 vCPU, 6.5 GB + "n1-highmem-2": (2000, 13312), # 2 vCPU, 13 GB + "n1-highmem-4": (4000, 26624), # 4 vCPU, 26 GB + "n1-highmem-8": (8000, 53248), # 8 vCPU, 52 GB + "n1-highmem-16": (16000, 106496), # 16 vCPU, 104 GB + "n1-highmem-32": (32000, 212992), # 32 vCPU, 208 GB + "n1-highmem-64": (64000, 425984), # 64 vCPU, 416 GB + "n1-highmem-96": (96000, 638976), # 96 vCPU, 624 GB + # N1 High-CPU series + "n1-highcpu-2": (2000, 1843), # 2 vCPU, 1.8 GB + "n1-highcpu-4": (4000, 3686), # 4 vCPU, 3.6 GB + "n1-highcpu-8": (8000, 7373), # 8 vCPU, 7.2 GB + "n1-highcpu-16": (16000, 14746), # 16 vCPU, 14.4 GB + "n1-highcpu-32": (32000, 29491), # 32 vCPU, 28.8 GB + "n1-highcpu-64": (64000, 58982), # 64 vCPU, 57.6 GB + "n1-highcpu-96": (96000, 88474), # 96 vCPU, 86.4 GB + # N2 Standard series + "n2-standard-2": (2000, 8192), # 2 vCPU, 8 GB + "n2-standard-4": (4000, 16384), # 4 vCPU, 16 GB + "n2-standard-8": (8000, 32768), # 8 vCPU, 32 GB + "n2-standard-16": (16000, 65536), # 16 vCPU, 64 GB + "n2-standard-32": (32000, 131072), # 32 vCPU, 128 GB + "n2-standard-48": (48000, 196608), # 48 vCPU, 192 GB + "n2-standard-64": (64000, 262144), # 64 vCPU, 256 GB + "n2-standard-80": (80000, 327680), # 80 vCPU, 320 GB + "n2-standard-96": (96000, 393216), # 96 vCPU, 384 GB + "n2-standard-128": (128000, 524288), # 128 vCPU, 512 GB + # N2 High-memory series + "n2-highmem-2": (2000, 16384), # 2 vCPU, 16 GB + "n2-highmem-4": (4000, 32768), # 4 vCPU, 32 GB + "n2-highmem-8": (8000, 65536), # 8 vCPU, 64 GB + "n2-highmem-16": (16000, 131072), # 16 vCPU, 128 GB + "n2-highmem-32": (32000, 262144), # 32 vCPU, 256 GB + "n2-highmem-48": (48000, 393216), # 48 vCPU, 384 GB + "n2-highmem-64": (64000, 524288), # 64 vCPU, 512 GB + "n2-highmem-80": (80000, 655360), # 80 vCPU, 640 GB + "n2-highmem-96": (96000, 786432), # 96 vCPU, 768 GB + "n2-highmem-128": (128000, 884736), # 128 vCPU, 864 GB + # N2 High-CPU series + "n2-highcpu-2": (2000, 2048), # 2 vCPU, 2 GB + "n2-highcpu-4": (4000, 4096), # 4 vCPU, 4 GB + "n2-highcpu-8": (8000, 8192), # 8 vCPU, 8 GB + "n2-highcpu-16": (16000, 16384), # 16 vCPU, 16 GB + "n2-highcpu-32": (32000, 32768), # 32 vCPU, 32 GB + "n2-highcpu-48": (48000, 49152), # 48 vCPU, 48 GB + "n2-highcpu-64": (64000, 65536), # 64 vCPU, 64 GB + "n2-highcpu-80": (80000, 81920), # 80 vCPU, 80 GB + "n2-highcpu-96": (96000, 98304), # 96 vCPU, 96 GB + # C2 High-CPU series + "c2-standard-4": (4000, 16384), # 4 vCPU, 16 GB + "c2-standard-8": (8000, 32768), # 8 vCPU, 32 GB + "c2-standard-16": (16000, 65536), # 16 vCPU, 64 GB + "c2-standard-30": (30000, 122880), # 30 vCPU, 120 GB + "c2-standard-60": (60000, 245760), # 60 vCPU, 240 GB + # E2 series + "e2-standard-2": (2000, 8192), # 2 vCPU, 8 GB + "e2-standard-4": (4000, 16384), # 4 vCPU, 16 GB + "e2-standard-8": (8000, 32768), # 8 vCPU, 32 GB + "e2-standard-16": (16000, 65536), # 16 vCPU, 64 GB + "e2-standard-32": (32000, 131072), # 32 vCPU, 128 GB + # A2 GPU-optimized series + "a2-highgpu-1g": (12000, 87040), # 12 vCPU, 85 GB, 1 A100 + "a2-highgpu-2g": (24000, 174080), # 24 vCPU, 170 GB, 2 A100 + "a2-highgpu-4g": (48000, 348160), # 48 vCPU, 340 GB, 4 A100 + "a2-highgpu-8g": (96000, 696320), # 96 vCPU, 680 GB, 8 A100 + "a2-megagpu-16g": (96000, 1392640), # 96 vCPU, 1360 GB, 16 A100 + # G2 GPU-optimized series + "g2-standard-4": (4000, 16384), # 4 vCPU, 16 GB, 1 L4 + "g2-standard-8": (8000, 32768), # 8 vCPU, 32 GB, 1 L4 + "g2-standard-12": (12000, 49152), # 12 vCPU, 48 GB, 1 L4 + "g2-standard-16": (16000, 65536), # 16 vCPU, 64 GB, 1 L4 + "g2-standard-24": (24000, 98304), # 24 vCPU, 96 GB, 2 L4 + "g2-standard-32": (32000, 131072), # 32 vCPU, 128 GB, 1 L4 + "g2-standard-48": (48000, 196608), # 48 vCPU, 192 GB, 4 L4 + "g2-standard-96": (96000, 393216), # 96 vCPU, 384 GB, 8 L4 + } + + if machine_type in machine_type_specs: + cpu_milli, memory_mib = machine_type_specs[machine_type] + print( + f" Machine type {machine_type}: {cpu_milli // 1000} vCPU, {memory_mib // 1024:.1f} GB" + ) + return cpu_milli, memory_mib + else: + # For unknown machine types, use conservative defaults + print(f" Unknown machine type '{machine_type}', using defaults: 2 vCPU, 2 GB") + print( + " Consider adding this machine type to the mapping for optimal resource allocation" + ) + return 2000, 2048 # 2 vCPU, 2 GB + + +def get_train_op_requirements() -> list[str]: + """Return the packages required for train_op execution.""" + return [ + "gcsfs", + "tensorboard", + "psutil", + "ruamel.yaml", + "pyarrow", + "fastparquet", + "owlready2", + "networkx", + ] + + +def _load_script_as_string(script_name: str) -> str: + """Load a Python script from the scripts directory as a string.""" + scripts_dir = Path(__file__).parent / "scripts" + script_path = scripts_dir / f"{script_name}.py" + + if not script_path.exists(): + raise FileNotFoundError(f"Script not found: {script_path}") + + with open(script_path, "r") as f: + content = f.read() + + # Remove the docstring and any module-level comments for cleaner execution + lines = content.split("\n") + # Skip lines that are just comments or docstrings at the top + start_idx = 0 + for i, line in enumerate(lines): + stripped = line.strip() + if ( + stripped + and not stripped.startswith("#") + and not stripped.startswith('"""') + and not stripped.startswith("'''") + ): + start_idx = i + break + + return "\n".join(lines[start_idx:]) + + +def _load_bash_script_as_string(script_name: str) -> str: + """Load a bash script from the scripts directory as a string.""" + scripts_dir = Path(__file__).parent / "scripts" + script_path = scripts_dir / f"{script_name}.sh" + + if not script_path.exists(): + raise FileNotFoundError(f"Script not found: {script_path}") + + with open(script_path, "r") as f: + content = f.read() + + return content + + +def get_pytorch_setup_code() -> str: + """Returns the PyTorch setup code as a string.""" + return _load_script_as_string("pytorch_setup") + + +def get_git_install_code() -> str: + """Returns the git installation code as a string.""" + return _load_script_as_string("git_install") + + +def get_data_download_code() -> str: + """Returns the data download code as a string.""" + return _load_script_as_string("data_download") + + +def get_cellarium_cli_code() -> str: + """Returns the cellarium CLI execution code as a string.""" + return _load_script_as_string("cellarium_cli") + + +def _get_train_op_text(copy_data_to_local_disk: bool = True) -> str: + """ + Returns the ground-truth train_op implementation as a string. + + This is the single source of truth for the train_op logic. + + Args: + copy_data_to_local_disk: Whether to include data download code + + Returns: + Complete train_op implementation as a string with inlined script content + """ + # Get the actual script content and inline it + git_install_code = get_git_install_code() + pytorch_setup_code = get_pytorch_setup_code() + cellarium_cli_code = get_cellarium_cli_code() + + if copy_data_to_local_disk: + data_download_code = get_data_download_code() + data_download_section = f""" +# optionally copy data from GCS to local disk +if copy_data_to_local_disk: +{_indent_code(data_download_code, 4)} +""" + else: + data_download_section = "" + + return f"""# re-install cellarium-ml if a git sha is provided +{git_install_code} +{data_download_section}# set up PyTorch environment +{pytorch_setup_code} + +# run the cellarium CLI +{cellarium_cli_code}""".strip() + + +def _indent_code(code: str, spaces: int) -> str: + """Indent each line of code by the specified number of spaces.""" + indent = " " * spaces + return "\n".join( + indent + line if line.strip() else line for line in code.split("\n") + ) + + +def create_train_op_function(copy_data_to_local_disk: bool = True): + """ + Returns a train_op function that can be used for both local and Vertex AI execution. + + Args: + copy_data_to_local_disk: Whether to copy GCS data to local disk + + Returns: + A function that executes the training pipeline + """ + train_op_text = _get_train_op_text(copy_data_to_local_disk) + + def train_op( + tool: str, + subcommand: str, + config: str, + git_sha: str = "", + ) -> None: + # Create execution context with all necessary variables and functions + exec_globals = { + "tool": tool, + "subcommand": subcommand, + "config": config, + "git_sha": git_sha, + "copy_data_to_local_disk": copy_data_to_local_disk, + "get_git_install_code": get_git_install_code, + "get_data_download_code": get_data_download_code, + "get_pytorch_setup_code": get_pytorch_setup_code, + "get_cellarium_cli_code": get_cellarium_cli_code, + } + # Execute the ground-truth train_op implementation + exec(train_op_text, exec_globals) + + return train_op + + +def get_train_op_code(copy_data_to_local_disk: bool = True) -> str: + """ + Returns the complete train_op code as a string for use in dsl.component. + + Args: + copy_data_to_local_disk: Whether to include data download code + + Returns: + Complete train_op implementation as a string + """ + return _get_train_op_text(copy_data_to_local_disk) + + +def get_batch_setup_script() -> str: + """Returns the batch setup script as a string.""" + return _load_bash_script_as_string("batch_setup") + + +def create_post_upload_runnable(ssd_task_dir: str, gcs_dest: str): + """Return a host-VM bash runnable that bulk-uploads task outputs to GCS. + + Runs *after* the Docker container exits, on the bare VM where ``gsutil`` is + available. Uses ``gsutil -m cp -r`` for parallel upload — orders of magnitude + faster than per-file gcsfs uploads for large prediction outputs. + + Args: + ssd_task_dir: Absolute path on the VM's local SSD where the container + wrote task output files (e.g. ``/mnt/disks/local-ssd/task_outputs``). + gcs_dest: GCS destination prefix, e.g. ``gs://my-bucket/run/task_outputs``. + """ + import textwrap + from google.cloud import batch_v1 as _batch_v1 + + script = textwrap.dedent(f""" + #!/bin/bash + echo "Post-run task output upload: {ssd_task_dir} -> {gcs_dest}" + if [ ! -d "{ssd_task_dir}" ]; then + echo " Task output dir not found, nothing to upload." + exit 0 + fi + shopt -s nullglob + files=("{ssd_task_dir}"/*) + shopt -u nullglob + if [ ${{#files[@]}} -eq 0 ]; then + echo " Task output dir is empty, nothing to upload." + exit 0 + fi + echo " Uploading ${{#files[@]}} top-level entries..." + gsutil -m cp -r "{ssd_task_dir}"/* "{gcs_dest}/" + echo " Task output upload complete." + """).strip() + + runnable = _batch_v1.Runnable() + runnable.script = _batch_v1.Runnable.Script() + runnable.script.text = script + return runnable + + +def create_batch_script( + tool: str, + subcommand: str, + config: str, + git_sha: str, + copy_data_to_local_disk: bool, + capture_logs_to_gcs: bool = False, + output_gcs_bucket: str = "", +) -> str: + """ + Create a batch script for Google Cloud Batch execution. + + This generates the complete bash script that sets up the environment + and executes the train_op code. + + Args: + tool: Cellarium tool to run + subcommand: Subcommand (fit/predict) + config: Path to config file + git_sha: Git SHA for cellarium-ml + copy_data_to_local_disk: Whether to copy data locally + capture_logs_to_gcs: Whether to capture logs to GCS + + Returns: + Complete bash script as a string + """ + train_op_code = get_train_op_code(copy_data_to_local_disk) + train_op_requirements = get_train_op_requirements() + batch_setup_script = get_batch_setup_script() + + return f'''#!/bin/bash +set -euxo pipefail +export PYTHONFAULTHANDLER=1 # better logging for python crashes + +# Set up environment variables for the setup script +export TOOL="{tool}" +export SUBCOMMAND="{subcommand}" +export CONFIG="{config}" +export GIT_SHA="{git_sha}" +export COPY_DATA_TO_LOCAL_DISK="{copy_data_to_local_disk}" +export CELLARIUM_CAPTURE_LOGS="{str(capture_logs_to_gcs).lower()}" +export TRAIN_OP_REQUIREMENTS="{" ".join(train_op_requirements)}" +export MOUNTED_GCS_PATH="{output_gcs_bucket}" + +# Set up local output directory for training artifacts (TensorBoard, checkpoints, custom outputs). +# The rsync sidecar below keeps this directory synced to GCS every 60 s so that TensorBoard +# remains watchable in real time - no GCS FUSE mount required. +# Use the local SSD only if it is actually mounted (not just a directory on the boot disk). +if mountpoint -q /mnt/disks/local-ssd 2>/dev/null; then + LOCAL_OUTPUT_DIR="{local_output_path}" + mkdir -p "$LOCAL_OUTPUT_DIR" + echo " Using local SSD for output directory: $LOCAL_OUTPUT_DIR" +else + LOCAL_OUTPUT_DIR="/tmp/run_outputs" + mkdir -p "$LOCAL_OUTPUT_DIR" + echo " Local SSD not mounted, using boot disk fallback: $LOCAL_OUTPUT_DIR" +fi + +# Run the batch setup script first — this installs gcsfs, which the sync sidecar needs. +cat > /tmp/batch_setup.sh << 'SETUP_EOF' +{batch_setup_script} +SETUP_EOF + +chmod +x /tmp/batch_setup.sh +/tmp/batch_setup.sh + +# Write a Python GCS sync helper using gcsfs (gsutil is not present in the ML container image). +# Incremental: tracks (size, mtime) of already-uploaded files in /tmp/gcs_sync_state.json +# so that unchanged files are skipped on subsequent runs. +cat > /tmp/gcs_sync.py << 'SYNC_EOF' +#!/usr/bin/env python3 +import json, os, sys, gcsfs +local_dir, gcs_dest = sys.argv[1], sys.argv[2].rstrip("/") +if not os.path.isdir(local_dir): + sys.exit(0) +state_file = "/tmp/gcs_sync_state.json" +try: + with open(state_file) as _f: + state = json.load(_f) +except (FileNotFoundError, json.JSONDecodeError): + state = {{}} +fs = gcsfs.GCSFileSystem() +synced = skipped = 0 +for root, _dirs, files in os.walk(local_dir): + for fname in files: + src = os.path.join(root, fname) + rel = os.path.relpath(src, local_dir) + st = os.stat(src) + key = rel + if state.get(key) == [st.st_size, st.st_mtime]: + skipped += 1 + continue + dst = f"{{gcs_dest}}/{{rel}}" + try: + fs.put(src, dst) + state[key] = [st.st_size, st.st_mtime] + synced += 1 + except Exception as e: + print(f" Warning: could not sync {{rel}}: {{e}}", file=sys.stderr) +try: + with open(state_file, "w") as _f: + json.dump(state, _f) +except Exception as e: + print(f" Warning: could not write sync state: {{e}}", file=sys.stderr) +print(f" Synced {{synced}} file(s), skipped {{skipped}} unchanged from {{local_dir}} to {{gcs_dest}}") +SYNC_EOF + +# Start background GCS sync sidecar (non-fatal; training continues regardless of errors) +RSYNC_PID="" +if [ -n "$MOUNTED_GCS_PATH" ]; then + echo " Starting sync sidecar: $LOCAL_OUTPUT_DIR -> $MOUNTED_GCS_PATH (every 5 mins)" + while true; do + python3 /tmp/gcs_sync.py "$LOCAL_OUTPUT_DIR" "$MOUNTED_GCS_PATH" 2>/dev/null || true + sleep 300 + done & + RSYNC_PID=$! +fi + +# Create Python wrapper script that properly handles environment variables and config path replacement +cat > /tmp/train_op_wrapper.py << 'WRAPPER_EOF' +import os +import tempfile +import yaml +import shutil +from pathlib import Path + +# Get environment variables +tool = os.environ.get('TOOL') +subcommand = os.environ.get('SUBCOMMAND') +config = os.environ.get('CONFIG') +git_sha = os.environ.get('GIT_SHA') +copy_data_to_local_disk = os.environ.get('COPY_DATA_TO_LOCAL_DISK', 'false').lower() == 'true' + +# Apply config path substitution: replace the GCS run directory with the local +# output directory so training writes artifacts to local SSD. The rsync sidecar +# in the outer bash script keeps GCS up to date - no FUSE mount required. +original_config = config +mounted_gcs_path = os.environ.get('MOUNTED_GCS_PATH', '') +local_output_path = '{local_output_path}' +# Mirror the bash logic: use the local SSD path only if /mnt/disks/local-ssd is +# a real mount point, not just a directory on the boot disk. +import subprocess as _sp +_ssd_mounted = _sp.run( + ['mountpoint', '-q', '/mnt/disks/local-ssd'], + capture_output=True +).returncode == 0 +if not _ssd_mounted: + local_output_path = '/tmp/run_outputs' + os.makedirs(local_output_path, exist_ok=True) + +if mounted_gcs_path: + # Load the original config - handle both local and GCS paths + if config.startswith('gs://'): + print(f" Downloading config from GCS: {{config}}") + import gcsfs + fs = gcsfs.GCSFileSystem() + with fs.open(config, 'r') as f: + config_content = f.read() + else: + print(f" Reading local config: {{config}}") + with open(config, 'r') as f: + config_content = f.read() + + # Normalise to gs:// form + if not mounted_gcs_path.startswith('gs://'): + mounted_gcs_path = 'gs://' + mounted_gcs_path + + # Also build the /gcs/ form for configs that use that convention + gcs_local_path = '/gcs/' + mounted_gcs_path[5:] # Remove 'gs://' and add '/gcs/' + + # Rewrite ONLY the trainer.default_root_dir key so Lightning writes artifacts + # to local SSD (the rsync sidecar then pushes them to GCS). + # All other GCS paths — ckpt_path, filenames, !FileLoader file_path, etc. — + # are intentionally left as gs:// so Lightning/gcsfs can stream them directly. + # (data.dadc.init_args.filenames is rewritten separately by data_download.py + # when copy_data_to_local_disk=True.) + import re as _re + modified_content = _re.sub( + r'(?m)(^\s*default_root_dir:\s*)' + _re.escape(mounted_gcs_path), + lambda m: m.group(1) + local_output_path, + config_content, + ) + modified_content = _re.sub( + r'(?m)(^\s*default_root_dir:\s*)' + _re.escape(gcs_local_path), + lambda m: m.group(1) + local_output_path, + modified_content, + ) + + replaced = local_output_path in modified_content + if replaced: + print(f" Redirected default_root_dir -> {{local_output_path}}") + else: + print(" default_root_dir did not match mounted_gcs_path — no substitution made") + + # Write modified config to a temporary file + temp_config = '/tmp/modified_config.yaml' + with open(temp_config, 'w') as f: + f.write(modified_content) + + config = temp_config + os.environ["CONFIG"] = config + print(f" Created modified config with updated paths: {{config}}") +else: + print(" No GCS output path configured - skipping config path substitution") + +# Execute train_op code with proper variable scope +{train_op_code} +WRAPPER_EOF + +# Cleanup trap: stop the rsync sidecar and do a blocking final sync when the +# script exits (whether training succeeded or failed). +_cleanup() {{ + local _exit=$? + if [ -n "${{RSYNC_PID:-}}" ]; then + echo " Stopping rsync sidecar (PID $RSYNC_PID)" + kill "$RSYNC_PID" 2>/dev/null || true + wait "$RSYNC_PID" 2>/dev/null || true + fi + if [ -n "$MOUNTED_GCS_PATH" ]; then + echo " Final sync: $LOCAL_OUTPUT_DIR -> $MOUNTED_GCS_PATH" + python3 /tmp/gcs_sync.py "$LOCAL_OUTPUT_DIR" "$MOUNTED_GCS_PATH" || true + fi + exit $_exit +}} +trap _cleanup EXIT + +# Execute the training operation +echo " Starting training operation..." +python3 /tmp/train_op_wrapper.py +''' + + +def create_vertex_ai_train_op_component(base_image: str = ""): + """ + Creates a dsl.component decorated train_op function for Vertex AI execution. + + This passes the train_op_code as a parameter so it gets serialized properly. + + Args: + base_image: The base image for the component + + Returns: + A dsl.component decorated function ready for Vertex AI + """ + from kfp import dsl + + @dsl.component( + packages_to_install=get_train_op_requirements(), + base_image=base_image, + ) + def train_op( + tool: str, + subcommand: str, + config: str, + train_op_code: str, # Pass the code as a parameter + git_sha: str = "", + copy_data_to_local_disk: bool = True, + ) -> None: + # Create execution context with all necessary variables + exec_globals = { + "tool": tool, + "subcommand": subcommand, + "config": config, + "git_sha": git_sha, + "copy_data_to_local_disk": copy_data_to_local_disk, + } + # Execute the train_op code with proper context + exec(train_op_code, exec_globals) + + return train_op + + +# Re-export utility functions for backward compatibility +def get_current_google_user(): + """Get the current Google user from credentials.""" + try: + from google.auth import default + from google.auth.transport.requests import Request + import jwt + + credentials, _ = default() + credentials.refresh(Request()) + id_token = credentials.id_token + decoded_token = jwt.decode(id_token, options={"verify_signature": False}) + return decoded_token.get("email").split("@")[0] + except Exception as e: + print( + "NOTE: unable to prepend google user name to pipeline name. " + f"This is purely cosmetic. Continuing. Error was:\n{e}" + ) + return None + + +def fetch_url_with_retries(url, retries=3, delay=1): + """Fetch a URL with retries.""" + import requests + import time + + for attempt in range(retries): + try: + response = requests.get(url) + response.raise_for_status() + return response + except requests.exceptions.RequestException as e: + if attempt < retries - 1: + time.sleep(delay) + else: + raise e + + +def get_allowed_cli_tool_names(url: str): + """ + Parse python code at a given URL to obtain a list of allowed cellarium-ml CLI tool names. + + Args: + url: URL to fetch the python code from. + + Returns: + List of allowed CLI tool names, or None if the URL could not be fetched. + """ + import ast + import requests + + try: + response = fetch_url_with_retries(url) + content = response.text + module = ast.parse(content) + cli_tool_names = [ + node.name + for node in module.body + if isinstance(node, ast.FunctionDef) + and any( + isinstance(decorator, ast.Name) and decorator.id == "register_model" + for decorator in node.decorator_list + ) + ] + return cli_tool_names + except requests.exceptions.RequestException as e: + print( + f"WARNING:\nAttempted to fetch URL {url} to look up allowed CLI tool names.\n" + "This URL was inferred from the --base-image tag and assumes the tag matches a git SHA for cellarium-ml.\n" + f"Request returned:\n{e}\n" + "NOTE: The input --tool cannot be validated. Double check tool name!\n" + ) + return None diff --git a/cellarium/workflows/submit_batch_component.py b/cellarium/workflows/submit_batch_component.py new file mode 100644 index 0000000..23134e0 --- /dev/null +++ b/cellarium/workflows/submit_batch_component.py @@ -0,0 +1,579 @@ +"""Submit a single component cellarium-ml job to Google Cloud Batch.""" + +import re +import uuid +from datetime import datetime + +import click +from google.cloud import batch_v1 + +from .shared_components import ( + _assert_gcs_bucket_exists, + get_current_google_user, + get_allowed_cli_tool_names, + create_batch_script, + create_post_upload_runnable, + get_machine_type_resources, + extract_output_gcs_bucket_from_config, + extract_data_gcs_bucket_from_config, + assert_gcs_bucket_accessible_as_service_account, + assert_data_first_file_exists, + extract_ckpt_path_from_config, + assert_ckpt_path_exists, + prepare_config_with_overrides, +) + + +def create_batch_job_spec( + job_name: str, + project: str, + location: str, + tool: str, + subcommand: str, + config: str, + git_sha: str, + copy_data_to_local_disk: bool, + base_image: str, + machine_type: str = "n1-standard-4", + accelerator_type: str = "nvidia-tesla-t4", + accelerator_count: int = 1, + max_run_duration: str = "21600s", + capture_logs_to_gcs: bool = False, + output_gcs_bucket: str = None, + local_ssd_size_gb: int = 375, + network: str = "default-vpc", + subnetwork: str = "", +) -> batch_v1.Job: + """ + Create a Google Cloud Batch job specification. + + Args: + job_name: Name of the batch job + project: Google Cloud project ID + location: Google Cloud location + tool: Cellarium tool to run + subcommand: Subcommand (fit/predict) + config: Path to config file + git_sha: Git SHA for cellarium-ml + copy_data_to_local_disk: Whether to copy data locally + base_image: Container image to use + machine_type: Machine type for the job + accelerator_type: GPU type + accelerator_count: Number of GPUs + max_run_duration: Maximum runtime in seconds format + capture_logs_to_gcs: Capture stdout/stderr to files for GCS sync + output_gcs_bucket: GCS destination for outputs (synced via rsync sidecar) + local_ssd_size_gb: Size of Local SSD in GB (set to 0 to disable) + network: VPC network name or full resource URL + subnetwork: Subnet name or full resource URL (defaults to same name as network for AUTO-mode VPCs) + + Returns: + Google Cloud Batch job specification + """ + + # Create the batch script using the shared function + batch_script = create_batch_script( + tool=tool, + subcommand=subcommand, + config=config, + git_sha=git_sha, + copy_data_to_local_disk=copy_data_to_local_disk, + capture_logs_to_gcs=capture_logs_to_gcs, + output_gcs_bucket=output_gcs_bucket, + ) + + # Create environment variables for the container + env_vars = { + "TOOL": tool, + "SUBCOMMAND": subcommand, + "CONFIG": config, + "GIT_SHA": git_sha, + "COPY_DATA_TO_LOCAL_DISK": str(copy_data_to_local_disk).lower(), + "CELLARIUM_CAPTURE_LOGS": str(capture_logs_to_gcs).lower(), + "CLOUDSDK_PYTHON": "/usr/bin/python3", + } + + # Define the task specification + task_spec = batch_v1.TaskSpec() + + task_spec_runnables = [] + + # Resolve machine resources up front so shm-size can be derived from total RAM. + cpu_milli, memory_mib = get_machine_type_resources(machine_type) + + # Configure the container runnable + container = batch_v1.Runnable.Container() + container.image_uri = base_image + container.commands = ["/bin/bash", "-c", batch_script] + + # Configure shared memory for PyTorch DataLoader workers. + # Use ~75% of total RAM so workers can buffer freely without crowding out + # the training process. Outside a container /dev/shm is unbounded, which + # is why bus errors only appear when running containerised. + shm_gb = max(1, (memory_mib * 3) // (1024 * 4)) + container.options = f"--shm-size={shm_gb}g" + print( + f" Configured container with shared memory size: {shm_gb}GB (~75% of {memory_mib // 1024}GB RAM)" + ) + + # GPU access is automatically configured by Google Cloud Batch when GPUs are allocated + if accelerator_count > 0: + print(" GPU access will be automatically configured by Google Cloud Batch") + + runnable = batch_v1.Runnable() + runnable.container = container + + # Set environment variables on the runnable + runnable.environment = batch_v1.Environment() + for key, value in env_vars.items(): + runnable.environment.variables[key] = value + + task_spec_runnables.append(runnable) + + # Post-container runnable: host-VM bash that bulk-uploads task outputs via gsutil. + # Runs after the Docker container exits so it can use the VM's native gsutil + # for fast parallel upload (gsutil -m cp -r) instead of per-file gcsfs inside + # the container. Skipped when no SSD or no output bucket is configured. + if local_ssd_size_gb > 0 and output_gcs_bucket: + ssd_task_dir = "/mnt/disks/local-ssd/task_outputs" + task_output_gcs_dest = f"{output_gcs_bucket.rstrip('/')}/task_outputs" + post_runnable = create_post_upload_runnable(ssd_task_dir, task_output_gcs_dest) + task_spec_runnables.append(post_runnable) + print( + f" Added post-container upload runnable: gsutil -m cp -r {ssd_task_dir} -> {task_output_gcs_dest}/" + ) + + task_spec.runnables = task_spec_runnables + + # Task volumes (local SSD only — GCS sync is handled by the rsync sidecar in the batch script) + task_volumes = [] + + if output_gcs_bucket: + print( + f" GCS output destination: {output_gcs_bucket} (synced via rsync sidecar)" + ) + else: + print(" No GCS output bucket specified, outputs will remain on local disk") + + # Local SSD must be added as a Volume so Batch formats and mounts it into the container + # at the specified mount_path before the container starts. + if local_ssd_size_gb > 0: + ssd_volume = batch_v1.Volume() + ssd_volume.device_name = "local-ssd" # must match attached_disk.device_name + ssd_volume.mount_path = "/mnt/disks/local-ssd" + ssd_volume.mount_options = ["rw,async"] + task_volumes.append(ssd_volume) + print(f" Local SSD configured: {local_ssd_size_gb}GB -> /mnt/disks/local-ssd") + else: + print(" Local SSD disabled, using boot disk only") + + # Set volumes on task spec + if task_volumes: + task_spec.volumes = task_volumes + + # Set compute resources based on machine type (already resolved above) + compute_resource = batch_v1.ComputeResource() + compute_resource.cpu_milli = cpu_milli + compute_resource.memory_mib = memory_mib + task_spec.compute_resource = compute_resource + + # Set maximum run duration + task_spec.max_run_duration = {"seconds": int(max_run_duration.rstrip("s"))} + + # Create task groups + group = batch_v1.TaskGroup() + group.task_count = 1 + group.task_spec = task_spec + + # Create allocation policy for machine type + allocation_policy = batch_v1.AllocationPolicy() + instance_policy = batch_v1.AllocationPolicy.InstancePolicy() + instance_policy.machine_type = machine_type + + # Add GPU to instance policy if specified + if accelerator_count > 0 and accelerator_type: + print(" GPU Configuration:") + print(f" Type: {accelerator_type}") + print(f" Count: {accelerator_count}") + print(f" Formatted type: {accelerator_type.lower().replace('_', '-')}") + + # For Batch, GPUs are configured via accelerators in the instance policy + accelerator = batch_v1.AllocationPolicy.Accelerator() + accelerator.type_ = accelerator_type.lower().replace("_", "-") + accelerator.count = accelerator_count + instance_policy.accelerators = [accelerator] + + print(" Added GPU to job specification") + else: + print( + f" No GPU configured (count: {accelerator_count}, type: '{accelerator_type}')" + ) + + # Add Local SSD configuration — must be done BEFORE assigning instance_policy to + # instance_policy_or_template, because proto-plus copies the message on assignment; + # any mutations to instance_policy after that point are silently ignored. + if local_ssd_size_gb > 0: + attached_disk = batch_v1.AllocationPolicy.AttachedDisk() + attached_disk.new_disk = batch_v1.AllocationPolicy.Disk() + attached_disk.new_disk.type_ = "local-ssd" + attached_disk.new_disk.size_gb = local_ssd_size_gb + attached_disk.device_name = "local-ssd" + instance_policy.disks = [attached_disk] + + instance_policy_or_template = batch_v1.AllocationPolicy.InstancePolicyOrTemplate() + instance_policy_or_template.policy = instance_policy + if accelerator_count > 0 and accelerator_type: + instance_policy_or_template.install_gpu_drivers = True + + allocation_policy.instances = [instance_policy_or_template] + + # Configure VPC network + if network: + # Build full resource URLs if short names were given + net_url = ( + network + if "/" in network + else f"projects/{project}/global/networks/{network}" + ) + # For AUTO-mode VPCs the subnet name matches the network name + resolved_subnet = subnetwork or network + sub_url = ( + resolved_subnet + if "/" in resolved_subnet + else f"projects/{project}/regions/{location}/subnetworks/{resolved_subnet}" + ) + network_interface = batch_v1.AllocationPolicy.NetworkInterface() + network_interface.network = net_url + network_interface.subnetwork = sub_url + network_policy = batch_v1.AllocationPolicy.NetworkPolicy() + network_policy.network_interfaces = [network_interface] + allocation_policy.network = network_policy + print(f" Network: {net_url}") + print(f" Subnetwork: {sub_url}") + job = batch_v1.Job() + job.task_groups = [group] + job.allocation_policy = allocation_policy + job.logs_policy = batch_v1.LogsPolicy() + + # Set log destination based on capture_logs_to_gcs setting + if capture_logs_to_gcs: + # Logs saved to local SSD (synced to GCS by the rsync sidecar) to avoid Cloud Logging costs + if local_ssd_size_gb > 0: + job.logs_policy.destination = batch_v1.LogsPolicy.Destination.PATH + job.logs_policy.logs_path = "/mnt/disks/local-ssd/job_logs" + print( + " Logs will be saved to Local SSD and synced to GCS (Cloud Logging disabled to save costs)" + ) + else: + job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING + print( + " capture_logs_to_gcs requested but local_ssd_size_gb=0; falling back to Cloud Logging" + ) + else: + # Default behavior - all logs go to Cloud Logging + job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING + print(" Logs will be sent to Cloud Logging") + + return job + + +@click.command(short_help="Submit a single-component job to Google Cloud Batch.") +@click.option( + "--tool", + required=True, + help="Tool to run, e.g. 'onepass_mean_var_std'.", +) +@click.option( + "--subcommand", + required=True, + type=click.Choice(["fit", "predict"]), + help="Subcommand to run, either 'fit' or 'predict'.", +) +@click.option( + "--config", + required=True, + help="GCS path to the training config YAML file.", +) +@click.option( + "--copy-data-to-local-disk", + default=True, + type=bool, + help="True copies GCS data to local disk fully (once) before training. False is ephemeral.", +) +@click.option( + "--project", + default="dsp-cellarium", + help="Google Cloud project ID.", +) +@click.option( + "--location", + default="us-central1", + help="Google Cloud location, e.g. 'us-central1'.", +) +@click.option( + "--job-name", + default="", + help="Job name, defaults to f'{user}_{tool}_{subcommand}'.", +) +@click.option( + "--machine-type", + default="n1-standard-16", + help="Machine type for the training job, e.g. 'n1-standard-16'.", +) +@click.option( + "--accelerator-type", + default="nvidia-tesla-t4", + help="Type of accelerator (gpu), e.g. 'nvidia-tesla-t4'.", +) +@click.option( + "--accelerator-count", + default=1, + type=int, + help="Number of GPUs.", +) +@click.option( + "--max-run-duration", + default="604800s", + help="Maximum runtime in seconds format, max 7 days, e.g. '604800s'.", +) +@click.option( + "--git-sha", + default="main", + type=str, + help="Cellarium-ML git SHA to install (if provided).", +) +@click.option( + "--base-image", + default="us-central1-docker.pkg.dev/broad-dsde-methods/cellarium-ai/cellarium-ml:0.0.8", + help="Base image for the component.", +) +@click.option( + "--capture-logs-to-gcs", + default=False, + is_flag=True, + help="Capture stdout/stderr to files and sync to GCS instead of using Cloud Logging.", +) +@click.option( + "--local-ssd-size-gb", + default=750, + type=int, + help="Size of Local SSD in GB in 375 increments (375, 750, 1125, etc.). Set to 0 to disable Local SSD.", +) +@click.option( + "--dry-run", + default=False, + is_flag=True, + help="Build and validate the job spec without submitting to Google Cloud Batch.", +) +@click.option( + "--network", + default="default", + help="VPC network name or full resource URL. Defaults to 'default'.", +) +@click.option( + "--subnetwork", + default="", + help="Subnet name or full resource URL. Defaults to the network name (valid for AUTO-mode VPCs).", +) +@click.option( + "--extract-bucket", + default=None, + help="GCS URI prefix containing extract_*.h5ad files, e.g. gs://my-bucket/my-prefix. Overwrites config yaml paths.", +) +@click.option( + "--staging-bucket", + default=None, + help="GCS URI prefix for staging local config files, e.g. gs://my-bucket/staging. Required when config is a local path and the config has no gs:// default_root_dir.", +) +def submit_batch_component( + project: str, + location: str, + config: str, + tool: str, + subcommand: str, + copy_data_to_local_disk: bool, + job_name: str, + machine_type: str, + accelerator_type: str, + accelerator_count: int, + max_run_duration: str, + git_sha: str, + base_image: str, + capture_logs_to_gcs: bool, + local_ssd_size_gb: int, + dry_run: bool = False, + network: str = "default-vpc", + subnetwork: str = "", + extract_bucket=None, + staging_bucket=None, +): + """ + Submit a single component cellarium-ml job to Google Cloud Batch. + """ + # Input validation and defaults + if job_name == "": + user = get_current_google_user() + base_name = f"{user}-{tool}-{subcommand}" if user else f"{tool}-{subcommand}" + + # Add timestamp and short UUID to ensure uniqueness + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + short_uuid = str(uuid.uuid4())[:8] + job_name = f"{base_name}-{timestamp}-{short_uuid}" + else: + # If user provided a custom job name, still add UUID to avoid conflicts + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + short_uuid = str(uuid.uuid4())[:8] + job_name = f"{job_name}-{timestamp}-{short_uuid}" + + # Ensure job name is valid for Batch (lowercase, hyphens only, max 63 chars) + job_name = job_name.lower().replace("_", "-") + job_name = job_name[:63].rstrip("-") + + if not re.match(r"^\d+s$", max_run_duration): + raise ValueError( + f"max_run_duration must be in 's' format, e.g. '3600s'. Got: '{max_run_duration}'" + ) + + if (git_sha == "") and (len(base_image.split(":")[-1]) > 0): + git_sha = base_image.split(":")[-1] + + config = prepare_config_with_overrides(config, extract_bucket) + + # Auto-detect output GCS bucket from config (done early so it can be used for local config staging) + output_gcs_bucket = extract_output_gcs_bucket_from_config(config) + if output_gcs_bucket: + print(f" Auto-detected output GCS bucket from config: {output_gcs_bucket}") + else: + print(" No GCS output bucket detected, using local storage only") + + if not dry_run: + # Verify the output bucket is reachable before uploading config or submitting + if output_gcs_bucket: + _assert_gcs_bucket_exists(output_gcs_bucket) + + # If config is a local path, upload it to GCS staging so the Batch VM can access it + if not config.startswith("gs://"): + bucket_for_staging = output_gcs_bucket or staging_bucket + if not bucket_for_staging: + raise ValueError( + "Local config file provided but no GCS bucket is available for staging. " + "Either add a gs:// default_root_dir to your config or pass " + "--staging-bucket gs://my-bucket/path." + ) + import gcsfs as _gcsfs + + fs = _gcsfs.GCSFileSystem() + staged_config_path = ( + f"{bucket_for_staging.rstrip('/')}/staging/configs/{job_name}.yaml" + ) + fs.put(config, staged_config_path) + print(f" Uploaded local config to GCS staging: {staged_config_path}") + config = staged_config_path + + # Pre-flight: check that the Batch SA can access the data bucket, and the first file exists + data_bucket = extract_data_gcs_bucket_from_config(config) + if data_bucket: + assert_gcs_bucket_accessible_as_service_account(project, data_bucket) + assert_data_first_file_exists(config) + + # Pre-flight: verify the checkpoint file exists in GCS if ckpt_path is set + ckpt_path = extract_ckpt_path_from_config(config) + if ckpt_path: + assert_ckpt_path_exists(ckpt_path) + + # Validate tool name against upstream CLI registry + url = f"https://raw.githubusercontent.com/cellarium-ai/cellarium-ml/{git_sha}/cellarium/ml/cli.py" + cli_tool_names = get_allowed_cli_tool_names(url) + if cli_tool_names is not None: + if tool not in cli_tool_names: + raise ValueError( + f"Tool '{tool}' not found in allowed CLI tools at {url}.\n" + f"Allowed tool names:\n{cli_tool_names}" + ) + + # Handle GPU settings + if accelerator_count == 0: + accelerator_type = "" + accelerator_count = 0 + + print(f"Submitting job '{job_name}' to Google Cloud Batch...") + print(f"Project: {project}") + print(f"Location: {location}") + print(f"Tool: {tool}") + print(f"Subcommand: {subcommand}") + print(f"Config: {config}") + print(f"Machine type: {machine_type}") + print( + f"Accelerator: {accelerator_count}x {accelerator_type}" + if accelerator_count > 0 + else "No accelerator" + ) + print(f"Max runtime: {max_run_duration}") + print(f"Local SSD size: {local_ssd_size_gb}GB") + + # Create the batch job specification + job_spec = create_batch_job_spec( + job_name=job_name, + project=project, + location=location, + tool=tool, + subcommand=subcommand, + config=config, + git_sha=git_sha, + copy_data_to_local_disk=copy_data_to_local_disk, + base_image=base_image, + machine_type=machine_type, + accelerator_type=accelerator_type, + accelerator_count=accelerator_count, + max_run_duration=max_run_duration, + capture_logs_to_gcs=capture_logs_to_gcs, + output_gcs_bucket=output_gcs_bucket, + local_ssd_size_gb=local_ssd_size_gb, + network=network, + subnetwork=subnetwork, + ) + + # Dry-run: validate and print job spec without submitting + if dry_run: + print(" Dry run: job spec built successfully, skipping submission.") + print(job_spec) + return + + # Submit the job + client = batch_v1.BatchServiceClient() + parent = f"projects/{project}/locations/{location}" + + try: + request = batch_v1.CreateJobRequest() + request.parent = parent + request.job_id = job_name + request.job = job_spec + + result = client.create_job(request=request) + print(f" Job '{job_name}' submitted successfully!") + print(f"Job resource name: {result.name}") + print(f"Job UID: {result.uid}") + print(f"Job state: {result.status.state.name}") + + # Print monitoring information + print("\nMonitoring commands:") + print( + f" gcloud batch jobs describe {job_name} --location={location} --project={project}" + ) + print(f" gcloud batch jobs list --location={location} --project={project}") + print( + f' gcloud logging read \'resource.type="gce_instance" AND resource.labels.job_id="{job_name}"\' --project={project}' + ) + print( + f"\nConsole URL:\n https://console.cloud.google.com/batch/jobs?project={project}" + ) + + return result + + except Exception as e: + print(f" Failed to submit job: {e}") + raise + + +if __name__ == "__main__": + submit_batch_component() diff --git a/cellarium/workflows/submit_batch_pipeline.py b/cellarium/workflows/submit_batch_pipeline.py new file mode 100644 index 0000000..056a290 --- /dev/null +++ b/cellarium/workflows/submit_batch_pipeline.py @@ -0,0 +1,532 @@ +"""Submit a multi-component cellarium-ml pipeline to Google Cloud Batch.""" + +import yaml +import uuid +from datetime import datetime +from typing import List, Dict, Any + +import click +from google.cloud import batch_v1 + +from .shared_components import ( + get_current_google_user, + get_allowed_cli_tool_names, + create_batch_script, + create_post_upload_runnable, + get_machine_type_resources, + extract_output_gcs_bucket_from_config, + prepare_config_with_overrides, +) + + +def parse_pipeline_yaml(config: str) -> tuple[str, list[dict]]: + """Parse pipeline YAML configuration - same as submit_pipeline.py""" + with open(config) as f: + config_contents = yaml.safe_load(f) + top_level_keys = list(config_contents.keys()) + assert len(top_level_keys) == 1, ( + "Pipeline YAML config error: The top level of the config file must be the display_name of the pipeline. Only one top level key is allowed." + ) + display_name = list(config_contents.keys())[0] + component_definitions = config_contents[display_name] + assert isinstance(component_definitions, list), ( + "Pipeline YAML config error: The value of the top level key must be a list of component definition dictionaries." + ) + for item in component_definitions: + assert isinstance(item, dict), ( + "Pipeline YAML config error: Each component definition in the list must be a dictionary." + ) + assert "tool" in item, ( + "Pipeline YAML config error: Each component definition must have a 'tool' key." + ) + assert "subcommand" in item, ( + "Pipeline YAML config error: Each component definition must have a 'subcommand' key." + ) + assert item["subcommand"] in ["fit", "predict"], ( + "Pipeline YAML config error: The 'subcommand' key's value must be either 'fit' or 'predict'." + ) + assert "config" in item, ( + "Pipeline YAML config error: Each component definition must have a 'config' key." + ) + return display_name, component_definitions + + +def create_batch_pipeline_jobs( + pipeline_name: str, + project: str, + location: str, + component_definitions: List[Dict[str, Any]], + copy_data_to_local_disk: bool, + base_image: str, + git_sha: str = "", + default_machine_type: str = "n1-standard-4", + default_accelerator_type: str = "nvidia-tesla-t4", + default_accelerator_count: int = 0, + capture_logs_to_gcs: bool = False, + local_ssd_size_gb: int = 375, +) -> List[batch_v1.Job]: + """ + Create a list of Google Cloud Batch jobs for a pipeline. + + Note: Google Batch doesn't have built-in pipeline orchestration like Vertex AI, + so we create individual jobs that can be submitted sequentially or managed externally. + """ + jobs = [] + + for i, component_def in enumerate(component_definitions): + # Create unique job name with timestamp and UUID + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + short_uuid = str(uuid.uuid4())[:8] + job_name = f"{pipeline_name}-{i}-{component_def['tool']}-{component_def['subcommand']}-{timestamp}-{short_uuid}" + job_name = job_name.lower().replace("_", "-") + + # Auto-detect output bucket from config file + config_path = component_def["config"] + detected_bucket = extract_output_gcs_bucket_from_config(config_path) + + # Create batch script for this component using the shared function + batch_script = create_batch_script( + tool=component_def["tool"], + subcommand=component_def["subcommand"], + config=component_def["config"], + git_sha=component_def.get("git_sha", git_sha), + copy_data_to_local_disk=copy_data_to_local_disk, + capture_logs_to_gcs=capture_logs_to_gcs, + output_gcs_bucket=detected_bucket, + ) + + # Create environment variables + env_vars = { + "TOOL": component_def["tool"], + "SUBCOMMAND": component_def["subcommand"], + "CONFIG": component_def["config"], + "GIT_SHA": component_def.get("git_sha", git_sha), + "COPY_DATA_TO_LOCAL_DISK": str(copy_data_to_local_disk).lower(), + "CELLARIUM_CAPTURE_LOGS": str(capture_logs_to_gcs).lower(), + "CLOUDSDK_PYTHON": "/usr/bin/python3", + } + + # Define the task specification + task_spec = batch_v1.TaskSpec() + + # Set up volumes (local SSD only — GCS sync is handled by the rsync sidecar) + task_volumes = [] + + if detected_bucket: + print( + f" GCS output destination for {job_name}: {detected_bucket} (synced via rsync sidecar)" + ) + else: + print( + f" No output GCS bucket detected for job {job_name}, using local storage" + ) + + # Local SSD must be added as a Volume so Batch formats and mounts it at the + # specified mount_path before any runnable starts (including host-VM bash runnables). + if local_ssd_size_gb > 0: + ssd_volume = batch_v1.Volume() + ssd_volume.device_name = "local-ssd" # must match attached_disk.device_name + ssd_volume.mount_path = "/mnt/disks/local-ssd" + ssd_volume.mount_options = ["rw,async"] + task_volumes.append(ssd_volume) + print( + f" Local SSD configured for job {job_name}: {local_ssd_size_gb}GB -> /mnt/disks/local-ssd" + ) + else: + print(f" Local SSD disabled for job {job_name}, using boot disk only") + + # Set volumes on task spec + if task_volumes: + task_spec.volumes = task_volumes + + # Resolve machine resources up front so shm-size can be derived from total RAM. + current_machine_type = component_def.get("machine_type", default_machine_type) + cpu_milli, memory_mib = get_machine_type_resources(current_machine_type) + + # Configure the container runnable + container = batch_v1.Runnable.Container() + container.image_uri = base_image + container.commands = ["/bin/bash", "-c", batch_script] + + # Configure shared memory for PyTorch DataLoader workers. + # Use ~75% of total RAM so workers can buffer freely without crowding out + # the training process. Outside a container /dev/shm is unbounded, which + # is why bus errors only appear when running containerised. + shm_gb = max(1, (memory_mib * 3) // (1024 * 4)) + container.options = f"--shm-size={shm_gb}g" + print( + f" Configured container with shared memory size: {shm_gb}GB (~75% of {memory_mib // 1024}GB RAM) for job {job_name}" + ) + + # GPU access is automatically configured by Google Cloud Batch when GPUs are allocated + accelerator_count = component_def.get( + "accelerator_count", default_accelerator_count + ) + if accelerator_count > 0: + print( + f" GPU access will be automatically configured by Google Cloud Batch for job {job_name}" + ) + + runnable = batch_v1.Runnable() + runnable.container = container + + # Set environment variables on the runnable + runnable.environment = batch_v1.Environment() + for key, value in env_vars.items(): + runnable.environment.variables[key] = value + + task_spec.runnables = [runnable] + + # Post-container runnable: host-VM bash that bulk-uploads task outputs via gsutil. + # Runs after the Docker container exits so it can use the VM's native gsutil + # for fast parallel upload (gsutil -m cp -r) instead of per-file gcsfs. + if local_ssd_size_gb > 0 and detected_bucket: + ssd_task_dir = "/mnt/disks/local-ssd/task_outputs" + task_output_gcs_dest = f"{detected_bucket.rstrip('/')}/task_outputs" + post_runnable = create_post_upload_runnable( + ssd_task_dir, task_output_gcs_dest + ) + task_spec.runnables = [runnable, post_runnable] + print( + f" Added post-container upload runnable for {job_name}: " + f"gsutil -m cp -r {ssd_task_dir} -> {task_output_gcs_dest}/" + ) + + # Set compute resources based on machine type (already resolved above) + compute_resource = batch_v1.ComputeResource() + compute_resource.cpu_milli = cpu_milli + compute_resource.memory_mib = memory_mib + task_spec.compute_resource = compute_resource + + # Set maximum run duration + max_duration = component_def.get("max_run_duration", "3600s") + task_spec.max_run_duration = {"seconds": int(max_duration.rstrip("s"))} + + # Create task groups + group = batch_v1.TaskGroup() + group.task_count = 1 + group.task_spec = task_spec + + # Create allocation policy + allocation_policy = batch_v1.AllocationPolicy() + instance_policy = batch_v1.AllocationPolicy.InstancePolicy() + instance_policy.machine_type = component_def.get( + "machine_type", default_machine_type + ) + + # Add GPU if specified + accelerator_count = component_def.get( + "accelerator_count", default_accelerator_count + ) + accelerator_type = component_def.get( + "accelerator_type", default_accelerator_type + ) + + if accelerator_count > 0 and accelerator_type: + print(f" GPU Configuration for {job_name}:") + print(f" Type: {accelerator_type}") + print(f" Count: {accelerator_count}") + print(f" Formatted type: {accelerator_type.lower().replace('_', '-')}") + + # For Batch, GPUs are configured via accelerators in the instance policy + accelerator = batch_v1.AllocationPolicy.Accelerator() + accelerator.type_ = accelerator_type.lower().replace("_", "-") + accelerator.count = accelerator_count + instance_policy.accelerators = [accelerator] + + print(f" Added GPU to job specification for {job_name}") + else: + print( + f" No GPU configured for {job_name} (count: {accelerator_count}, type: '{accelerator_type}')" + ) + + instance_policy_or_template = ( + batch_v1.AllocationPolicy.InstancePolicyOrTemplate() + ) + instance_policy_or_template.policy = instance_policy + if accelerator_count > 0 and accelerator_type: + instance_policy_or_template.install_gpu_drivers = True + + # Add Local SSD configuration + if local_ssd_size_gb > 0: + attached_disk = batch_v1.AllocationPolicy.AttachedDisk() + attached_disk.new_disk = batch_v1.AllocationPolicy.Disk() + attached_disk.new_disk.type_ = "local-ssd" + attached_disk.new_disk.size_gb = local_ssd_size_gb + attached_disk.device_name = "local-ssd" + instance_policy.disks = [attached_disk] + + allocation_policy.instances = [instance_policy_or_template] + + # Create the job + job = batch_v1.Job() + job.task_groups = [group] + job.allocation_policy = allocation_policy + job.logs_policy = batch_v1.LogsPolicy() + + # Set log destination based on capture_logs_to_gcs setting + if capture_logs_to_gcs: + # Logs saved to local SSD (synced to GCS by the rsync sidecar) to avoid Cloud Logging costs + job.logs_policy.destination = batch_v1.LogsPolicy.Destination.PATH + job.logs_policy.logs_path = "/mnt/disks/local-ssd/job_logs" + else: + # Default behavior - all logs go to Cloud Logging + job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING + + jobs.append((job_name, job)) + + return jobs + + +@click.command(short_help="Submit a multi-step pipeline of jobs to Google Cloud Batch.") +@click.option( + "--config", + required=True, + help="Local path to the pipeline config YAML file.", +) +@click.option( + "--copy-data-to-local-disk", + default=True, + type=bool, + help="True copies GCS data to local disk fully (once) before training. False is ephemeral.", +) +@click.option( + "--project", + default="dsp-cell-annotation-service", + help="Google Cloud project ID.", +) +@click.option( + "--location", + default="us-central1", + help="Google Cloud location, e.g. 'us-central1'.", +) +@click.option( + "--pipeline-name", + default="", + help="Pipeline name, defaults to f'{user}_{display_name}'.", +) +@click.option( + "--git-sha", + default="", + type=str, + help="Cellarium-ML git SHA to install (if provided).", +) +@click.option( + "--base-image", + default="us-central1-docker.pkg.dev/broad-dsde-methods/cellarium-ai/cellarium-ml:cellarium-gpt-cstorch", + help="Base image for the component.", +) +@click.option( + "--submit-sequentially", + default=False, + type=bool, + help="Submit jobs one at a time (sequential) rather than all at once (parallel).", +) +@click.option( + "--wait-between-jobs", + default=0, + type=int, + help="Seconds to wait between job submissions when using --submit-sequentially.", +) +@click.option( + "--default-machine-type", + default="n1-standard-4", + help="Default machine type for components that don't specify one.", +) +@click.option( + "--default-accelerator-type", + default="nvidia-tesla-t4", + help="Default GPU type for components that don't specify one.", +) +@click.option( + "--default-accelerator-count", + default=0, + type=int, + help="Default number of GPUs for components that don't specify one.", +) +@click.option( + "--capture-logs-to-gcs", + default=False, + is_flag=True, + help="Capture stdout/stderr to files and sync to GCS instead of using Cloud Logging.", +) +@click.option( + "--local-ssd-size-gb", + default=375, + type=int, + help="Size of Local SSD in GB (375, 750, 1125, etc.). Set to 0 to disable Local SSD and use boot disk only.", +) +@click.option( + "--extract-bucket", + default=None, + help="GCS URI prefix containing extract_*.h5ad files, e.g. gs://my-bucket/my-prefix.", +) +def submit_batch_pipeline( + project: str, + location: str, + config: str, + copy_data_to_local_disk: bool, + pipeline_name: str, + git_sha: str, + base_image: str, + submit_sequentially: bool, + wait_between_jobs: int, + default_machine_type: str, + default_accelerator_type: str, + default_accelerator_count: int, + capture_logs_to_gcs: bool, + local_ssd_size_gb: int, + extract_bucket=None, +): + """ + Submit a multi-component cellarium-ml pipeline to Google Cloud Batch. + + Note: Unlike Vertex AI Pipelines, Google Batch doesn't have built-in + pipeline orchestration. This submits individual Batch jobs that can + be run in parallel or sequentially. + """ + # Parse the pipeline configuration + display_name, component_definitions = parse_pipeline_yaml(config) + + # Apply extract_bucket overrides to each component's config + if extract_bucket is not None: + for component in component_definitions: + component["config"] = prepare_config_with_overrides( + component["config"], extract_bucket + ) + + # Set pipeline name + if pipeline_name == "": + user = get_current_google_user() + base_name = f"{user}-{display_name}" if user else display_name + + # Add timestamp to ensure uniqueness (individual jobs will get their own UUIDs) + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + pipeline_name = f"{base_name}-{timestamp}" + else: + # If user provided a custom pipeline name, still add timestamp to avoid conflicts + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + pipeline_name = f"{pipeline_name}-{timestamp}" + + pipeline_name = pipeline_name.lower().replace("_", "-") + + if (git_sha == "") and (len(base_image.split(":")[-1]) > 0): + git_sha = base_image.split(":")[-1] + + print( + f"Preparing pipeline '{pipeline_name}' with {len(component_definitions)} components..." + ) + print(f"Project: {project}") + print(f"Location: {location}") + print(f"Sequential submission: {submit_sequentially}") + print(f"Local SSD size: {local_ssd_size_gb}GB") + + # Validate all tools + url = f"https://raw.githubusercontent.com/cellarium-ai/cellarium-ml/{git_sha}/cellarium/ml/cli.py" + cli_tool_names = get_allowed_cli_tool_names(url) + if cli_tool_names is not None: + for component_def in component_definitions: + tool = component_def["tool"] + if tool not in cli_tool_names: + raise ValueError( + f"Tool '{tool}' not found in allowed CLI tools at {url}.\n" + f"Allowed tool names:\n{cli_tool_names}" + ) + + # Create batch jobs + batch_jobs = create_batch_pipeline_jobs( + pipeline_name=pipeline_name, + project=project, + location=location, + component_definitions=component_definitions, + copy_data_to_local_disk=copy_data_to_local_disk, + base_image=base_image, + git_sha=git_sha, + default_machine_type=default_machine_type, + default_accelerator_type=default_accelerator_type, + default_accelerator_count=default_accelerator_count, + capture_logs_to_gcs=capture_logs_to_gcs, + local_ssd_size_gb=local_ssd_size_gb, + ) + + # Submit jobs + client = batch_v1.BatchServiceClient() + parent = f"projects/{project}/locations/{location}" + submitted_jobs = [] + + print(f"\nSubmitting {len(batch_jobs)} jobs...") + + for i, (job_name, job_spec) in enumerate(batch_jobs): + try: + print(f" Submitting job {i + 1}/{len(batch_jobs)}: {job_name}") + + request = batch_v1.CreateJobRequest() + request.parent = parent + request.job_id = job_name + request.job = job_spec + + operation = client.create_job(request=request) + try: + result = operation.result() + submitted_jobs.append((job_name, result)) + print(f" Job '{job_name}' submitted successfully (UID: {result.uid})") + except AttributeError: + # Handle case where operation.result() doesn't work as expected + result = operation + submitted_jobs.append((job_name, result)) + print( + f" Job '{job_name}' submitted successfully (Operation: {operation.name})" + ) + + # Try to get the job details directly + try: + job_resource = client.get_job( + name=f"projects/{project}/locations/{location}/jobs/{job_name}" + ) + submitted_jobs[-1] = ( + job_name, + job_resource, + ) # Update with actual job + print(f" Job UID: {job_resource.uid}") + except Exception as e: + print(f" Note: Could not fetch job details immediately: {e}") + + if ( + submit_sequentially + and i < len(batch_jobs) - 1 + and wait_between_jobs > 0 + ): + import time + + print(f" Waiting {wait_between_jobs} seconds before next job...") + time.sleep(wait_between_jobs) + + except Exception as e: + print(f" Failed to submit job '{job_name}': {e}") + if submit_sequentially: + print(" Stopping sequential submission due to error.") + break + continue + + print("\n Pipeline submission complete!") + print(f"Submitted {len(submitted_jobs)}/{len(batch_jobs)} jobs successfully.") + + if submitted_jobs: + print("\nMonitoring commands:") + print(f" gcloud batch jobs list --location={location} --project={project}") + for job_name, _ in submitted_jobs: + print( + f" gcloud batch jobs describe {job_name} --location={location} --project={project}" + ) + + print("\nView logs:") + for job_name, _ in submitted_jobs: + print( + f' gcloud logging read \'resource.type="gce_instance" AND resource.labels.job_id="{job_name}"\' --project={project}' + ) + + return submitted_jobs + + +if __name__ == "__main__": + submit_batch_pipeline() diff --git a/cellarium/workflows/submit_vertex_component.py b/cellarium/workflows/submit_vertex_component.py new file mode 100644 index 0000000..0ca07a3 --- /dev/null +++ b/cellarium/workflows/submit_vertex_component.py @@ -0,0 +1,185 @@ +import tempfile + +import click +from google.cloud import aiplatform +from google_cloud_pipeline_components.v1.custom_job import ( + create_custom_training_job_from_component, +) +from kfp import compiler, dsl + +from .shared_components import ( + get_current_google_user, + get_allowed_cli_tool_names, + get_train_op_code, + create_vertex_ai_train_op_component, + prepare_config_with_overrides, +) + + +@click.command(short_help="Submit a single-component job to Vertex AI Pipelines.") +@click.option( + "--tool", + required=True, + help="Tool to run, e.g. 'onepass_mean_var_std'.", +) +@click.option( + "--subcommand", + required=True, + type=click.Choice(["fit", "predict"]), + help="Subcommand to run, either 'fit' or 'predict'.", +) +@click.option( + "--config", + required=True, + help="GCS path to the training config YAML file.", +) +@click.option( + "--copy-data-to-local-disk", + default=True, + type=bool, + help="True copies GCS data to local disk fully (once) before training. False is ephemeral.", +) +@click.option( + "--project", + default="dsp-cell-annotation-service", + help="Google Cloud project ID.", +) +@click.option( + "--location", + default="us-central1", + help="Google Cloud location, e.g. 'us-central1'.", +) +@click.option( + "--pipeline-name", + default="", + help="Pipeline name, defaults to f'{user}_{tool}_{subcommand}'.", +) +@click.option( + "--machine-type", + default="n1-standard-8", + help="Machine type for the training job, e.g. 'n1-standard-16'.", +) +@click.option( + "--replica-count", + default=1, + type=int, + help="Number of replicas (nodes) for training.", +) +@click.option( + "--accelerator-type", + default="NVIDIA_TESLA_T4", + help="Type of accelerator (gpu), e.g. 'NVIDIA_TESLA_T4'.", +) +@click.option( + "--accelerator-count", + default=1, + type=int, + help="Number of GPUs.", +) +@click.option( + "--git-sha", + default="", + type=str, + help="Cellarium-ML git SHA to install (if provided).", +) +@click.option( + "--base-image", + default="us-central1-docker.pkg.dev/broad-dsde-methods/cellarium-ai/cellarium-ml:cellarium-gpt-cstorch", + help="Base image for the component.", +) +@click.option( + "--extract-bucket", + default=None, + help="GCS URI prefix containing extract_*.h5ad files, e.g. gs://my-bucket/my-prefix.", +) +def submit_single_component_pipeline( + project: str, + location: str, + config: str, + tool: str, + subcommand: str, + copy_data_to_local_disk: bool, + pipeline_name: str, + machine_type: str, + replica_count: int, + accelerator_type: str, + accelerator_count: int, + git_sha: str, + base_image: str, + extract_bucket=None, +): + """ + Submit a single component cellarium-ml pipeline to Vertex AI Pipelines. + """ + # input validation and defaults + display_name = f"{tool}_{subcommand}" + if pipeline_name == "": + user = get_current_google_user() + if user is not None: + pipeline_name = f"{user}_{display_name}" + else: + pipeline_name = display_name + if (git_sha == "") and (len(base_image.split(":")[-1]) > 0): + git_sha = base_image.split(":")[-1] + config = prepare_config_with_overrides(config, extract_bucket) + url = f"https://raw.githubusercontent.com/cellarium-ai/cellarium-ml/{git_sha}/cellarium/ml/cli.py" + cli_tool_names = get_allowed_cli_tool_names(url) + if cli_tool_names is not None: + if tool not in cli_tool_names: + raise ValueError( + f"Tool '{tool}' not found in allowed CLI tools at {url}.\n" + f"Allowed tool names:\n{cli_tool_names}" + ) + if ( + (accelerator_count is None) + or (accelerator_type is None) + or (accelerator_count == 0) + ): + # vertex ai wants None for both inputs if one of them is None + accelerator_count = None + accelerator_type = None + + aiplatform.init(project=project, location=location) + + # Create the train_op component using our dynamic creator + train_op = create_vertex_ai_train_op_component(base_image) + + # Get the train_op code that will be passed as a parameter + train_op_code = get_train_op_code(copy_data_to_local_disk) + + custom_training_job = create_custom_training_job_from_component( + train_op, + display_name=display_name, + replica_count=replica_count, + machine_type=machine_type, + accelerator_type=accelerator_type, + accelerator_count=accelerator_count, + ) + + @dsl.pipeline(name=pipeline_name, description=f"cellarium-ml {tool} {subcommand}") + def pipeline(): + custom_training_job( + project=project, + location=location, + tool=tool, + subcommand=subcommand, + config=config, + train_op_code=train_op_code, + git_sha=git_sha, + copy_data_to_local_disk=copy_data_to_local_disk, + ).set_display_name(display_name) + + with tempfile.NamedTemporaryFile(suffix=".yaml") as f: + compiler.Compiler().compile(pipeline_func=pipeline, package_path=f.name) + + job = aiplatform.PipelineJob( + display_name=display_name, + template_path=f.name, + enable_caching=False, # by default this is True + ) + + job.submit() + + +if __name__ == "__main__": + submit_single_component_pipeline() diff --git a/cellarium/workflows/submit_vertex_pipeline.py b/cellarium/workflows/submit_vertex_pipeline.py new file mode 100644 index 0000000..9fec933 --- /dev/null +++ b/cellarium/workflows/submit_vertex_pipeline.py @@ -0,0 +1,224 @@ +import tempfile +import yaml + +import click +from google.cloud import aiplatform +from google_cloud_pipeline_components.v1.custom_job import ( + create_custom_training_job_from_component, +) +from kfp import compiler, dsl + +from .shared_components import ( + get_current_google_user, + get_allowed_cli_tool_names, + get_train_op_code, + create_vertex_ai_train_op_component, + prepare_config_with_overrides, +) + + +def parse_pipeline_yaml(config: str) -> tuple[str, list[dict]]: + with open(config) as f: + config_contents = yaml.safe_load(f) + top_level_keys = list(config_contents.keys()) + assert len(top_level_keys) == 1, ( + "Pipeline YAML config error: The top level of the config file must be the display_name of the pipeline. Only one top level key is allowed." + ) + display_name = list(config_contents.keys())[0] + component_definitions = config_contents[display_name] + assert isinstance(component_definitions, list), ( + "Pipeline YAML config error: The value of the top level key must be a list of component definition dictionaries." + ) + for item in component_definitions: + assert isinstance(item, dict), ( + "Pipeline YAML config error: Each component definition in the list must be a dictionary." + ) + assert "tool" in item, ( + "Pipeline YAML config error: Each component definition must have a 'tool' key." + ) + assert "subcommand" in item, ( + "Pipeline YAML config error: Each component definition must have a 'subcommand' key." + ) + assert item["subcommand"] in ["fit", "predict"], ( + "Pipeline YAML config error: The 'subcommand' key's value must be either 'fit' or 'predict'." + ) + assert "config" in item, ( + "Pipeline YAML config error: Each component definition must have a 'config' key." + ) + return display_name, component_definitions + + +@click.command( + short_help="Submit a multi-step sequential pipeline to Vertex AI Pipelines." +) +@click.option( + "--pipeline-config", + required=True, + help="Local path to the pipeline config YAML file.", +) +@click.option( + "--project", + default="dsp-cell-annotation-service", + help="Google Cloud project ID.", +) +@click.option( + "--location", + default="us-central1", + help="Google Cloud location, e.g. 'us-central1'.", +) +@click.option( + "--pipeline-name", + default="", + help="Pipeline name, defaults to f'{user}_{tool}_{subcommand}'.", +) +@click.option( + "--copy-data-to-local-disk", + default=True, + type=bool, + help="True copies GCS data to local disk fully (once) before training. False is ephemeral.", +) +@click.option( + "--base-image", + default="us-central1-docker.pkg.dev/broad-dsde-methods/cellarium-ai/cellarium-ml:cellarium-gpt-cstorch", + help="Base image for the component.", +) +@click.option( + "--extract-bucket", + default=None, + help="GCS URI prefix containing extract_*.h5ad files, e.g. gs://my-bucket/my-prefix.", +) +def submit_sequential_pipeline( + project: str, + location: str, + pipeline_config: str, + pipeline_name: str, + copy_data_to_local_disk: bool, + base_image: str, + extract_bucket=None, +): + """ + Submit a pipeline of sequential cellarium-ml tools to Vertex AI Pipelines. + + Example contents of pipeline-config: + + .. code-block:: yaml + + scvi_vanilla_with_full_latent_batch: + + - tool: scvi + + subcommand: fit + + config: gs://cellarium-human-primary-data/curriculum/human_all_primary_20241108/configs/20241120_scvi_train_config.yaml + + machine_type: n1-standard-16 + + accelerator_type: NVIDIA_TESLA_T4 + + accelerator_count: 4 + + git_sha: c14705370d2a7a805286fa3dd0e4795c10e6cefd + + """ + # parse pipeline config + display_name, component_definitions = parse_pipeline_yaml(pipeline_config) + + # Apply extract_bucket overrides to each component's config + if extract_bucket is not None: + for component in component_definitions: + component["config"] = prepare_config_with_overrides( + component["config"], extract_bucket + ) + + # input validation and defaults + if pipeline_name == "": + user = get_current_google_user() + if user is not None: + pipeline_name = f"{user}_{display_name}" + else: + pipeline_name = display_name + + for t, sha in [(c["tool"], c.get("git_sha", "")) for c in component_definitions]: + if (sha == "") and (len(base_image.split(":")[-1]) > 0): + sha = base_image.split(":")[-1] + url = f"https://raw.githubusercontent.com/cellarium-ai/cellarium-ml/{sha}/cellarium/ml/cli.py" + cli_tool_names = get_allowed_cli_tool_names(url) + if cli_tool_names is not None: + if t not in cli_tool_names: + raise ValueError( + f"Tool '{t}' not found in allowed CLI tools at {url}.\n" + f"Allowed tool names:\n{cli_tool_names}" + ) + for subcommand in [c["subcommand"] for c in component_definitions]: + if subcommand not in ["fit", "predict"]: + raise ValueError( + f"Subcommand '{subcommand}' not recognized. Must be either 'fit' or 'predict'." + ) + for i, c in enumerate(component_definitions): + accelerator_type = c.get("accelerator_type", None) + accelerator_count = c.get("accelerator_count", None) + if ( + (accelerator_count is None) + or (accelerator_type is None) + or (accelerator_count == 0) + ): + # vertex ai wants None for both inputs if one of them is None + component_definitions[i]["accelerator_count"] = None + component_definitions[i]["accelerator_type"] = None + + aiplatform.init(project=project, location=location) + + # Create the train_op component using our dynamic creator + train_op = create_vertex_ai_train_op_component(base_image) + + # Get the train_op code that will be passed as a parameter + train_op_code = get_train_op_code(copy_data_to_local_disk) + + # create component definitions + custom_training_jobs = [ + create_custom_training_job_from_component( + train_op, + display_name=f"{i}__{c['tool']}_{c['subcommand']}", + replica_count=c.get("replica_count", 1), + machine_type=c.get("machine_type", None), + accelerator_type=c.get("accelerator_type", None), + accelerator_count=c.get("accelerator_count", None), + ) + for i, c in enumerate(component_definitions) + ] + + @dsl.pipeline(name=pipeline_name, description="cellarium-ml sequence") + def pipeline(): + tasks = [] + for i, (component_definition, custom_training_job) in enumerate( + zip(component_definitions, custom_training_jobs) + ): + task = custom_training_job( + project=project, + location=location, + tool=component_definition["tool"], + subcommand=component_definition["subcommand"], + config=component_definition["config"], + train_op_code=train_op_code, + git_sha=component_definition.get("git_sha", ""), + copy_data_to_local_disk=copy_data_to_local_disk, + ).set_display_name( + f"{i}__{component_definition['tool']}_{component_definition['subcommand']}" + ) + if tasks: # Set dependency if there's a previous task + task.after(tasks[-1]) + tasks.append(task) + + with tempfile.NamedTemporaryFile(suffix=".yaml") as f: + compiler.Compiler().compile(pipeline_func=pipeline, package_path=f.name) + + job = aiplatform.PipelineJob( + display_name=display_name, + template_path=f.name, + ) + + job.submit() + + +if __name__ == "__main__": + submit_sequential_pipeline() diff --git a/default_configs/hvg_seurat_v3_train.yaml b/default_configs/hvg_seurat_v3_train.yaml new file mode 100644 index 0000000..d05228f --- /dev/null +++ b/default_configs/hvg_seurat_v3_train.yaml @@ -0,0 +1,46 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + init_args: + broadcast_buffers: false + max_epochs: 2 + default_root_dir: gs://cellarium-dev-central/workflows/20260519_scvi_seurat_v3_4k_batch_suspension +model: + model: + class_path: cellarium.ml.models.HVGSeuratV3 + init_args: + n_top_genes: 4000 + use_batch_key: true + flavor: seurat_v3_paper +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_human_primary_gr300genes_fullschema/extract_files/extract_{000000..000050}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 2 + cache_size_strictly_enforced: true + indices_strict: true + obs_columns_to_validate: + - assay_suspension_type + # - donor_id_experiment + # - suspension_type + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var_names + batch_index_n: + attr: obs + key: assay_suspension_type + # key: donor_id_experiment + # key: suspension_type + convert_fn: cellarium.ml.utilities.data.categories_to_codes + batch_size: 5000 + num_workers: 8 + prefetch_factor: 2 +ckpt_path: null diff --git a/default_configs/imputation_scvi_train.yaml b/default_configs/imputation_scvi_train.yaml new file mode 100644 index 0000000..fb1d369 --- /dev/null +++ b/default_configs/imputation_scvi_train.yaml @@ -0,0 +1,94 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + logger: null + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + every_n_epochs: 1 + max_epochs: 25 + log_every_n_steps: 100 + gradient_clip_algorithm: norm + gradient_clip_val: 0.5 + inference_mode: false + default_root_dir: gs://cellarium-dev-central/workflows/20260407_human_scvi_imputation_default +model: + cpu_transforms: + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: gs://cellarium-dev-central/workflows/20260331_czi_human_primary_20251108_gte300genes_hvg_seurat_v3/hvg_seurat_v3_output__hvg_only.csv + loader_fn: pandas.read_csv + attr: "feature_id" + convert_fn: pandas.Series.to_list + model: + class_path: cellarium.ml.models.ImputationModel + init_args: + var_names_g: null + n_batch: null + masking_probability: 0.5 + noise2self_ratio: 0.5 + # n_latent_batch: 128 + batch_embedded: false + batch_representation_sampled: false + n_cats_per_cov: null + kl_warmup_epochs: 20 + batch_kl_weight_max: 0.0 + use_batch_norm: both + encoder: + hidden_layers: + - class_path: torch.nn.Linear + init_args: + out_features: 256 + final_layer: + class_path: torch.nn.Linear + init_args: {} + decoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 256 + label_to_bias_hidden_layers: [] + final_layer: + class_path: torch.nn.Linear + init_args: {} + final_additive_bias: false + n_latent: 50 + optim_fn: torch.optim.AdamW + optim_kwargs: + lr: 1e-4 +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_human_primary_gr300genes/extract_files/extract_{000000..000050}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 2 + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + key: feature_id + batch_index_n: + attr: obs + key: assay_dataset_id_donor_id_suspension_type + convert_fn: cellarium.ml.utilities.data.categories_to_codes + # categorical_covariate_index_nd: + # attr: obs + # key: + # - suspension_type + # - assay + # convert_fn: cellarium.ml.utilities.data.categories_to_codes + batch_size: 1024 + shuffle: true + num_workers: 8 + prefetch_factor: 4 + persistent_workers: false +ckpt_path: null diff --git a/default_configs/incremental_pca_predict.yaml b/default_configs/incremental_pca_predict.yaml new file mode 100644 index 0000000..57d6148 --- /dev/null +++ b/default_configs/incremental_pca_predict.yaml @@ -0,0 +1,75 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + callbacks: + - class_path: cellarium.ml.callbacks.PredictionWriter + init_args: + output_dir: . + max_epochs: 1 + logger: null + default_root_dir: gs://cellarium-dev-central/workflows/20260410_cas_pca_predict_vsindex +model: + transforms: + - cellarium.ml.transforms.NormalizeTotal + - cellarium.ml.transforms.Log1p + - class_path: cellarium.ml.transforms.ZScore + init_args: + mean_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260403_cas_onepass/lightning_logs/version_0/checkpoints/epoch=0-step=17681.ckpt + attr: model.mean_g + convert_fn: torch.Tensor.cpu + std_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260403_cas_onepass/lightning_logs/version_0/checkpoints/epoch=0-step=17681.ckpt + attr: model.std_g + convert_fn: torch.Tensor.cpu + var_names_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260403_cas_onepass/lightning_logs/version_0/checkpoints/epoch=0-step=17681.ckpt + attr: model.var_names_g + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + # file_path: gs://cellarium-dev-central/workflows/20260403_cas_hvg_seurat_v3_8k_batch_suspension/hvg_seurat_v3_output__hvg_only.csv + file_path: gs://cellarium-file-system/curriculum/human_10x_ebd_dfr_gt_3_extract/models/shared_meta/mean-var-std-highly-variable-genes-log1p-default-cutoffs.csv + loader_fn: pandas.read_csv + # attr: "feature_id" + attr: "original_feature_id" + convert_fn: pandas.Series.to_list + ordering: false + model: + class_path: cellarium.ml.models.IncrementalPCA + init_args: + var_names_g: null + n_components: 64 +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_vsindex_10x/extract_files/extract_{000000..000200}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 2 + obs_columns_to_validate: [] + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + key: feature_id + obs_names_n: + attr: obs_names + batch_size: 10000 + shuffle: false + num_workers: 4 # limited by memory on n1-standard-16 + prefetch_factor: 4 + persistent_workers: false +return_predictions: false +ckpt_path: gs://cellarium-dev-central/workflows/20260410_cas_pca_oldgenelist/lightning_logs/version_0/checkpoints/epoch=0-step=8841.ckpt diff --git a/default_configs/incremental_pca_train.yaml b/default_configs/incremental_pca_train.yaml new file mode 100644 index 0000000..b04e084 --- /dev/null +++ b/default_configs/incremental_pca_train.yaml @@ -0,0 +1,67 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + max_epochs: 1 + logger: null + default_root_dir: gs://cellarium-dev-central/workflows/20260410_cas_pca_oldgenelist +model: + transforms: + - cellarium.ml.transforms.NormalizeTotal + - cellarium.ml.transforms.Log1p + - class_path: cellarium.ml.transforms.ZScore + init_args: + mean_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260403_cas_onepass/lightning_logs/version_0/checkpoints/epoch=0-step=17681.ckpt + attr: model.mean_g + convert_fn: torch.Tensor.cpu + std_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260403_cas_onepass/lightning_logs/version_0/checkpoints/epoch=0-step=17681.ckpt + attr: model.std_g + convert_fn: torch.Tensor.cpu + var_names_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260403_cas_onepass/lightning_logs/version_0/checkpoints/epoch=0-step=17681.ckpt + attr: model.var_names_g + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + # file_path: gs://cellarium-dev-central/workflows/20260403_cas_hvg_seurat_v3_8k_batch_suspension/hvg_seurat_v3_output__hvg_only.csv + file_path: gs://cellarium-file-system/curriculum/human_10x_ebd_dfr_gt_3_extract/models/shared_meta/mean-var-std-highly-variable-genes-log1p-default-cutoffs.csv + loader_fn: pandas.read_csv + # attr: "feature_id" + attr: "original_feature_id" + convert_fn: pandas.Series.to_list + ordering: false + model: + class_path: cellarium.ml.models.IncrementalPCA + init_args: + var_names_g: null + n_components: 64 +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files/extract_{000000..000050}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 2 + obs_columns_to_validate: [] + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + key: feature_id + batch_size: 5000 + num_workers: 2 # limited by memory on n1-standard-16 + prefetch_factor: 4 + persistent_workers: false +ckpt_path: null diff --git a/default_configs/nmf_train.yaml b/default_configs/nmf_train.yaml new file mode 100644 index 0000000..f5919e2 --- /dev/null +++ b/default_configs/nmf_train.yaml @@ -0,0 +1,56 @@ +# lightning.pytorch==2.4.0 +seed_everything: true +trainer: + num_nodes: 1 + max_epochs: 10 +model: + cpu_transforms: + - class_path: cellarium.ml.transforms.DivideByScale + init_args: + var_names_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260402_human_onepass/lightning_logs/version_0/checkpoints/epoch=1-step=37788.ckpt + attr: model.var_names_g + scale_g: + !CheckpointLoader + file_path: gs://cellarium-dev-central/workflows/20260402_human_onepass/lightning_logs/version_0/checkpoints/epoch=1-step=37788.ckpt + attr: model.scale_g + eps: 1e-4 + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: gs://cellarium-dev-central/workflows/20260331_czi_human_primary_20251108_gte300genes_hvg_seurat_v3/hvg_seurat_v3_output__hvg_only.csv + loader_fn: pandas.read_csv + attr: "feature_id" + convert_fn: pandas.Series.to_list + model: + class_path: cellarium.ml.models.OnlineNonNegativeMatrixFactorization + init_args: + var_names_g: null + k_values: [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30] + r: 100 + algorithm: nmf_torch_hals + n_cells_total: 500_000 + is_initialized: false +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/bican_freeze_1_full_extract/extract_files/extract_{000000..000050}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 2 + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + key: feature_id + batch_size: 1024 + shuffle: true + num_workers: 12 + prefetch_factor: 4 + persistent_workers: false +ckpt_path: null diff --git a/default_configs/onepass_mean_var_std_train.yaml b/default_configs/onepass_mean_var_std_train.yaml new file mode 100644 index 0000000..5f0997e --- /dev/null +++ b/default_configs/onepass_mean_var_std_train.yaml @@ -0,0 +1,42 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + max_epochs: 1 + logger: null + default_root_dir: gs://cellarium-dev-central/workflows/20260403_cas_onepass +model: + transforms: + - cellarium.ml.transforms.Densify + - class_path: cellarium.ml.transforms.NormalizeTotal + init_args: + target_count: 10_000 + - cellarium.ml.transforms.Log1p + model: + class_path: cellarium.ml.models.OnePassMeanVarStd + init_args: + algorithm: shifted_data + output_path: onepass_output.csv +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files/extract_{000000..000050}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 3 + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.to_torch_sparse_csr + var_names_g: + attr: var + key: feature_id + batch_size: 4096 + num_workers: 12 + prefetch_factor: 2 + persistent_workers: false +ckpt_path: null diff --git a/default_configs/raw_onepass_mean_var_std_train.yaml b/default_configs/raw_onepass_mean_var_std_train.yaml new file mode 100644 index 0000000..00410c0 --- /dev/null +++ b/default_configs/raw_onepass_mean_var_std_train.yaml @@ -0,0 +1,37 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + max_epochs: 1 + logger: null + default_root_dir: gs://cellarium-dev-central/workflows/20260403_bican_raw_onepass +model: + model: + class_path: cellarium.ml.models.OnePassMeanVarStd + init_args: + algorithm: shifted_data +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_human_primary_gr300genes/extract_files/extract_{000000..000050}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 2 + obs_columns_to_validate: [] + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + # key: feature_id + key: gene_id + batch_size: 5000 + num_workers: 8 # limited by memory on n1-standard-16 + prefetch_factor: 4 + persistent_workers: false +ckpt_path: null diff --git a/default_configs/scvi_reconstruct.yaml b/default_configs/scvi_reconstruct.yaml new file mode 100644 index 0000000..c1f70d4 --- /dev/null +++ b/default_configs/scvi_reconstruct.yaml @@ -0,0 +1,114 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + logger: + callbacks: + - class_path: cellarium.ml.callbacks.PredictionWriter + init_args: + output_dir: reconstructions + prediction_size: null + max_epochs: 1 + log_every_n_steps: 100 + inference_mode: true + default_root_dir: + gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco_reconstruction +model: + cpu_transforms: + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: + gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco/20260514_hvg_plus_glyco_genes.csv + loader_fn: pandas.read_csv + attr: "feature_id" + convert_fn: pandas.Series.to_list + transforms: + - cellarium.ml.transforms.Densify + model: + class_path: cellarium.ml.models.SingleCellVariationalInference + init_args: + + # reconstruction + reconstruct_counts_on_predict: true + reconstruction_transform_batch: 0 + reconstruction_transform_categorical_covariates: null + reconstruction_var_names_g: + !FileLoader + file_path: + gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco/20260514_hvg_plus_glyco_genes.csv + loader_fn: pandas.read_csv + attr: "feature_id" + convert_fn: pandas.Series.to_list + + var_names_g: null + n_batch: null + n_cats_per_cov: null + + # n_latent_batch: 128 + batch_embedded: false + batch_representation_sampled: false + kl_warmup_epochs: 20 + batch_kl_weight_max: 0.0 + use_batch_norm: both + encoder: + hidden_layers: + - class_path: torch.nn.Linear + init_args: + out_features: 512 + - class_path: torch.nn.Linear + init_args: + out_features: 512 + final_layer: + class_path: torch.nn.Linear + init_args: {} + decoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 512 + label_to_bias_hidden_layers: [] + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 512 + label_to_bias_hidden_layers: [] + final_layer: + class_path: torch.nn.Linear + init_args: {} + final_additive_bias: false + n_latent: 256 + optim_fn: torch.optim.AdamW + optim_kwargs: + lr: 1e-4 +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: + gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_human_primary_gr300genes_fullschema/extract_files/extract_{000000..000100}.h5ad + shard_size: 10000 + # last_shard_size: 6707 + max_cache_size: 2 + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.keep_sparse + var_names_g: + attr: var + key: feature_id + batch_index_n: + attr: obs + key: assay_dataset_id_donor_id_suspension_type + convert_fn: cellarium.ml.utilities.data.categories_to_codes + # categorical_covariate_index_nd: + # attr: obs + # key: + # - suspension_type + # - assay + # convert_fn: cellarium.ml.utilities.data.categories_to_codes + batch_size: 1024 + shuffle: true + num_workers: 8 + prefetch_factor: 4 + persistent_workers: false +ckpt_path: gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco/lightning_logs/version_0/checkpoints/epoch=91-step=8487276.ckpt diff --git a/default_configs/scvi_train.yaml b/default_configs/scvi_train.yaml new file mode 100644 index 0000000..db13c71 --- /dev/null +++ b/default_configs/scvi_train.yaml @@ -0,0 +1,126 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + accelerator: gpu + devices: 4 + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + find_unused_parameters: false + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + every_n_epochs: 1 + max_epochs: 20 + log_every_n_steps: 1000 + val_check_interval: 10000 + gradient_clip_algorithm: norm + gradient_clip_val: 1.0 + inference_mode: false + default_root_dir: gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco_validation_runs +model: + cpu_transforms: + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco/20260514_hvg_plus_glyco_genes.csv + loader_fn: pandas.read_csv + attr: "feature_id" + convert_fn: pandas.Series.to_list + model: + class_path: cellarium.ml.models.SingleCellVariationalInference + init_args: + var_names_g: null + n_batch: null + n_latent_batch: 64 + batch_embedded: true + batch_representation_sampled: true + n_cats_per_cov: null + kl_warmup_epochs: null + kl_warmup_steps: 200_000 + batch_kl_weight_max: 0.0 + use_batch_norm: both + + # if using a validation data set + cell_type_categories: null + ontology_distance_matrix: + !FileLoader + file_path: https://github.com/obophenotype/cell-ontology/releases/download/v2025-07-30/cl-basic.owl + loader_fn: cellarium.ml.utilities.data.compute_cl_distance_matrix + val_cell_type_classifier_reservoir_size: 200_000 + + encoder: + hidden_layers: + - class_path: torch.nn.Linear + init_args: + out_features: 512 + # - class_path: torch.nn.Linear + # init_args: + # out_features: 512 + final_layer: + class_path: torch.nn.Linear + init_args: {} + decoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 512 + label_to_bias_hidden_layers: [] + # - class_path: cellarium.ml.models.scvi.LinearWithBatch + # init_args: + # out_features: 512 + # label_to_bias_hidden_layers: [] + final_layer: + class_path: torch.nn.Linear + init_args: {} + final_additive_bias: false + n_latent: 64 + optim_fn: torch.optim.AdamW + optim_kwargs: + lr: 1e-3 + scheduler_fn: torch.optim.lr_scheduler.LinearLR + scheduler_kwargs: + start_factor: 0.01 + total_iters: 100000 +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_human_primary_gr300genes_fullschema/extract_files/extract_{000000..001000}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 3 + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + key: feature_id + batch_index_n: + attr: obs + key: assay_dataset_id_donor_id_suspension_type + convert_fn: cellarium.ml.utilities.data.categories_to_codes + # categorical_covariate_index_nd: + # attr: obs + # key: + # - suspension_type + # - assay + # convert_fn: cellarium.ml.utilities.data.categories_to_codes + # only used for validation data + validation_cell_type_index_n: + attr: obs + key: cell_type_ontology_term_id + convert_fn: cellarium.ml.utilities.data.categories_to_codes + batch_size: 4096 + shuffle: true + val_size: 0.02 + num_workers: 2 + prefetch_factor: 2 + persistent_workers: false +# ckpt_path: gs://cellarium-dev-central/workflows/20260401_human_scvi_our_hvg_copy_czi/lightning_logs/version_0/checkpoints/epoch=6-step=645771.ckpt diff --git a/pyproject.toml b/pyproject.toml index ee78166..84e7c47 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,14 +18,43 @@ classifiers = [ "Programming Language :: Python :: 3", "Programming Language :: Python :: Implementation :: CPython", ] -dynamic = ["dependencies", "optional-dependencies", "version", "readme"] +dynamic = ["version", "readme"] +dependencies = [ + "google-cloud-aiplatform==1.68.0", + "google-cloud-pipeline-components==2.17.0", + "google-cloud-batch==0.17.36", + "kfp==2.7.0", + "click==8.1.7", + "pyjwt==2.9.0", + "gcsfs", + "h5py", + "ruamel.yaml", +] + +[project.optional-dependencies] +dev = [ + "tox~=4.6", + "ruff", + "pytest~=7.3", +] +train_op = [ + "gcsfs", + "tensorboard", + "psutil", + "ruamel.yaml", +] [tool.setuptools.packages.find] -include = ["cellarium.workflows"] +include = ["cellarium*"] + +[tool.setuptools.package-data] +"cellarium.workflows.scripts" = ["*.sh"] +"cellarium.workflows.example" = ["*.yaml"] + +[project.scripts] +cellarium-workflow = "cellarium.workflows.cli:cli" [tool.setuptools.dynamic] -dependencies = { file = ["requirements/base.txt"] } -optional-dependencies = { docs = { file = ["requirements/docs.txt"] }, test = { file = ["requirements/test.txt"] }, vis = { file = ["requirements/vis.txt"] } } readme = { file = ["README.md"], content-type = "text/markdown" } [project.urls] diff --git a/requirements/base.txt b/requirements/base.txt deleted file mode 100644 index 0088d2c..0000000 --- a/requirements/base.txt +++ /dev/null @@ -1,4 +0,0 @@ -google-cloud-aiplatform==1.68.0 -google-cloud-pipeline-components==2.17.0 -kfp==2.7.0 -click==8.1.7 \ No newline at end of file diff --git a/requirements/dev.txt b/requirements/dev.txt deleted file mode 100644 index 2d54ff9..0000000 --- a/requirements/dev.txt +++ /dev/null @@ -1 +0,0 @@ -tox~=4.6 \ No newline at end of file diff --git a/requirements/test.txt b/requirements/test.txt deleted file mode 100644 index e382152..0000000 --- a/requirements/test.txt +++ /dev/null @@ -1,5 +0,0 @@ -pytest~=7.3 -coverage~=4.5 -click~=8.0 -mockito>=1.5.0 -parameterized>=0.9.0 \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_prepare_config_with_overrides.py b/tests/test_prepare_config_with_overrides.py new file mode 100644 index 0000000..a903358 --- /dev/null +++ b/tests/test_prepare_config_with_overrides.py @@ -0,0 +1,221 @@ +""" +Tests for prepare_config_with_overrides in shared_components. + +The local extract_*.h5ad files in the repo root (10 files, 10 000 obs each) +are used as fixtures. gcsfs is patched so no real GCS calls are made; the +mock maps expected GCS paths straight back to the local file paths. +""" + +import os +import re +import textwrap +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from ruamel.yaml import YAML + +from cellarium.workflows.shared_components import prepare_config_with_overrides + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +REPO_ROOT = Path(__file__).parent.parent +EXTRACT_FILES = sorted( + REPO_ROOT.glob("extract_*.h5ad"), + key=lambda p: int(re.search(r"extract_(\d+)", p.name).group(1)), +) +N_FILES = len(EXTRACT_FILES) # 10 +OBS_PER_FILE = 10_000 # confirmed from local files +FAKE_BUCKET = "gs://fake-bucket/fake-prefix" + + +def _make_fake_fs(): + """ + Return a mock that mimics the gcsfs.GCSFileSystem API used by + prepare_config_with_overrides: + - .glob(pattern) -> list of GCS-style paths + - .open(path, mode) -> file-like object backed by the local h5ad + """ + fake_fs = MagicMock() + + # glob: return one fake GCS path per local extract file + fake_gcs_paths = [ + f"fake-bucket/fake-prefix/extract_{i}.h5ad" for i in range(N_FILES) + ] + fake_fs.glob.return_value = fake_gcs_paths + + def _open(gcs_path, mode="rb"): + # Map e.g. "fake-bucket/fake-prefix/extract_3.h5ad" -> local file + filename = gcs_path.split("/")[-1] + local_path = REPO_ROOT / filename + return open(local_path, "rb") + + fake_fs.open.side_effect = _open + return fake_fs + + +def _load_yaml(path: str) -> dict: + yaml = YAML() + with open(path) as f: + return yaml.load(f) + + +def _base_config_yaml() -> str: + return textwrap.dedent("""\ + data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://old-bucket/old-prefix/extract_{0..4}.h5ad + shard_size: 500 + last_shard_size: 123 + max_cache_size: 2 + """) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestNoExtractBucket: + """When extract_bucket is None the original path is returned unchanged.""" + + def test_returns_original_path(self, tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + result = prepare_config_with_overrides(str(cfg), extract_bucket=None) + assert result == str(cfg) + + def test_file_is_not_modified(self, tmp_path): + cfg = tmp_path / "config.yaml" + original = _base_config_yaml() + cfg.write_text(original) + prepare_config_with_overrides(str(cfg), extract_bucket=None) + assert cfg.read_text() == original + + +class TestExtractBucketPatching: + """With extract_bucket supplied, filenames / shard_size / last_shard_size + are updated to match the discovered shards.""" + + def test_filenames_pattern(self, tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + out = prepare_config_with_overrides(str(cfg), FAKE_BUCKET) + doc = _load_yaml(out) + filenames = doc["data"]["dadc"]["init_args"]["filenames"] + assert filenames == f"{FAKE_BUCKET}/extract_{{0..{N_FILES - 1}}}.h5ad" + + def test_shard_size_from_first_file(self, tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + out = prepare_config_with_overrides(str(cfg), FAKE_BUCKET) + doc = _load_yaml(out) + assert doc["data"]["dadc"]["init_args"]["shard_size"] == OBS_PER_FILE + + def test_last_shard_size_from_last_file(self, tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + out = prepare_config_with_overrides(str(cfg), FAKE_BUCKET) + doc = _load_yaml(out) + assert doc["data"]["dadc"]["init_args"]["last_shard_size"] == OBS_PER_FILE + + def test_returns_different_path_from_input(self, tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + out = prepare_config_with_overrides(str(cfg), FAKE_BUCKET) + assert out != str(cfg) + + def test_original_config_file_unchanged(self, tmp_path): + cfg = tmp_path / "config.yaml" + original = _base_config_yaml() + cfg.write_text(original) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + prepare_config_with_overrides(str(cfg), FAKE_BUCKET) + assert cfg.read_text() == original + + def test_output_is_valid_yaml_file(self, tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + out = prepare_config_with_overrides(str(cfg), FAKE_BUCKET) + assert os.path.isfile(out) + doc = _load_yaml(out) + assert "data" in doc + + def test_trailing_slash_on_bucket_stripped(self, tmp_path): + """A trailing slash on the bucket prefix should not produce double slashes.""" + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + out = prepare_config_with_overrides(str(cfg), FAKE_BUCKET + "/") + doc = _load_yaml(out) + filenames = doc["data"]["dadc"]["init_args"]["filenames"] + assert "//" not in filenames.replace("gs://", "") + + +class TestErrorCases: + def test_no_matching_files_raises(self, tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text(_base_config_yaml()) + fake_fs = MagicMock() + fake_fs.glob.return_value = [] + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + with pytest.raises( + FileNotFoundError, match="No extract_.*\\.h5ad files found" + ): + prepare_config_with_overrides(str(cfg), FAKE_BUCKET) + + def test_missing_dadc_key_raises(self, tmp_path): + bad_yaml = textwrap.dedent("""\ + data: + something_else: + init_args: {} + """) + cfg = tmp_path / "config.yaml" + cfg.write_text(bad_yaml) + fake_fs = _make_fake_fs() + with patch( + "cellarium.workflows.shared_components.gcsfs.GCSFileSystem", + return_value=fake_fs, + ): + with pytest.raises(KeyError): + prepare_config_with_overrides(str(cfg), FAKE_BUCKET) diff --git a/tests/test_shared_components.py b/tests/test_shared_components.py new file mode 100644 index 0000000..83477fa --- /dev/null +++ b/tests/test_shared_components.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Test script to verify the shared components work correctly.""" + +import sys +import os + +sys.path.insert(0, os.path.dirname(__file__)) + +from cellarium.workflows.shared_components import ( + get_pytorch_setup_code, + get_git_install_code, + get_data_download_code, + get_cellarium_cli_code, + get_current_google_user, + create_train_op_function, + get_train_op_code, +) + + +def test_code_generation(): + """Test that all code generation functions work.""" + print("Testing code generation functions...") + + # Test git install code + git_code = get_git_install_code() + print(f"Git install code length: {len(git_code)}") + assert "pip install" in git_code + assert "cellarium-ml" in git_code + + # Test pytorch setup code + pytorch_code = get_pytorch_setup_code() + print(f"PyTorch setup code length: {len(pytorch_code)}") + assert "psutil" in pytorch_code + assert "torch" in pytorch_code + assert "OMP_NUM_THREADS" in pytorch_code + + # Test data download code + data_code = get_data_download_code() + print(f"Data download code length: {len(data_code)}") + assert "gcsfs" in data_code + assert "download_file" in data_code + + # Test cellarium CLI code + cli_code = get_cellarium_cli_code() + print(f"Cellarium CLI code length: {len(cli_code)}") + assert "cellarium_ml_cli" in cli_code + + print("All code generation tests passed!") + + +def test_train_op_functions(): + """Test train_op function creation and code generation.""" + print("Testing train_op functions...") + + # Test creating train_op function + train_op = create_train_op_function(copy_data_to_local_disk=True) + assert callable(train_op) + print("✓ create_train_op_function works") + + # Test with copy_data_to_local_disk=False + train_op_no_copy = create_train_op_function(copy_data_to_local_disk=False) + assert callable(train_op_no_copy) + print("✓ create_train_op_function works with copy_data_to_local_disk=False") + + # Test get_train_op_code + code_with_download = get_train_op_code(copy_data_to_local_disk=True) + assert len(code_with_download) > 0 + assert "import os" in code_with_download # Should contain actual script content + assert "cellarium_ml_cli" in code_with_download # Should contain CLI execution + assert ( + "copy_data_to_local_disk" in code_with_download + ) # Should contain data download check + print("✓ get_train_op_code works with data download") + + code_without_download = get_train_op_code(copy_data_to_local_disk=False) + assert len(code_without_download) > 0 + assert "import os" in code_without_download # Should contain actual script content + assert "cellarium_ml_cli" in code_without_download # Should contain CLI execution + assert ( + "copy_data_to_local_disk" not in code_without_download + ) # Should not contain data download check + print("✓ get_train_op_code works without data download") + + print("Train_op function tests passed!") + + +def test_imports(): + """Test that all workflow files can be imported.""" + print("Testing workflow file imports...") + + try: + print("✓ submit_single_component imports successfully") + except Exception as e: + print(f"✗ submit_single_component import failed: {e}") + raise + + try: + print("✓ submit_pipeline imports successfully") + except Exception as e: + print(f"✗ submit_pipeline import failed: {e}") + raise + + try: + print("✓ local_single_component imports successfully") + except Exception as e: + print(f"✗ local_single_component import failed: {e}") + raise + + print("All import tests passed!") + + +def test_utility_functions(): + """Test utility functions.""" + print("Testing utility functions...") + + # Test get_current_google_user (may fail if not authenticated) + try: + user = get_current_google_user() + print(f"Current Google user: {user}") + except Exception as e: + print(f"Google user test skipped (expected): {e}") + + print("Utility function tests completed!") + + +if __name__ == "__main__": + test_code_generation() + test_train_op_functions() + test_imports() + test_utility_functions() + print("All tests completed successfully!") diff --git a/tmp/bican_cnmfpreprocess_onepass_mean_var_std_train.yaml b/tmp/bican_cnmfpreprocess_onepass_mean_var_std_train.yaml new file mode 100644 index 0000000..7c15ea0 --- /dev/null +++ b/tmp/bican_cnmfpreprocess_onepass_mean_var_std_train.yaml @@ -0,0 +1,89 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + accelerator: auto + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + init_args: + accelerator: null + parallel_devices: null + cluster_environment: null + checkpoint_io: null + precision_plugin: null + ddp_comm_state: null + ddp_comm_hook: null + ddp_comm_wrapper: null + model_averaging_period: null + process_group_backend: null + timeout: 0:30:00 + start_method: popen + dict_kwargs: + broadcast_buffers: false + devices: auto + num_nodes: 1 + precision: null + callbacks: null + fast_dev_run: false + max_epochs: 1 + min_epochs: null + max_steps: -1 + min_steps: null + max_time: null + limit_train_batches: null + limit_val_batches: null + limit_test_batches: null + limit_predict_batches: null + overfit_batches: 0.0 + val_check_interval: null + check_val_every_n_epoch: 1 + num_sanity_val_steps: null + log_every_n_steps: null + enable_checkpointing: null + enable_progress_bar: null + enable_model_summary: null + accumulate_grad_batches: 1 + gradient_clip_val: null + gradient_clip_algorithm: null + deterministic: null + benchmark: null + inference_mode: true + use_distributed_sampler: true + profiler: null + detect_anomaly: false + barebones: false + plugins: null + sync_batchnorm: false + reload_dataloaders_every_n_epochs: 0 + default_root_dir: gs://cellarium-dev-central/workflows/20260605_bican_allfreeze1_onepass_for_cnmf +model: + transforms: + - cellarium.ml.transforms.Densify + - class_path: cellarium.ml.transforms.NormalizeTotal + init_args: + target_count: 1_000_000 + - cellarium.ml.transforms.Log1p + model: + class_path: cellarium.ml.models.OnePassMeanVarStd + init_args: + algorithm: shifted_data + output_path: onepass_output.csv +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/bican_freeze_1_full_extract/extract_files/extract_{000000..000050}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 3 + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + key: gene_id + batch_size: 2048 + num_workers: 0 + prefetch_factor: null + persistent_workers: false +ckpt_path: null