Merge pull request #110 from swisstopo/repo-in-mlflow

Tag git branch and last commit in MLFlow
swisstopo · Jan 22, 2025 · b5ef183 · b5ef183 · github-actions · Jan 22, 2025
2 parents e7f76aa + b823ca0
commit b5ef183
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -45,7 +45,8 @@ lint = [
     "pre-commit==3.6.2",
 ]
 experiment-tracking = [
-    "mlflow==2.11.1"
+    "mlflow==2.11.1",
+    "pygit2"
 ]
 visualize = [
     "matplotlib==3.8.0"

diff --git a/src/stratigraphy/main.py b/src/stratigraphy/main.py
@@ -29,6 +29,7 @@
 mlflow_tracking = os.getenv("MLFLOW_TRACKING") == "True"  # Checks whether MLFlow tracking is enabled
 if mlflow_tracking:
     import mlflow
+    import pygit2
 
 logging.basicConfig(format="%(asctime)s %(levelname)-8s %(message)s", level=logging.INFO, datefmt="%Y-%m-%d %H:%M:%S")
 logger = logging.getLogger(__name__)
@@ -147,7 +148,7 @@ def setup_mlflow_tracking(
     out_directory: Path = None,
     predictions_path: Path = None,
     metadata_path: Path = None,
-    experiment_name: str = "Boreholes Stratigraphy",
+    experiment_name: str = "Boreholes data extraction",
 ):
     """Set up MLFlow tracking."""
     mlflow.set_experiment(experiment_name)
@@ -163,6 +164,12 @@ def setup_mlflow_tracking(
     mlflow.log_params(flatten(line_detection_params))
     mlflow.log_params(flatten(matching_params))
 
+    repo = pygit2.Repository(".")
+    commit = repo[repo.head.target]
+    mlflow.set_tag("git_branch", repo.head.shorthand)
+    mlflow.set_tag("git_commit_message", commit.message)
+    mlflow.set_tag("git_commit_sha", commit.id)
+
 
 def start_pipeline(
     input_directory: Path,
File	Stmts	Miss	Cover	Missing
src/stratigraphy
__init__.py	8	1	88%	11
extract.py	167	167	0%	3–446
get_files.py	19	19	0%	3–47
main.py	126	126	0%	3–333
src/stratigraphy/benchmark
metrics.py	59	42	29%	22–25, 29–32, 36–39, 46–49, 53–54, 58, 65–74, 78–91, 96–133
src/stratigraphy/data_extractor
data_extractor.py	76	4	95%	32, 45, 120, 164
utility.py	6	3	50%	28–36
src/stratigraphy/depth
a_to_b_interval_extractor.py	37	15	59%	41–60, 79, 92
depthcolumnentry_extractor.py	23	2	91%	45–46
interval.py	101	52	49%	25–28, 33–36, 42, 48, 52, 91–137, 158, 164–180
src/stratigraphy/depths_materials_column_pairs
bounding_boxes.py	30	10	67%	23, 32, 50, 60, 72–78
material_description_rect_with_sidebar.py	18	8	56%	27–40
src/stratigraphy/evaluation
evaluation_dataclasses.py	49	11	78%	52, 71–74, 90, 104, 125–131, 147
groundwater_evaluator.py	48	1	98%	77
layer_evaluator.py	66	46	30%	29–30, 35–39, 47, 69–95, 105–113, 128–149
metadata_evaluator.py	37	14	62%	46–65, 86–93
utility.py	16	7	56%	43–52
src/stratigraphy/groundwater
groundwater_extraction.py	146	90	38%	52, 94, 137–148, 180–184, 199–215, 226–314, 335–363
utility.py	42	36	14%	10–17, 30–50, 62–76, 91–105
src/stratigraphy/layer
layer.py	37	13	65%	26, 29, 37, 52–72
src/stratigraphy/lines
geometric_line_utilities.py	86	2	98%	81, 131
line.py	51	4	92%	25, 50, 60, 110
linesquadtree.py	46	1	98%	75
src/stratigraphy/metadata
coordinate_extraction.py	106	4	96%	29, 93–94, 106
elevation_extraction.py	90	60	33%	34–39, 47, 55, 63, 79–87, 124–138, 150–153, 165–197, 212–220, 234–238
language_detection.py	18	13	28%	17–23, 37–45
metadata.py	66	24	64%	27, 83, 101–127, 146–155, 195–198, 206
src/stratigraphy/sidebar
a_above_b_sidebar.py	96	39	59%	40, 43, 58–66, 71, 78, 99, 104–111, 141–142, 184–225
a_above_b_sidebar_validator.py	41	19	54%	48, 58, 61, 63, 83–86, 111–129, 141–145
a_to_b_sidebar.py	43	14	67%	36, 49–50, 67, 95–108
layer_identifier_sidebar.py	46	29	37%	46–65, 81–97, 109, 122
layer_identifier_sidebar_extractor.py	29	20	31%	31–41, 55–75
sidebar.py	38	3	92%	41, 46, 87
sidebarentry.py	20	3	85%	27, 31, 43
src/stratigraphy/text
description_block_splitter.py	70	2	97%	24, 139
extract_text.py	29	3	90%	19, 53–54
find_description.py	41	8	80%	26–34, 111–114
textblock.py	90	11	88%	22, 27, 39, 44, 71, 79, 104, 116, 139, 160, 189
src/stratigraphy/util
dataclasses.py	32	3	91%	37–39
predictions.py	72	34	53%	72, 95–115, 143–187
util.py	34	12	65%	69–76, 90–92, 116–117
TOTAL	2386	975	59%