diff --git a/tests/backend/test_ui_docker_bridge.py b/tests/backend/test_ui_docker_bridge.py deleted file mode 100644 index 4778240e..00000000 --- a/tests/backend/test_ui_docker_bridge.py +++ /dev/null @@ -1,732 +0,0 @@ -import argparse -import contextlib -import io -import os -import sys -import unittest -from pathlib import Path -from unittest.mock import patch, MagicMock - -from weightslab.ui_docker_bridge import ( - _check_docker, - _clean_stale_docker_resources, - _compose_cmd, - _ensure_certificates, - _ensure_scripts_executable, - _generate_certs_with_fallback, - _get_example_dir, - _install_example_requirements, - _make_executable, - _remove_docker_image, - _strip_derived_deploy_env, - _DERIVED_DEPLOY_ENV_VARS, - _FRONTEND_IMAGE, - _STACK_CONTAINERS, - example_start, - main, - ui_launch, - ui_secure_environment, -) - - -class TestCheckDocker(unittest.TestCase): - @patch("weightslab.ui_docker_bridge.shutil.which", return_value=None) - def test_exits_when_docker_not_found(self, _mock_which): - with self.assertRaises(SystemExit) as ctx: - _check_docker() - self.assertEqual(ctx.exception.code, 1) - - @patch( - "weightslab.ui_docker_bridge.subprocess.run", - side_effect=__import__("subprocess").CalledProcessError(1, "docker info"), - ) - @patch("weightslab.ui_docker_bridge.shutil.which", return_value="/usr/bin/docker") - def test_exits_when_daemon_not_running(self, _mock_which, _mock_run): - with self.assertRaises(SystemExit) as ctx: - _check_docker() - self.assertEqual(ctx.exception.code, 1) - - @patch("weightslab.ui_docker_bridge.subprocess.run") - @patch("weightslab.ui_docker_bridge.shutil.which", return_value="/usr/bin/docker") - def test_passes_when_docker_available(self, _mock_which, mock_run): - _check_docker() - mock_run.assert_called_once() - - -class TestComposeCmd(unittest.TestCase): - @patch("weightslab.ui_docker_bridge._compose_base_cmd", return_value=["docker", "compose"]) - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_runs_docker_compose_with_env(self, mock_run, _mock_base): - mock_run.return_value = MagicMock(stdout="", returncode=0) - _compose_cmd("/path/to/compose.yml", "/path/to/envoy.yaml", ["up", "-d"]) - mock_run.assert_called_once() - args, kwargs = mock_run.call_args - self.assertEqual( - args[0], - ["docker", "compose", "-f", "/path/to/compose.yml", "up", "-d"], - ) - self.assertEqual(kwargs["env"]["WS_ENVOY_CONFIG"], "/path/to/envoy.yaml") - self.assertTrue(kwargs["stdout"]) - self.assertTrue(kwargs["text"]) - - @patch("weightslab.ui_docker_bridge._compose_base_cmd", return_value=["docker-compose"]) - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_v1_translates_up_pull_into_pull_then_up(self, mock_run, _mock_base): - """Compose v1 has no `up --pull`; it must `pull` first, then `up` without the flag.""" - mock_run.return_value = MagicMock(stdout="", returncode=0) - _compose_cmd("/c.yml", "/e.yaml", ["up", "-d", "--pull", "always"]) - cmds = [c.args[0] for c in mock_run.call_args_list] - # A separate `docker-compose -f /c.yml pull` ran first. - self.assertIn(["docker-compose", "-f", "/c.yml", "pull"], cmds) - # The final `up` carries no --pull flag (v1 would reject it). - up_cmd = cmds[-1] - self.assertEqual(up_cmd, ["docker-compose", "-f", "/c.yml", "up", "-d"]) - self.assertNotIn("--pull", up_cmd) - - @patch("weightslab.ui_docker_bridge._compose_base_cmd", return_value=None) - def test_exits_when_no_compose_cli(self, _mock_base): - with self.assertRaises(SystemExit) as ctx: - _compose_cmd("/c.yml", "/e.yaml", ["up", "-d"]) - self.assertEqual(ctx.exception.code, 1) - - -class TestComposeDetection(unittest.TestCase): - """Prefer Compose v2 (`docker compose`), fall back to v1 (`docker-compose`).""" - - @staticmethod - def _fake_run(ok_keys): - """Return a subprocess.run stub: rc 0 when ' '.join(cmd[:2]) is in ok_keys.""" - def run(cmd, *a, **k): - return MagicMock(returncode=0 if " ".join(cmd[:2]) in ok_keys else 1) - return run - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_prefers_v2(self, mock_run): - from weightslab.ui_docker_bridge import _detect_compose_cmd - mock_run.side_effect = self._fake_run({"docker compose"}) - self.assertEqual(_detect_compose_cmd(), ["docker", "compose"]) - - @patch("weightslab.ui_docker_bridge.shutil.which", return_value="/usr/bin/docker-compose") - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_falls_back_to_v1_when_v2_absent(self, mock_run, _which): - from weightslab.ui_docker_bridge import _detect_compose_cmd - # v2 probe fails (rc 1); v1 probe (`docker-compose version`) succeeds. - mock_run.side_effect = self._fake_run({"docker-compose version"}) - self.assertEqual(_detect_compose_cmd(), ["docker-compose"]) - - @patch("weightslab.ui_docker_bridge.shutil.which", return_value=None) - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_returns_none_when_neither_available(self, mock_run, _which): - from weightslab.ui_docker_bridge import _detect_compose_cmd - mock_run.side_effect = self._fake_run(set()) - self.assertIsNone(_detect_compose_cmd()) - - -class TestScriptsExecutable(unittest.TestCase): - """Bundled .sh scripts are made executable so users skip the manual `chmod +x`.""" - - @unittest.skipIf(sys.platform == "win32", "execute bit is POSIX-only") - def test_make_executable_adds_exec_bits(self): - import stat as _stat - import tempfile - with tempfile.NamedTemporaryFile(suffix=".sh", delete=False) as f: - path = f.name - try: - os.chmod(path, 0o644) - _make_executable(path) - mode = os.stat(path).st_mode - self.assertTrue(mode & _stat.S_IXUSR) - self.assertTrue(mode & _stat.S_IXGRP) - self.assertTrue(mode & _stat.S_IXOTH) - finally: - os.unlink(path) - - def test_make_executable_is_noop_on_windows(self): - with patch("weightslab.ui_docker_bridge._is_windows", return_value=True): - with patch("weightslab.ui_docker_bridge.os.chmod") as mock_chmod: - _make_executable("/whatever/path.sh") - mock_chmod.assert_not_called() - - def test_make_executable_swallows_oserror(self): - # A non-chmod-able path (e.g. root-owned system install) must not raise. - with patch("weightslab.ui_docker_bridge.os.stat", side_effect=OSError("denied")): - _make_executable("/root/owned.sh") # should not raise - - @unittest.skipIf(sys.platform == "win32", "execute bit is POSIX-only") - def test_ensure_scripts_executable_marks_bundled_scripts(self): - import stat as _stat - from weightslab.ui_docker_bridge import _get_bootstrap_script - _ensure_scripts_executable() - bootstrap = _get_bootstrap_script() - self.assertTrue(bootstrap.exists(), bootstrap) - self.assertTrue(os.stat(bootstrap).st_mode & _stat.S_IXUSR) - - def test_ensure_scripts_executable_noop_on_windows(self): - with patch("weightslab.ui_docker_bridge._is_windows", return_value=True): - with patch("weightslab.ui_docker_bridge._make_executable") as mock_mk: - _ensure_scripts_executable() - mock_mk.assert_not_called() - - -class TestUiLaunch(unittest.TestCase): - """ui_launch: unsecured by default (no cert gen); --certs opts into TLS.""" - - @patch("weightslab.ui_docker_bridge.CertAuthManager") - @patch("weightslab.ui_docker_bridge._get_bootstrap_script", return_value="/does/not/exist.sh") - @patch("weightslab.ui_docker_bridge._run_shell_script", return_value=0) - @patch("weightslab.ui_docker_bridge._ensure_certificates", return_value=True) - @patch("weightslab.ui_docker_bridge._clean_stale_docker_resources") - @patch("weightslab.ui_docker_bridge._compose_cmd") - @patch("weightslab.ui_docker_bridge._check_docker") - @patch("weightslab.ui_docker_bridge._get_envoy_config", return_value="/fake/envoy.yaml") - @patch("weightslab.ui_docker_bridge._get_compose_file", return_value="/fake/docker-compose.yml") - def test_launch_default_no_cert_gen_cleans_and_launches_unsecured( - self, _gc, _ge, mock_check, mock_compose, mock_clean, mock_ensure, - _mock_shell, _gb, mock_mgr, - ): - mgr = MagicMock() - mgr.has_valid_certs.return_value = False # no certs on disk -> unsecured - mock_mgr.from_env_or_default.return_value = mgr - with patch.dict(os.environ, {}, clear=False): - os.environ.pop("VITE_PORT", None) - with self.assertLogs("weightslab.ui_docker_bridge", level="INFO") as log_context: - ui_launch(argparse.Namespace()) - mock_check.assert_called_once() - mock_ensure.assert_not_called() # certs NOT generated by default - mock_clean.assert_called_once() # stale cleanup ran - mock_compose.assert_called_once_with( - "/fake/docker-compose.yml", - "/fake/envoy.yaml", - ["up", "-d", "--pull", "always"], - ) - self.assertTrue(any("http://localhost:5173" in msg for msg in log_context.output)) - self.assertFalse(any("https://localhost:5173" in msg for msg in log_context.output)) - - @patch("weightslab.ui_docker_bridge.CertAuthManager") - @patch("weightslab.ui_docker_bridge._get_bootstrap_script", return_value="/does/not/exist.sh") - @patch("weightslab.ui_docker_bridge._run_shell_script", return_value=0) - @patch("weightslab.ui_docker_bridge._ensure_certificates", return_value=True) - @patch("weightslab.ui_docker_bridge._clean_stale_docker_resources") - @patch("weightslab.ui_docker_bridge._compose_cmd") - @patch("weightslab.ui_docker_bridge._check_docker") - @patch("weightslab.ui_docker_bridge._get_envoy_config", return_value="/fake/envoy.yaml") - @patch("weightslab.ui_docker_bridge._get_compose_file", return_value="/fake/docker-compose.yml") - def test_launch_respects_custom_port( - self, _gc, _ge, _mock_check, _mock_compose, _mock_clean, _mock_ensure, - _mock_shell, _gb, mock_mgr, - ): - mgr = MagicMock(certs_dir="/fake/certs") - mgr.has_valid_certs.return_value = True - mock_mgr.from_env_or_default.return_value = mgr - with patch.dict(os.environ, {"VITE_PORT": "3000"}): - with self.assertLogs("weightslab.ui_docker_bridge", level="INFO") as log_context: - ui_launch(argparse.Namespace(certs=True)) - self.assertTrue(any("https://localhost:3000" in msg for msg in log_context.output)) - - @patch("weightslab.ui_docker_bridge.CertAuthManager") - @patch("weightslab.ui_docker_bridge._get_bootstrap_script", return_value="/does/not/exist.sh") - @patch("weightslab.ui_docker_bridge._run_shell_script", return_value=0) - @patch("weightslab.ui_docker_bridge._ensure_certificates", return_value=True) - @patch("weightslab.ui_docker_bridge._clean_stale_docker_resources") - @patch("weightslab.ui_docker_bridge._compose_cmd") - @patch("weightslab.ui_docker_bridge._check_docker") - @patch("weightslab.ui_docker_bridge._get_envoy_config", return_value="/fake/envoy.yaml") - @patch("weightslab.ui_docker_bridge._get_compose_file", return_value="/fake/docker-compose.yml") - def test_launch_certs_flag_generates_and_runs_secured( - self, _gc, _ge, _mock_check, _mock_compose, _mock_clean, mock_ensure, - _mock_shell, _gb, mock_mgr, - ): - mgr = MagicMock(certs_dir="/fake/certs") - mgr.has_valid_certs.return_value = True # certs present after generation - mock_mgr.from_env_or_default.return_value = mgr - with patch.dict(os.environ, {}, clear=False): - os.environ.pop("VITE_PORT", None) - with self.assertLogs("weightslab.ui_docker_bridge", level="INFO") as log_context: - ui_launch(argparse.Namespace(certs=True)) - mock_ensure.assert_called_once() # --certs generates certs - self.assertTrue(any("https://localhost:5173" in msg for msg in log_context.output)) - - -class TestEnsureCertificates(unittest.TestCase): - """_ensure_certificates only generates files; it never exports TLS/auth env.""" - - @patch("weightslab.ui_docker_bridge._install_ca_trust") - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback") - def test_uses_existing_certs_without_generating(self, mock_gen, _mock_trust): - manager = MagicMock() - # Gate is has_any_credentials(); existing creds short-circuit generation. - manager.has_any_credentials.return_value = True - manager.has_valid_certs.return_value = True - result = _ensure_certificates(manager, force_certs=False) - self.assertTrue(result) - mock_gen.assert_not_called() - manager.get_or_create_auth_token.assert_called_once() - # Derived TLS env must NOT be set here — the deploy pipeline derives it. - manager.setup_tls_environment.assert_not_called() - - @patch("weightslab.ui_docker_bridge._install_ca_trust") - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback", return_value=0) - def test_generates_when_missing_and_forwards_certs_dir(self, mock_gen, _mock_trust): - manager = MagicMock() - # No credentials at the gate -> generate; certs valid afterwards. - manager.has_any_credentials.return_value = False - manager.has_valid_certs.return_value = True - result = _ensure_certificates(manager, force_certs=False) - self.assertTrue(result) - mock_gen.assert_called_once_with(force_certs=False, certs_dir=manager.certs_dir) - manager.setup_tls_environment.assert_not_called() - - @patch("weightslab.ui_docker_bridge._install_ca_trust") - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback", return_value=0) - def test_force_regenerates_even_when_present(self, mock_gen, _mock_trust): - manager = MagicMock() - manager.has_any_credentials.return_value = True - manager.has_valid_certs.return_value = True - _ensure_certificates(manager, force_certs=True) - mock_gen.assert_called_once_with(force_certs=True, certs_dir=manager.certs_dir) - - @patch("weightslab.ui_docker_bridge._install_ca_trust") - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback", return_value=1) - def test_returns_false_on_generation_failure(self, mock_gen, _mock_trust): - manager = MagicMock() - manager.has_any_credentials.return_value = False - manager.has_valid_certs.return_value = False - result = _ensure_certificates(manager, force_certs=False) - self.assertFalse(result) - - -class TestRemoveDockerImage(unittest.TestCase): - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_removes_when_present(self, mock_run): - mock_run.side_effect = [MagicMock(stdout="abc123\nabc123\ndef456\n"), MagicMock()] - _remove_docker_image(_FRONTEND_IMAGE) - self.assertEqual(mock_run.call_count, 2) - rmi_call = mock_run.call_args_list[1].args[0] - self.assertEqual(rmi_call[:3], ["docker", "rmi", "-f"]) - self.assertIn("abc123", rmi_call) - self.assertIn("def456", rmi_call) - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_noop_when_absent(self, mock_run): - mock_run.return_value = MagicMock(stdout="") - _remove_docker_image(_FRONTEND_IMAGE) - mock_run.assert_called_once() # only the 'docker images -q' query, no rmi - - -class TestCleanStaleDockerResources(unittest.TestCase): - @patch("weightslab.ui_docker_bridge._remove_docker_image") - @patch("weightslab.ui_docker_bridge.subprocess.run") - @patch("weightslab.ui_docker_bridge._compose_cmd") - @patch("weightslab.ui_docker_bridge._get_envoy_config", return_value="/fake/envoy.yaml") - @patch("weightslab.ui_docker_bridge._get_compose_file", return_value="/fake/docker-compose.yml") - def test_clean_tears_down_and_removes_image(self, _gc, _ge, mock_compose, mock_run, mock_rmimg): - _clean_stale_docker_resources() - # 1. compose down --remove-orphans --volumes - mock_compose.assert_called_once_with( - "/fake/docker-compose.yml", - "/fake/envoy.yaml", - ["down", "--remove-orphans", "--volumes"], - ) - # 2. docker rm -f for each known stack container - removed = [c.args[0] for c in mock_run.call_args_list] - for container in _STACK_CONTAINERS: - self.assertTrue( - any(call[:3] == ["docker", "rm", "-f"] and container in call for call in removed), - f"expected 'docker rm -f {container}'", - ) - # 3. cached frontend image removed - mock_rmimg.assert_called_once_with(_FRONTEND_IMAGE) - - @patch("weightslab.ui_docker_bridge._remove_docker_image") - @patch("weightslab.ui_docker_bridge.subprocess.run") - @patch("weightslab.ui_docker_bridge._compose_cmd", - side_effect=__import__("subprocess").CalledProcessError(1, "down")) - @patch("weightslab.ui_docker_bridge._get_envoy_config", return_value="/fake/envoy.yaml") - @patch("weightslab.ui_docker_bridge._get_compose_file", return_value="/fake/docker-compose.yml") - def test_clean_tolerates_compose_down_failure(self, _gc, _ge, _mc, _mr, mock_rmimg): - # Should not raise even when 'compose down' returns non-zero (nothing to remove). - _clean_stale_docker_resources() - mock_rmimg.assert_called_once_with(_FRONTEND_IMAGE) - - -class TestUiSecureEnvironment(unittest.TestCase): - @patch("weightslab.ui_docker_bridge.CertAuthManager") - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback", return_value=0) - def test_ui_secure_environment_success(self, mock_gen_certs, mock_cert_manager): - """`weightslab se`: generate certs + token, export WEIGHTSLAB_CERTS_DIR.""" - mock_manager_instance = MagicMock() # certs_dir is a MagicMock (supports .mkdir) - mock_manager_instance.get_or_create_auth_token.return_value = "fake_token" - mock_cert_manager.return_value = mock_manager_instance - - args = argparse.Namespace(force_certs=False) - - with self.assertLogs("weightslab.ui_docker_bridge", level="INFO") as log_context: - ui_secure_environment(args) - - self.assertTrue(any("Certificates generated successfully" in msg for msg in log_context.output)) - self.assertTrue(any("gRPC auth token created" in msg for msg in log_context.output)) - # Generation is pointed at the chosen certs dir (single source of truth). - mock_gen_certs.assert_called_once_with(force_certs=False, certs_dir=mock_manager_instance.certs_dir) - mock_manager_instance.certs_dir.mkdir.assert_called_once() - # se exports WEIGHTSLAB_CERTS_DIR for the process. - self.assertTrue(any("WEIGHTSLAB_CERTS_DIR exported" in msg for msg in log_context.output)) - - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback", return_value=0) - def test_ui_secure_environment_force_certs(self, mock_gen_certs): - """`weightslab se --force-certs` forwards force_certs to generation.""" - args = argparse.Namespace(force_certs=True) - with patch.dict(os.environ, {}, clear=False): - ui_secure_environment(args) - self.assertTrue(mock_gen_certs.call_args.kwargs["force_certs"]) - - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback", return_value=1) - def test_ui_secure_environment_cert_failure(self, mock_gen_certs): - """Certificate generation failure exits non-zero.""" - args = argparse.Namespace(force_certs=False) - with self.assertRaises(SystemExit) as ctx: - ui_secure_environment(args) - self.assertEqual(ctx.exception.code, 1) - - -class TestMainCLI(unittest.TestCase): - """The CLI exposes: se, launch (+ bare start), start example, cli, tunnel, - help; the legacy `ui launch` still dispatches to the same handler.""" - - @patch("weightslab.ui_docker_bridge.ui_secure_environment") - def test_main_dispatches_se(self, mock_se): - with patch("sys.argv", ["weightslab", "se"]): - main() - mock_se.assert_called_once() - - @patch("weightslab.ui_docker_bridge.ui_launch") - def test_main_dispatches_ui_launch(self, mock_launch): - with patch("sys.argv", ["weightslab", "ui", "launch"]): - main() - mock_launch.assert_called_once() - - @patch("weightslab.ui_docker_bridge.example_start") - def test_main_dispatches_start_example(self, mock_example): - with patch("sys.argv", ["weightslab", "start", "example"]): - main() - mock_example.assert_called_once() - - def test_main_ui_without_action_does_not_crash(self): - with patch("sys.argv", ["weightslab", "ui"]): - main() # should print ui help, not raise - - @patch("weightslab.ui_docker_bridge.ui_launch") - def test_main_dispatches_launch(self, mock_launch): - with patch("sys.argv", ["weightslab", "launch"]): - main() - mock_launch.assert_called_once() - - @patch("weightslab.ui_docker_bridge.ui_launch") - def test_main_bare_start_launches_ui(self, mock_launch): - # bare `weightslab start` is an alias for `weightslab launch` - with patch("sys.argv", ["weightslab", "start"]): - main() - mock_launch.assert_called_once() - - def test_main_help_does_not_crash(self): - with patch("sys.argv", ["weightslab", "help"]): - main() # should not raise - - def test_main_no_args_does_not_crash(self): - with patch("sys.argv", ["weightslab"]): - main() # should not raise - - -class TestUserOnboardingFlow(unittest.TestCase): - """Integration-like test of the full onboarding flow, fully hermetic.""" - - @patch("weightslab.ui_docker_bridge.subprocess.run") - @patch("weightslab.ui_docker_bridge._compose_cmd") - @patch("weightslab.ui_docker_bridge._check_docker") - @patch("weightslab.ui_docker_bridge._run_shell_script", return_value=0) - @patch("weightslab.ui_docker_bridge._clean_stale_docker_resources") - @patch("weightslab.ui_docker_bridge._ensure_certificates", return_value=True) - @patch("weightslab.ui_docker_bridge._generate_certs_with_fallback", return_value=0) - @patch("weightslab.ui_docker_bridge.CertAuthManager") - def test_complete_onboarding_workflow( - self, mock_cert_manager, mock_gen, mock_ensure, mock_clean, - mock_shell, mock_check, mock_compose, mock_subproc, - ): - """se -> ui launch -> start example, with no real Docker/daemon/subprocess.""" - mgr = MagicMock() - mock_cert_manager.return_value = mgr - mock_cert_manager.from_env_or_default.return_value = mgr - mock_subproc.return_value = MagicMock(returncode=0) - - with patch.dict(os.environ, {}, clear=False): - try: - ui_secure_environment(argparse.Namespace(force_certs=False)) - ui_launch(argparse.Namespace()) - example_start(argparse.Namespace()) - except Exception as e: - self.fail(f"Onboarding workflow failed: {e}") - - @patch("sys.argv", ["weightslab", "se", "--force-certs"]) - @patch("weightslab.ui_docker_bridge.ui_secure_environment") - def test_cli_se_force_certs(self, mock_se): - main() - mock_se.assert_called_once() - self.assertTrue(mock_se.call_args.args[0].force_certs) - - -class TestBackendConnectionDetection(unittest.TestCase): - """Test backend connection detection utility.""" - - @patch("socket.socket") - def test_backend_connection_success(self, mock_socket_class): - from weightslab.ui_docker_bridge import _test_backend_connection - mock_socket = MagicMock() - mock_socket.connect_ex.return_value = 0 - mock_socket_class.return_value = mock_socket - self.assertTrue(_test_backend_connection()) - - @patch("socket.socket") - def test_backend_connection_failure(self, mock_socket_class): - from weightslab.ui_docker_bridge import _test_backend_connection - mock_socket = MagicMock() - mock_socket.connect_ex.return_value = 1 - mock_socket_class.return_value = mock_socket - self.assertFalse(_test_backend_connection()) - - @patch("socket.socket") - def test_backend_connection_timeout(self, mock_socket_class): - from weightslab.ui_docker_bridge import _test_backend_connection - mock_socket = MagicMock() - mock_socket.connect_ex.side_effect = Exception("Connection timeout") - mock_socket_class.return_value = mock_socket - self.assertFalse(_test_backend_connection()) - - -class TestPathConversion(unittest.TestCase): - """Test Windows path to Git Bash conversion.""" - - def test_windows_path_conversion(self): - from weightslab.ui_docker_bridge import _convert_to_git_bash_path - win_path = r"C:\Users\testuser\.weightslab-certs" - bash_path = _convert_to_git_bash_path(win_path) - self.assertEqual(bash_path, "/mnt/c/Users/testuser/.weightslab-certs") - - def test_unix_path_passthrough(self): - from weightslab.ui_docker_bridge import _convert_to_git_bash_path - unix_path = "/home/testuser/.weightslab-certs" - self.assertEqual(_convert_to_git_bash_path(unix_path), unix_path) - - -class TestSingleSourceOfTruth(unittest.TestCase): - """WEIGHTSLAB_CERTS_DIR is the only input; everything else derives from it.""" - - def test_strip_derived_deploy_env_removes_all(self): - sentinel = {k: "stale" for k in _DERIVED_DEPLOY_ENV_VARS} - with patch.dict(os.environ, sentinel, clear=False): - _strip_derived_deploy_env() - for key in _DERIVED_DEPLOY_ENV_VARS: - self.assertNotIn(key, os.environ, f"{key} should have been stripped") - - @patch("weightslab.ui_docker_bridge._run_shell_script", return_value=0) - def test_generate_certs_forwards_certs_dir(self, mock_shell): - rc = _generate_certs_with_fallback(force_certs=False, certs_dir="/custom/certs") - self.assertEqual(rc, 0) - env_vars = mock_shell.call_args.args[2] - self.assertIsInstance(env_vars, dict) - self.assertIn("WEIGHTSLAB_CERTS_DIR", env_vars) - self.assertIn("custom/certs", env_vars["WEIGHTSLAB_CERTS_DIR"].replace("\\", "/")) - - @patch("weightslab.ui_docker_bridge._run_shell_script", return_value=0) - def test_generate_certs_without_dir_passes_no_env(self, mock_shell): - _generate_certs_with_fallback(force_certs=False) - self.assertIsNone(mock_shell.call_args.args[2]) - - @patch("weightslab.ui_docker_bridge.CertAuthManager") - @patch("weightslab.ui_docker_bridge._get_bootstrap_script", return_value="/does/not/exist.sh") - @patch("weightslab.ui_docker_bridge._ensure_certificates", return_value=True) - @patch("weightslab.ui_docker_bridge._clean_stale_docker_resources") - @patch("weightslab.ui_docker_bridge._compose_cmd") - @patch("weightslab.ui_docker_bridge._check_docker") - @patch("weightslab.ui_docker_bridge._get_envoy_config", return_value="/fake/envoy.yaml") - @patch("weightslab.ui_docker_bridge._get_compose_file", return_value="/fake/docker-compose.yml") - def test_ui_launch_strips_derived_env( - self, _gc, _ge, _mock_check, _mock_compose, _mock_clean, _mock_ensure, _gb, mock_mgr, - ): - mock_mgr.from_env_or_default.return_value = MagicMock() - # Simulate a stale derived env var leaking in (e.g. from import-time check). - with patch.dict(os.environ, {"ENVOY_DOWNSTREAM_TLS": "on", "VITE_SERVER_PROTOCOL": "https"}, clear=False): - ui_launch(argparse.Namespace()) - self.assertNotIn("ENVOY_DOWNSTREAM_TLS", os.environ) - self.assertNotIn("VITE_SERVER_PROTOCOL", os.environ) - - @patch("weightslab.ui_docker_bridge.CertAuthManager") - @patch("weightslab.ui_docker_bridge._get_bootstrap_script", return_value="/does/not/exist.sh") - @patch("weightslab.ui_docker_bridge._ensure_certificates", return_value=True) - @patch("weightslab.ui_docker_bridge._clean_stale_docker_resources") - @patch("weightslab.ui_docker_bridge._compose_cmd") - @patch("weightslab.ui_docker_bridge._check_docker") - @patch("weightslab.ui_docker_bridge._get_envoy_config", return_value="/fake/envoy.yaml") - @patch("weightslab.ui_docker_bridge._get_compose_file", return_value="/fake/docker-compose.yml") - def test_ui_launch_url_https_only_when_certs_present( - self, _gc, _ge, _mock_check, _mock_compose, _mock_clean, _mock_ensure, _gb, mock_mgr, - ): - # Certs absent on disk -> URL must be http (default unsecured launch). - mgr = MagicMock() - mgr.has_valid_certs.return_value = False - mock_mgr.from_env_or_default.return_value = mgr - with patch.dict(os.environ, {}, clear=False): - os.environ.pop("VITE_PORT", None) - with self.assertLogs("weightslab.ui_docker_bridge", level="INFO") as log_context: - ui_launch(argparse.Namespace()) - self.assertTrue(any("http://localhost:5173" in m for m in log_context.output)) - self.assertFalse(any("https://localhost:5173" in m for m in log_context.output)) - - -class TestExampleStart(unittest.TestCase): - """`weightslab start example` runs the bundled classification example.""" - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_example_start_runs_classification(self, mock_run): - mock_run.return_value = MagicMock(returncode=0) - with self.assertLogs("weightslab.ui_docker_bridge", level="INFO") as log_context: - example_start(argparse.Namespace()) - mock_run.assert_called_once() - cmd = mock_run.call_args.args[0] - self.assertEqual(cmd[0], sys.executable) - main_py = cmd[1].replace("\\", "/") - self.assertTrue(main_py.endswith("examples/PyTorch/wl-classification/main.py"), main_py) - cwd = mock_run.call_args.kwargs["cwd"].replace("\\", "/") - self.assertTrue(cwd.endswith("examples/PyTorch/wl-classification"), cwd) - self.assertTrue(any("classification (cls) example" in m for m in log_context.output)) - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_example_start_propagates_nonzero_exit(self, mock_run): - mock_run.return_value = MagicMock(returncode=3) - with self.assertRaises(SystemExit) as ctx: - example_start(argparse.Namespace()) - self.assertEqual(ctx.exception.code, 3) - - @patch("weightslab.ui_docker_bridge._get_example_dir", return_value=Path("/does/not/exist")) - def test_example_start_errors_when_missing(self, _mock_dir): - with self.assertRaises(SystemExit) as ctx: - example_start(argparse.Namespace()) - self.assertEqual(ctx.exception.code, 1) - - def test_example_dir_points_at_bundled_example(self): - # The bundled classification example must actually ship with the package. - self.assertTrue((_get_example_dir("wl-classification") / "main.py").exists()) - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_example_start_seg_runs_segmentation(self, mock_run): - mock_run.return_value = MagicMock(returncode=0) - with self.assertLogs("weightslab.ui_docker_bridge", level="INFO") as log_context: - example_start(argparse.Namespace(example_kind="seg")) - main_py = mock_run.call_args.args[0][1].replace("\\", "/") - self.assertTrue(main_py.endswith("examples/PyTorch/wl-segmentation/main.py"), main_py) - self.assertTrue(any("segmentation (seg) example" in m for m in log_context.output)) - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_example_start_defaults_to_cls_when_flag_absent(self, mock_run): - # A Namespace without example_kind (e.g. older call sites) still defaults to cls. - mock_run.return_value = MagicMock(returncode=0) - example_start(argparse.Namespace()) - main_py = mock_run.call_args.args[0][1].replace("\\", "/") - self.assertTrue(main_py.endswith("examples/PyTorch/wl-classification/main.py"), main_py) - - -class TestInstallExampleRequirements(unittest.TestCase): - """Requirements install is non-interactive and only runs when a file is present.""" - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_installs_requirements_non_interactively_when_present(self, mock_run): - import tempfile - mock_run.return_value = MagicMock(returncode=0) - with tempfile.TemporaryDirectory() as tmp: - req = Path(tmp) / "requirements.txt" - req.write_text("numpy\n") - _install_example_requirements(Path(tmp)) - mock_run.assert_called_once() - cmd = mock_run.call_args.args[0] - self.assertEqual(cmd[:5], [sys.executable, "-m", "pip", "install", "-r"]) - self.assertIn("--no-input", cmd) # never prompts - self.assertTrue(mock_run.call_args.kwargs.get("check")) - - @patch("weightslab.ui_docker_bridge.subprocess.run") - def test_skips_when_no_requirements_file(self, mock_run): - import tempfile - with tempfile.TemporaryDirectory() as tmp: - _install_example_requirements(Path(tmp)) - mock_run.assert_not_called() - - -class TestBannerAndHelp(unittest.TestCase): - """`weightslab`, `weightslab help`, `-h`, `--help` show banner + the command set.""" - - @staticmethod - def _capture_main(argv): - buf = io.StringIO() - with patch("sys.argv", argv): - with contextlib.redirect_stdout(buf): - try: - main() - except SystemExit: - pass # argparse -h/--help exits 0 - return buf.getvalue() - - def test_dash_h_shows_banner_and_command_reference(self): - out = self._capture_main(["weightslab", "-h"]) - self.assertIn("WeightsLab", out) # tagline from description - self.assertIn("ui launch", out) - self.assertIn("--certs", out) - self.assertIn("start example", out) - - def test_long_help_flag_shows_command_reference(self): - out = self._capture_main(["weightslab", "--help"]) - self.assertIn("ui launch", out) - self.assertIn("--force-certs", out) - - def test_help_subcommand_shows_command_reference(self): - out = self._capture_main(["weightslab", "help"]) - self.assertIn("se", out) - self.assertIn("ui launch", out) - self.assertIn("start example", out) - - def test_no_args_shows_command_reference(self): - out = self._capture_main(["weightslab"]) - self.assertIn("ui launch", out) - - def test_help_mentions_new_launch_command(self): - out = self._capture_main(["weightslab", "help"]) - self.assertIn("weightslab launch", out) - - def test_help_does_not_mention_removed_commands(self): - out = self._capture_main(["weightslab", "help"]) - for removed in ("ui stop", "ui drop", "ui docker", "--no-auth", "--no-clean"): - self.assertNotIn(removed, out, f"help should not mention removed '{removed}'") - - -class TestLaunchCliFlags(unittest.TestCase): - """`ui launch` accepts only --certs (TLS is opt-in; default unsecured).""" - - @patch("weightslab.ui_docker_bridge.ui_launch") - def test_launch_certs_flag_parsed(self, mock_launch): - with patch("sys.argv", ["weightslab", "ui", "launch", "--certs"]): - main() - mock_launch.assert_called_once() - self.assertTrue(mock_launch.call_args.args[0].certs) - - @patch("weightslab.ui_docker_bridge.ui_launch") - def test_launch_default_certs_false(self, mock_launch): - with patch("sys.argv", ["weightslab", "ui", "launch"]): - main() - self.assertFalse(mock_launch.call_args.args[0].certs) - - def test_launch_removed_flags_are_rejected(self): - for flag in ("--no-auth", "--force-certs", "--no-clean", "--dev"): - with patch("sys.argv", ["weightslab", "ui", "launch", flag]): - with self.assertRaises(SystemExit) as ctx: - main() - self.assertNotEqual(ctx.exception.code, 0, flag) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/components/test_checkpoint_workflow.py b/tests/components/test_checkpoint_workflow.py index e4cb5e99..6bbef8e9 100644 --- a/tests/components/test_checkpoint_workflow.py +++ b/tests/components/test_checkpoint_workflow.py @@ -1298,33 +1298,15 @@ def test_logger_queue_saved_with_weights(self): self.chkpt_manager.update_experiment_hash(force=False, dump_immediately=False) self.chkpt_manager.save_pending_changes(force=True) - snapshot_dir = Path(self.chkpt_manager.loggers_dir) - manifest_path = snapshot_dir / "loggers.manifest.json" - legacy_snapshot_path = snapshot_dir / "loggers.json" - self.assertTrue( - manifest_path.exists() or legacy_snapshot_path.exists(), - "Logger snapshot should be saved with checkpoint" - ) - - if manifest_path.exists(): - with open(manifest_path, "r") as f: - manifest = json.load(f) - chunk_files = manifest.get("chunks", []) - self.assertGreaterEqual(len(chunk_files), 1, "Chunked logger snapshot should contain at least one chunk") - self.assertTrue( - all((snapshot_dir / chunk_name).exists() for chunk_name in chunk_files), - "All logger snapshot chunks referenced by manifest should exist" - ) - self.assertTrue(self.chkpt_manager.load_logger_snapshot()) - lg = ledgers.get_logger('main') - loggers = {'main': {'signal_history': lg.get_signal_history() if hasattr(lg, 'get_signal_history') else []}} - else: - with open(legacy_snapshot_path, "r") as f: - snapshot = json.load(f) - loggers = snapshot.get("loggers", {}) - - self.assertIn('main', loggers, "Logger entry should be present") - signals = loggers['main'].get("signal_history", []) + # Logger history is persisted to an on-disk DuckDB file alongside the + # checkpoint (the chunked-zstd snapshot was replaced by DuckDB in #257). + db_path = Path(self.chkpt_manager.get_logger_db_path()) + self.assertTrue(db_path.exists(), "Logger DuckDB should be saved with checkpoint") + + # The registered logger carries the persisted signal history. + lg = ledgers.get_logger('main') + self.assertIsNotNone(lg, "Logger entry should be present") + signals = lg.get_signal_history() if hasattr(lg, 'get_signal_history') else [] if isinstance(signals, dict): total_signals = 0 for experiments in signals.values(): diff --git a/tests/general/test_logger_snapshot_rotation.py b/tests/general/test_logger_snapshot_rotation.py deleted file mode 100644 index f0a0001e..00000000 --- a/tests/general/test_logger_snapshot_rotation.py +++ /dev/null @@ -1,102 +0,0 @@ -import json -import tempfile -import unittest -from pathlib import Path - -from weightslab.backend import ledgers -from weightslab.components.checkpoint_manager import CheckpointManager -from weightslab.backend.logger import LoggerQueue - - -class LoggerSnapshotRotationTests(unittest.TestCase): - def setUp(self): - ledgers.clear_all() - - def tearDown(self): - ledgers.clear_all() - - def test_save_logger_snapshot_writes_chunked_zstd_and_can_reload(self): - with tempfile.TemporaryDirectory(prefix="weightslab_logger_snapshot_") as temp_dir: - manager = CheckpointManager(root_log_dir=temp_dir, load_model=False, load_config=False, load_data=False) - exp_hash = "12345678abcdef0199aabbcc" - manager.current_exp_hash = exp_hash - - logger_queue = LoggerQueue(register=False) - logger_queue.load_snapshot( - { - "graph_names": ["train/loss"], - "signal_history": [ - { - "experiment_name": "train/loss", - "model_age": 1, - "metric_name": "train/loss", - "metric_value": 0.42, - "experiment_hash": exp_hash, - } - ], - "signal_history_per_sample": {}, - } - ) - ledgers.register_logger(logger_queue, name="main") - - saved_path = manager.save_logger_snapshot() - self.assertIsNotNone(saved_path, "save_logger_snapshot should return a saved path") - self.assertTrue(saved_path.exists(), "Logger snapshot manifest should exist") - - snapshot_dir = Path(manager.loggers_dir) - manifest_path = snapshot_dir / "loggers.manifest.json" - self.assertTrue(manifest_path.exists(), "Manifest should be written for chunked snapshot format") - - with open(manifest_path, "r", encoding="utf-8") as f: - manifest = json.load(f) - chunk_names = manifest.get("chunks", []) - self.assertGreaterEqual(len(chunk_names), 1, "At least one compressed chunk should be created") - for chunk_name in chunk_names: - self.assertTrue((snapshot_dir / chunk_name).exists(), f"Missing expected chunk file: {chunk_name}") - - ledgers.clear_all() - loaded = manager.load_logger_snapshot() - self.assertTrue(loaded, "Chunked snapshot should load successfully") - restored_logger = ledgers.get_logger("main") - self.assertTrue(hasattr(restored_logger, "get_signal_history"), "Restored logger should support signal history") - self.assertGreaterEqual(len(restored_logger.get_signal_history()), 1, "Restored logger should contain signals") - - def test_load_logger_snapshot_supports_legacy_json(self): - with tempfile.TemporaryDirectory(prefix="weightslab_logger_snapshot_legacy_") as temp_dir: - manager = CheckpointManager(root_log_dir=temp_dir, load_model=False, load_config=False, load_data=False) - exp_hash = "abcdef0199aabbcc12345678" - manager.current_exp_hash = exp_hash - - snapshot_dir = Path(manager.loggers_dir) - snapshot_dir.mkdir(parents=True, exist_ok=True) - legacy_payload = { - "exp_hash": exp_hash, - "timestamp": "2026-02-24T00:00:00", - "loggers": { - "main": { - "graph_names": ["train/acc"], - "signal_history": [ - { - "experiment_name": "train/acc", - "model_age": 2, - "metric_name": "train/acc", - "metric_value": 0.9, - "experiment_hash": exp_hash, - } - ], - "signal_history_per_sample": {}, - } - }, - } - with open(snapshot_dir / "loggers.json", "w", encoding="utf-8") as f: - json.dump(legacy_payload, f) - - loaded = manager.load_logger_snapshot() - self.assertTrue(loaded, "Legacy JSON logger snapshot should still load") - restored_logger = ledgers.get_logger("main") - self.assertTrue(hasattr(restored_logger, "get_signal_history"), "Restored logger should support signal history") - self.assertGreaterEqual(len(restored_logger.get_signal_history()), 1, "Legacy snapshot should restore signals") - - -if __name__ == "__main__": - unittest.main() diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py index 4da02b6f..8b7069e1 100755 --- a/weightslab/trainer/services/data_service.py +++ b/weightslab/trainer/services/data_service.py @@ -3453,6 +3453,48 @@ def _is_metadata_only_request(self, request) -> bool: except Exception: return False + @staticmethod + def _downsample_curve(vals, max_points=32): + """Uniformly downsample a series to <= max_points floats (endpoints kept).""" + if len(vals) <= max_points: + return [float(v) for v in vals] + idx = np.linspace(0, len(vals) - 1, max_points).round().astype(int) + return [float(vals[j]) for j in idx] + + def _loss_trajectory_curves(self, sample_ids): + """Downsampled per-sample loss curve for the on-cell sparkline. + + Returns ``{str(sample_id): [float, ...]}`` for samples with history; an + empty dict when there is no logger or no per-sample loss signal. Read-only — + shape classification lives on the write path (enable_loss_shape_signal), + not here. Signal name overridable via WL_LOSS_TRAJ_SIGNAL (default "loss"). + """ + from weightslab.backend import ledgers + + logger_q = ledgers.get_logger() + if logger_q is None: + return {} + + signal = os.environ.get("WL_LOSS_TRAJ_SIGNAL", "loss") + try: + rows = logger_q.query_per_sample(signal, sample_ids=sample_ids) + except Exception: + logger.debug("loss_trajectory query skipped", exc_info=True) + return {} + + by_sid = {} + for sid, step, val, _ in rows: + by_sid.setdefault(str(sid), []).append((int(step), float(val))) + + curves = { + sid: self._downsample_curve([v for _, v in sorted(pts)]) + for sid, pts in by_sid.items() + } + if curves: + logger.info("[loss_traj] page_ids=%d with_history=%d signal=%s", + len(sample_ids), len(curves), signal) + return curves + def _build_metadata_only_response(self, df_slice: pd.DataFrame, requested_cols=None): """Build a DataSamplesResponse of metadata DataRecords from dataframe columns only. @@ -3571,6 +3613,20 @@ def _build_metadata_only_response(self, df_slice: pd.DataFrame, requested_cols=N _DataStat(name=col, type="string", shape=[1], value_string="1") ) + # -- Per-sample loss trajectory (the on-cell sparkline; read-only) ----- + # Ship each visible sample's loss curve as an 'array' DataStat for the grid + # sparkline. The loss-SHAPE haze is NOT computed here: classification runs on + # the write/train path (enable_loss_shape_signal / write_loss_shapes) which + # maintains the tag:loss_shape column, so this path does no per-request work. + loss_curves = self._loss_trajectory_curves(sample_ids) + for i, sid in enumerate(sample_ids): + curve = loss_curves.get(str(sid)) + if not curve: + continue + row_stats[i].append( + _DataStat(name="loss_trajectory", type="array", + shape=[len(curve)], value=curve)) + # -- Build DataRecord list in one comprehension ----------------------- _DataRecord = pb2.DataRecord data_records = [ diff --git a/weightslab/trainer/services/experiment_service.py b/weightslab/trainer/services/experiment_service.py index 3de658a5..b7313c4e 100644 --- a/weightslab/trainer/services/experiment_service.py +++ b/weightslab/trainer/services/experiment_service.py @@ -792,11 +792,11 @@ def ExperimentCommand(self, request, context): checkpoint_manager = ledgers.get_checkpoint_manager() except Exception: checkpoint_manager = None - if checkpoint_manager is not None and hasattr(checkpoint_manager, "save_logger_snapshot"): + if checkpoint_manager is not None and hasattr(checkpoint_manager, "flush_logger_to_disk"): try: - checkpoint_manager.save_logger_snapshot() + checkpoint_manager.flush_logger_to_disk() except Exception: - logger.debug("Could not persist logger snapshot after note update", exc_info=True) + logger.debug("Could not persist logger history after note update", exc_info=True) # Log to audit trail self._log_audit( diff --git a/weightslab/ui_docker_bridge.py b/weightslab/ui_docker_bridge.py deleted file mode 100644 index 11a1e826..00000000 --- a/weightslab/ui_docker_bridge.py +++ /dev/null @@ -1,1319 +0,0 @@ -"""Weights Studio Docker management with secure TLS.""" - -import argparse -import os -import re -import shutil -import stat -import subprocess -import sys -import socket -import logging -from pathlib import Path - -from weightslab.security import CertAuthManager -from weightslab.tunnel import DEFAULT_LISTEN_PORT - -logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') -logger = logging.getLogger(__name__) - -# Capture whether WEIGHTSLAB_CERTS_DIR was already in the shell env when this -# process started (set by the user) vs. injected later by our own code. -_CERTS_DIR_IN_ORIGINAL_ENV: bool = "WEIGHTSLAB_CERTS_DIR" in os.environ - -# Docker resources owned by the bundled stack. Cleanup is scoped strictly to -# these names — we never run a global `docker system prune`. -_FRONTEND_IMAGE = "graybx/weightslab" -_STACK_CONTAINERS = ("weights_studio_envoy", "weights_studio_frontend") - -# Default frontend image repo (no tag) and tag. Overridable per-launch via the -# `--image/-i` and `--version/-v` flags on `weightslab ui launch`, which are -# passed to docker compose through the WS_FRONTEND_IMAGE env var (see the -# `${WS_FRONTEND_IMAGE:-...}` substitution in docker-compose.yml). -_DEFAULT_FRONTEND_REPO = "graybx/weightslab" -_DEFAULT_FRONTEND_TAG = "latest" - - -def _resolve_frontend_image(image_arg=None, version_arg=None) -> str: - """Return the full ``repo:tag`` frontend image ref from the CLI flags. - - ``image_arg`` is the repo (e.g. ``guillaumep2705/weightslab``) and may itself - carry a tag; ``version_arg`` is the tag (e.g. ``latest`` or ``v1.2.3``). An - explicit ``--version`` always wins over any tag embedded in ``--image``. - Falls back to ``graybx/weightslab:latest`` when neither is given. - """ - image = image_arg or _DEFAULT_FRONTEND_REPO - # A ':' in the last path segment is a tag (not a registry port like host:5000/img). - last_segment = image.rsplit("/", 1)[-1] - has_tag = ":" in last_segment - if version_arg: - repo = image.rsplit(":", 1)[0] if has_tag else image - return f"{repo}:{version_arg}" - if has_tag: - return image - return f"{image}:{_DEFAULT_FRONTEND_TAG}" - -# Cached Docker Compose base command. Resolved once per process to either the -# v2 plugin (``["docker", "compose"]``) or the legacy v1 standalone binary -# (``["docker-compose"]``) — see _detect_compose_cmd(). None until first probe. -_COMPOSE_BASE_CMD = None - -# TLS/auth env vars that are *derived* from cert-file presence in -# WEIGHTSLAB_CERTS_DIR. WEIGHTSLAB_CERTS_DIR is the single source of truth; these -# are computed by the deploy pipeline (build-and-deploy.sh + the compose `auto` -# logic) from the actual files. They are stripped before launching so a stale or -# pre-set value can never override the file-based decision. -_DERIVED_DEPLOY_ENV_VARS = ( - "GRPC_TLS_ENABLED", - "GRPC_TLS_REQUIRE_CLIENT_AUTH", - "GRPC_TLS_CERT_FILE", - "GRPC_TLS_KEY_FILE", - "GRPC_TLS_CA_FILE", - "ENVOY_UPSTREAM_TLS", - "ENVOY_DOWNSTREAM_TLS", - "WS_SERVER_PROTOCOL", - "VITE_SERVER_PROTOCOL", - "VITE_DEV_SERVER_HTTPS", - "WL_ENABLE_GRPC_AUTH_TOKEN", - "VITE_WL_ENABLE_GRPC_AUTH_TOKEN", - "GRPC_AUTH_TOKEN", - "VITE_GRPC_AUTH_TOKEN", -) - - -def _persist_certs_dir(certs_dir_str: str) -> None: - """Persist WEIGHTSLAB_CERTS_DIR so future terminals and the training backend find it. - - Windows — runs `setx` (permanent user env) and prints the PS one-liner for - the current session. - Linux/macOS — appends an export line to ~/.bashrc (idempotent) and prints the - source command for the current session. - """ - export_line = f'export WEIGHTSLAB_CERTS_DIR="{certs_dir_str}"' - if _is_windows(): - result = subprocess.run( - ["setx", "WEIGHTSLAB_CERTS_DIR", certs_dir_str], - stdout=subprocess.PIPE, stderr=subprocess.PIPE, - ) - if result.returncode == 0: - logger.info(" WEIGHTSLAB_CERTS_DIR saved permanently via setx (new terminals will have it)") - else: - logger.warning(f"setx failed — set it manually: setx WEIGHTSLAB_CERTS_DIR \"{certs_dir_str}\"") - logger.info(f" Current terminal (PowerShell): $env:WEIGHTSLAB_CERTS_DIR = \"{certs_dir_str}\"") - else: - bashrc = Path.home() / ".bashrc" - try: - existing = bashrc.read_text(encoding="utf-8") if bashrc.exists() else "" - if export_line not in existing: - with open(bashrc, "a", encoding="utf-8") as f: - f.write(f"\n# Added by weightslab\n{export_line}\n") - logger.info(f" WEIGHTSLAB_CERTS_DIR appended to {bashrc} (new terminals will have it)") - else: - logger.info(f" WEIGHTSLAB_CERTS_DIR already in {bashrc}") - except OSError as e: - logger.warning(f"Could not write to {bashrc}: {e}") - logger.info(f" Add manually: {export_line}") - logger.info(f" Current terminal: source ~/.bashrc (or open a new terminal)") - - -def _strip_derived_deploy_env() -> None: - """Drop derived TLS/auth env vars so the deploy pipeline decides solely from - cert-file presence in WEIGHTSLAB_CERTS_DIR (the single source of truth).""" - for key in _DERIVED_DEPLOY_ENV_VARS: - os.environ.pop(key, None) - - -def _banner() -> str: - """Return the WeightsLab ASCII banner, or a plain title if art is unavailable.""" - try: - from weightslab.art import _BANNER - return _BANNER - except Exception: - return "WeightsLab" - - -_DESCRIPTION = ( - _banner() - + "\nWeightsLab — Inspect, Edit, and Evolve Neural Networks\n" - + "Manage the Weights Studio UI, its Docker stack, and the secure " - + "(TLS + gRPC auth) environment." -) - -_EPILOG = """\ -commands: - se Set up the secure environment: generate TLS - certificates + a gRPC auth token in - ~/.weightslab-certs. Then set WEIGHTSLAB_CERTS_DIR - (the single source of truth) so the backend + new - shells find them. - --force-certs regenerate even if certs exist - - launch Purge stale weightslab/weights_studio Docker - resources, build & start the UI stack, then open it - in your browser. UNSECURED (HTTP) by default — no - certs generated. Also: bare `weightslab start`. - (Legacy alias: `weightslab ui launch`.) - --certs generate (if missing) + use TLS - certs + gRPC auth (HTTPS) - - start example Run a bundled PyTorch example (foreground; stop with - Ctrl+C). Installs the example's requirements first, - without prompting. Defaults to classification: - --cls classification example (default) - --seg segmentation example - --det detection example - --clus clustering example - --gen generation example - --3d_det 3D LiDAR point-cloud detection example - --2d_det 2D LiDAR point-cloud detection example - - cli Open an interactive terminal connected to a - currently-running experiment (pause/resume, status, - evaluate, agent query, etc.). Auto-discovers the - running experiment; the experiment must be serving - the CLI (e.g. wl.serve(serving_cli=True)). - --port PORT connect to a specific CLI port - --host HOST connect to a specific host (default: localhost) - - tunnel Forward a REMOTE gRPC backend (e.g. a Colab run - behind `ngrok tcp 50051`) to a LOCAL port so the UI - stack — whose Envoy dials localhost:50051 — reaches - it. Run alongside `weightslab launch`. Raw TCP, so - the backend must be plaintext (default launch). - ENDPOINT remote host:port (e.g. bore.pub:12345) - (default: $WEIGHTSLAB_TUNNEL_ENDPOINT) - --listen-port N local port to expose (default 50051) - --listen-host H interface to bind (default: auto) - --remote-port N remote port, if not in ENDPOINT - -examples: - weightslab se # one-time secure setup (then export WEIGHTSLAB_CERTS_DIR) - weightslab se --force-certs # regenerate the certs - weightslab launch # clean + launch UI, open browser (unsecured HTTP, default) - weightslab start # same as `weightslab launch` - weightslab launch --certs # secured launch (HTTPS + gRPC auth) - weightslab launch -i guillaumep2705/weightslab # pull frontend from a custom repo (latest) - weightslab launch -i guillaumep2705/weightslab -v v1.2.3 # pin a specific version/tag - weightslab ui launch # legacy alias for `weightslab launch` - weightslab start example # run the classification demo (default) - weightslab start example --seg # run the segmentation demo - weightslab start example --det # run the detection demo - weightslab start example --3d_det # run the 3D LiDAR detection demo - weightslab start example --2d_det # run the 2D LiDAR detection demo - weightslab cli # connect a terminal to the running experiment - weightslab cli --port 60000 # connect to a specific CLI port - weightslab tunnel bore.pub:12345 # expose a remote (Colab) backend at localhost:50051 for the UI -""" - - -def _get_compose_file(): - """Return the path to the bundled docker-compose.yml.""" - # return files("weightslab.docker.docker") / "docker-compose.yml" - return Path(__file__).parent / 'docker' / 'docker' / 'docker-compose.yml' - - -def _get_envoy_config(): - """Return the path to the bundled envoy.yaml.""" - # return files("weightslab.docker.envoy") / "envoy.yaml" - return Path(__file__).parent / 'docker' / 'docker' / 'envoy.yaml' - - -def _get_bootstrap_script() -> Path: - """Get the bootstrap-secure.ps1 script path.""" - return Path(__file__).parent / 'docker' / 'docker' / 'utils' / 'build-and-deploy.sh' - - -def _get_cert_script() -> Path: - """Get the generate-certs-auth-token.sh script path.""" - return Path(__file__).parent / 'docker' / 'docker' / 'utils' / 'generate-certs-auth-token.sh' - - -def _get_cert_script_ps1() -> Path: - """Get the generate-certs-auth-token.ps1 script path.""" - return Path(__file__).parent / 'docker' / 'docker' / 'utils' / 'generate-certs-auth-token.ps1' - - -def _is_windows() -> bool: - """Check if running on Windows.""" - return sys.platform == 'win32' - - -def _make_executable(path) -> None: - """Add the execute bit to a file so it can be run directly (POSIX only). - - pip installs the bundled ``.sh`` scripts as package data *without* the execute - bit, and they are committed/shipped non-executable. ``_run_shell_script`` runs - them as a command (``bash -c "VAR=... '/path/script.sh'"``), which requires - the execute bit — otherwise bash fails with "Permission denied" until the user - ``chmod +x`` by hand. This grants u+x,g+x,o+x (e.g. 0644 -> 0755) and is a - best-effort no-op on Windows, where the POSIX execute bit is not used. - """ - if _is_windows(): - return - try: - mode = os.stat(path).st_mode - os.chmod(path, mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) - except OSError as exc: - # Non-fatal: e.g. a root-owned system install the user can't chmod. They - # would not have been able to chmod it manually either; surface as debug. - logger.debug(f"Could not set execute bit on {path}: {exc}") - - -def _ensure_scripts_executable() -> None: - """Make every bundled shell script executable (POSIX only). - - Covers build-and-deploy.sh, generate-certs-auth-token.sh and the other ``.sh`` - files under ``weightslab/ui`` so a freshly pip-installed package can run them - without the user having to ``chmod +x`` first. No-op on Windows. - """ - if _is_windows(): - return - ui_dir = Path(__file__).parent / 'docker' - try: - scripts = list(ui_dir.rglob('*.sh')) - except OSError as exc: - logger.debug(f"Could not enumerate bundled scripts under {ui_dir}: {exc}") - return - for script in scripts: - _make_executable(script) - - -def _detect_compose_cmd(): - """Return the base command for Docker Compose, preferring v2 over v1. - - Docker Compose ships in two forms: - * v2 — the ``docker compose`` CLI plugin (bundled with Docker Desktop and - the recommended install). Invoked as two words: ``docker compose``. - * v1 — the legacy standalone ``docker-compose`` binary (one word, hyphen). - - We probe v2 first (``docker compose version``) then fall back to v1 - (``docker-compose version``) so ``weightslab ui launch`` works with either. - Returns ``["docker", "compose"]``, ``["docker-compose"]``, or None if neither - is available. - """ - # v2: the `docker compose` plugin. - try: - result = subprocess.run( - ["docker", "compose", "version"], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - ) - if result.returncode == 0: - return ["docker", "compose"] - except FileNotFoundError: - # docker itself is missing; _check_docker() surfaces the user-facing error. - pass - - # v1: the legacy standalone `docker-compose` binary. - if shutil.which("docker-compose") is not None: - try: - result = subprocess.run( - ["docker-compose", "version"], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - ) - if result.returncode == 0: - return ["docker-compose"] - except FileNotFoundError: - pass - - return None - - -def _compose_base_cmd(): - """Cached Docker Compose base command (v2 ``docker compose``, else v1 ``docker-compose``).""" - global _COMPOSE_BASE_CMD - if _COMPOSE_BASE_CMD is None: - _COMPOSE_BASE_CMD = _detect_compose_cmd() - return _COMPOSE_BASE_CMD - - -def _check_docker(): - """Verify that docker is installed and the daemon is running.""" - if shutil.which("docker") is None: - logger.error( - "Docker is required but not found on your PATH.\n" - "Install it from: https://docs.docker.com/get-docker/" - ) - sys.exit(1) - - try: - subprocess.run( - ["docker", "info"], - check=True, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - except subprocess.CalledProcessError: - logger.error( - "\n" - "=============================================================" - "=============================================================" - "Docker is installed but the daemon is not running.\n" - "Start it with: Docker Desktop or 'sudo systemctl start docker'." - "=============================================================" - "\n" - ) - sys.exit(1) - - -def _compose_cmd(compose_file, envoy_config, action): - """Build and run a docker compose command.""" - env = os.environ.copy() - env["WS_ENVOY_CONFIG"] = str(envoy_config) - - # If WEIGHTSLAB_CERTS_DIR isn't set, fall back to the default certs dir when - # it actually holds a full cert set + gRPC token. This gives docker compose a - # valid host-native bind-mount source (so Envoy gets its certs) even when the - # user never exported the var — and covers compose calls that run before - # ui_launch sets it (e.g. the stale-resource cleanup `compose down`). - if not env.get("WEIGHTSLAB_CERTS_DIR"): - default_mgr = CertAuthManager() - if default_mgr.has_valid_certs() and default_mgr.token_file.exists(): - host_path = _to_docker_host_path(default_mgr.certs_dir) - env["WEIGHTSLAB_CERTS_DIR"] = host_path - logger.info( - f"WEIGHTSLAB_CERTS_DIR not set — using default certs dir for docker: {host_path}" - ) - - base = _compose_base_cmd() - if base is None: - logger.error( - "Docker Compose is required but was not found.\n" - "Install Compose v2 (recommended) — the `docker compose` CLI plugin, " - "bundled with Docker Desktop: https://docs.docker.com/compose/install/\n" - "The legacy v1 `docker-compose` binary also works if it is on your PATH." - ) - sys.exit(1) - - action = list(action) - # Compose v1 (`docker-compose`) has no `up --pull ` flag (it is - # v2-only). Emulate it: `pull` first (best-effort — some services only build - # locally), then `up` without the flag. v2 supports it inline, so leave it. - if base == ["docker-compose"] and action and action[0] == "up" and "--pull" in action: - i = action.index("--pull") - del action[i:i + 2] # drop '--pull' and its policy value (e.g. 'always') - logger.info("Docker Compose v1 detected — pulling images before 'up'...") - pull_result = subprocess.run( - base + ["-f", str(compose_file), "pull"], - env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, - ) - if pull_result.stdout: - for line in pull_result.stdout.splitlines(): - logger.info(line) - - cmd = base + ["-f", str(compose_file)] + action - result = subprocess.run(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) - - if result.stdout: - for line in result.stdout.splitlines(): - logger.info(line) - - if result.returncode != 0: - raise subprocess.CalledProcessError(result.returncode, cmd) - - -def _test_backend_connection(host: str = '127.0.0.1', port: int = 50051, timeout: float = 5.0) -> bool: - """Test if backend gRPC server is reachable.""" - try: - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(timeout) - result = sock.connect_ex((host, port)) - sock.close() - return result == 0 - except Exception as e: - logger.debug(f"Backend connection test failed: {e}") - return False - - -def _run_powershell_script(script_path: str, args: list = None, env_vars: dict = None) -> int: - """Run a PowerShell script and return exit code.""" - if not _is_windows(): - logger.error("Secure launch requires Windows with PowerShell") - return 1 - - cmd = [ - 'powershell', - '-NoProfile', - '-ExecutionPolicy', 'Bypass', - '-File', script_path - ] - - if args: - cmd.extend(args) - - try: - env = os.environ.copy() - if env_vars: - env.update(env_vars) - result = subprocess.run(cmd, env=env) - return result.returncode - except Exception as e: - logger.error(f"Failed to run script: {e}") - return 1 - - -def _convert_to_git_bash_path(win_path: str) -> str: - """Convert Windows path to Git Bash compatible format.""" - # Normalize backslashes to forward slashes explicitly: Path(...).as_posix() - # does NOT convert '\' on POSIX hosts (e.g. Linux CI runners), so a Windows - # input path would keep its separators there. - p = str(win_path).replace("\\", "/") - # Convert C:/Users/... to /mnt/c/Users/... for Git Bash - if len(p) > 1 and p[1] == ':': - drive = p[0].lower() - rest = p[2:] - return f"/mnt/{drive}{rest}" - return p - - -def _to_docker_host_path(path) -> str: - """Return a path that Docker Desktop can bind-mount. - - On Windows, normalize WSL (/mnt/c/...) and Git Bash (/c/...) forms — and - native paths that Path() mangled into \\mnt\\c\\... — to C:/... form. - Docker Desktop cannot resolve /mnt-style sources and silently mounts an - empty directory, which crashes Envoy on missing certs. On Linux/macOS the - path is returned unchanged. - """ - p = str(path).replace('\\', '/') - if _is_windows(): - m = re.match(r'^/mnt/([a-zA-Z])/(.*)$', p) or re.match(r'^/([a-zA-Z])/(.*)$', p) - if m: - return f"{m.group(1).upper()}:/{m.group(2)}" - return p - - -def _run_shell_script(script_path: str, args: list = None, env_vars: dict = None) -> int: - """Run a shell script using bash with proper environment variable passing.""" - - try: - # Fix line endings in the file before running - with open(script_path, 'rb') as f: - script_bytes = f.read() - - # Ensure Unix line endings - fixed_bytes = script_bytes.replace(b'\r\n', b'\n').replace(b'\r', b'\n') - - # Write back if needed - if fixed_bytes != script_bytes: - with open(script_path, 'wb') as f: - f.write(fixed_bytes) - - # We invoke the script as a command below (bash -c "VAR=... 'script'"), - # which needs the execute bit. pip installs it without one, so add it here - # (no-op on Windows). Uses the on-disk path before any Git Bash conversion. - _make_executable(script_path) - - env = os.environ.copy() - if env_vars: - env.update(env_vars) - # Debug: log all environment variables being passed - for key, value in env_vars.items(): - logger.info(f"Passing env var: {key}={value}") - - # Build bash command - pass Windows path directly, script will handle conversion - # # Process path to ensure it's compatible with bash, especially on Windows - if _is_windows() and '\\' in script_path: - script_path = script_path.replace("\\", "/") # Ensure path is Unix-style for bash - script_path = _convert_to_git_bash_path(script_path) - logger.info(f"Converted script path for bash: {script_path}") - logger.info(f"Running shell script: {script_path} with args: {args} and env_vars: {env_vars}") - - # Build environment variable assignments for bash command - env_assignments = ' '.join([f"{k}='{v}'" for k, v in env_vars.items()]) if env_vars else "" - - if env_assignments: - # Pass env vars directly in bash command using -c flag - # This works on Windows (Git Bash), Linux, and macOS - logger.info('Using env assignments') - bash_script_cmd = f"{env_assignments} '{script_path}'" - if args: - bash_script_cmd += " " + " ".join(f"'{arg}'" for arg in args) - bash_cmd = ['bash', '-c', bash_script_cmd] - else: - logger.info('Not using env assignments') - bash_cmd = ['bash', '-x', str(script_path)] - if args: - bash_cmd.extend(args) - - result = subprocess.run(bash_cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) - if result.stdout: - for line in result.stdout.splitlines(): - logger.info(line) - return result.returncode - except FileNotFoundError: - logger.error(f"Script file not found: {script_path}") - return 1 - except Exception as e: - logger.error(f"Failed to run script: {e}") - return 1 - - -def _generate_certs_with_fallback(force_certs: bool = False, certs_dir=None) -> int: - """Try shell script first, fall back to PowerShell on Windows if it fails. - - ``certs_dir`` is forwarded to the generation scripts as ``WEIGHTSLAB_CERTS_DIR`` - so certs land in the single source-of-truth directory (the scripts default to - ``~/.weightslab-certs`` when it is not provided). The shell script receives a - POSIX absolute path (``/mnt/c/...`` on Windows/WSL); PowerShell receives the - host-native path (``C:/...``). - """ - env_vars = None - if certs_dir is not None: - # Shell scripts (bash/WSL) need a POSIX-style absolute path. - # _to_docker_host_path gives C:/... which WSL bash treats as a relative - # path, creating the certs dir under cwd instead of ~/.weightslab-certs. - # Use the /mnt/c/... form so the path is absolute inside WSL. - if _is_windows(): - bash_certs_dir = _convert_to_git_bash_path(str(certs_dir)) - else: - bash_certs_dir = str(certs_dir) - env_vars = {'WEIGHTSLAB_CERTS_DIR': bash_certs_dir} - - cert_script = str(_get_cert_script()) - if not Path(cert_script).exists(): - logger.warning(f"Shell script not found: {cert_script}") - else: - script_args = [] - if force_certs: - script_args.append('--force-create-certs') - - logger.info("Attempting certificate generation with shell script...") - exit_code = _run_shell_script(cert_script, script_args, env_vars) - if exit_code == 0: - return 0 - logger.warning(f"Shell script failed (exit code {exit_code})") - - # Fallback to PowerShell on Windows - if _is_windows(): - logger.info("Falling back to PowerShell for certificate generation...") - cert_script_ps1 = str(_get_cert_script_ps1()) - if not Path(cert_script_ps1).exists(): - logger.error(f"PowerShell script not found: {cert_script_ps1}") - return 1 - - script_args = [] - if force_certs: - script_args.append('-ForceCreateCerts') - - exit_code = _run_powershell_script(cert_script_ps1, script_args, env_vars) - return exit_code - else: - logger.error("Neither shell nor PowerShell script could generate certificates") - return 1 - - -_DEV_CA_SUBJECT = "weightslab-dev-ca" - - -def _install_ca_trust(ca_file: Path) -> None: - """Install the dev CA into the OS trust store so browsers trust the HTTPS UI. - - Idempotent and safe to call on every launch. Platform behavior: - * Windows — adds to the CurrentUser\\Root store via the .NET X509Store API - (silent, no prompt). - * macOS — adds to the login keychain (may show a one-time auth prompt). - * Linux — installs into the system trust store via sudo (one-time prompt) - and, best-effort, the user's NSS DB so Chrome/Firefox trust it too. - - A failure here is non-fatal: TLS still works, the browser just shows a - self-signed warning until the CA is trusted manually. - """ - if not ca_file.exists(): - logger.warning(f"CA file not found, skipping trust install: {ca_file}") - return - - if _is_windows(): - ps = ( - "$ErrorActionPreference='Stop';" - f"$caPath='{ca_file}';" - "$certObj=New-Object System.Security.Cryptography.X509Certificates.X509Certificate2($caPath);" - "$store=New-Object System.Security.Cryptography.X509Certificates.X509Store('Root','CurrentUser');" - "$store.Open('ReadWrite');" - "try{" - # Already trusted (same thumbprint)? Nothing to do — avoids a fragile Remove. - "$match=$store.Certificates|Where-Object{$_.Thumbprint -eq $certObj.Thumbprint};" - "if(-not $match){" - # Drop stale same-subject CAs from a previous rotation (best-effort). - f"$stale=$store.Certificates|Where-Object{{$_.Subject -eq 'CN={_DEV_CA_SUBJECT}'}};" - "foreach($c in $stale){try{$store.Remove($c)}catch{}};" - "$store.Add($certObj);" - "}" - "}finally{$store.Close()}" - ) - result = subprocess.run( - ["powershell", "-NoProfile", "-Command", ps], - stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, - ) - if result.returncode == 0: - logger.info(" Dev CA trusted in Windows CurrentUser\\Root store (restart browser to apply)") - else: - logger.warning(f"Could not auto-trust dev CA: {result.stderr.strip()}") - return - - if sys.platform == "darwin": - check = subprocess.run( - ["security", "find-certificate", "-c", _DEV_CA_SUBJECT], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - ) - if check.returncode == 0: - logger.info(" Dev CA already trusted (macOS keychain)") - return - logger.info("Installing dev CA into macOS login keychain (may prompt)...") - subprocess.run( - ["security", "add-trusted-cert", "-r", "trustRoot", - "-k", os.path.expanduser("~/Library/Keychains/login.keychain-db"), str(ca_file)], - ) - return - - # Linux: system trust store (curl/openssl) + best-effort NSS (browsers). - system_ca = Path("/usr/local/share/ca-certificates/weightslab-dev-ca.crt") - try: - already = system_ca.exists() and system_ca.read_bytes() == ca_file.read_bytes() - except OSError: - already = False - if not already: - logger.info("Installing dev CA into the Linux system trust store (may prompt for sudo)...") - subprocess.run(["sudo", "cp", str(ca_file), str(system_ca)]) - subprocess.run(["sudo", "update-ca-certificates"]) - else: - logger.info(" Dev CA already in Linux system trust store") - - # Browsers use their own NSS DB; add it there too if certutil is available. - if shutil.which("certutil"): - nssdb = Path.home() / ".pki" / "nssdb" - nssdb.mkdir(parents=True, exist_ok=True) - subprocess.run( - ["certutil", "-A", "-n", _DEV_CA_SUBJECT, "-t", "C,,", - "-i", str(ca_file), "-d", f"sql:{nssdb}"], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - ) - - -def _ensure_certificates(manager: CertAuthManager, force_certs: bool = False) -> bool: - """Generate certs + auth token in ``manager.certs_dir`` if missing (or forced). - - Does NOT export any TLS/auth env vars: the launch pipeline derives TLS purely - from cert-file presence in ``WEIGHTSLAB_CERTS_DIR`` (the single source of - truth). Returns True if certs are present afterwards, False otherwise. - """ - if manager.has_any_credentials() and not force_certs: - logger.info(f" Using existing credentials in {manager.certs_dir}") - manager.get_or_create_auth_token() - # Ensure the CA is trusted even when reusing certs from a prior run that - # was generated via bash (which does not install OS trust). - _install_ca_trust(manager.ca_file) - return manager.has_valid_certs() - - logger.info( - "Regenerating certificates (--force-certs)..." - if force_certs - else "No certificates found — generating them now..." - ) - Path(manager.certs_dir).mkdir(parents=True, exist_ok=True) - exit_code = _generate_certs_with_fallback(force_certs=force_certs, certs_dir=manager.certs_dir) - if exit_code != 0: - logger.warning("Certificate generation failed — continuing in unsecured mode") - return False - - manager.get_or_create_auth_token() - _install_ca_trust(manager.ca_file) - logger.info(f" Certificates ready in {manager.certs_dir}") - return manager.has_valid_certs() - - -def _remove_docker_image(image: str) -> None: - """Remove a Docker image (all tags) if present; no-op when absent.""" - result = subprocess.run( - ["docker", "images", "-q", image], - stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, - ) - image_ids = sorted(set(result.stdout.split())) - if not image_ids: - return - logger.info(f"Removing cached image '{image}' ({len(image_ids)} ref(s))...") - subprocess.run( - ["docker", "rmi", "-f", *image_ids], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - ) - - -def _clean_stale_docker_resources(frontend_image: str = _FRONTEND_IMAGE) -> None: - """Remove leftover weightslab / weights_studio Docker state before a launch. - - Stale containers, the compose network, anonymous volumes, a cached frontend - image, and a leftover generated ``.env`` can each silently break a fresh - launch (an old image served instead of a rebuild, or an empty cert mount - that crashes Envoy). This is scoped STRICTLY to weightslab/weights_studio - resources — it never runs a global ``docker system prune``. - - ``frontend_image`` is the repo (optionally ``repo:tag``) whose cached copy is - dropped so ``up --pull always`` fetches a fresh one — pass the resolved value - when ``--image/--version`` point at a non-default repo. - """ - logger.info("Cleaning stale Docker resources (weightslab/weights_studio only)...") - compose_file = _get_compose_file() - envoy_config = _get_envoy_config() - - # 1. Tear down any existing stack: containers + default network + anon volumes. - try: - _compose_cmd(compose_file, envoy_config, ["down", "--remove-orphans", "--volumes"]) - except subprocess.CalledProcessError as exc: - logger.debug(f"'compose down' returned non-zero (nothing to remove?): {exc}") - - # 2. Force-remove leftover named containers started outside this compose project. - for container in _STACK_CONTAINERS: - subprocess.run( - ["docker", "rm", "-f", container], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - ) - - # 3. Drop the cached frontend image so a fresh one is built/pulled. - _remove_docker_image(frontend_image) - - # 4. Remove the generated .env so stale values don't leak into compose. - env_file = Path(compose_file).parent / ".env" - if env_file.exists(): - try: - env_file.unlink() - logger.info(f"Removed stale env file: {env_file}") - except OSError as exc: - logger.debug(f"Could not remove {env_file}: {exc}") - - -def _open_ui_in_browser(url: str) -> None: - """Best-effort: open the Studio UI in the default browser after launch. - - Skipped when WEIGHTSLAB_NO_BROWSER is set, in CI, or on a headless Linux box - (no DISPLAY / Wayland). Never raises — a failed open must not fail the launch. - """ - if (os.environ.get("WEIGHTSLAB_NO_BROWSER") - or os.environ.get("CI") - or os.environ.get("PYTEST_CURRENT_TEST")): - return - if sys.platform.startswith("linux") and not ( - os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY") - ): - logger.info("Headless environment — open the UI manually at %s", url) - return - try: - import webbrowser - if webbrowser.open_new_tab(url): - logger.info("Opening the Weights Studio UI in your browser…") - except Exception as e: # noqa: BLE001 — auto-open is best-effort only - logger.debug("Could not auto-open the browser: %s", e) - - -def ui_launch(args): - """Purge stale Docker state, then build & start the UI. - - By default the UI launches UNSECURED (HTTP, no gRPC auth) — no certificates - are generated. Pass ``--certs`` to generate (if missing) and use TLS certs + - a gRPC auth token. Existing certs in WEIGHTSLAB_CERTS_DIR are always honored - (file presence is the single source of truth) and are never deleted here. - - Flags (all optional, read defensively so legacy callers still work): - --certs generate (if missing) and use TLS certs + gRPC auth (HTTPS) - --force-certs with --certs, regenerate certificates even if they exist - --no-clean skip the stale Docker resource cleanup step - --dev use the dev compose overlay - -i/--image frontend image repo to run/pull (default: graybx/weightslab) - -v/--version frontend image tag/version to pull (default: latest) - """ - try: - from weightslab.utils.telemetry import ping_ui_launch - from weightslab import __version__ as _wl_version - ping_ui_launch(_wl_version) - except Exception: - pass - _check_docker() - # pip installs the bundled .sh scripts without the execute bit; make them - # runnable so the user never has to `chmod +x` before `weightslab ui launch`. - _ensure_scripts_executable() - - use_certs = getattr(args, "certs", False) - no_auth = getattr(args, "no_auth", False) - force_certs = getattr(args, "force_certs", False) - no_clean = getattr(args, "no_clean", False) - certs_dir_arg = getattr(args, "certs_dir", None) - - # Resolve which frontend image to run. docker compose reads it via the - # WS_FRONTEND_IMAGE env var (see `${WS_FRONTEND_IMAGE:-...}` in the compose - # file); export it so both `up --pull always` below and the stale-image - # cleanup target the same repo/tag. - frontend_image = _resolve_frontend_image( - getattr(args, "image", None), getattr(args, "version", None) - ) - os.environ["WS_FRONTEND_IMAGE"] = frontend_image - frontend_repo = frontend_image.rsplit(":", 1)[0] if ":" in frontend_image.rsplit("/", 1)[-1] else frontend_image - logger.info(f"Using frontend image: {frontend_image}") - - # A custom certs dir (positional arg) takes precedence over the env var / - # default. Otherwise fall back to $WEIGHTSLAB_CERTS_DIR or ~/.weightslab-certs. - if certs_dir_arg: - # Resolve to an absolute path so Windows Python, WSL bash and docker all - # agree on the location (a relative path resolves differently per shell). - certs_dir_arg = str(Path(certs_dir_arg).resolve()) - manager = CertAuthManager(certs_dir=certs_dir_arg, enable_auth=not no_auth) - logger.info(f"Using custom certs directory: {manager.certs_dir}") - else: - manager = CertAuthManager.from_env_or_default(enable_auth=not no_auth) - - if use_certs: - _ensure_certificates(manager, force_certs=force_certs) - else: - # Default: do NOT generate certs. Existing certs in WEIGHTSLAB_CERTS_DIR - # are still honored (file-presence is the single source of truth); this - # never deletes anything. Ensure the dir exists for the compose bind-mount. - logger.info("Launching WITHOUT cert generation (default; HTTP). " - "Pass --certs for secured HTTPS + gRPC auth.") - try: - manager.certs_dir.mkdir(parents=True, exist_ok=True) - except (OSError, AttributeError): - # AttributeError: certs_dir is a plain str (e.g. in unit tests). - try: - os.makedirs(str(manager.certs_dir), exist_ok=True) - except OSError: - logger.warning('Fail to create weightslab-certs directory!') - - # Single source of truth: TLS is on iff cert files exist in the dir. - secured = manager.has_valid_certs() - - # Set the correct certs path in the process env NOW so that _compose_cmd - # passes it to docker compose — this overrides any stale/malformed value - # in a leftover .env file and prevents the "too many colons" bind-mount error. - os.environ["WEIGHTSLAB_CERTS_DIR"] = _to_docker_host_path(manager.certs_dir) - - # Remove stale weightslab/weights_studio Docker resources that could prevent - # a clean launch. Scoped strictly to our own resources (see helper docstring). - if no_clean: - logger.info("Skipping stale Docker resource cleanup (--no-clean)") - else: - _clean_stale_docker_resources(frontend_repo) - - # Single source of truth: the deploy pipeline (build-and-deploy.sh + the - # compose `auto` logic) derives TLS/auth solely from cert-file presence in - # WEIGHTSLAB_CERTS_DIR. Strip any pre-set/derived TLS env so it cannot - # override that file-based decision. - _strip_derived_deploy_env() - - # Run bootstrap script to setup environment - bootstrap_script = str(_get_bootstrap_script()) - if Path(bootstrap_script).exists(): - logger.info("Running bootstrap script...") - logger.info(f"Bootstrap script path: {bootstrap_script}") - # Convert Windows path to Unix-style for bash - certs_dir_str = str(manager.certs_dir) - if _is_windows() and '\\' in certs_dir_str: - certs_dir_str = _convert_to_git_bash_path(certs_dir_str) - logger.info(f"Converted path to Unix-style: {certs_dir_str}") - - # Calculate WEIGHTSLAB_ROOT (parent of weightslab package) - weightslab_root = str(Path(__file__).parent) - if _is_windows() and '\\' in weightslab_root: - weightslab_root = _convert_to_git_bash_path(weightslab_root) - - # Docker Desktop (Windows) bind mounts need a host-native path - # (e.g. C:/Users/...), NOT the /mnt/c Unix path used for the bash - # script's own file checks. as_posix() yields C:/... on Windows and a - # normal /abs/path on Linux/macOS — correct for docker compose on every - # platform. The bash script writes this into .env for the compose mount. - certs_dir_host = _to_docker_host_path(manager.certs_dir) - - # Always pass the real certs dir so the compose bind-mount has a valid - # source. When there are no certs we add --unsecure, which forces - # Envoy/nginx to plaintext (mounted files, if any, are ignored) without - # leaving the mount source empty. - bootstrap_env_vars = { - 'WEIGHTSLAB_CERTS_DIR': certs_dir_str, - 'WEIGHTSLAB_CERTS_DIR_HOST': certs_dir_host, - 'WEIGHTSLAB_ROOT': weightslab_root - } - # On Windows the bootstrap runs in WSL/Git-Bash where `docker` usually - # isn't on PATH; the compose up below (run from Windows Python) handles - # the deploy. Tell the script to only write .env and skip docker ops so - # it doesn't emit a spurious "build failed" error. - if _is_windows(): - bootstrap_env_vars['WEIGHTSLAB_SKIP_DOCKER_OPS'] = '1' - script_args = [] - if not secured: - script_args.append('--unsecure') - if getattr(args, "dev", False): - script_args.append('--dev') - logger.info(f"WEIGHTSLAB_CERTS_DIR={bootstrap_env_vars['WEIGHTSLAB_CERTS_DIR']}") - logger.info(f"WEIGHTSLAB_CERTS_DIR_HOST={bootstrap_env_vars['WEIGHTSLAB_CERTS_DIR_HOST']}") - logger.info(f"WEIGHTSLAB_ROOT={bootstrap_env_vars['WEIGHTSLAB_ROOT']}") - exit_code = _run_shell_script(bootstrap_script, script_args, bootstrap_env_vars) - if exit_code != 0: - logger.warning(f"Bootstrap script exited with code {exit_code}, continuing anyway...") - else: - logger.warning(f"Bootstrap script not found: {bootstrap_script}") - - # docker compose gives an exported env var precedence over .env. Force the - # host-native certs path here (on every path, secured or not) so our - # compose call below bind-mounts a real folder — a stray /mnt/c or empty - # value would mount an empty/invalid dir and can crash Envoy. - os.environ["WEIGHTSLAB_CERTS_DIR"] = _to_docker_host_path(manager.certs_dir) - - _compose_cmd( - _get_compose_file(), - _get_envoy_config(), - ["up", "-d", "--pull", "always"], - ) - port = os.environ.get("VITE_PORT", "5173") - # HTTPS iff cert files actually exist in WEIGHTSLAB_CERTS_DIR — the same - # file-presence rule the deploy pipeline applies. - secured = manager.has_valid_certs() - protocol = "https" if secured else "http" - url = f"{protocol}://localhost:{port}" - logger.info(f"Weights Studio UI is running at: {url}") - # Auto-open the UI (best-effort; set WEIGHTSLAB_NO_BROWSER to disable). - _open_ui_in_browser(url) - if secured: - certs_dir_str = str(manager.certs_dir) - # Persist first (it logs its own lines), so the export/setx reminder - # below stays the FINAL output and can't be missed. - # Persist when a custom dir was given (it won't match the default, so the - # backend must be told where to look) or when the var isn't already set. - if certs_dir_arg or not _CERTS_DIR_IN_ORIGINAL_ENV: - _persist_certs_dir(certs_dir_str) - # The backend and any new shell must point at the same certs dir, or - # they'll mismatch the UI's TLS/auth. Keep this the last thing printed. - logger.warning("") - logger.warning(" ACTION REQUIRED — TLS is ON. Set WEIGHTSLAB_CERTS_DIR so the " - "training backend and new terminals use the same certificates:") - logger.warning(f" (bash) export WEIGHTSLAB_CERTS_DIR=\"{certs_dir_str}\"") - logger.warning(f" (Windows) setx WEIGHTSLAB_CERTS_DIR \"{certs_dir_str}\"") - else: - logger.info("UI is running UNSECURED (HTTP, no gRPC auth). " - "Re-run with `weightslab launch --certs` for TLS.") - - -def ui_secure_environment(args): - """`weightslab se`: create a certs directory with certs + gRPC token. - - The directory is the single source of truth — WEIGHTSLAB_CERTS_DIR is exported - for this process and the user is asked to export it globally. Everything else - (TLS on/off, auth on/off) is derived from the files in that directory by the - backend and the deploy pipeline, so this command does not set any other env. - """ - logger.info("Setting up secure environment...") - # Bundled .sh scripts ship without the execute bit (pip strips it); make them - # runnable so the cert-generation script below doesn't fail on "Permission denied". - _ensure_scripts_executable() - - force_certs = getattr(args, "force_certs", False) - no_auth = getattr(args, "no_auth", False) - certs_dir = getattr(args, "certs_dir", None) - if certs_dir: - # Absolute path so Windows Python, WSL bash and docker agree on location. - certs_dir = str(Path(certs_dir).resolve()) - - # Resolve the target directory first (no filesystem work in __init__), so we - # can point the generation scripts at it via WEIGHTSLAB_CERTS_DIR. - manager = CertAuthManager(certs_dir=certs_dir, enable_auth=not no_auth) - - exit_code = _generate_certs_with_fallback(force_certs=force_certs, certs_dir=manager.certs_dir) - if exit_code != 0: - logger.error("Certificate generation failed") - sys.exit(1) - - manager.certs_dir.mkdir(parents=True, exist_ok=True) - manager.get_or_create_auth_token() - - # Export ONLY the single source of truth for this process. - os.environ["WEIGHTSLAB_CERTS_DIR"] = str(manager.certs_dir) - - logger.info(" Certificates generated successfully") - logger.info(" gRPC auth token created") - logger.info(f" Certs and token stored in: {manager.certs_dir}") - logger.info(f" WEIGHTSLAB_CERTS_DIR exported for this process: {manager.certs_dir}") - logger.info("Then launch the secured UI with: weightslab launch --certs") - # Keep this the FINAL output so the user can't miss the action they must take. - logger.warning("") - logger.warning(" ACTION REQUIRED — set WEIGHTSLAB_CERTS_DIR globally so new shells " - "and the training backend find these certs (single source of truth):") - logger.warning(f" (bash) echo 'export WEIGHTSLAB_CERTS_DIR=\"{manager.certs_dir}\"' >> ~/.bashrc && source ~/.bashrc") - logger.warning(f" (Windows) setx WEIGHTSLAB_CERTS_DIR \"{manager.certs_dir}\"") - - -# Bundled PyTorch examples, keyed by the CLI flag (e.g. --cls -> wl-classification). -# Each value is (directory name under examples/PyTorch, human-readable label). -# kind -> (dir_name, label, category) where category is the examples/ subfolder. -_EXAMPLES = { - "cls": ("wl-classification", "classification", "PyTorch"), - "seg": ("wl-segmentation", "segmentation", "PyTorch"), - "det": ("wl-detection", "detection", "PyTorch"), - "clus": ("wl-clustering", "clustering", "PyTorch"), - "gen": ("wl-generation", "generation", "PyTorch"), - "3d_det": ("wl-3d-lidar-detection", "3D LiDAR detection", "Usecases"), - "2d_det": ("wl-2d-lidar-detection", "2D LiDAR detection", "Usecases"), -} -_DEFAULT_EXAMPLE = "cls" - - -def _get_example_dir(name: str = "wl-classification", category: str = "PyTorch") -> Path: - """Path to a bundled example directory (under examples//).""" - return Path(__file__).parent / 'examples' / category / name - - -def _install_example_requirements(example_dir: Path) -> None: - """Install an example's requirements non-interactively, if a file is present. - - Looks for requirements.txt (then requirements.in) in the example directory and - runs `pip install -r` with the current interpreter and `--no-input` so it never - prompts. Non-fatal: a failure is logged and the example is still attempted, so a - transient install hiccup doesn't block a run where deps are already satisfied. - """ - for fname in ("requirements.txt", "requirements.in"): - req = example_dir / fname - if not req.exists(): - continue - logger.info(f"Installing example requirements (non-interactive): {req}") - try: - subprocess.run( - [sys.executable, "-m", "pip", "install", "-r", str(req), - "--no-input", "--disable-pip-version-check"], - check=True, - ) - except subprocess.CalledProcessError as exc: - logger.warning( - f"Failed to install requirements ({req}): {exc}. " - "Continuing — the example may still run if deps are already installed." - ) - return # only the first matching requirements file is used - - -def example_start(args): - """`weightslab start example [--cls|--seg|--clus|--gen]`: run a bundled example. - - Defaults to the classification (cls) example. First installs the example's - requirements (if a requirements file is present) without prompting, then runs - its main.py with the current Python interpreter from its own directory so it - resolves its sibling config.yaml. Runs in the foreground (serves until Ctrl+C). - """ - kind = getattr(args, "example_kind", None) or _DEFAULT_EXAMPLE - dir_name, label, category = _EXAMPLES.get(kind, _EXAMPLES[_DEFAULT_EXAMPLE]) - - example_dir = _get_example_dir(dir_name, category) - main_py = example_dir / "main.py" - if not main_py.exists(): - logger.error(f"{label.capitalize()} example not found: {main_py}") - sys.exit(1) - - # Install the example's own requirements first, without any interaction. - _install_example_requirements(example_dir) - - logger.info(f"Starting the WeightsLab {label} ({kind}) example...") - logger.info(f" {main_py}") - logger.info("In another terminal, launch the UI with: weightslab launch") - logger.info(f"Then open http://localhost:5173 — stop the example with Ctrl+C.") - if not _CERTS_DIR_IN_ORIGINAL_ENV: - manager = CertAuthManager.from_env_or_default() - if manager.has_valid_certs(): - logger.warning( - "WEIGHTSLAB_CERTS_DIR is not set in your shell environment. " - "TLS will work this session (certs found at default location) " - "but may not persist across terminals." - ) - _persist_certs_dir(str(manager.certs_dir)) - try: - env = os.environ.copy() - env['WEIGHTSLAB_SUPPRESS_BANNER'] = '1' - result = subprocess.run([sys.executable, str(main_py)], cwd=str(example_dir), env=env) - except KeyboardInterrupt: - logger.info("Example stopped.") - return - if result.returncode != 0: - sys.exit(result.returncode) - - -def cli_connect(args): - """`weightslab cli [--port N] [--host H]`: open an interactive terminal - connected to a currently-running experiment's CLI server. - - With no --port, auto-discovers the running experiment (the backend advertises - its actual port on startup). Pass --port to target a specific server. - """ - try: - import weightslab.backend.cli as cli_backend - except Exception as exc: - logger.error(f"Could not load the WeightsLab CLI client: {exc}") - sys.exit(1) - - port = getattr(args, "port", None) - host = getattr(args, "host", None) - exit_code = cli_backend.cli_connect(cli_port=port, cli_host=host) - sys.exit(exit_code) - - -def _add_example_kind_flags(p: argparse.ArgumentParser) -> None: - """Attach the mutually-exclusive example-kind flags (default: classification).""" - group = p.add_mutually_exclusive_group() - group.add_argument("--cls", action="store_const", dest="example_kind", const="cls", - help="Run the classification example (default)") - group.add_argument("--seg", action="store_const", dest="example_kind", const="seg", - help="Run the segmentation example") - group.add_argument("--det", action="store_const", dest="example_kind", const="det", - help="Run the detection example") - group.add_argument("--clus", action="store_const", dest="example_kind", const="clus", - help="Run the clustering example") - group.add_argument("--gen", action="store_const", dest="example_kind", const="gen", - help="Run the generation example") - group.add_argument("--3d_det", action="store_const", dest="example_kind", const="3d_det", - help="Run the 3D LiDAR point-cloud detection example") - group.add_argument("--2d_det", action="store_const", dest="example_kind", const="2d_det", - help="Run the 2D LiDAR point-cloud detection example") - p.set_defaults(example_kind=_DEFAULT_EXAMPLE) - - -def _add_launch_flags(p, with_certs_dir: bool = True) -> None: - """Shared flags for every UI-launch entry point: `weightslab launch`, the bare - `weightslab start`, and the legacy `weightslab ui launch`.""" - p.add_argument('--certs', action='store_true', - help='Generate (if missing) and use TLS certs + gRPC auth token (secured HTTPS). Default: unsecured HTTP.') - p.add_argument('-i', '--image', default=None, - help='Frontend image repo to run/pull (default: graybx/weightslab). e.g. guillaumep2705/weightslab') - p.add_argument('-v', '--version', default=None, - help='Frontend image tag/version to pull (default: latest). e.g. v1.2.3') - if with_certs_dir: - p.add_argument('certs_dir', nargs='?', default=None, - help='Custom directory for certs/token (default: $WEIGHTSLAB_CERTS_DIR or ~/.weightslab-certs)') - - -def _build_parser() -> argparse.ArgumentParser: - """Build the top-level argument parser (banner + detailed command reference). - - The CLI is intentionally minimal — exactly these commands: - weightslab --help | -h | help - weightslab se [--force-certs] - weightslab launch [--certs] (also: bare `weightslab start`) - weightslab start example [--cls|--seg|--det|--clus|--gen|--3d_det|--2d_det] - weightslab cli [--port PORT] [--host HOST] - weightslab ui launch [--certs] (legacy alias of `launch`) - """ - parser = argparse.ArgumentParser( - prog="weightslab", - description=_DESCRIPTION, - epilog=_EPILOG, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - # metavar lists only the documented commands; the `example` alias is accepted - # but intentionally omitted here (and help=SUPPRESS'd below) so it stays hidden. - sub = parser.add_subparsers(dest="command", metavar="{se,launch,start,ui,cli,tunnel,help}") - - # weightslab se [--force-certs] [certs_dir] - se_parser = sub.add_parser("se", help="Set up the secure environment (TLS certs + gRPC auth token)") - se_parser.add_argument('--force-certs', action='store_true', help='Regenerate certificates even if they already exist') - se_parser.add_argument('certs_dir', nargs='?', default=None, - help='Custom directory for certs/token (default: $WEIGHTSLAB_CERTS_DIR or ~/.weightslab-certs)') - - # weightslab launch [--certs] [-i IMAGE] [-v VERSION] [certs_dir] - # Primary command to bring up the Weights Studio UI (also reachable as the - # bare `weightslab start`, and the legacy `weightslab ui launch`). - launch_parser = sub.add_parser( - "launch", - help="Clean stale Docker state, launch the UI, and open it in the browser " - "(unsecured by default; --certs for TLS)") - _add_launch_flags(launch_parser) - - # weightslab ui launch [...] — LEGACY alias, kept working for backward - # compatibility (existing notebooks/docs). Prefer `weightslab launch`. - # (help=SUPPRESS hides the `ui` command from the top-level --help listing.) - ui_parser = sub.add_parser("ui", help=argparse.SUPPRESS) - ui_sub = ui_parser.add_subparsers(dest="action") - launch_ui_parser = ui_sub.add_parser( - "launch", help="(legacy) alias for `weightslab launch`") - _add_launch_flags(launch_ui_parser) - - # weightslab cli [--port N] [--host H] - cli_parser = sub.add_parser( - "cli", help="Open an interactive terminal connected to the running experiment") - cli_parser.add_argument('--port', type=int, default=None, - help='CLI server port (default: auto-discover the running experiment)') - cli_parser.add_argument('--host', default=None, - help='CLI server host (default: localhost)') - - # weightslab tunnel ENDPOINT [--listen-port N] [--listen-host H] [--remote-port N] - tunnel_parser = sub.add_parser( - "tunnel", - help="Forward a remote gRPC backend (e.g. a Colab run via `ngrok tcp 50051`) " - "to a local port for the UI") - tunnel_parser.add_argument( - 'endpoint', nargs='?', default=None, - help="Remote backend endpoint host:port (e.g. bore.pub:12345). " - "A tcp:// prefix is accepted. Default: $WEIGHTSLAB_TUNNEL_ENDPOINT.") - tunnel_parser.add_argument( - '--listen-port', '-p', type=int, default=DEFAULT_LISTEN_PORT, - help=f"Local port to expose (default: {DEFAULT_LISTEN_PORT} — the port the UI's Envoy dials)") - tunnel_parser.add_argument( - '--listen-host', default=None, - help="Local interface to bind (default: 127.0.0.1 on Windows/macOS, 0.0.0.0 on Linux)") - tunnel_parser.add_argument( - '--remote-port', type=int, default=None, - help="Remote port, if not included in ENDPOINT") - - # weightslab start [--certs ...] → launch the UI (alias of `launch`) - # weightslab start example [--cls|...] → run a bundled PyTorch example - start_parser = sub.add_parser( - "start", - help="Launch the Weights Studio UI (bare `start`), or `start example` to run a demo") - _add_launch_flags(start_parser, with_certs_dir=False) - start_sub = start_parser.add_subparsers(dest="start_target") - example_parser = start_sub.add_parser( - "example", help="Start a bundled PyTorch example (default: classification)") - _add_example_kind_flags(example_parser) - - # Tolerate the swapped order: `weightslab example start [flags]` (and bare - # `weightslab example`) behave exactly like `weightslab start example`. Hidden - # from --help on purpose (argparse.SUPPRESS) — it's a forgiving fallback, not a - # documented command. - example_alias = sub.add_parser("example", help=argparse.SUPPRESS) - example_alias_sub = example_alias.add_subparsers(dest="example_action") - example_alias_start = example_alias_sub.add_parser( - "start", help="Start a bundled PyTorch example (default: classification)") - _add_example_kind_flags(example_alias_start) - - sub.add_parser("help", help="Show this help message") - - return parser, ui_parser, start_parser, cli_parser - - -def main(): - parser, ui_parser, start_parser, cli_parser = _build_parser() - args = parser.parse_args() - - if args.command == "help" or args.command is None: - parser.print_help() - elif args.command == "cli": - cli_connect(args) - elif args.command == "tunnel": - from weightslab.tunnel import tunnel_connect - tunnel_connect(args) - elif args.command == "se": - ui_secure_environment(args) - elif args.command == "launch": - # Primary UI-launch command. - ui_launch(args) - elif args.command == "ui": - # Legacy: `weightslab ui launch` still works (prefer `weightslab launch`). - if getattr(args, "action", None) == "launch": - ui_launch(args) - else: - ui_parser.print_help() - elif args.command == "start": - # `weightslab start example ...` runs a demo; bare `weightslab start` - # launches the UI (alias of `weightslab launch`). - if getattr(args, "start_target", None) == "example": - example_start(args) - else: - ui_launch(args) - elif args.command == "example": - # Alias for `start example` — tolerate the swapped subcommand order - # (`weightslab example start [flags]`) and the bare `weightslab example`. - example_start(args) - else: - parser.print_help() - - -if __name__ == "__main__": - main()