Repository: Galileo-Galilei/kedro-mlflow (python)

Issue:
PipelineML breaks on kedro-viz 9.2.0
## Description

PipelineML breaks on kedro-viz 9.2.0

## Context

The new kedro-viz 9.2.0 release heavily refactors graph generation for namespaced pipelines - https://github.com/kedro-org/kedro-viz/pull/1897. It now uses `only_nodes_with_namespace` and `__sub__` pipeline methods, but PipelineML doesn't have them implemented. 

## Steps to Reproduce

1. Have a PipelineML specified in kedro project (by using `pipeline_ml_factory`)
2. Run kedro viz

## Expected Result

Kedro viz should run, and pipeline should display training nodes

## Actual Result

Kedro viz cannot run at all

```
  File "repo/venv/lib/python3.10/site-packages/kedro_viz/data_access/repositories/modular_pipelines.py", line 61, in populate_tree
    sub_pipeline = pipeline.only_nodes_with_namespace(modular_pipeline_id)
  File "repo/venv/lib/python3.10/site-packages/kedro_mlflow/pipeline/pipeline_ml.py", line 168, in only_nodes_with_namespace
    raise NotImplementedError(MSG_NOT_IMPLEMENTED)
NotImplementedError: This method is not implemented because it does not make sense for 'PipelineML'. Manipulate directly the training pipeline and recreate the 'PipelineML' with 'pipeline_ml_factory' factory.

```

```
  File "repo/venv/lib/python3.10/site-packages/kedro_viz/data_access/repositories/modular_pipelines.py", line 62, in populate_tree
    rest_of_the_pipeline = pipeline - sub_pipeline
  File "repo/venv/lib/python3.10/site-packages/kedro_mlflow/pipeline/pipeline_ml.py", line 239, in __sub__
    raise NotImplementedError(MSG_NOT_IMPLEMENTED)
NotImplementedError: This method is not implemented because it does not make sense for 'PipelineML'. Manipulate directly the training pipeline and recreate the 'PipelineML' with 'pipeline_ml_factory' factory.

```


## Workaround

Quick and dirty implementation in PipelineML class:

```py
def only_nodes_with_namespace(
    self, node_namespace: str
) -> "Pipeline":  # pragma: no cover
    return self.training.only_nodes_with_namespace(node_namespace)
```

```py
def __sub__(self, other):  # pragma: no cover
      return self.training - other
```

Can't use `self._turn_pipeline_to_ml(...)` here as the pipelines are not consistent anymore

## Your Environment

<!-- Include as many relevant details about the environment in which you experienced the bug: -->

* `kedro` and `kedro-mlflow` : 0.19.6 and 0.12.2
* Python version used (`python -V`): 3.10.11
* Operating system and version: M2 Sonoma 14.5

## Does the bug also happen with the last version on master?

Probably yes

Pull request:
Allow only_nodes_with_namespace in PipelineML for kedro-viz compatibility
## Description
Close #569

## Development notes

- Add ``only_nodes_with_namespace`` and ``__sub__`` methods for ``PipelineML`` object
- I have tested manually but kedro-viz 9.2.0 does not fail with a simple PipelineML in spaceflights-pandas. I need 


## Checklist

