fix: lineage to the root respects experiment (#14)

pitercl · web-flow · commit 1bd2f7ec9c12 · 2025-08-25T12:21:16.000Z
* fix: lineage to the root respects experiment

* Add tests

* Fix unit test

* Fix import

* Add fix for fetch_series + test

* Import fix

* Improve series test

* Improve metrics test

* Fix tests
diff --git a/src/neptune_query/internal/composition/fetch_metrics.py b/src/neptune_query/internal/composition/fetch_metrics.py
@@ -171,6 +171,7 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
                         run_attribute_definitions=run_attribute_definitions_split,
                         include_inherited=lineage_to_the_root,
                         include_preview=include_point_previews,
+                        container_type=container_type,
                         step_range=step_range,
                         tail_limit=tail_limit,
                     )
diff --git a/src/neptune_query/internal/composition/fetch_series.py b/src/neptune_query/internal/composition/fetch_series.py
@@ -127,6 +127,7 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
                             client=client,
                             run_attribute_definitions=run_attribute_definitions_split,
                             include_inherited=lineage_to_the_root,
+                            container_type=container_type,
                             step_range=step_range,
                             tail_limit=tail_limit,
                         ),
diff --git a/src/neptune_query/internal/retrieval/metrics.py b/src/neptune_query/internal/retrieval/metrics.py
@@ -33,6 +33,7 @@
     retry,
     util,
 )
+from .search import ContainerType
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +54,7 @@ def fetch_multiple_series_values(
     client: AuthenticatedClient,
     run_attribute_definitions: list[identifiers.RunAttributeDefinition],
     include_inherited: bool,
+    container_type: ContainerType,
     include_preview: bool,
     step_range: tuple[Union[float, None], Union[float, None]] = (None, None),
     tail_limit: Optional[int] = None,
@@ -81,6 +83,7 @@ def fetch_multiple_series_values(
                     },
                     "attribute": run_attribute.attribute_definition.name,
                     "lineage": "FULL" if include_inherited else "NONE",
+                    "lineageEntityType": "EXPERIMENT" if container_type == ContainerType.EXPERIMENT else "RUN",
                     "includePreview": include_preview,
                 },
             }
diff --git a/src/neptune_query/internal/retrieval/series.py b/src/neptune_query/internal/retrieval/series.py
@@ -43,6 +43,7 @@
     File,
     Histogram,
 )
+from .search import ContainerType
 
 SeriesValue = NamedTuple("SeriesValue", [("step", float), ("value", Any), ("timestamp_millis", float)])
 
@@ -51,6 +52,7 @@ def fetch_series_values(
     client: AuthenticatedClient,
     run_attribute_definitions: Iterable[RunAttributeDefinition],
     include_inherited: bool,
+    container_type: ContainerType,
     step_range: Tuple[Union[float, None], Union[float, None]] = (None, None),
     tail_limit: Optional[int] = None,
 ) -> Generator[util.Page[tuple[RunAttributeDefinition, list[SeriesValue]]], None, None]:
@@ -75,6 +77,7 @@ def fetch_series_values(
                     },
                     "attribute": run_definition.attribute_definition.name,
                     "lineage": "FULL" if include_inherited else "NONE",
+                    "lineageEntityType": "EXPERIMENT" if container_type == ContainerType.EXPERIMENT else "RUN",
                 },
             }
             for request_id, run_definition in request_id_to_run_attr_definition.items()
diff --git a/tests/e2e/internal/retrieval/test_series.py b/tests/e2e/internal/retrieval/test_series.py
@@ -3,6 +3,7 @@
 import pytest
 
 from neptune_query.internal.identifiers import AttributeDefinition
+from neptune_query.internal.retrieval.search import ContainerType
 from neptune_query.internal.retrieval.series import (
     RunAttributeDefinition,
     fetch_series_values,
@@ -27,6 +28,7 @@ def test_fetch_series_values_does_not_exist(client, project, experiment_identifi
             client,
             [run_definition],
             include_inherited=False,
+            container_type=ContainerType.EXPERIMENT,
         )
     )
 
@@ -59,6 +61,7 @@ def test_fetch_series_values_single_series(
             client,
             [run_definition],
             include_inherited=False,
+            container_type=ContainerType.EXPERIMENT,
         )
     )
 
