altar-chrombpnet 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- altar_chrombpnet-0.1.0/.gitignore +32 -0
- altar_chrombpnet-0.1.0/LICENSE +21 -0
- altar_chrombpnet-0.1.0/PKG-INFO +62 -0
- altar_chrombpnet-0.1.0/README.md +37 -0
- altar_chrombpnet-0.1.0/altar_chrombpnet/__init__.py +65 -0
- altar_chrombpnet-0.1.0/altar_chrombpnet/layouts.py +108 -0
- altar_chrombpnet-0.1.0/altar_chrombpnet/local_scoring.py +194 -0
- altar_chrombpnet-0.1.0/altar_chrombpnet/plugin.py +762 -0
- altar_chrombpnet-0.1.0/altar_chrombpnet/py.typed +0 -0
- altar_chrombpnet-0.1.0/pyproject.toml +55 -0
- altar_chrombpnet-0.1.0/tests/test_binding.py +29 -0
- altar_chrombpnet-0.1.0/tests/test_interpretation_verification.py +120 -0
- altar_chrombpnet-0.1.0/tests/test_local_scoring_run.py +165 -0
- altar_chrombpnet-0.1.0/tests/test_plans.py +708 -0
- altar_chrombpnet-0.1.0/tests/test_plugin.py +150 -0
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
*.pyc
|
|
2
|
+
*~
|
|
3
|
+
**/__pycache__/*
|
|
4
|
+
*.swp
|
|
5
|
+
.vscode/
|
|
6
|
+
.idea/
|
|
7
|
+
.DS_Store
|
|
8
|
+
.env*
|
|
9
|
+
.mypy_cache/
|
|
10
|
+
.ruff_cache/
|
|
11
|
+
.pytest_cache/
|
|
12
|
+
.coverage
|
|
13
|
+
htmlcov/
|
|
14
|
+
coverage.xml
|
|
15
|
+
pytest.xml
|
|
16
|
+
.hypothesis/
|
|
17
|
+
.python-version
|
|
18
|
+
.venv/
|
|
19
|
+
.venv-*/
|
|
20
|
+
*.egg-info/
|
|
21
|
+
build/
|
|
22
|
+
dist/
|
|
23
|
+
site/
|
|
24
|
+
*.sqlite3
|
|
25
|
+
*.log
|
|
26
|
+
|
|
27
|
+
# Generated variant indexes and datasets. The small canonical gene table is tracked.
|
|
28
|
+
altar/altar/variants/data/ccres.dnatree
|
|
29
|
+
altar/altar/variants/data/region_annotations.parquet
|
|
30
|
+
altar/altar/variants/data/variants.pkl.gz
|
|
31
|
+
altar/altar/variants/data/raw/
|
|
32
|
+
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Riya Sinha
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: altar-chrombpnet
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: ChromBPNet model binding for large-scale Altar variant scoring
|
|
5
|
+
Project-URL: Documentation, https://kundajelab.github.io/altar/
|
|
6
|
+
Project-URL: Issues, https://github.com/kundajelab/altar/issues
|
|
7
|
+
Project-URL: Repository, https://github.com/kundajelab/altar
|
|
8
|
+
Author: Riya Sinha
|
|
9
|
+
License-Expression: MIT
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
17
|
+
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
|
|
18
|
+
Classifier: Typing :: Typed
|
|
19
|
+
Requires-Python: >=3.12
|
|
20
|
+
Requires-Dist: altar<0.2,>=0.1
|
|
21
|
+
Provides-Extra: test
|
|
22
|
+
Requires-Dist: pytest-asyncio>=0.24; extra == 'test'
|
|
23
|
+
Requires-Dist: pytest>=8; extra == 'test'
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
|
|
26
|
+
# Altar ChromBPNet binding
|
|
27
|
+
|
|
28
|
+
`altar-chrombpnet` is the lightweight control-plane binding between Altar and the independently locked
|
|
29
|
+
ChromBPNet runtime under `runtimes/chrombpnet`. Installing it registers `CHROMBPNET` through Altar's standard
|
|
30
|
+
model-plugin entry point; the base `altar` distribution does not contain ChromBPNet-specific science.
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
pip install altar-chrombpnet
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
from altar.models import get_model_plugin
|
|
38
|
+
from altar_chrombpnet import ChromBPNetConfiguration
|
|
39
|
+
|
|
40
|
+
plugin = get_model_plugin("CHROMBPNET")
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
The binding owns ChromBPNet configuration, score schemas, prioritization, preprocessing, scoring, and
|
|
44
|
+
interpretation plans. It depends only on Altar's public façades and Pydantic. TensorFlow, model code, weights,
|
|
45
|
+
reference genomes, peaks, and provider SDKs remain outside the binding.
|
|
46
|
+
|
|
47
|
+
The binding also owns ChromBPNet's logical paths (`altar_chrombpnet.layouts`): the derived per-fold peak
|
|
48
|
+
distributions and interval index, the staged peaks and motif set, and the interpretation run tree. Altar core
|
|
49
|
+
provides only the model-neutral `ModelArtifactLayout` and `ScoringLayout`; every model-owned ChromBPNet file is
|
|
50
|
+
composed from `model_resource_file`, so preparation and scoring name the same files under one layout. Staged
|
|
51
|
+
file names come from the declared `weights_format` and `peaks_compression`, never from a resource URI.
|
|
52
|
+
|
|
53
|
+
Preprocessing can emit a manifest that records the ordered weights, peaks, genome, strand policy, producer
|
|
54
|
+
image, and output digests. That provenance is additive: the normal scoring plan consumes the conventional
|
|
55
|
+
peak-distribution and safe interval-index paths directly and does not require a manifest, an Altar-owned
|
|
56
|
+
directory, or an adoption step. Standalone runtime callers may supply both `--preprocessing-manifest` and
|
|
57
|
+
`--preprocessing-identity` to opt into strict verification. Historical scores and externally managed
|
|
58
|
+
artifacts are not rewritten or recomputed.
|
|
59
|
+
|
|
60
|
+
Existing callers migrating from Altar core must install this package and replace imports such as
|
|
61
|
+
`from altar.models import ChromBPNetPlugin` with `from altar_chrombpnet import ChromBPNetPlugin`. Registry-based
|
|
62
|
+
callers need only install the package; the `CHROMBPNET` key is unchanged.
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
# Altar ChromBPNet binding
|
|
2
|
+
|
|
3
|
+
`altar-chrombpnet` is the lightweight control-plane binding between Altar and the independently locked
|
|
4
|
+
ChromBPNet runtime under `runtimes/chrombpnet`. Installing it registers `CHROMBPNET` through Altar's standard
|
|
5
|
+
model-plugin entry point; the base `altar` distribution does not contain ChromBPNet-specific science.
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install altar-chrombpnet
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
```python
|
|
12
|
+
from altar.models import get_model_plugin
|
|
13
|
+
from altar_chrombpnet import ChromBPNetConfiguration
|
|
14
|
+
|
|
15
|
+
plugin = get_model_plugin("CHROMBPNET")
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
The binding owns ChromBPNet configuration, score schemas, prioritization, preprocessing, scoring, and
|
|
19
|
+
interpretation plans. It depends only on Altar's public façades and Pydantic. TensorFlow, model code, weights,
|
|
20
|
+
reference genomes, peaks, and provider SDKs remain outside the binding.
|
|
21
|
+
|
|
22
|
+
The binding also owns ChromBPNet's logical paths (`altar_chrombpnet.layouts`): the derived per-fold peak
|
|
23
|
+
distributions and interval index, the staged peaks and motif set, and the interpretation run tree. Altar core
|
|
24
|
+
provides only the model-neutral `ModelArtifactLayout` and `ScoringLayout`; every model-owned ChromBPNet file is
|
|
25
|
+
composed from `model_resource_file`, so preparation and scoring name the same files under one layout. Staged
|
|
26
|
+
file names come from the declared `weights_format` and `peaks_compression`, never from a resource URI.
|
|
27
|
+
|
|
28
|
+
Preprocessing can emit a manifest that records the ordered weights, peaks, genome, strand policy, producer
|
|
29
|
+
image, and output digests. That provenance is additive: the normal scoring plan consumes the conventional
|
|
30
|
+
peak-distribution and safe interval-index paths directly and does not require a manifest, an Altar-owned
|
|
31
|
+
directory, or an adoption step. Standalone runtime callers may supply both `--preprocessing-manifest` and
|
|
32
|
+
`--preprocessing-identity` to opt into strict verification. Historical scores and externally managed
|
|
33
|
+
artifacts are not rewritten or recomputed.
|
|
34
|
+
|
|
35
|
+
Existing callers migrating from Altar core must install this package and replace imports such as
|
|
36
|
+
`from altar.models import ChromBPNetPlugin` with `from altar_chrombpnet import ChromBPNetPlugin`. Registry-based
|
|
37
|
+
callers need only install the package; the `CHROMBPNET` key is unchanged.
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
# Copyright 2026 Riya Sinha
|
|
2
|
+
"""ChromBPNet model binding for Altar."""
|
|
3
|
+
|
|
4
|
+
from altar_chrombpnet.layouts import (
|
|
5
|
+
MOTIF_SET_FILENAME,
|
|
6
|
+
PEAK_INTERVAL_INDEX_FILENAME,
|
|
7
|
+
PREPROCESSING_MANIFEST_FILENAME,
|
|
8
|
+
ChromBPNetInterpretationLayout,
|
|
9
|
+
PeaksCompression,
|
|
10
|
+
WeightsFormat,
|
|
11
|
+
fold_weights_file,
|
|
12
|
+
model_dir,
|
|
13
|
+
motif_set_file,
|
|
14
|
+
peak_distribution_file,
|
|
15
|
+
peak_interval_index_file,
|
|
16
|
+
peaks_file,
|
|
17
|
+
preprocessing_manifest_file,
|
|
18
|
+
)
|
|
19
|
+
from altar_chrombpnet.plugin import (
|
|
20
|
+
CHROMBPNET_DEFAULT_IMAGE,
|
|
21
|
+
FINEMO_IMAGE,
|
|
22
|
+
KIND_INTERPRETATION_AVERAGE,
|
|
23
|
+
KIND_INTERPRETATION_FOLDS,
|
|
24
|
+
KIND_INTERPRETATION_MOTIF,
|
|
25
|
+
KIND_INTERPRETATION_PLOT,
|
|
26
|
+
KIND_MODEL_PREPROCESSING,
|
|
27
|
+
KIND_SCORING_FOLDS,
|
|
28
|
+
KIND_SCORING_SUMMARIZATION,
|
|
29
|
+
NUM_FOLDS,
|
|
30
|
+
ChromBPNetConfiguration,
|
|
31
|
+
ChromBPNetPlugin,
|
|
32
|
+
ChromBPNetRunInputs,
|
|
33
|
+
ShardKind,
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
__all__ = [
|
|
38
|
+
"CHROMBPNET_DEFAULT_IMAGE",
|
|
39
|
+
"FINEMO_IMAGE",
|
|
40
|
+
"KIND_INTERPRETATION_AVERAGE",
|
|
41
|
+
"KIND_INTERPRETATION_FOLDS",
|
|
42
|
+
"KIND_INTERPRETATION_MOTIF",
|
|
43
|
+
"KIND_INTERPRETATION_PLOT",
|
|
44
|
+
"KIND_MODEL_PREPROCESSING",
|
|
45
|
+
"KIND_SCORING_FOLDS",
|
|
46
|
+
"KIND_SCORING_SUMMARIZATION",
|
|
47
|
+
"MOTIF_SET_FILENAME",
|
|
48
|
+
"NUM_FOLDS",
|
|
49
|
+
"PEAK_INTERVAL_INDEX_FILENAME",
|
|
50
|
+
"PREPROCESSING_MANIFEST_FILENAME",
|
|
51
|
+
"ChromBPNetConfiguration",
|
|
52
|
+
"ChromBPNetInterpretationLayout",
|
|
53
|
+
"ChromBPNetPlugin",
|
|
54
|
+
"ChromBPNetRunInputs",
|
|
55
|
+
"PeaksCompression",
|
|
56
|
+
"ShardKind",
|
|
57
|
+
"WeightsFormat",
|
|
58
|
+
"fold_weights_file",
|
|
59
|
+
"model_dir",
|
|
60
|
+
"motif_set_file",
|
|
61
|
+
"peak_distribution_file",
|
|
62
|
+
"peak_interval_index_file",
|
|
63
|
+
"peaks_file",
|
|
64
|
+
"preprocessing_manifest_file",
|
|
65
|
+
]
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
# Copyright 2026 Riya Sinha
|
|
2
|
+
"""ChromBPNet-owned logical paths for preparation, scoring, and interpretation.
|
|
3
|
+
|
|
4
|
+
Altar core supplies only model-neutral paths (`ModelArtifactLayout` and `ScoringLayout`): the staged genome,
|
|
5
|
+
model-owned resource files, variant inputs, and result tables. The files below are ChromBPNet science: the
|
|
6
|
+
per-fold peak-activity distributions and peak interval index that preparation derives, the matched peak set,
|
|
7
|
+
the TF-MoDISco motif set used for hit calling, and the interpretation run tree. Each model-owned path is built
|
|
8
|
+
through the caller's `ModelArtifactLayout.model_resource_file`, so a scoring plan names exactly the files that
|
|
9
|
+
the preparation plan wrote under the same layout.
|
|
10
|
+
|
|
11
|
+
Staged file names come from declared configuration (`weights_format`, `peaks_compression`), never from a
|
|
12
|
+
resource URI: a content-addressed URI need not carry a file extension.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
import posixpath
|
|
17
|
+
from dataclasses import dataclass
|
|
18
|
+
from typing import TYPE_CHECKING, Literal
|
|
19
|
+
|
|
20
|
+
from altar.models import ModelArtifactLayout
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
if TYPE_CHECKING:
|
|
24
|
+
from collections.abc import Mapping
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
WeightsFormat = Literal["h5"]
|
|
28
|
+
"""Serialized fold-weight formats the ChromBPNet runtime loads. Keras loads a single-file model by suffix."""
|
|
29
|
+
|
|
30
|
+
PeaksCompression = Literal["gzip", "none"]
|
|
31
|
+
"""Compression of the narrowPeak-schema peak set. The runtime infers decompression from the staged suffix."""
|
|
32
|
+
|
|
33
|
+
PEAK_INTERVAL_INDEX_FILENAME = "peaks.intervals.jsonl"
|
|
34
|
+
MOTIF_SET_FILENAME = "motifs.h5"
|
|
35
|
+
PREPROCESSING_MANIFEST_FILENAME = "preprocessing.manifest.json"
|
|
36
|
+
|
|
37
|
+
_WEIGHTS_EXTENSIONS: Mapping[str, str] = {"h5": ".h5"}
|
|
38
|
+
_PEAKS_FILENAMES: Mapping[str, str] = {"gzip": "peaks.bed.gz", "none": "peaks.bed"}
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def fold_weights_file(layout: ModelArtifactLayout, model_id: str, fold: int, weights_format: WeightsFormat) -> str:
|
|
42
|
+
"""Return one fold's staged weights, `models/{model_id}/fold_{fold}_model.h5` for `weights_format="h5"`."""
|
|
43
|
+
return layout.fold_model_file(model_id, fold, _WEIGHTS_EXTENSIONS[weights_format])
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def peaks_file(layout: ModelArtifactLayout, model_id: str, compression: PeaksCompression) -> str:
|
|
47
|
+
"""Return the staged raw peak set, `models/{model_id}/peaks.bed.gz` or `.../peaks.bed`."""
|
|
48
|
+
return layout.model_resource_file(model_id, _PEAKS_FILENAMES[compression])
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def peak_distribution_file(layout: ModelArtifactLayout, model_id: str, fold: int) -> str:
|
|
52
|
+
"""Return one fold's derived peak-activity distribution, `models/{model_id}/fold_{fold}_peak_distribution.npy`."""
|
|
53
|
+
return layout.model_resource_file(model_id, f"fold_{fold}_peak_distribution.npy")
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def peak_interval_index_file(layout: ModelArtifactLayout, model_id: str) -> str:
|
|
57
|
+
"""Return the safe serialized peak interval index used to rebuild the in-memory peak tree."""
|
|
58
|
+
return layout.model_resource_file(model_id, PEAK_INTERVAL_INDEX_FILENAME)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def preprocessing_manifest_file(layout: ModelArtifactLayout, model_id: str) -> str:
|
|
62
|
+
"""Return the optional preparation provenance manifest written beside the derived artifacts."""
|
|
63
|
+
return layout.model_resource_file(model_id, PREPROCESSING_MANIFEST_FILENAME)
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def motif_set_file(layout: ModelArtifactLayout, model_id: str) -> str:
|
|
67
|
+
"""Return the staged TF-MoDISco motif set, `models/{model_id}/motifs.h5`.
|
|
68
|
+
|
|
69
|
+
The motif set is a model-owned configured resource, so it lives beside that model's weights rather than at
|
|
70
|
+
a shared storage-root path that two models could overwrite.
|
|
71
|
+
"""
|
|
72
|
+
return layout.model_resource_file(model_id, MOTIF_SET_FILENAME)
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def model_dir(layout: ModelArtifactLayout, model_id: str) -> str:
|
|
76
|
+
"""Return the directory that holds model-owned files: preparation's `-o` output directory."""
|
|
77
|
+
return posixpath.dirname(peak_interval_index_file(layout, model_id))
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
@dataclass(frozen=True)
|
|
81
|
+
class ChromBPNetInterpretationLayout(ModelArtifactLayout):
|
|
82
|
+
"""Paths for a ChromBPNet interpretation run, one method per file or directory the stage commands name.
|
|
83
|
+
|
|
84
|
+
Inherits the shared genome and model-file paths from `ModelArtifactLayout`. Several outputs are directories
|
|
85
|
+
rather than single files: a fold's contribution scores, the averaged run directory, and the per-allele hit
|
|
86
|
+
directories. The model runtime and finemo write a tree of files into each, so those stages do not declare
|
|
87
|
+
file `Transfer`s the way scoring does.
|
|
88
|
+
"""
|
|
89
|
+
|
|
90
|
+
def run_dir(self, run_id: str, model_id: str) -> str:
|
|
91
|
+
return posixpath.join("interpretations", run_id, model_id)
|
|
92
|
+
|
|
93
|
+
def model_variants_file(self, run_id: str, model_id: str) -> str:
|
|
94
|
+
return posixpath.join(self.run_dir(run_id, model_id), "variants.tsv")
|
|
95
|
+
|
|
96
|
+
def fold_results_dir(self, run_id: str, model_id: str, fold: int) -> str:
|
|
97
|
+
return posixpath.join(self.run_dir(run_id, model_id), f"fold_{fold}")
|
|
98
|
+
|
|
99
|
+
def variants_with_plots_file(self, run_id: str, model_id: str) -> str:
|
|
100
|
+
return posixpath.join(self.run_dir(run_id, model_id), "variants_w_interpret_plots.tsv")
|
|
101
|
+
|
|
102
|
+
def hitcaller_region_file(self, run_id: str, model_id: str, allele: str) -> str:
|
|
103
|
+
"""Return the per-allele finemo input the average step writes, `…/{allele}_finemo_input.npz`."""
|
|
104
|
+
return posixpath.join(self.run_dir(run_id, model_id), f"{allele}_finemo_input.npz")
|
|
105
|
+
|
|
106
|
+
def hits_dir(self, run_id: str, model_id: str, allele: str) -> str:
|
|
107
|
+
"""Return the per-allele finemo call-hits output directory, `…/ref_hits` or `…/alt_hits`."""
|
|
108
|
+
return posixpath.join(self.run_dir(run_id, model_id), f"{allele}_hits")
|
|
@@ -0,0 +1,194 @@
|
|
|
1
|
+
# Copyright 2026 Riya Sinha
|
|
2
|
+
"""Helpers for running a ChromBPNet scoring plan on your own machine.
|
|
3
|
+
|
|
4
|
+
`examples/scoring/run.py` and the local run tests both use these, so the logic for turning a plan into
|
|
5
|
+
container submissions — build the plan, submit each shard, fan in the folds, collect the output — lives in
|
|
6
|
+
one place.
|
|
7
|
+
|
|
8
|
+
The helpers need no cloud infrastructure. They combine the reference `ChromBPNetPlugin`, any
|
|
9
|
+
`ExecutionBackend`, and the reference reconcile loop. A caller supplies a backend and a `prep` function that
|
|
10
|
+
maps each planned spec to the spec actually submitted: the identity function for a real run, or `stubify`
|
|
11
|
+
for a run with no container image and no GPU.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
import asyncio
|
|
16
|
+
import dataclasses
|
|
17
|
+
import os
|
|
18
|
+
from collections.abc import Callable
|
|
19
|
+
from pathlib import Path
|
|
20
|
+
|
|
21
|
+
from altar.execution import (
|
|
22
|
+
ContainerTaskSpec,
|
|
23
|
+
ExecutionBackend,
|
|
24
|
+
GroupSpec,
|
|
25
|
+
InMemoryTaskLedger,
|
|
26
|
+
PathResolver,
|
|
27
|
+
ReconcileResult,
|
|
28
|
+
TaskHandle,
|
|
29
|
+
all_folds,
|
|
30
|
+
group_by,
|
|
31
|
+
run_reconcile_loop,
|
|
32
|
+
single,
|
|
33
|
+
)
|
|
34
|
+
from altar.models import (
|
|
35
|
+
ContainerScoringPlan,
|
|
36
|
+
ResourceReference,
|
|
37
|
+
ScoringLayout,
|
|
38
|
+
ScoringRequest,
|
|
39
|
+
)
|
|
40
|
+
from altar.testing import make_stub_sources, source_transfers, stubify
|
|
41
|
+
|
|
42
|
+
from altar_chrombpnet.plugin import (
|
|
43
|
+
KIND_SCORING_FOLDS,
|
|
44
|
+
KIND_SCORING_SUMMARIZATION,
|
|
45
|
+
ChromBPNetConfiguration,
|
|
46
|
+
ChromBPNetPlugin,
|
|
47
|
+
ChromBPNetRunInputs,
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
# Real runs must provide an immutable published digest.
|
|
52
|
+
CHROMBPNET_RUNTIME_IMAGE_ENV = "ALTAR_CHROMBPNET_IMAGE"
|
|
53
|
+
_TEST_CHROMBPNET_RUNTIME_IMAGE = "test/chrombpnet@sha256:" + "a" * 64
|
|
54
|
+
|
|
55
|
+
# Prep maps a planned spec to the spec to submit: identity for a real run, `stubify` for a stub run.
|
|
56
|
+
Prep = Callable[[ContainerTaskSpec], ContainerTaskSpec]
|
|
57
|
+
|
|
58
|
+
__all__ = [
|
|
59
|
+
"CHROMBPNET_RUNTIME_IMAGE_ENV",
|
|
60
|
+
"Prep",
|
|
61
|
+
"build_scoring_plan",
|
|
62
|
+
"drive_scoring_run",
|
|
63
|
+
"make_stub_sources",
|
|
64
|
+
"remap_chrombpnet_runtime_image",
|
|
65
|
+
"resolve_chrombpnet_runtime_image",
|
|
66
|
+
"stubify",
|
|
67
|
+
]
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def resolve_chrombpnet_runtime_image() -> str:
|
|
71
|
+
"""Return the configured ChromBPNet runtime image."""
|
|
72
|
+
image = os.environ.get(CHROMBPNET_RUNTIME_IMAGE_ENV)
|
|
73
|
+
if image is None:
|
|
74
|
+
msg = f"set {CHROMBPNET_RUNTIME_IMAGE_ENV} to an immutable ChromBPNet image digest"
|
|
75
|
+
raise ValueError(msg)
|
|
76
|
+
return image
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def remap_chrombpnet_runtime_image(spec: ContainerTaskSpec, image: str | None) -> ContainerTaskSpec:
|
|
80
|
+
"""Replace the plan's ChromBPNet image while preserving other stage images."""
|
|
81
|
+
if image and spec.kind in {KIND_SCORING_FOLDS, KIND_SCORING_SUMMARIZATION}:
|
|
82
|
+
return dataclasses.replace(spec, image=image)
|
|
83
|
+
return spec
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def build_scoring_plan( # noqa: PLR0913
|
|
87
|
+
resolver: PathResolver,
|
|
88
|
+
*,
|
|
89
|
+
job_id: str = "example-job",
|
|
90
|
+
model_id: str = "example-model",
|
|
91
|
+
genome: str = "hg38",
|
|
92
|
+
accession_pattern: str = "chrombpnet.fold_{fold}.h5",
|
|
93
|
+
image: str = _TEST_CHROMBPNET_RUNTIME_IMAGE,
|
|
94
|
+
) -> ContainerScoringPlan:
|
|
95
|
+
"""Build a five-fold ChromBPNet plan plus its summarize shard against `resolver`'s mount."""
|
|
96
|
+
plugin = ChromBPNetPlugin()
|
|
97
|
+
configuration = ChromBPNetConfiguration(
|
|
98
|
+
weights=tuple(
|
|
99
|
+
ResourceReference(
|
|
100
|
+
uri=accession_pattern.format(fold=fold),
|
|
101
|
+
digest=f"sha256:{fold + 1:064x}",
|
|
102
|
+
)
|
|
103
|
+
for fold in range(5)
|
|
104
|
+
),
|
|
105
|
+
peaks=ResourceReference(uri="peaks.bed.gz", digest=f"sha256:{15:064x}"),
|
|
106
|
+
image=image,
|
|
107
|
+
)
|
|
108
|
+
return plugin.build_scoring_plan(
|
|
109
|
+
ScoringRequest(
|
|
110
|
+
job_id=job_id,
|
|
111
|
+
model_id=model_id,
|
|
112
|
+
configuration=configuration,
|
|
113
|
+
run_inputs=ChromBPNetRunInputs(
|
|
114
|
+
genome_build=genome,
|
|
115
|
+
genome=ResourceReference(
|
|
116
|
+
uri=ScoringLayout().genome_fasta(genome),
|
|
117
|
+
digest=f"sha256:{'f' * 64}",
|
|
118
|
+
),
|
|
119
|
+
),
|
|
120
|
+
output=plugin.default_output,
|
|
121
|
+
resolver=resolver,
|
|
122
|
+
layout=ScoringLayout(),
|
|
123
|
+
)
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
async def drive_scoring_run( # noqa: PLR0913
|
|
128
|
+
backend: ExecutionBackend,
|
|
129
|
+
plan: ContainerScoringPlan,
|
|
130
|
+
*,
|
|
131
|
+
prep: Prep,
|
|
132
|
+
source_dir: str,
|
|
133
|
+
results_dir: str,
|
|
134
|
+
interval_s: float = 1.0,
|
|
135
|
+
max_ticks: int = 180,
|
|
136
|
+
max_in_flight: int | None = None,
|
|
137
|
+
) -> tuple[str, list[ReconcileResult]]:
|
|
138
|
+
"""Run `plan` to completion on `backend` and return `(final_result_host_path, tick_results)`.
|
|
139
|
+
|
|
140
|
+
It stages and submits each shard, runs the label-based fan-in that waits for all five folds before
|
|
141
|
+
submitting summarize, and collects the summarize output. `examples/scoring/run.py` runs this same drive
|
|
142
|
+
with progress printing added.
|
|
143
|
+
|
|
144
|
+
`prep` maps each planned spec to the spec actually submitted: the identity function for a real run, or
|
|
145
|
+
`stubify` for a stub run.
|
|
146
|
+
|
|
147
|
+
`max_in_flight` caps how many fold shards run at once. Without a cap, every shard is submitted as soon as
|
|
148
|
+
it is staged. On a host with a single GPU that oversubscribes the card: each ChromBPNet fold container
|
|
149
|
+
claims most of the GPU, so five running at once fail with `CUDA_ERROR_OUT_OF_MEMORY`. Deciding how much
|
|
150
|
+
compute to use is the caller's job, not the core library's, so this knob lives on the driver. Pass `1` on a
|
|
151
|
+
one-GPU host; leave it `None` where the backend autoscales, such as Modal or a Kubernetes cluster.
|
|
152
|
+
"""
|
|
153
|
+
num_folds = ChromBPNetPlugin().num_folds
|
|
154
|
+
ledger = InMemoryTaskLedger()
|
|
155
|
+
summarize = plan.summarize
|
|
156
|
+
if summarize is None:
|
|
157
|
+
msg = "scoring plan has no summarize shard to fan into"
|
|
158
|
+
raise ValueError(msg)
|
|
159
|
+
|
|
160
|
+
in_flight: list[TaskHandle] = []
|
|
161
|
+
|
|
162
|
+
async def _await_slot() -> None:
|
|
163
|
+
"""Wait until fewer than `max_in_flight` submitted folds are still running."""
|
|
164
|
+
if not max_in_flight:
|
|
165
|
+
return
|
|
166
|
+
while len(in_flight) >= max_in_flight:
|
|
167
|
+
statuses = await backend.poll(in_flight)
|
|
168
|
+
done = {s.handle.external_id for s in statuses if s.state.is_terminal}
|
|
169
|
+
if done:
|
|
170
|
+
in_flight[:] = [h for h in in_flight if h.external_id not in done]
|
|
171
|
+
else:
|
|
172
|
+
await asyncio.sleep(interval_s)
|
|
173
|
+
|
|
174
|
+
for shard in plan.shards:
|
|
175
|
+
await _await_slot()
|
|
176
|
+
await backend.stage_inputs(source_transfers(shard.inputs, source_dir))
|
|
177
|
+
handle = await backend.submit(prep(shard))
|
|
178
|
+
await ledger.record(handle, kind=shard.kind)
|
|
179
|
+
in_flight.append(handle)
|
|
180
|
+
|
|
181
|
+
async def on_folds_ready(_group_key: str, _handles: list[TaskHandle]) -> None:
|
|
182
|
+
handle = await backend.submit(prep(summarize))
|
|
183
|
+
await ledger.record(handle, kind=summarize.kind)
|
|
184
|
+
|
|
185
|
+
async def on_summarize_ready(_group_key: str, _handles: list[TaskHandle]) -> None:
|
|
186
|
+
await backend.collect_outputs(source_transfers(summarize.outputs, results_dir))
|
|
187
|
+
|
|
188
|
+
specs = [
|
|
189
|
+
GroupSpec(KIND_SCORING_FOLDS, group_by("job_id", "model_id"), all_folds(num_folds), on_folds_ready),
|
|
190
|
+
GroupSpec(KIND_SCORING_SUMMARIZATION, group_by("job_id", "model_id"), single(), on_summarize_ready),
|
|
191
|
+
]
|
|
192
|
+
results = await run_reconcile_loop(backend, ledger, specs, interval_s=interval_s, max_ticks=max_ticks)
|
|
193
|
+
final = Path(results_dir) / summarize.outputs[0].logical_path
|
|
194
|
+
return str(final), results
|