diff --git a/skore/src/skore/cli/cli.py b/skore/src/skore/cli/cli.py index 3d59bc05a..d5bdf531e 100644 --- a/skore/src/skore/cli/cli.py +++ b/skore/src/skore/cli/cli.py @@ -6,7 +6,7 @@ from skore.cli.color_format import ColorArgumentParser from skore.project import open from skore.project._create import _create -from skore.project._launch import _launch +from skore.project._launch import _cleanup_potential_zombie_process, _launch from skore.project._load import _load @@ -45,6 +45,15 @@ def cli(args: list[str]): "project_name", help="the name or path of the project to open", ) + parser_launch.add_argument( + "--keep-alive", + action=argparse.BooleanOptionalAction, + help=( + "whether to keep the server alive once the main process finishes " + "(default: %(default)s)" + ), + default="auto", + ) parser_launch.add_argument( "--port", type=int, @@ -66,6 +75,16 @@ def cli(args: list[str]): help="increase logging verbosity", ) + # cleanup potential zombie processes + parser_cleanup = subparsers.add_parser( + "cleanup", help="Clean up all UI running processes" + ) + parser_cleanup.add_argument( + "--verbose", + action="store_true", + help="increase logging verbosity", + ) + # open a skore project parser_open = subparsers.add_parser( "open", help='Create a "project.skore" file and start the UI' @@ -93,6 +112,15 @@ def cli(args: list[str]): help=("whether to serve the project (default: %(default)s)"), default=True, ) + parser_open.add_argument( + "--keep-alive", + action=argparse.BooleanOptionalAction, + help=( + "whether to keep the server alive once the main process finishes " + "(default: %(default)s)" + ), + default="auto", + ) parser_open.add_argument( "--port", type=int, @@ -116,6 +144,7 @@ def cli(args: list[str]): elif parsed_args.subcommand == "launch": _launch( project=_load(project_name=parsed_args.project_name), + keep_alive=parsed_args.keep_alive, port=parsed_args.port, open_browser=parsed_args.open_browser, verbose=parsed_args.verbose, @@ -126,8 +155,11 @@ def cli(args: list[str]): create=parsed_args.create, overwrite=parsed_args.overwrite, serve=parsed_args.serve, + keep_alive=parsed_args.keep_alive, port=parsed_args.port, verbose=parsed_args.verbose, ) + elif parsed_args.subcommand == "cleanup": + _cleanup_potential_zombie_process() else: parser.print_help() diff --git a/skore/src/skore/cli/color_format.py b/skore/src/skore/cli/color_format.py index 6a3e3a960..eeba9994c 100644 --- a/skore/src/skore/cli/color_format.py +++ b/skore/src/skore/cli/color_format.py @@ -79,7 +79,7 @@ def format_help(self): # Color the subcommands in cyan help_text = re.sub( - r"(?<=\s)(launch|create|open)(?=\s+)", + r"(?<=\s)(launch|create|open|cleanup)(?=\s+)", r"[cyan bold]\1[/cyan bold]", help_text, ) diff --git a/skore/src/skore/project/_launch.py b/skore/src/skore/project/_launch.py index 2f586eb4f..40c60f4e3 100644 --- a/skore/src/skore/project/_launch.py +++ b/skore/src/skore/project/_launch.py @@ -17,6 +17,7 @@ from fastapi import FastAPI from skore.project.project import Project, logger +from skore.utils._environment import is_environment_notebook_like from skore.utils._logger import logger_context @@ -214,8 +215,27 @@ def cleanup_server(project: Project, timeout: float = 5.0) -> bool: return True +def _cleanup_potential_zombie_process(): + state_path = platformdirs.user_state_path(appname="skore") + for pid_file in state_path.glob("skore-server-*.json"): + try: + pid_info = json.load(pid_file.open()) + try: + process = psutil.Process(pid_info["pid"]) + process.terminate() + try: + process.wait(timeout=5.0) + except psutil.TimeoutExpired: + process.kill() + except psutil.NoSuchProcess: + pass + finally: + pid_file.unlink() + + def _launch( project: Project, + keep_alive: Union[str, bool] = "auto", port: Union[int, None] = None, open_browser: bool = True, verbose: bool = False, @@ -224,14 +244,20 @@ def _launch( Parameters ---------- - project : Project - The project to launch. - port : int, optional - The port to use for the server, by default None. - open_browser : bool, optional - Whether to open the browser, by default True. - verbose : bool, optional - Whether to print verbose output, by default False. + project : Project + The project to launch. + keep_alive : Union[str, bool], default="auto" + Whether to keep the server alive once the main process finishes. When False, + the server will be terminated when the main process finishes. When True, + the server will be kept alive and thus block the main process from exiting. + When `"auto"`, the server will be kept alive if the current execution context + is not a notebook-like environment. + port : int, optional + The port to use for the server, by default None. + open_browser : bool, optional + Whether to open the browser, by default True. + verbose : bool, optional + Whether to print verbose output, by default False. """ from skore import console # avoid circular import @@ -253,11 +279,13 @@ def _launch( ready_event = multiprocessing.Event() + daemon = is_environment_notebook_like() if keep_alive == "auto" else not keep_alive + with logger_context(logger, verbose): process = multiprocessing.Process( target=run_server, args=(project, port, open_browser, ready_event), - daemon=True, + daemon=daemon, ) process.start() project._server_info = ServerInfo(project, port, process.pid) @@ -265,8 +293,16 @@ def _launch( ready_event.wait() # wait for server to been started console.rule("[bold cyan]skore-UI[/bold cyan]") - console.print( - f"Running skore UI from '{project.name}' at URL http://localhost:{port}" - ) - atexit.register(cleanup_server, project) + msg = f"Running skore UI from '{project.name}' at URL http://localhost:{port}" + if not daemon: + msg += " (Press CTRL+C to quit)" + console.print(msg) + + if not daemon: + try: + process.join() + except KeyboardInterrupt: + cleanup_server(project) + else: + atexit.register(cleanup_server, project) diff --git a/skore/src/skore/project/_open.py b/skore/src/skore/project/_open.py index 931fff794..3bd14ba19 100644 --- a/skore/src/skore/project/_open.py +++ b/skore/src/skore/project/_open.py @@ -12,6 +12,7 @@ def open( create: bool = True, overwrite: bool = False, serve: bool = True, + keep_alive: Union[str, bool] = "auto", port: Union[int, None] = None, verbose: bool = False, ) -> Project: @@ -34,6 +35,12 @@ def open( Has no effect otherwise. serve: bool, default=True Whether to launch the skore UI. + keep_alive : Union[str, bool], default="auto" + Whether to keep the server alive once the main process finishes. When False, + the server will be terminated when the main process finishes. When True, + the server will be kept alive and thus block the main process from exiting. + When `"auto"`, the server will be kept alive if the current execution context + is not a notebook-like environment. port: int, default=None Port at which to bind the UI server. If ``None``, the server will be bound to an available port. @@ -68,6 +75,6 @@ def open( raise if serve: - _launch(project, port=port, verbose=verbose) + _launch(project, keep_alive=keep_alive, port=port, verbose=verbose) return project diff --git a/skore/src/skore/utils/_environment.py b/skore/src/skore/utils/_environment.py new file mode 100644 index 000000000..cd91d72fd --- /dev/null +++ b/skore/src/skore/utils/_environment.py @@ -0,0 +1,52 @@ +import os +import sys + + +def get_environment_info(): + """Detect the current Python execution environment. + + Returns a dictionary with information about the environment. + """ + env_info = { + "is_jupyter": False, + "is_vscode": False, + "is_interactive": False, + "environment_name": "standard_python", + "details": {}, + } + + # Check for interactive mode + env_info["is_interactive"] = hasattr(sys, "ps1") + + # Check for Jupyter + try: + shell = get_ipython().__class__.__name__ # type: ignore + env_info["details"]["ipython_shell"] = shell + + if shell == "ZMQInteractiveShell": # Jupyter notebook/lab + env_info["is_jupyter"] = True + env_info["environment_name"] = "jupyter" + elif shell == "TerminalInteractiveShell": # IPython terminal + env_info["environment_name"] = "ipython_terminal" + except NameError: + pass + + # Check for VSCode + if "VSCODE_PID" in os.environ: + env_info["is_vscode"] = True + if env_info["is_interactive"]: + env_info["environment_name"] = "vscode_interactive" + else: + env_info["environment_name"] = "vscode_script" + + # Add additional environment details + env_info["details"]["python_executable"] = sys.executable + env_info["details"]["python_version"] = sys.version + + return env_info + + +def is_environment_notebook_like() -> bool: + """Return `True` if the execution context dcan render HTML. `False` otherwise.""" + info = get_environment_info() + return info["is_vscode"] or info["is_jupyter"] diff --git a/skore/tests/unit/cli/test_cli.py b/skore/tests/unit/cli/test_cli.py index 5af315ac6..94d6f9465 100644 --- a/skore/tests/unit/cli/test_cli.py +++ b/skore/tests/unit/cli/test_cli.py @@ -35,6 +35,7 @@ def test_cli_launch(tmp_project_path): [ "launch", str(tmp_project_path), + "--no-keep-alive", "--no-open-browser", "--verbose", ] @@ -44,7 +45,9 @@ def test_cli_launch(tmp_project_path): def test_cli_launch_no_project_name(): with pytest.raises(SystemExit): - cli(["launch", "--port", 0, "--no-open-browser", "--verbose"]) + cli( + ["launch", "--port", 0, "--no-keep-alive", "--no-open-browser", "--verbose"] + ) def test_cli_create(tmp_path): @@ -82,9 +85,9 @@ def test_cli_open(tmp_path, monkeypatch): assert (project_path / "items").exists() assert (project_path / "views").exists() - cli(["open", str(project_path), "--verbose"]) + cli(["open", str(project_path), "--no-keep-alive", "--verbose"]) close_project(project_path) - cli(["open", str(project_path), "--no-create"]) + cli(["open", str(project_path), "--no-keep-alive", "--no-create"]) close_project(project_path) with pytest.raises(FileNotFoundError): diff --git a/skore/tests/unit/cli/test_open.py b/skore/tests/unit/cli/test_open.py index b67a4ad19..ed74b99c1 100644 --- a/skore/tests/unit/cli/test_open.py +++ b/skore/tests/unit/cli/test_open.py @@ -20,8 +20,10 @@ def mock_launch(monkeypatch): """Fixture that mocks the _launch function and tracks calls to it.""" calls = [] - def _mock_launch(project, port=None, open_browser=True, verbose=False): - calls.append((project, port, open_browser, verbose)) + def _mock_launch( + project, keep_alive="auto", port=None, open_browser=True, verbose=False + ): + calls.append((project, keep_alive, port, open_browser, verbose)) monkeypatch.setattr("skore.project._launch._launch", _mock_launch) return calls @@ -38,6 +40,7 @@ def test_cli_open(tmp_project_path, mock_launch): "--port", "8000", "--serve", + "--no-keep-alive", "--verbose", ] ) @@ -49,7 +52,7 @@ def test_cli_open_creates_project(tmp_path, mock_launch): project_path = tmp_path / "new_project.skore" assert not project_path.exists() - cli(["open", str(project_path), "--create"]) + cli(["open", str(project_path), "--create", "--no-keep-alive"]) assert project_path.exists() assert len(mock_launch) == 1 @@ -59,7 +62,7 @@ def test_cli_open_no_create_fails(tmp_path, mock_launch): project_path = tmp_path / "nonexistent.skore" with pytest.raises(FileNotFoundError): - cli(["open", str(project_path), "--no-create"]) + cli(["open", str(project_path), "--no-create", "--no-keep-alive"]) assert len(mock_launch) == 0 @@ -67,10 +70,10 @@ def test_cli_open_overwrite(tmp_path, mock_launch): """Test that CLI open can overwrite existing project.""" project_path = tmp_path / "overwrite_test.skore" - cli(["open", str(project_path), "--create"]) + cli(["open", str(project_path), "--create", "--no-keep-alive"]) initial_time = os.path.getmtime(project_path) - cli(["open", str(project_path), "--create", "--overwrite"]) + cli(["open", str(project_path), "--create", "--overwrite", "--no-keep-alive"]) new_time = os.path.getmtime(project_path) assert new_time > initial_time assert len(mock_launch) == 2 diff --git a/skore/tests/unit/project/test_launch.py b/skore/tests/unit/project/test_launch.py index 1e2174e40..18229656a 100644 --- a/skore/tests/unit/project/test_launch.py +++ b/skore/tests/unit/project/test_launch.py @@ -50,7 +50,7 @@ def test_server_info(on_disk_project): def test_launch(capsys, tmp_path): """Check the general behaviour of the launch function.""" skore_project = _create(tmp_path / "test_project") - _launch(skore_project, open_browser=False, verbose=True) + _launch(skore_project, open_browser=False, keep_alive=False, verbose=True) assert "Running skore UI" in capsys.readouterr().out assert skore_project._server_info is not None @@ -64,8 +64,12 @@ def test_launch(capsys, tmp_path): output = capsys.readouterr().out assert "Server that was running" in output - _launch(skore_project, port=8000, open_browser=False, verbose=True) - _launch(skore_project, port=8000, open_browser=False, verbose=True) + _launch( + skore_project, port=8000, open_browser=False, keep_alive=False, verbose=True + ) + _launch( + skore_project, port=8000, open_browser=False, keep_alive=False, verbose=True + ) assert "Server is already running" in capsys.readouterr().out @@ -139,7 +143,9 @@ def mock_kill(pid, signal): monkeypatch.setattr(os, "kill", mock_kill) - _launch(skore_project, port=8001, open_browser=False, verbose=True) + _launch( + skore_project, port=8001, open_browser=False, keep_alive=False, verbose=True + ) assert skore_project._server_info is not None assert skore_project._server_info.port == 8001