@@ -110,7 +113,13 @@ def test_fetch_series_values_single_series_stop_range(
 
     #  when
     series = extract_pages(
-        fetch_series_values(client, [run_definition], include_inherited=False, step_range=step_range)
+        fetch_series_values(
+            client,
+            [run_definition],
+            include_inherited=False,
+            container_type=ContainerType.EXPERIMENT,
+            step_range=step_range,
+        )
     )
 
     # then
@@ -149,7 +158,13 @@ def test_fetch_series_values_single_series_tail_limit(
 
     #  when
     series = extract_pages(
-        fetch_series_values(client, [run_definition], include_inherited=False, tail_limit=tail_limit)
+        fetch_series_values(
+            client,
+            [run_definition],
+            include_inherited=False,
+            container_type=ContainerType.EXPERIMENT,
+            tail_limit=tail_limit,
+        )
     )
 
     # then
diff --git a/tests/e2e/internal/test_split.py b/tests/e2e/internal/test_split.py
@@ -17,6 +17,7 @@
     fetch_attribute_values,
 )
 from neptune_query.internal.retrieval.metrics import fetch_multiple_series_values
+from neptune_query.internal.retrieval.search import ContainerType
 from neptune_query.internal.retrieval.series import (
     SeriesValue,
     fetch_series_values,
@@ -186,7 +187,12 @@ def test_fetch_string_series_values_retrieval(client, project, experiment_identi
     try:
         result = extract_pages(
             fetch_series_values(
-                client, attribute_definitions, include_inherited=True, step_range=(None, None), tail_limit=None
+                client,
+                attribute_definitions,
+                include_inherited=True,
+                container_type=ContainerType.EXPERIMENT,
+                step_range=(None, None),
+                tail_limit=None,
             )
         )
     except (NeptuneRetryError, NeptuneUnexpectedResponseError) as e:
@@ -265,6 +271,7 @@ def test_fetch_float_series_values_retrieval(client, project, experiment_identif
             client,
             attribute_definitions,
             include_inherited=True,
+            container_type=ContainerType.EXPERIMENT,
             include_preview=False,
             step_range=(None, None),
             tail_limit=None,
diff --git a/tests/e2e/v1/generator.py b/tests/e2e/v1/generator.py
@@ -1,5 +1,8 @@
 import concurrent.futures
-from dataclasses import dataclass
+from dataclasses import (
+    dataclass,
+    field,
+)
 from datetime import (
     datetime,
     timedelta,
@@ -23,15 +26,16 @@
 class GeneratedRun:
     custom_run_id: str
     experiment_name: str
-    fork_run_id: Union[str, None]
-    fork_level: Optional[int]
-    fork_point: Optional[int]
-    configs: dict[AttributeName, Union[float, bool, int, str, datetime, list, set, tuple]]
-    metrics: dict[AttributeName, dict[Step, Value]]
-    tags: list[str]
+    fork_run_id: Union[str, None] = None
+    fork_level: Optional[int] = None
+    fork_point: Optional[int] = None
+    configs: dict[AttributeName, Union[float, bool, int, str, datetime, list, set, tuple]] = field(default_factory=dict)
+    metrics: dict[AttributeName, dict[Step, Value]] = field(default_factory=dict)
+    string_series: dict[AttributeName, dict[Step, str]] = field(default_factory=dict)
+    tags: list[str] = field(default_factory=list)
 
     def attributes(self):
-        return set().union(self.configs.keys(), self.metrics.keys())
+        return set().union(self.configs.keys(), self.metrics.keys(), self.string_series.keys())
 
     def metrics_values(self, name: AttributeName) -> list[tuple[Step, Value]]:
         return list(self.metrics[name].items())
@@ -51,9 +55,6 @@ def metrics_values(self, name: AttributeName) -> list[tuple[Step, Value]]:
     GeneratedRun(
         custom_run_id="linear_history_root",
         experiment_name=LINEAR_TREE_EXP_NAME,
-        fork_level=None,
-        fork_point=None,
-        fork_run_id=None,
         tags=["linear_root", "linear"],
         configs={
             "int-value": 1,
@@ -127,9 +128,6 @@ def metrics_values(self, name: AttributeName) -> list[tuple[Step, Value]]:
     GeneratedRun(
         custom_run_id="forked_history_root",
         experiment_name=FORKED_TREE_EXP_NAME,
-        fork_level=None,
-        fork_point=None,
-        fork_run_id=None,
         tags=["forked_history_root", "forked_history"],
         configs={
             "int-value": 1,
@@ -186,7 +184,54 @@ def metrics_values(self, name: AttributeName) -> list[tuple[Step, Value]]:
     ),
 ]
 
-ALL_STATIC_RUNS = LINEAR_HISTORY_TREE + FORKED_HISTORY_TREE
+# Tree structure:
+#
+# multi_experiment_history:
+# root (level: None, experiment: exp_with_multi_experiment_history_1)
+#   └── fork1 (level: 1, fork_point: 4, experiment: exp_with_multi_experiment_history_2)
+#         └── fork2 (level: 2, fork_point: 8, experiment: exp_with_multi_experiment_history_2)
+MULT_EXPERIMENT_HISTORY_EXP_1 = "exp_with_multi_experiment_history_1"
+MULT_EXPERIMENT_HISTORY_EXP_2 = "exp_with_multi_experiment_history_2"
+MULTI_EXPERIMENT_HISTORY = [
+    GeneratedRun(
+        custom_run_id="mult_exp_history_run_1",
+        experiment_name=MULT_EXPERIMENT_HISTORY_EXP_1,
+        metrics={
+            "metrics/m1": {step: step * 0.1 for step in range(0, 5)},
+        },
+        string_series={
+            "string_series/s1": {step: f"val_run1_{step}" for step in range(0, 5)},
+        },
+    ),
+    GeneratedRun(
+        custom_run_id="mult_exp_history_run_2",
+        experiment_name=MULT_EXPERIMENT_HISTORY_EXP_2,
+        fork_level=1,
+        fork_point=4,
+        fork_run_id="mult_exp_history_run_1",
+        metrics={
+            "metrics/m1": {step: step * 0.2 for step in range(5, 9)},
+        },
+        string_series={
+            "string_series/s1": {step: f"val_run2_{step}" for step in range(5, 9)},
+        },
+    ),
+    GeneratedRun(
+        custom_run_id="mult_exp_history_run_3",
+        experiment_name=MULT_EXPERIMENT_HISTORY_EXP_2,
+        fork_level=2,
+        fork_point=8,
+        fork_run_id="mult_exp_history_run_2",
+        metrics={
+            "metrics/m1": {step: step * 0.3 for step in range(9, 12)},
+        },
+        string_series={
+            "string_series/s1": {step: f"val_run3_{step}" for step in range(9, 12)},
+        },
+    ),
+]
+
+ALL_STATIC_RUNS = LINEAR_HISTORY_TREE + FORKED_HISTORY_TREE + MULTI_EXPERIMENT_HISTORY
 RUN_BY_ID = {run.custom_run_id: run for run in ALL_STATIC_RUNS}
 
 
@@ -209,6 +254,10 @@ def log_run(generated: GeneratedRun, api_token: str, e2e_alpha_project: str):
             for step, value in metric_values.items():
                 run.log_metrics(step=step, data={metric_name: value}, timestamp=timestamp_for_step(step))
 
+        for string_series_name, string_series_values in generated.string_series.items():
+            for step, value in string_series_values.items():
+                run.log_string_series(step=step, data={string_series_name: value}, timestamp=timestamp_for_step(step))
+
 
 def log_runs(api_token: str, e2e_alpha_project: str, runs: list[GeneratedRun]):
     max_level = max(run.fork_level or 0 for run in runs)
diff --git a/tests/e2e/v1/test_fetch_metrics.py b/tests/e2e/v1/test_fetch_metrics.py
@@ -34,6 +34,21 @@
     TEST_DATA,
     ExperimentData,
 )
+from tests.e2e.v1.generator import (
+    MULT_EXPERIMENT_HISTORY_EXP_2,
+    timestamp_for_step,
+)
+
+
+def _to_run_attribute_definition(project, run, metric_name):
+    return RunAttributeDefinition(
+        RunIdentifier(ProjectIdentifier(project), SysId(run)),
+        AttributeDefinition(metric_name, "float_series"),
+    )
+
+
+def _to_float_point_value(step, value):
+    return int(timestamp_for_step(step).timestamp() * 1000), step, value, False, 1.0
 
 
 def create_expected_data(
@@ -295,3 +310,39 @@ def test__fetch_metrics_unique__output_format_variants(
     assert result.columns.tolist() == columns
     assert result.index.names == ["experiment", "step"]
     assert {t[0] for t in result.index.tolist()} == filtred_exps
+
+
+@pytest.mark.parametrize(
+    "lineage_to_the_root,expected_values",
+    [
+        (
+            True,
+            [(step, step * 0.1) for step in range(0, 5)]
+            + [(step, step * 0.2) for step in range(5, 9)]
+            + [(step, step * 0.3) for step in range(9, 12)],
+        ),
+        (False, [(step, step * 0.2) for step in range(5, 9)] + [(step, step * 0.3) for step in range(9, 12)]),
+    ],
+)
+def test__fetch_metrics__lineage(new_project_id, lineage_to_the_root, expected_values):
+    df = fetch_metrics(
+        project=new_project_id,
+        experiments=[MULT_EXPERIMENT_HISTORY_EXP_2],
+        attributes=r"metrics/m1",
+        lineage_to_the_root=lineage_to_the_root,
+    )
+
+    expected = create_metrics_dataframe(
+        metrics_data={
+            _to_run_attribute_definition(new_project_id, MULT_EXPERIMENT_HISTORY_EXP_2, "metrics/m1"): [
+                _to_float_point_value(step, value) for step, value in expected_values
+            ]
+        },
+        sys_id_label_mapping={SysId(MULT_EXPERIMENT_HISTORY_EXP_2): MULT_EXPERIMENT_HISTORY_EXP_2},
+        type_suffix_in_column_names=False,
+        include_point_previews=False,
+        timestamp_column_name=None,
+        index_column_name="experiment",
+    )
+
+    pd.testing.assert_frame_equal(df, expected)
diff --git a/tests/e2e/v1/test_fetch_series.py b/tests/e2e/v1/test_fetch_series.py
diff --git a/tests/unit/v1/test_split.py b/tests/unit/v1/test_split.py

Original file line number	Diff line number	Diff line change
`@@ -171,6 +171,7 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:`
`171`	`171`	`run_attribute_definitions=run_attribute_definitions_split,`
`172`	`172`	`include_inherited=lineage_to_the_root,`
`173`	`173`	`include_preview=include_point_previews,`
	`174`	`+ container_type=container_type,`
`174`	`175`	`step_range=step_range,`
`175`	`176`	`tail_limit=tail_limit,`
`176`	`177`	`)`