- [X] Read the [contributing](https://github.com/Galileo-Galilei/kedro-mlflow/blob/master/CONTRIBUTING.md) guidelines
- [X] Open this PR as a 'Draft Pull Request' if it is work-in-progress
- [ ] Update the documentation to reflect the code changes
- [X] Add a description of this change and add your name to the list of supporting contributions in the [`CHANGELOG.md`](https://github.com/Galileo-Galilei/kedro-mlflow/blob/master/CHANGELOG.md) file. Please respect [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) guidelines.
- [ ] Add tests to cover your changes

## Notice

- [X] I acknowledge and agree that, by checking this box and clicking "Submit Pull Request":

- I submit this contribution under the [Apache 2.0 license](https://www.apache.org/licenses/LICENSE-2.0.txt) and represent that I am entitled to do so on behalf of myself, my employer, or relevant third parties, as applicable.
- I certify that (a) this contribution is my original creation and / or (b) to the extent it is not my original creation, I am authorised to submit this contribution on behalf of the original creator(s) or their licensees.
- I certify that the use of this contribution as authorised by the Apache 2.0 license does not violate the intellectual property rights of anyone else.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 889cc74..abf386f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,11 @@
 
 -   :memo: :loud_sound:  Various improvements to documentation and logging ([#549](https://github.com/Galileo-Galilei/kedro-mlflow/issues/549), [#554](https://github.com/Galileo-Galilei/kedro-mlflow/issues/554), [#567](https://github.com/Galileo-Galilei/kedro-mlflow/issues/567))
 
+### Fixed
+
+-   :bug: Implement ``only_nodes_with_namespace`` and ``__sub__`` methods in ``PipelineML`` for ``kedro-viz>=9.2.0`` compatibility ([#569](https://github.com/Galileo-Galilei/kedro-mlflow/issues/569))
+
+
 ## [0.12.2] - 2024-04-18
 
 ### Added
diff --git a/kedro_mlflow/pipeline/pipeline_ml.py b/kedro_mlflow/pipeline/pipeline_ml.py
index b8c17a6..bdd6f80 100644
--- a/kedro_mlflow/pipeline/pipeline_ml.py
+++ b/kedro_mlflow/pipeline/pipeline_ml.py
@@ -1,3 +1,4 @@
+from logging import Logger, getLogger
 from typing import Dict, Iterable, Optional, Union
 
 from kedro.pipeline import Pipeline
@@ -10,6 +11,14 @@ MSG_NOT_IMPLEMENTED = (
     "recreate the 'PipelineML' with 'pipeline_ml_factory' factory."
 )
 
+MSG_WARNING_KEDRO_VIZ = (
+    "BEWARE - This 'Pipeline' is no longer a 'PipelineML' object. "
+    "This method is only implemented for compatibility with kedro-viz "
+    "but should never be used directly.\nSee "
+    "https://github.com/Galileo-Galilei/kedro-mlflow/issues/569 "
+    " for more context. "
+)
+
 
 class PipelineML(Pipeline):
     """
@@ -86,6 +95,10 @@ class PipelineML(Pipeline):
         self.log_model_kwargs = {**self.LOG_MODEL_KWARGS_DEFAULT, **log_model_kwargs}
         self._check_consistency()
 
+    @property
+    def _logger(self) -> Logger:
+        return getLogger(__name__)
+
     @property
     def training(self) -> Pipeline:
         return Pipeline(self.nodes)
@@ -165,7 +178,8 @@ class PipelineML(Pipeline):
     def only_nodes_with_namespace(
         self, node_namespace: str
     ) -> "Pipeline":  # pragma: no cover
-        raise NotImplementedError(MSG_NOT_IMPLEMENTED)
+        self._logger.warning(MSG_WARNING_KEDRO_VIZ)
+        return self.training.only_nodes_with_namespace(node_namespace)
 
     def only_nodes_with_inputs(self, *inputs: str) -> "PipelineML":  # pragma: no cover
         raise NotImplementedError(MSG_NOT_IMPLEMENTED)
@@ -235,7 +249,8 @@ class PipelineML(Pipeline):
         raise NotImplementedError(MSG_NOT_IMPLEMENTED)
 
     def __sub__(self, other):  # pragma: no cover
-        raise NotImplementedError(MSG_NOT_IMPLEMENTED)
+        self._logger.warning(MSG_WARNING_KEDRO_VIZ)
+        return self.training - other
 
     def __and__(self, other):  # pragma: no cover
         # kept for compatibility with KedroContext _filter_pipelinefunction


Test patch:
diff --git a/tests/pipeline/test_pipeline_ml.py b/tests/pipeline/test_pipeline_ml.py
index 3774058..a56f8ff 100644
--- a/tests/pipeline/test_pipeline_ml.py
+++ b/tests/pipeline/test_pipeline_ml.py
@@ -1,6 +1,6 @@
 import pytest
 from kedro.io import DataCatalog, MemoryDataset
-from kedro.pipeline import Pipeline, node
+from kedro.pipeline import Pipeline, node, pipeline
 from kedro_datasets.pandas import CSVDataset
 
 from kedro_mlflow.pipeline import pipeline_ml_factory
@@ -119,6 +119,18 @@ def pipeline_ml_with_intermediary_artifacts():
     return pipeline_ml_with_intermediary_artifacts
 
 
+@pytest.fixture
+def pipeline_ml_with_namespace(pipeline_ml_with_tag):
+    pipeline_ml_with_namespace = pipeline_ml_factory(
+        training=pipeline(pipeline_ml_with_tag.training, namespace="tmp_namespace"),
+        inference=pipeline(
+            pipeline_ml_with_tag.inference, inputs={"model": "tmp_namespace.model"}
+        ),
+        input_name=pipeline_ml_with_tag.input_name,
+    )
+    return pipeline_ml_with_namespace
+
+
 @pytest.fixture
 def pipeline_ml_with_inputs_artifacts():
     full_pipeline = Pipeline(
@@ -246,6 +258,52 @@ def catalog_with_parameters():
     return catalog_with_parameters
 
 
+def test_pipeline_ml_only_nodes_with_namespace(
+    caplog,
+    pipeline_ml_with_namespace,
+):
+    """When the pipeline is filtered with only_nodes_with_namespace, we return only the training pipeline. This is for kedro viz compatibility"""
+
+    # pipeline_ml_with_namespace are fixture in conftest
+
+    # remember : the arguments are iterable, so do not pass string directly (e.g ["training"] rather than training)
+
+    filtered_pipeline_ml = pipeline_ml_with_namespace.only_nodes_with_namespace(
+        node_namespace="tmp_namespace"
+    )
+
+    # PipelineML class must be preserved when filtering
+    # inference should be unmodified
+    # training pipeline nodes must be identical to kedro filtering.
+    assert isinstance(filtered_pipeline_ml, Pipeline)
+    assert not isinstance(filtered_pipeline_ml, PipelineML)
+    assert str(filtered_pipeline_ml) == str(pipeline_ml_with_namespace.training)
+    assert "kedro-viz but should never be" in caplog.text
+
+
+def test_pipeline_ml_substraction(
+    caplog,
+    pipeline_ml_with_intermediary_artifacts,
+):
+    """When the pipeline is filtered with only_nodes_with_namespace, we return only the training pipeline. This is for kedro viz compatibility"""
+
+    # pipeline_ml_with_namespace are fixture in conftest
+
+    # remember : the arguments are iterable, so do not pass string directly (e.g ["training"] rather than training)
+
+    filtered_pipeline_ml = (
+        pipeline_ml_with_intermediary_artifacts
+        - pipeline_ml_with_intermediary_artifacts.inference
+    )
+
+    # PipelineML class must be preserved when filtering
+    # inference should be unmodified
+    # training pipeline nodes must be identical to kedro filtering.
+    assert isinstance(filtered_pipeline_ml, Pipeline)
+    assert not isinstance(filtered_pipeline_ml, PipelineML)
+    assert "kedro-viz but should never be" in caplog.text
+
+
 @pytest.mark.parametrize(
     "tags,from_nodes,to_nodes,node_names,from_inputs",
     [