From 99de3db7a5019f2482607beb0bfe4cc03dc00e28 Mon Sep 17 00:00:00 2001 From: Sajid Alam Date: Thu, 1 Feb 2024 10:03:39 +0000 Subject: [PATCH 1/7] Update jupyter.py Signed-off-by: Sajid Alam --- package/kedro_viz/launchers/jupyter.py | 33 ++++++++++++++++++++------ 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/package/kedro_viz/launchers/jupyter.py b/package/kedro_viz/launchers/jupyter.py index 407ee7c3c1..785871f01d 100644 --- a/package/kedro_viz/launchers/jupyter.py +++ b/package/kedro_viz/launchers/jupyter.py @@ -5,6 +5,7 @@ import logging import multiprocessing import os +import shlex import socket from contextlib import closing from typing import Any, Dict @@ -12,6 +13,7 @@ import IPython from IPython.display import HTML, display +from kedro_viz.launchers.cli import viz_cli from kedro_viz.launchers.utils import _check_viz_up, _wait_for from kedro_viz.server import DEFAULT_HOST, DEFAULT_PORT, run_server @@ -78,13 +80,13 @@ def _display_databricks_html(port: int): # pragma: no cover print(f"Kedro-Viz is available at {url}") -def run_viz(port: int = None, local_ns: Dict[str, Any] = None) -> None: +def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: """ Line magic function to start kedro viz. It calls a kedro viz in a process and displays it in the Jupyter notebook environment. Args: - port: TCP port that viz will listen to. Defaults to 4141. + args: String of arguments to pass to Kedro Viz. local_ns: Local namespace with local variables of the scope where the line magic is invoked. This argument must be in the signature, even though it is not used. This is because the Kedro IPython extension registers line magics with @@ -92,15 +94,32 @@ def run_viz(port: int = None, local_ns: Dict[str, Any] = None) -> None: https://ipython.readthedocs.io/en/stable/config/custommagics.html """ - port = port or DEFAULT_PORT # Default argument doesn't work in Jupyter line magic. - host = _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST - port = _allocate_port( - host, start_at=int(port) - ) # Line magic provides string arguments by default, so we need to convert to int. + # Parse the args string + parsed_args = shlex.split(args) + + # Use Click's context to parse + with viz_cli.make_context("viz", parsed_args) as ctx: + cli_args = ctx.params + + run_server_kwargs = { + "host": cli_args.get("host", DEFAULT_HOST), + "port": cli_args.get("port", DEFAULT_PORT), + "pipeline_name": cli_args.get("pipeline"), + "env": cli_args.get("env"), + "autoreload": cli_args.get("autoreload"), + "ignore_plugins": cli_args.get("ignore_plugins"), + "extra_params": cli_args.get("params"), + } + + # Ensure that the port is available + port = run_server_kwargs["port"] + port = _allocate_port(DEFAULT_HOST, start_at=int(port)) if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): _VIZ_PROCESSES[port].terminate() + host = run_server_kwargs["host"] + project_path = ( local_ns["context"].project_path if local_ns is not None and "context" in local_ns From 8861cf6b016608b80a5f4d6b704ce8d040d5a25b Mon Sep 17 00:00:00 2001 From: Sajid Alam Date: Thu, 1 Feb 2024 13:06:01 +0000 Subject: [PATCH 2/7] Update jupyter.py Signed-off-by: Sajid Alam --- package/kedro_viz/launchers/jupyter.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/package/kedro_viz/launchers/jupyter.py b/package/kedro_viz/launchers/jupyter.py index 785871f01d..347aa93b40 100644 --- a/package/kedro_viz/launchers/jupyter.py +++ b/package/kedro_viz/launchers/jupyter.py @@ -102,7 +102,7 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: cli_args = ctx.params run_server_kwargs = { - "host": cli_args.get("host", DEFAULT_HOST), + "host": cli_args.get("host", _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST), "port": cli_args.get("port", DEFAULT_PORT), "pipeline_name": cli_args.get("pipeline"), "env": cli_args.get("env"), @@ -113,13 +113,13 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: # Ensure that the port is available port = run_server_kwargs["port"] + host = run_server_kwargs["host"] + port = _allocate_port(DEFAULT_HOST, start_at=int(port)) if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): _VIZ_PROCESSES[port].terminate() - host = run_server_kwargs["host"] - project_path = ( local_ns["context"].project_path if local_ns is not None and "context" in local_ns From b77ae7bed36f7371a96fb256c054d429dd2f849b Mon Sep 17 00:00:00 2001 From: Sajid Alam Date: Mon, 5 Feb 2024 10:43:46 +0000 Subject: [PATCH 3/7] Update jupyter.py Signed-off-by: Sajid Alam --- package/kedro_viz/launchers/jupyter.py | 40 ++++++++++++-------------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/package/kedro_viz/launchers/jupyter.py b/package/kedro_viz/launchers/jupyter.py index 347aa93b40..4822d8daa7 100644 --- a/package/kedro_viz/launchers/jupyter.py +++ b/package/kedro_viz/launchers/jupyter.py @@ -83,10 +83,10 @@ def _display_databricks_html(port: int): # pragma: no cover def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: """ Line magic function to start kedro viz. It calls a kedro viz in a process and displays it in - the Jupyter notebook environment. + the Jupyter notebook environment. It can be used with or without additional arguments. Args: - args: String of arguments to pass to Kedro Viz. + args: String of arguments to pass to Kedro Viz. If empty, defaults will be used. local_ns: Local namespace with local variables of the scope where the line magic is invoked. This argument must be in the signature, even though it is not used. This is because the Kedro IPython extension registers line magics with @@ -95,15 +95,27 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: """ # Parse the args string - parsed_args = shlex.split(args) + parsed_args = shlex.split(args) if args else [] # Use Click's context to parse with viz_cli.make_context("viz", parsed_args) as ctx: cli_args = ctx.params + # Determine host and port, considering Databricks environment + host = cli_args.get("host", _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST) + port = cli_args.get("port", DEFAULT_PORT) + + # Allocate an available port if necessary + port = _allocate_port(host, start_at=int(port)) + + # Terminate existing process on the same port, if any + if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): + _VIZ_PROCESSES[port].terminate() + + # Set up other parameters run_server_kwargs = { - "host": cli_args.get("host", _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST), - "port": cli_args.get("port", DEFAULT_PORT), + "host": host, + "port": port, "pipeline_name": cli_args.get("pipeline"), "env": cli_args.get("env"), "autoreload": cli_args.get("autoreload"), @@ -111,26 +123,12 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: "extra_params": cli_args.get("params"), } - # Ensure that the port is available - port = run_server_kwargs["port"] - host = run_server_kwargs["host"] - - port = _allocate_port(DEFAULT_HOST, start_at=int(port)) - - if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): - _VIZ_PROCESSES[port].terminate() - - project_path = ( - local_ns["context"].project_path - if local_ns is not None and "context" in local_ns - else None - ) - + # Start Kedro-Viz server in a new process process_context = multiprocessing.get_context("spawn") viz_process = process_context.Process( target=run_server, daemon=True, - kwargs={"project_path": project_path, "host": host, "port": port}, + kwargs=run_server_kwargs, ) viz_process.start() From 85122f5c2006a5f8771fb158d199ef8f29cd54a7 Mon Sep 17 00:00:00 2001 From: Sajid Alam Date: Mon, 5 Feb 2024 15:03:38 +0000 Subject: [PATCH 4/7] manually add logic for args Signed-off-by: Sajid Alam --- package/kedro_viz/launchers/jupyter.py | 50 ++++++++++++++------------ 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/package/kedro_viz/launchers/jupyter.py b/package/kedro_viz/launchers/jupyter.py index 61167bdcc6..0096fe4063 100644 --- a/package/kedro_viz/launchers/jupyter.py +++ b/package/kedro_viz/launchers/jupyter.py @@ -9,6 +9,7 @@ import socket from contextlib import closing from typing import Any, Dict +from pathlib import Path import IPython from IPython.display import HTML, display @@ -82,8 +83,7 @@ def _display_databricks_html(port: int): # pragma: no cover def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: """ - Line magic function to start kedro viz. It calls a kedro viz in a process and displays it in - the Jupyter notebook environment. It can be used with or without additional arguments. + Line magic function to start Kedro Viz with optional arguments. Args: args: String of arguments to pass to Kedro Viz. If empty, defaults will be used. @@ -94,21 +94,24 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: https://ipython.readthedocs.io/en/stable/config/custommagics.html """ - # Parse the args string - parsed_args = shlex.split(args) if args else [] - - # Use Click's context to parse - with viz_cli.make_context("viz", parsed_args) as ctx: - cli_args = ctx.params - - # Determine host and port, considering Databricks environment - host = cli_args.get("host", _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST) - port = cli_args.get("port", DEFAULT_PORT) - - # Allocate an available port if necessary - port = _allocate_port(host, start_at=int(port)) - - # Terminate existing process on the same port, if any + # Parse arguments + parsed_args = shlex.split(args) + arg_dict = {arg.lstrip('-').split('=')[0]: arg.split('=')[1] if '=' in arg else True + for arg in parsed_args} + + host = arg_dict.get("host", _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST) + port = int(arg_dict.get("port", DEFAULT_PORT)) + load_file = arg_dict.get("load-file", None) + save_file = arg_dict.get("save-file", None) + pipeline = arg_dict.get("pipeline", None) + env = arg_dict.get("env", None) + ignore_plugins = arg_dict.get("ignore-plugins", False) + params = arg_dict.get("params", "") + + # Allocate port + port = _allocate_port(host, start_at=port) + + # Terminate existing process if necessary if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): _VIZ_PROCESSES[port].terminate() @@ -116,11 +119,12 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: run_server_kwargs = { "host": host, "port": port, - "pipeline_name": cli_args.get("pipeline"), - "env": cli_args.get("env"), - "autoreload": cli_args.get("autoreload"), - "ignore_plugins": cli_args.get("ignore_plugins"), - "extra_params": cli_args.get("params"), + "load_file": load_file, + "save_file": save_file, + "pipeline_name": pipeline, + "env": env, + "ignore_plugins": ignore_plugins, + "extra_params": params, } # Start Kedro-Viz server in a new process @@ -128,7 +132,7 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: viz_process = process_context.Process( target=run_server, daemon=True, - kwargs=run_server_kwargs, + kwargs=run_server_kwargs ) viz_process.start() From caef1ecced44622fde380da1901f2767e8a839be Mon Sep 17 00:00:00 2001 From: Sajid Alam Date: Mon, 5 Feb 2024 15:27:43 +0000 Subject: [PATCH 5/7] lint Signed-off-by: Sajid Alam --- package/kedro_viz/launchers/jupyter.py | 33 ++++++++++++++++++-------- 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/package/kedro_viz/launchers/jupyter.py b/package/kedro_viz/launchers/jupyter.py index 0096fe4063..94b092e8b1 100644 --- a/package/kedro_viz/launchers/jupyter.py +++ b/package/kedro_viz/launchers/jupyter.py @@ -9,12 +9,10 @@ import socket from contextlib import closing from typing import Any, Dict -from pathlib import Path import IPython from IPython.display import HTML, display -from kedro_viz.launchers.cli import viz_cli from kedro_viz.launchers.utils import _check_viz_up, _wait_for from kedro_viz.server import DEFAULT_HOST, DEFAULT_PORT, run_server @@ -81,7 +79,19 @@ def _display_databricks_html(port: int): # pragma: no cover print(f"Kedro-Viz is available at {url}") -def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: +def parse_args(args): # pragma: no cover + """Parse the argument string and return a dictionary of arguments.""" + parsed_args = shlex.split(args) + arg_dict = { + arg.lstrip("-").split("=")[0]: arg.split("=")[1] if "=" in arg else True + for arg in parsed_args + } + return arg_dict + + +def run_viz( # pylint: disable=too-many-locals + args: str = "", local_ns: Dict[str, Any] = None +) -> None: """ Line magic function to start Kedro Viz with optional arguments. @@ -94,10 +104,8 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: https://ipython.readthedocs.io/en/stable/config/custommagics.html """ - # Parse arguments - parsed_args = shlex.split(args) - arg_dict = {arg.lstrip('-').split('=')[0]: arg.split('=')[1] if '=' in arg else True - for arg in parsed_args} + # Parse arguments using the new function + arg_dict = parse_args(args) host = arg_dict.get("host", _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST) port = int(arg_dict.get("port", DEFAULT_PORT)) @@ -115,6 +123,12 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): _VIZ_PROCESSES[port].terminate() + project_path = ( + local_ns["context"].project_path + if local_ns is not None and "context" in local_ns + else None + ) + # Set up other parameters run_server_kwargs = { "host": host, @@ -125,14 +139,13 @@ def run_viz(args: str = "", local_ns: Dict[str, Any] = None) -> None: "env": env, "ignore_plugins": ignore_plugins, "extra_params": params, + "project_path": project_path, } # Start Kedro-Viz server in a new process process_context = multiprocessing.get_context("spawn") viz_process = process_context.Process( - target=run_server, - daemon=True, - kwargs=run_server_kwargs + target=run_server, daemon=True, kwargs=run_server_kwargs ) viz_process.start() From bdb3bcb7468227cd59822e0689a1b15edf7c503f Mon Sep 17 00:00:00 2001 From: Sajid Alam Date: Tue, 6 Feb 2024 15:34:36 +0000 Subject: [PATCH 6/7] fix tests Signed-off-by: Sajid Alam --- package/kedro_viz/launchers/jupyter.py | 30 ++++++++++++++------ package/tests/test_launchers/test_jupyter.py | 23 ++++++++++++++- 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/package/kedro_viz/launchers/jupyter.py b/package/kedro_viz/launchers/jupyter.py index 94b092e8b1..5c9da6dcd6 100644 --- a/package/kedro_viz/launchers/jupyter.py +++ b/package/kedro_viz/launchers/jupyter.py @@ -12,6 +12,7 @@ import IPython from IPython.display import HTML, display +from watchgod import RegExpWatcher, run_process from kedro_viz.launchers.utils import _check_viz_up, _wait_for from kedro_viz.server import DEFAULT_HOST, DEFAULT_PORT, run_server @@ -80,7 +81,7 @@ def _display_databricks_html(port: int): # pragma: no cover def parse_args(args): # pragma: no cover - """Parse the argument string and return a dictionary of arguments.""" + """Parses the args string and returns a dictionary of arguments.""" parsed_args = shlex.split(args) arg_dict = { arg.lstrip("-").split("=")[0]: arg.split("=")[1] if "=" in arg else True @@ -104,7 +105,7 @@ def run_viz( # pylint: disable=too-many-locals https://ipython.readthedocs.io/en/stable/config/custommagics.html """ - # Parse arguments using the new function + # Parse arguments arg_dict = parse_args(args) host = arg_dict.get("host", _DATABRICKS_HOST if _is_databricks() else DEFAULT_HOST) @@ -113,13 +114,14 @@ def run_viz( # pylint: disable=too-many-locals save_file = arg_dict.get("save-file", None) pipeline = arg_dict.get("pipeline", None) env = arg_dict.get("env", None) + autoreload = arg_dict.get("autoreload", False) ignore_plugins = arg_dict.get("ignore-plugins", False) params = arg_dict.get("params", "") # Allocate port port = _allocate_port(host, start_at=port) - # Terminate existing process if necessary + # Terminate existing process if needed if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): _VIZ_PROCESSES[port].terminate() @@ -129,7 +131,6 @@ def run_viz( # pylint: disable=too-many-locals else None ) - # Set up other parameters run_server_kwargs = { "host": host, "port": port, @@ -137,16 +138,27 @@ def run_viz( # pylint: disable=too-many-locals "save_file": save_file, "pipeline_name": pipeline, "env": env, + "autoreload": autoreload, "ignore_plugins": ignore_plugins, "extra_params": params, "project_path": project_path, } - - # Start Kedro-Viz server in a new process process_context = multiprocessing.get_context("spawn") - viz_process = process_context.Process( - target=run_server, daemon=True, kwargs=run_server_kwargs - ) + if autoreload: + run_process_kwargs = { + "path": project_path, + "target": run_server, + "kwargs": run_server_kwargs, + "watcher_cls": RegExpWatcher, + "watcher_kwargs": {"re_files": r"^.*(\.yml|\.yaml|\.py|\.json)$"}, + } + viz_process = process_context.Process( + target=run_process, daemon=False, kwargs={**run_process_kwargs} + ) + else: + viz_process = process_context.Process( + target=run_server, daemon=True, kwargs={**run_server_kwargs} + ) viz_process.start() _VIZ_PROCESSES[port] = viz_process diff --git a/package/tests/test_launchers/test_jupyter.py b/package/tests/test_launchers/test_jupyter.py index 60447a4550..02128424de 100644 --- a/package/tests/test_launchers/test_jupyter.py +++ b/package/tests/test_launchers/test_jupyter.py @@ -29,6 +29,13 @@ def test_run_viz(self, mocker, patched_check_viz_up): "project_path": None, "host": "127.0.0.1", "port": 4141, + "load_file": None, + "save_file": None, + "pipeline_name": None, + "env": None, + "autoreload": False, + "ignore_plugins": False, + "extra_params": "", }, ) mock_jupyter_display.assert_called_once() @@ -46,13 +53,20 @@ def test_run_viz(self, mocker, patched_check_viz_up): "project_path": None, "host": "127.0.0.1", "port": 4141, + "load_file": None, + "save_file": None, + "pipeline_name": None, + "env": None, + "autoreload": False, + "ignore_plugins": False, + "extra_params": "", }, ) assert set(_VIZ_PROCESSES.keys()) == {4141} def test_run_viz_invalid_port(self, mocker, patched_check_viz_up): with pytest.raises(ValueError): - run_viz(port=999999) + run_viz("--port=999999") def test_exception_when_viz_cannot_be_launched(self, mocker): mocker.patch( @@ -85,6 +99,13 @@ def test_run_viz_on_databricks(self, mocker, patched_check_viz_up, monkeypatch): "project_path": None, "host": "0.0.0.0", "port": 4141, + "load_file": None, + "save_file": None, + "pipeline_name": None, + "env": None, + "autoreload": False, + "ignore_plugins": False, + "extra_params": "", }, ) databricks_display.assert_called_once() From 98f02822d9a4aa73911332ffe7963f1ab4b6461d Mon Sep 17 00:00:00 2001 From: Sajid Alam Date: Tue, 6 Feb 2024 15:57:35 +0000 Subject: [PATCH 7/7] coverage Signed-off-by: Sajid Alam --- package/tests/test_launchers/test_jupyter.py | 26 ++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/package/tests/test_launchers/test_jupyter.py b/package/tests/test_launchers/test_jupyter.py index 02128424de..6a13309713 100644 --- a/package/tests/test_launchers/test_jupyter.py +++ b/package/tests/test_launchers/test_jupyter.py @@ -125,3 +125,29 @@ def test_run_viz_creates_correct_link(self, mocker, patched_check_viz_up): displayed_html = mock_display_html.call_args[0][0].data assert 'target="_blank"' in displayed_html assert "Open Kedro-Viz" in displayed_html + + def test_run_viz_with_autoreload(self, mocker, patched_check_viz_up): + mock_process_context = mocker.patch("multiprocessing.get_context") + mock_context_instance = mocker.Mock() + mock_process_context.return_value = mock_context_instance + mock_process = mocker.patch.object(mock_context_instance, "Process") + + run_viz("--autoreload", None) + + mock_process.assert_called_once_with( + target=mocker.ANY, + daemon=False, # No daemon for autoreload + kwargs=mocker.ANY, + ) + + def test_run_viz_without_autoreload(self, mocker, patched_check_viz_up): + mock_process_context = mocker.patch("multiprocessing.get_context") + mock_context_instance = mocker.Mock() + mock_process_context.return_value = mock_context_instance + mock_process = mocker.patch.object(mock_context_instance, "Process") + + run_viz("", None) + + mock_process.assert_called_once_with( + target=run_server, daemon=True, kwargs=mocker.ANY + )